diff --git a/ultraplot/axes/base.py b/ultraplot/axes/base.py index 8174d870b..f8373914f 100644 --- a/ultraplot/axes/base.py +++ b/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`. diff --git a/ultraplot/axes/plot.py b/ultraplot/axes/plot.py index 83eba0d93..48254187e 100644 --- a/ultraplot/axes/plot.py +++ b/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) diff --git a/ultraplot/figure.py b/ultraplot/figure.py index 5493755ab..33125de67 100644 --- a/ultraplot/figure.py +++ b/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, diff --git a/ultraplot/legend.py b/ultraplot/legend.py index c11d4e14b..4faa47f26 100644 --- a/ultraplot/legend.py +++ b/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,7 +1989,24 @@ 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)) @@ -1920,7 +2014,7 @@ def sizelegend( 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, diff --git a/ultraplot/tests/test_legend.py b/ultraplot/tests/test_legend.py index a41b2c1db..a382d0e8b 100644 --- a/ultraplot/tests/test_legend.py +++ b/ultraplot/tests/test_legend.py @@ -538,6 +538,131 @@ def test_sizelegend_marker_size_overrides_use_semantic_size_rules(): uplt.close(fig) +def test_sizelegend_area_levels_match_absolute_scatter_sizes(): + fig, ax = uplt.subplots() + scatter_size = 50 + scatter = ax.scatter( + [0], + [0], + marker=".", + s=scatter_size, + absolute_size=True, + ) + handles, labels = ax.sizelegend( + [scatter_size], + marker=".", + add=False, + ) + + assert labels == ["50"] + assert handles[0].get_marker() == "." + assert handles[0].get_markersize() == pytest.approx(scatter.get_sizes()[0] ** 0.5) + uplt.close(fig) + + +def test_sizelegend_values_follow_scatter_scaled_sizes(): + fig, ax = uplt.subplots() + scatter_values = np.array([10.0, 30.0, 50.0, 70.0, 90.0]) + legend_levels = scatter_values[1:-1] + scatter = ax.scatter( + np.arange(scatter_values.size), + np.zeros(scatter_values.size), + s=scatter_values, + smin=4, + smax=100, + absolute_size=False, + ) + handles, labels = ax.sizelegend( + legend_levels, + values=scatter_values, + smin=4, + smax=100, + marker=".", + add=False, + ) + + assert labels == ["30", "50", "70"] + assert [handle.get_marker() for handle in handles] == [".", ".", "."] + assert [handle.get_markersize() for handle in handles] == pytest.approx( + np.sqrt(scatter.get_sizes()[1:-1]) + ) + uplt.close(fig) + + +def test_sizelegend_infers_scatter_scaled_sizes_by_default(): + fig, ax = uplt.subplots() + scatter_values = np.array([10.0, 30.0, 50.0, 70.0, 90.0]) + legend_levels = scatter_values[1:-1] + scatter = ax.scatter( + np.arange(scatter_values.size), + np.zeros(scatter_values.size), + s=scatter_values, + smin=4, + smax=100, + absolute_size=False, + ) + handles, labels = ax.sizelegend( + legend_levels, + marker=".", + add=False, + ) + + assert labels == ["30", "50", "70"] + assert [handle.get_markersize() for handle in handles] == pytest.approx( + np.sqrt(scatter.get_sizes()[1:-1]) + ) + uplt.close(fig) + + +def test_sizelegend_explicit_area_false_skips_scatter_inference(): + fig, ax = uplt.subplots() + ax.scatter([0, 1, 2], [0, 0, 0], s=[10, 50, 90], smin=4, smax=100) + handles, labels = ax.sizelegend([10], area=False, add=False) + + assert labels == ["10"] + assert handles[0].get_markersize() == pytest.approx(10) + uplt.close(fig) + + +def test_sizelegend_skips_incompatible_scatter_inference(): + fig, ax = uplt.subplots() + ax.scatter([0, 1, 2], [0, 0, 0], s=[10, 50, 90], smin=4, smax=100) + handles, labels = ax.sizelegend([200], add=False) + + assert labels == ["200"] + assert handles[0].get_markersize() == pytest.approx(np.sqrt(200)) + uplt.close(fig) + + +def test_sizelegend_area_false_follows_scatter_radius_scaling(): + fig, ax = uplt.subplots() + scatter_values = np.array([10.0, 50.0, 90.0]) + scatter = ax.scatter( + np.arange(scatter_values.size), + np.zeros(scatter_values.size), + s=scatter_values, + smin=2, + smax=10, + area_size=False, + absolute_size=False, + ) + handles, labels = ax.sizelegend( + scatter_values, + values=scatter_values, + area=False, + smin=2, + smax=10, + marker=".", + add=False, + ) + + assert labels == ["10", "50", "90"] + assert [handle.get_markersize() for handle in handles] == pytest.approx( + np.sqrt(scatter.get_sizes()) + ) + uplt.close(fig) + + def test_legend_single_point_plot_matches_marker_only_artist(): fig, ax = uplt.subplots() ax.plot([0], [0], marker="o", label="point")