Automatic inference for semantic legends for marker/scatter size

ultraplot/axes/base.py

@@ -3835,6 +3835,27 @@ def sizelegend(self, levels, **kwargs):
             Treat ``levels`` as marker areas (``True``, default) or
             diameters (``False``). Areas are converted with
             ``ms = sqrt(level) * scale``. Falls back to :rc:`legend.size.area`.
+        values : array-like, optional
+            Full scatter-size data used to infer the scaling range for
+            ``levels``. When provided, or when any of ``vmin``, ``vmax``,
+            ``smin``, ``smax``, ``area_size``, or ``absolute_size`` are
+            provided, ``levels`` are transformed with the same size scaling
+            rules used by :meth:`~ultraplot.axes.PlotAxes.scatter` while
+            labels remain based on the original ``levels``. When these options
+            are omitted and a compatible UltraPlot scatter artist already exists
+            on the axes, its size scale is inferred automatically.
+        vmin, vmax : float, optional
+            Explicit data range for scatter-style size scaling. Defaults to the
+            finite range of ``values`` or ``levels``.
+        smin, smax : float, optional
+            Minimum and maximum scaled marker sizes, with the same meaning as
+            in :meth:`~ultraplot.axes.PlotAxes.scatter`.
+        area_size, absolute_size : bool, optional
+            Scatter-style size scaling switches. Defaults match
+            :meth:`~ultraplot.axes.PlotAxes.scatter` when scatter-style scaling
+            is active. When scatter-style scaling is active and ``area_size`` is
+            omitted, an explicit ``area=False`` is treated like
+            ``area_size=False``.
         scale : float, optional
             Multiplier applied after area/diameter conversion.
             Falls back to :rc:`legend.size.scale`.

ultraplot/axes/plot.py

@@ -5618,7 +5618,12 @@ def _apply_scatter(self, xs, ys, ss, cc, *, vert=True, **kwargs):
         kw, extents = self._inbounds_extent(inbounds=inbounds, **kw)
         xs, ys, kw = self._parse_1d_args(xs, ys, vert=vert, autoreverse=False, **kw)
         ys, kw = inputs._dist_reduce(ys, **kw)
-        ss, kw = self._parse_markersize(ss, **kw)  # parse 's'
+        size_scale = {
+            key: kw.get(key, None)
+            for key in ("smin", "smax", "area_size", "absolute_size")
+        }
+        ss_source = inputs._to_numpy_array(ss) if ss is not None else None
+        ss, kw = self._parse_markersize(ss_source, **kw)  # parse 's'
 
         # Only parse color if explicitly provided
         infer_rgb = True
@@ -5645,7 +5650,9 @@ def _apply_scatter(self, xs, ys, ss, cc, *, vert=True, **kwargs):
         # Create the cycler object by manually cycling and sanitzing the inputs
         guide_kw = _pop_params(kw, self._update_guide)
         objs = []
-        for _, n, x, y, s, c, kw in self._iter_arg_cols(xs, ys, ss, cc, **kw):
+        for _, n, x, y, s, c, s_source, kw in self._iter_arg_cols(
+            xs, ys, ss, cc, ss_source, **kw
+        ):
             # Cycle s and c as they are in cycle_manually
             # Note: they could be None
             kw["s"], kw["c"] = s, c
@@ -5655,6 +5662,11 @@ def _apply_scatter(self, xs, ys, ss, cc, *, vert=True, **kwargs):
             if not vert:
                 x, y = y, x
             obj = self._call_native("scatter", x, y, **kw)
+            if s_source is not None:
+                obj._ultraplot_size_scale = {
+                    "values": inputs._to_numpy_array(s_source),
+                    **size_scale,
+                }
             self._inbounds_xylim(extents, x, y)
             objs.append((*eb, *es, obj) if eb or es else obj)
 

ultraplot/figure.py

@@ -473,6 +473,10 @@
 ----------
 levels
     Numeric levels used to generate marker-size entries.
+values, vmin, vmax, smin, smax, area_size, absolute_size
+    Optional scatter-style size scaling controls forwarded to
+    `~ultraplot.axes.Axes.sizelegend`. When omitted, a compatible UltraPlot
+    scatter artist can be used to infer the size scale automatically.
 
 Other parameters
 ----------------
