Support custom labels in sizelegend (#629)

* Add custom labels to sizelegend

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: cvanelteren

docs/colorbars_legends.py

@@ -479,18 +479,154 @@
 # standalone semantic keys (categories, size scales, color levels, or geometry types).
 # UltraPlot provides helper methods that build these entries directly:
 #
+# * :meth:`~ultraplot.axes.Axes.entrylegend`
 # * :meth:`~ultraplot.axes.Axes.catlegend`
 # * :meth:`~ultraplot.axes.Axes.sizelegend`
 # * :meth:`~ultraplot.axes.Axes.numlegend`
 # * :meth:`~ultraplot.axes.Axes.geolegend`
+#
+# These helpers are useful whenever the legend should describe an encoding rather than
+# mirror artists that already happen to be drawn. In practice there are two distinct
+# workflows:
+#
+# * Use :meth:`~ultraplot.axes.Axes.legend` when you already have artists and want to
+#   reuse their labels or lightly restyle the legend handles.
+# * Use the semantic helpers when you want to define the legend from meaning-first
+#   inputs such as categories, numeric size levels, numeric color levels, or geometry
+#   types, even if no matching exemplar artist exists on the axes.
+#
+# Choosing a helper
+# ~~~~~~~~~~~~~~~~~
+#
+# * :meth:`~ultraplot.axes.Axes.entrylegend` is the most general helper. Use it when
+#   you want explicit labels, mixed line and marker entries, or fully custom legend
+#   rows that are not easily described by a single category or numeric scale.
+# * :meth:`~ultraplot.axes.Axes.catlegend` is for discrete categories mapped to colors,
+#   markers, and optional line styles. Labels come from the category names.
+# * :meth:`~ultraplot.axes.Axes.sizelegend` is for marker-size semantics. Labels are
+#   derived from the numeric levels by default, can be formatted with ``fmt=``, and
+#   can now be overridden directly with ``labels=[...]`` or ``labels={level: label}``.
+# * :meth:`~ultraplot.axes.Axes.numlegend` is for numeric color encodings rendered as
+#   discrete patches without requiring a pre-existing mappable.
+# * :meth:`~ultraplot.axes.Axes.geolegend` is for shapes and map-like semantics. It can
+#   mix named symbols, Shapely geometries, and country shorthands in one legend.
+#
+# The helpers are intentionally composable. Each one accepts ``add=False`` and returns
+# ``(handles, labels)`` so you can merge semantic sections and pass the result through
+# :meth:`~ultraplot.axes.Axes.legend` yourself.
+#
+# .. code-block:: python
+#
+#    # Reuse plotted artists when they already exist.
+#    hs = ax.plot(data, labels=["control", "treatment"])
+#    ax.legend(hs, loc="r")
+#
+#    # Build a category key without plotting one exemplar artist per category.
+#    ax.catlegend(
+#        ["Control", "Treatment"],
+#        colors={"Control": "blue7", "Treatment": "red7"},
+#        markers={"Control": "o", "Treatment": "^"},
+#        loc="r",
+#    )
+#
+#    # Build fully custom entries with explicit labels and mixed semantics.
+#    ax.entrylegend(
+#        [
+#            {
+#                "label": "Observed samples",
+#                "line": False,
+#                "marker": "o",
+#                "markersize": 8,
+#                "markerfacecolor": "blue7",
+#                "markeredgecolor": "black",
+#            },
+#            {
+#                "label": "Model fit",
+#                "line": True,
+#                "color": "black",
+#                "linewidth": 2.5,
+#                "linestyle": "--",
+#            },
+#        ],
+#        title="Entry styles",
+#        loc="l",
+#    )
+#
+#    # Size legends can format labels automatically or accept explicit labels.
+#    ax.sizelegend(
+#        [10, 50, 200],
+#        labels=["Small", "Medium", "Large"],
+#        title="Population",
+#        loc="ur",
+#    )
+#
+#    # Numeric color legends are discrete color keys decoupled from a mappable.
+#    ax.numlegend(vmin=0, vmax=1, n=5, cmap="viko", fmt="{:.2f}", loc="ll")
+#
+#    # Geometry legends can mix named shapes, Shapely geometries, and country codes.
+#    ax.geolegend([("Triangle", "triangle"), ("Australia", "country:AU")], loc="r")
+#
+# .. code-block:: python
+#
+#    # Compose multiple semantic helpers into one legend.
+#    size_handles, size_labels = ax.sizelegend(
+#        [10, 50, 200],
+#        labels=["Small", "Medium", "Large"],
+#        add=False,
+#    )
+#    entry_handles, entry_labels = ax.entrylegend(
+#        [
+#            {
+#                "label": "Observed",
+#                "line": False,
+#                "marker": "o",
+#                "markerfacecolor": "blue7",
+#            },
+#            {
+#                "label": "Fit",
+#                "line": True,
+#                "color": "black",
+#                "linewidth": 2,
+#            },
+#        ],
+#        add=False,
+#    )
+#    ax.legend(
+#        size_handles + entry_handles,
+#        size_labels + entry_labels,
+#        loc="r",
+#        title="Combined semantic key",
+#    )
 
 # %%
 import cartopy.crs as ccrs
 import shapely.geometry as sg
 
-fig, ax = uplt.subplots(refwidth=4.2)
+fig, ax = uplt.subplots(refwidth=5.0)
 ax.format(title="Semantic legend helpers", grid=False)
 
+ax.entrylegend(
+    [
+        {
+            "label": "Observed samples",
+            "line": False,
+            "marker": "o",
+            "markersize": 8,
+            "markerfacecolor": "blue7",
+            "markeredgecolor": "black",
+        },
+        {
+            "label": "Model fit",
+            "line": True,
+            "color": "black",
+            "linewidth": 2.5,
+            "linestyle": "--",
+        },
+    ],
+    loc="l",
+    title="Entry styles",
+    frameon=False,
+)
 ax.catlegend(
     ["A", "B", "C"],
     colors={"A": "red7", "B": "green7", "C": "blue7"},
@@ -500,6 +636,7 @@
 )
 ax.sizelegend(
     [10, 50, 200],
+    labels=["Small", "Medium", "Large"],
     loc="upper right",
     title="Population",
     ncols=1,

docs/examples/legends_colorbars/03_semantic_legends.py

@@ -7,10 +7,11 @@
 Why UltraPlot here?
 -------------------
 UltraPlot adds semantic legend helpers directly on axes:
-``catlegend``, ``sizelegend``, ``numlegend``, and ``geolegend``.
-These are useful when you want legend meaning decoupled from plotted handles.
+``entrylegend``, ``catlegend``, ``sizelegend``, ``numlegend``, and ``geolegend``.
+These are useful when you want legend meaning decoupled from plotted handles, or
+when you want a standalone semantic key that describes an encoding directly.
 
-Key functions: :py:meth:`ultraplot.axes.Axes.catlegend`, :py:meth:`ultraplot.axes.Axes.sizelegend`, :py:meth:`ultraplot.axes.Axes.numlegend`, :py:meth:`ultraplot.axes.Axes.geolegend`.
+Key functions: :py:meth:`ultraplot.axes.Axes.entrylegend`, :py:meth:`ultraplot.axes.Axes.catlegend`, :py:meth:`ultraplot.axes.Axes.sizelegend`, :py:meth:`ultraplot.axes.Axes.numlegend`, :py:meth:`ultraplot.axes.Axes.geolegend`.
 
 See also
 --------
@@ -19,21 +20,35 @@
 
 # %%
 import cartopy.crs as ccrs
-import numpy as np
 import shapely.geometry as sg
-from matplotlib.path import Path
 
 import ultraplot as uplt
 
-np.random.seed(0)
-data = np.random.randn(2, 100)
-sizes = np.random.randint(10, 512, data.shape[1])
-colors = np.random.rand(data.shape[1])
-
-fig, ax = uplt.subplots()
-ax.scatter(*data, color=colors, s=sizes, cmap="viko")
+fig, ax = uplt.subplots(refwidth=5.0)
 ax.format(title="Semantic legend helpers")
 
+ax.entrylegend(
+    [
+        {
+            "label": "Observed samples",
+            "line": False,
+            "marker": "o",
+            "markersize": 8,
+            "markerfacecolor": "blue7",
+            "markeredgecolor": "black",
+        },
+        {
+            "label": "Model fit",
+            "line": True,
+            "color": "black",
+            "linewidth": 2.5,
+            "linestyle": "--",
+        },
+    ],
+    loc="l",
+    title="Entry styles",
+    frameon=False,
+)
 ax.catlegend(
     ["A", "B", "C"],
     colors={"A": "red7", "B": "green7", "C": "blue7"},
@@ -43,6 +58,7 @@
 )
 ax.sizelegend(
     [10, 50, 200],
+    labels=["Small", "Medium", "Large"],
     loc="upper right",
     title="Population",
     ncols=1,
@@ -88,4 +104,5 @@
     frameon=False,
     country_reso="10m",
 )
+ax.axis("off")
 fig.show()

ultraplot/axes/base.py

@@ -3555,6 +3555,8 @@ def sizelegend(self, levels, **kwargs):
             Numeric levels used to generate marker-size entries.
         **kwargs
             Forwarded to `ultraplot.legend.UltraLegend.sizelegend`.
+            Pass ``labels=[...]`` or ``labels={level: label}`` to override the
+            generated labels.
             Pass ``add=False`` to return ``(handles, labels)`` without drawing.
         """
         return plegend.UltraLegend(self).sizelegend(levels, **kwargs)

ultraplot/legend.py

@@ -1119,6 +1119,7 @@ def _entry_legend_entries(
 def _size_legend_entries(
     levels: Iterable[float],
     *,
+    labels=None,
     color="0.35",
     marker="o",
     area=True,
@@ -1142,7 +1143,21 @@ def _size_legend_entries(
     else:
         ms = np.abs(values)
     ms = np.maximum(ms * scale, minsize)
-    labels = [_format_label(value, fmt) for value in values]
+    if labels is None:
+        label_list = [_format_label(value, fmt) for value in values]
+    elif isinstance(labels, Mapping):
+        label_list = []
+        for value in values:
+            key = float(value)
+            if key not in labels:
+                raise ValueError(
+                    "sizelegend labels mapping must include a label for every level."
+                )
+            label_list.append(str(labels[key]))
+    else:
+        label_list = [str(label) for label in labels]
+        if len(label_list) != len(values):
+            raise ValueError("sizelegend labels must have the same length as levels.")
     base_styles = {
         "line": False,
         "alpha": alpha,
@@ -1152,7 +1167,7 @@ def _size_legend_entries(
     }
     base_styles.update(entry_kwargs)
     handles = []
-    for idx, (value, label, size) in enumerate(zip(values, labels, ms)):
+    for idx, (value, label, size) in enumerate(zip(values, label_list, ms)):
         styles = _resolve_style_values(base_styles, float(value), idx)
         color_value = _style_lookup(color, float(value), idx, default="0.35")
         marker_value = _style_lookup(marker, float(value), idx, default="o")
@@ -1171,7 +1186,7 @@ def _size_legend_entries(
                 **styles,
             )
         )
-    return handles, labels
+    return handles, label_list
 
 
 def _num_legend_entries(
@@ -1561,6 +1576,7 @@ def sizelegend(
         self,
         levels: Iterable[float],
         *,
+        labels=None,
         color=None,
         marker=None,
         area: Optional[bool] = None,
@@ -1603,6 +1619,7 @@ def sizelegend(
         )
         handles, labels = _size_legend_entries(
             levels,
+            labels=labels,
             color=color,
             marker=marker,
             area=area,

ultraplot/tests/test_legend.py

@@ -469,6 +469,44 @@ def test_sizelegend_handle_kw_accepts_line_scatter_aliases():
     uplt.close(fig)
 
 
+def test_sizelegend_supports_custom_labels_sequence():
+    fig, ax = uplt.subplots()
+    handles, labels = ax.sizelegend(
+        [10, 50, 200],
+        labels=["small", "medium", "large"],
+        add=False,
+    )
+    assert labels == ["small", "medium", "large"]
+    assert [handle.get_label() for handle in handles] == labels
+    uplt.close(fig)
+
+
+def test_sizelegend_supports_custom_labels_mapping():
+    fig, ax = uplt.subplots()
+    handles, labels = ax.sizelegend(
+        [10, 50, 200],
+        labels={10: "small", 50: "medium", 200: "large"},
+        add=False,
+    )
+    assert labels == ["small", "medium", "large"]
+    assert [handle.get_label() for handle in handles] == labels
+    uplt.close(fig)
+
+
+def test_sizelegend_custom_labels_validate_length():
+    fig, ax = uplt.subplots()
+    with pytest.raises(ValueError, match="same length as levels"):
+        ax.sizelegend([10, 50], labels=["small"], add=False)
+    uplt.close(fig)
+
+
+def test_sizelegend_custom_labels_mapping_must_cover_levels():
+    fig, ax = uplt.subplots()
+    with pytest.raises(ValueError, match="include a label for every level"):
+        ax.sizelegend([10, 50], labels={10: "small"}, add=False)
+    uplt.close(fig)
+
+
 def test_numlegend_handle_kw_accepts_patch_aliases():
     fig, ax = uplt.subplots()
     handles, labels = ax.numlegend(
@@ -585,7 +623,6 @@ def test_semantic_legend_rejects_label_kwarg(builder, args, kwargs):
     (
         ("entrylegend", (["A", "B"],), {}),
         ("catlegend", (["A", "B"],), {}),
-        ("sizelegend", ([10, 50],), {}),
         ("numlegend", tuple(), {"vmin": 0, "vmax": 1}),
     ),
 )