@@ -3336,6 +3340,13 @@ def sizelegend(
         color=None,
         marker=None,
         area=None,
+        values=None,
+        vmin=None,
+        vmax=None,
+        smin=None,
+        smax=None,
+        area_size=None,
+        absolute_size=None,
         scale=None,
         minsize=None,
         fmt=None,
@@ -3359,6 +3370,13 @@ def sizelegend(
             color=color,
             marker=marker,
             area=area,
+            values=values,
+            vmin=vmin,
+            vmax=vmax,
+            smin=smin,
+            smax=smax,
+            area_size=area_size,
+            absolute_size=absolute_size,
             scale=scale,
             minsize=minsize,
             fmt=fmt,

ultraplot/legend.py

@@ -16,7 +16,7 @@
 from matplotlib.markers import MarkerStyle
 
 from .config import rc
-from .internals import _not_none, _pop_props, docstring, guides, rcsetup
+from .internals import _not_none, _pop_props, docstring, guides, inputs, rcsetup
 from .utils import _fontsize_to_pt, units
 
 try:
@@ -1396,6 +1396,7 @@ def _entry_legend_entries(
 def _size_legend_entries(
     levels: Iterable[float],
     *,
+    label_values=None,
     labels=None,
     color="0.35",
     marker="o",
@@ -1415,16 +1416,21 @@ def _size_legend_entries(
     values = np.asarray(list(levels), dtype=float)
     if values.size == 0:
         return [], []
+    label_values = (
+        values if label_values is None else np.asarray(label_values, dtype=float)
+    )
+    if label_values.size != values.size:
+        raise ValueError("sizelegend label values must have the same length as levels.")
     if area:
         ms = np.sqrt(np.clip(values, 0, None))
     else:
         ms = np.abs(values)
     ms = np.maximum(ms * scale, minsize)
     if labels is None:
-        label_list = [_format_label(value, fmt) for value in values]
+        label_list = [_format_label(value, fmt) for value in label_values]
     elif isinstance(labels, Mapping):
         label_list = []
-        for value in values:
+        for value in label_values:
             key = float(value)
             if key not in labels:
                 raise ValueError(
@@ -1444,13 +1450,13 @@ def _size_legend_entries(
     }
     base_styles.update(entry_kwargs)
     handles = []
-    for idx, (value, label, size) in enumerate(zip(values, label_list, ms)):
-        styles = _resolve_style_values(base_styles, float(value), idx)
+    for idx, (label_value, label, size) in enumerate(zip(label_values, label_list, ms)):
+        styles = _resolve_style_values(base_styles, float(label_value), idx)
         color_value = _style_lookup(
-            color, float(value), idx, default="0.35", prop="color"
+            color, float(label_value), idx, default="0.35", prop="color"
         )
         marker_value = _style_lookup(
-            marker, float(value), idx, default="o", prop="marker"
+            marker, float(label_value), idx, default="o", prop="marker"
         )
         line_value = bool(styles.pop("line", False))
         if line_value and marker_value in ("", None):
@@ -1470,6 +1476,70 @@ def _size_legend_entries(
     return handles, label_list
 
 
+def _scale_size_legend_values(
+    values,
+    *,
+    source=None,
+    vmin=None,
+    vmax=None,
+    smin=None,
+    smax=None,
+    area_size=True,
+    absolute_size=None,
+):
+    """
+    Transform semantic size values with the same rules used by scatter().
+    """
+    values = np.asarray(values, dtype=float)
+    source = values if source is None else inputs._to_numpy_array(source)
+    if absolute_size is None:
+        absolute_size = np.size(source) == 1
+    if not absolute_size or smin is not None or smax is not None:
+        smin = _not_none(smin, 1)
+        smax = _not_none(smax, rc["lines.markersize"] ** (1, 2)[area_size])
+        dmin, dmax = inputs._safe_range(source)
+        dmin = _not_none(vmin, dmin)
+        dmax = _not_none(vmax, dmax)
+        if dmin is not None and dmax is not None and dmin != dmax:
+            values = smin + (smax - smin) * (values - dmin) / (dmax - dmin)
+    areas = values ** (2, 1)[area_size]
+    return np.sqrt(np.clip(areas, 0, None))
+
+
+def _infer_size_legend_scale(axes, values):
+    """
+    Infer scatter-style size scaling from the latest compatible scatter artist.
+    """
+    values = np.asarray(values, dtype=float)
+    finite_values = values[np.isfinite(values)]
+    candidates = []
+    for artist in getattr(axes, "collections", ()):
+        metadata = getattr(artist, "_ultraplot_size_scale", None)
+        if metadata and metadata.get("values", None) is not None:
+            candidates.append(metadata)
+    if not candidates:
+        return None
+
+    def _covers(metadata):
+        if finite_values.size == 0:
+            return False
+        try:
+            dmin, dmax = inputs._safe_range(metadata["values"])
+            return (
+                dmin is not None
+                and dmax is not None
+                and np.nanmin(finite_values) >= dmin
+                and np.nanmax(finite_values) <= dmax
+            )
+        except Exception:
+            return False
+
+    compatible = [metadata for metadata in candidates if _covers(metadata)]
+    if not compatible:
+        return None
+    return dict(compatible[-1])
+
+
 def _num_legend_entries(
     levels=None,
     *,
@@ -1901,6 +1971,13 @@ def sizelegend(
         color=None,
         marker=None,
         area: Optional[bool] = None,
+        values=None,
+        vmin: Optional[float] = None,
+        vmax: Optional[float] = None,
+        smin: Optional[float] = None,
+        smax: Optional[float] = None,
+        area_size: Optional[bool] = None,
+        absolute_size: Optional[bool] = None,
         scale: Optional[float] = None,
         minsize: Optional[float] = None,
         fmt=None,
@@ -1912,15 +1989,32 @@ def sizelegend(
         Build size legend entries and optionally draw a legend.
         Public docs live on :meth:`Axes.sizelegend`.
         """
+        area_input = area
+        levels = np.asarray(list(levels), dtype=float)
+        explicit_scaled = any(
+            value is not None
+            for value in (values, vmin, vmax, smin, smax, area_size, absolute_size)
+        )
+        auto_scale = None
+        if not explicit_scaled and area_input is None:
+            auto_scale = _infer_size_legend_scale(self.axes, levels)
+        if auto_scale is not None:
+            values = auto_scale.get("values", values)
+            smin = auto_scale.get("smin", smin)
+            smax = auto_scale.get("smax", smax)
+            area_size = auto_scale.get("area_size", area_size)
+            absolute_size = auto_scale.get("absolute_size", absolute_size)
+        scaled = explicit_scaled or auto_scale is not None
         area = _not_none(area, rc["legend.size.area"])
+        area_size = _not_none(area_size, area_input if scaled else None, True)
         styles = {}
         if handle_kw:
             styles.update(_pop_entry_props(handle_kw))
         styles.update(_pop_entry_props(kwargs))
         color = _not_none(color, styles.pop("color", None), rc["legend.size.color"])
         marker = _not_none(marker, styles.pop("marker", None), rc["legend.size.marker"])
         scale = _not_none(scale, rc["legend.size.scale"])
-        minsize = _not_none(minsize, rc["legend.size.minsize"])
+        minsize = _not_none(minsize, 0.0 if scaled else rc["legend.size.minsize"])
         fmt = _not_none(fmt, rc["legend.size.format"])
         alpha = _not_none(styles.pop("alpha", None), rc["legend.size.alpha"])
         markeredgecolor = _not_none(
@@ -1930,8 +2024,23 @@ def sizelegend(
             styles.pop("markeredgewidth", None), rc["legend.size.markeredgewidth"]
         )
         markerfacecolor = _not_none(styles.pop("markerfacecolor", None), None)
+        if scaled:
+            visual_sizes = _scale_size_legend_values(
+                levels,
+                source=values,
+                vmin=vmin,
+                vmax=vmax,
+                smin=smin,
+                smax=smax,
+                area_size=area_size,
+                absolute_size=absolute_size,
+            )
+            area = False
+        else:
+            visual_sizes = levels
         handles, labels = _size_legend_entries(
-            levels,
+            visual_sizes,
+            label_values=levels if scaled else None,
             labels=labels,
             color=color,
             marker=marker,