Add ref argument to fig.legend, support 1D slicing, and intelligent placement inference

Author: cvanelteren

docs/colorbars_legends.py

@@ -469,3 +469,44 @@
 ax = axs[1]
 ax.legend(hs2, loc="b", ncols=3, center=True, title="centered rows")
 axs.format(xlabel="xlabel", ylabel="ylabel", suptitle="Legend formatting demo")
+# %% [raw] raw_mimetype="text/restructuredtext"
+# .. _ug_guides_decouple:
+#
+# Decoupling legend content and location
+# --------------------------------------
+#
+# Sometimes you may want to generate a legend using handles from specific axes
+# but place it relative to other axes. In UltraPlot, you can achieve this by passing
+# both the `ax` and `ref` keywords to :func:`~ultraplot.figure.Figure.legend`
+# (or :func:`~ultraplot.figure.Figure.colorbar`). The `ax` keyword specifies the
+# axes used to generate the legend handles, while the `ref` keyword specifies the
+# reference axes used to determine the legend location.
+#
+# For example, to draw a legend based on the handles in the second row of subplots
+# but place it below the first row of subplots, you can use
+# ``fig.legend(ax=axs[1, :], ref=axs[0, :], loc='bottom')``. If ``ref`` is a list
+# of axes, UltraPlot intelligently infers the span (width or height) and anchors
+# the legend to the appropriate outer edge (e.g., the bottom-most axis for ``loc='bottom'``
+# or the right-most axis for ``loc='right'``).
+
+# %%
+import numpy as np
+
+import ultraplot as uplt
+
+fig, axs = uplt.subplots(nrows=2, ncols=2, refwidth=2, share=False)
+axs.format(abc="A.", suptitle="Decoupled legend location demo")
+
+# Plot data on all axes
+state = np.random.RandomState(51423)
+data = (state.rand(20, 4) - 0.5).cumsum(axis=0)
+for ax in axs:
+    ax.plot(data, cycle="mplotcolors", labels=list("abcd"))
+
+# Legend 1: Content from Row 2 (ax=axs[1, :]), Location below Row 1 (ref=axs[0, :])
+# This places a legend describing the bottom row data underneath the top row.
+fig.legend(ax=axs[1, :], ref=axs[0, :], loc="bottom", title="Data from Row 2")
+
+# Legend 2: Content from Row 1 (ax=axs[0, :]), Location below Row 2 (ref=axs[1, :])
+# This places a legend describing the top row data underneath the bottom row.
+fig.legend(ax=axs[0, :], ref=axs[1, :], loc="bottom", title="Data from Row 1")

ultraplot/figure.py

@@ -1417,18 +1417,6 @@ def _add_axes_panel(
         if span_override is not None:
             kw["span_override"] = span_override
 
-        # Check for position override (row for horizontal panels, col for vertical panels)
-        pos_override = None
-        if side in ("left", "right"):
-            if _not_none(cols, col) is not None:
-                pos_override = _not_none(cols, col)
-        else:
-            if _not_none(rows, row) is not None:
-                pos_override = _not_none(rows, row)
-
-        if pos_override is not None:
-            kw["pos_override"] = pos_override
-
         ss, share = gs._insert_panel_slot(side, ax, **kw)
         # Guard: GeoAxes with non-rectilinear projections cannot share with panels
         if isinstance(ax, paxes.GeoAxes) and not ax._is_rectilinear():
@@ -2712,27 +2700,125 @@ def legend(
         matplotlib.axes.Axes.legend
         """
         ax = kwargs.pop("ax", None)
+        ref = kwargs.pop("ref", None)
+        loc_ax = ref if ref is not None else ax
+
         # Axes panel legend
-        if ax is not None:
+        if loc_ax is not None:
+            content_ax = ax if ax is not None else loc_ax
             # Check if span parameters are provided
             has_span = _not_none(span, row, col, rows, cols) is not None
-            # Extract a single axes from array if span is provided
-            # Otherwise, pass the array as-is for normal legend behavior
-            # Automatically collect handles and labels from spanned axes if not provided
-            if has_span and np.iterable(ax) and not isinstance(ax, (str, maxes.Axes)):
-                # Auto-collect handles and labels if not explicitly provided
-                if handles is None and labels is None:
-                    handles, labels = [], []
-                    for axi in ax:
+
+            # Automatically collect handles and labels from content axes if not provided
+            # Case 1: content_ax is a list (we must auto-collect)
+            # Case 2: content_ax != loc_ax (we must auto-collect because loc_ax.legend won't find content_ax handles)
+            must_collect = (
+                np.iterable(content_ax)
+                and not isinstance(content_ax, (str, maxes.Axes))
+            ) or (content_ax is not loc_ax)
+
+            if must_collect and handles is None and labels is None:
+                handles, labels = [], []
+                # Handle list of axes
+                if np.iterable(content_ax) and not isinstance(
+                    content_ax, (str, maxes.Axes)
+                ):
+                    for axi in content_ax:
                         h, l = axi.get_legend_handles_labels()
                         handles.extend(h)
                         labels.extend(l)
-                try:
-                    ax_single = next(iter(ax))
-                except (TypeError, StopIteration):
-                    ax_single = ax
+                # Handle single axis
+                else:
+                    handles, labels = content_ax.get_legend_handles_labels()
+
+            # Infer span from loc_ax if it is a list and no span provided
+            if (
+                not has_span
+                and np.iterable(loc_ax)
+                and not isinstance(loc_ax, (str, maxes.Axes))
+            ):
+                loc_trans = _translate_loc(loc, "legend", default=rc["legend.loc"])
+                side = (
+                    loc_trans
+                    if loc_trans in ("left", "right", "top", "bottom")
+                    else None
+                )
+
+                if side:
+                    r_min, r_max = float("inf"), float("-inf")
+                    c_min, c_max = float("inf"), float("-inf")
+                    valid_ax = False
+                    for axi in loc_ax:
+                        if not hasattr(axi, "get_subplotspec"):
+                            continue
+                        ss = axi.get_subplotspec().get_topmost_subplotspec()
+                        r1, r2, c1, c2 = ss._get_rows_columns()
+                        r_min = min(r_min, r1)
+                        r_max = max(r_max, r2)
+                        c_min = min(c_min, c1)
+                        c_max = max(c_max, c2)
+                        valid_ax = True
+
+                    if valid_ax:
+                        if side in ("left", "right"):
+                            rows = (r_min + 1, r_max + 1)
+                        else:
+                            cols = (c_min + 1, c_max + 1)
+                        has_span = True
+
+            # Extract a single axes from array if span is provided (or if ref is a list)
+            # Otherwise, pass the array as-is for normal legend behavior (only if loc_ax is list)
+            if (
+                has_span
+                and np.iterable(loc_ax)
+                and not isinstance(loc_ax, (str, maxes.Axes))
+            ):
+                # Pick the best axis to anchor to based on the legend side
+                loc_trans = _translate_loc(loc, "legend", default=rc["legend.loc"])
+                side = (
+                    loc_trans
+                    if loc_trans in ("left", "right", "top", "bottom")
+                    else None
+                )
+
+                best_ax = None
+                best_coord = float("-inf")
+
+                # If side is determined, search for the edge axis
+                if side:
+                    for axi in loc_ax:
+                        if not hasattr(axi, "get_subplotspec"):
+                            continue
+                        ss = axi.get_subplotspec().get_topmost_subplotspec()
+                        r1, r2, c1, c2 = ss._get_rows_columns()
+
+                        if side == "right":
+                            val = c2  # Maximize column index
+                        elif side == "left":
+                            val = -c1  # Minimize column index
+                        elif side == "bottom":
+                            val = r2  # Maximize row index
+                        elif side == "top":
+                            val = -r1  # Minimize row index
+                        else:
+                            val = 0
+
+                        if val > best_coord:
+                            best_coord = val
+                            best_ax = axi
+
+                # Fallback to first axis if no best axis found (or side is None)
+                if best_ax is None:
+                    try:
+                        ax_single = next(iter(loc_ax))
+                    except (TypeError, StopIteration):
+                        ax_single = loc_ax
+                else:
+                    ax_single = best_ax
+
             else:
-                ax_single = ax
+                ax_single = loc_ax
+
             leg = ax_single.legend(
                 handles,
                 labels,

ultraplot/gridspec.py

@@ -601,7 +601,6 @@ def _parse_panel_arg_with_span(
         side: str,
         ax: "paxes.Axes",
         span_override: Optional[Union[int, Tuple[int, int]]],
-        pos_override: Optional[Union[int, Tuple[int, int]]] = None,
     ) -> Tuple[str, int, slice]:
         """
         Parse panel arg with span override. Uses ax for position, span for extent.
@@ -614,8 +613,6 @@ def _parse_panel_arg_with_span(
             The axes to position the panel relative to
         span_override : int or tuple
             The span extent (1-indexed like subplot numbers)
-        pos_override : int or tuple, optional
-            The row or column index (1-indexed like subplot numbers)
 
         Returns
         -------
@@ -630,20 +627,6 @@ def _parse_panel_arg_with_span(
         ss = ax.get_subplotspec().get_topmost_subplotspec()
         row1, row2, col1, col2 = ss._get_rows_columns()
 
-        # Override axes position if requested
-        if pos_override is not None:
-            if isinstance(pos_override, Integral):
-                pos1, pos2 = pos_override - 1, pos_override - 1
-            else:
-                pos_override = np.atleast_1d(pos_override)
-                pos1, pos2 = pos_override[0] - 1, pos_override[-1] - 1
-
-            # NOTE: We only need the relevant coordinate (row or col)
-            if side in ("left", "right"):
-                col1, col2 = pos1, pos2
-            else:
-                row1, row2 = pos1, pos2
-
         # Determine slot and index based on side
         slot = side[0]
         offset = len(ax._panel_dict[side]) + 1
@@ -686,7 +669,6 @@ def _insert_panel_slot(
         pad: Optional[Union[float, str]] = None,
         filled: bool = False,
         span_override: Optional[Union[int, Tuple[int, int]]] = None,
-        pos_override: Optional[Union[int, Tuple[int, int]]] = None,
     ):
         """
         Insert a panel slot into the existing gridspec. The `side` is the panel side
@@ -700,9 +682,7 @@ def _insert_panel_slot(
             raise ValueError(f"Invalid side {side}.")
         # Use span override if provided
         if span_override is not None:
-            slot, idx, span = self._parse_panel_arg_with_span(
-                side, arg, span_override, pos_override=pos_override
-            )
+            slot, idx, span = self._parse_panel_arg_with_span(side, arg, span_override)
         else:
             slot, idx, span = self._parse_panel_arg(side, arg)
         pad = units(pad, "em", "in")

ultraplot/tests/test_gridspec.py

@@ -1,5 +1,6 @@
-import ultraplot as uplt
 import pytest
+
+import ultraplot as uplt
 from ultraplot.gridspec import SubplotGrid
 
 
@@ -72,3 +73,56 @@ def test_tight_layout_disabled():
     gs = ax.get_subplotspec().get_gridspec()
     with pytest.raises(RuntimeError):
         gs.tight_layout(fig)
+
+
+def test_gridspec_slicing():
+    """
+    Test various slicing methods on SubplotGrid, including 1D list/array indexing.
+    """
+    import numpy as np
+
+    fig, axs = uplt.subplots(nrows=4, ncols=4)
+
+    # Test 1D integer indexing
+    assert axs[0].number == 1
+    assert axs[15].number == 16
+
+    # Test 1D slice indexing
+    subset = axs[0:2]
+    assert isinstance(subset, SubplotGrid)
+    assert len(subset) == 2
+    assert subset[0].number == 1
+    assert subset[1].number == 2
+
+    # Test 1D list indexing (Fix #1)
+    subset_list = axs[[0, 5]]
+    assert isinstance(subset_list, SubplotGrid)
+    assert len(subset_list) == 2
+    assert subset_list[0].number == 1
+    assert subset_list[1].number == 6
+
+    # Test 1D array indexing
+    subset_array = axs[np.array([0, 5])]
+    assert isinstance(subset_array, SubplotGrid)
+    assert len(subset_array) == 2
+    assert subset_array[0].number == 1
+    assert subset_array[1].number == 6
+
+    # Test 2D slicing (tuple of slices)
+    # axs[0:2, :] -> Rows 0 and 1, all cols
+    subset_2d = axs[0:2, :]
+    assert isinstance(subset_2d, SubplotGrid)
+    # 2 rows * 4 cols = 8 axes
+    assert len(subset_2d) == 8
+
+    # Test 2D mixed slicing (list in one dim) (Fix #2 related to _encode_indices)
+    # axs[[0, 1], :] -> Row indices 0 and 1, all cols
+    subset_mixed = axs[[0, 1], :]
+    assert isinstance(subset_mixed, SubplotGrid)
+    assert len(subset_mixed) == 8
+
+    # Verify content
+    # subset_mixed[0] -> Row 0, Col 0 -> Number 1
+    # subset_mixed[4] -> Row 1, Col 0 -> Number 5 (since 4 cols per row)
+    assert subset_mixed[0].number == 1
+    assert subset_mixed[4].number == 5

ultraplot/tests/test_legend.py

@@ -529,3 +529,43 @@ def test_legend_explicit_handles_labels_override_auto_collection():
     assert leg is not None
     assert len(leg.get_texts()) == 1
     assert leg.get_texts()[0].get_text() == "custom_label"
+
+
+import numpy as np
+
+import ultraplot as uplt
+
+
+def test_legend_ref_argument():
+    """Test using 'ref' to decouple legend location from content axes."""
+    fig, axs = uplt.subplots(nrows=2, ncols=2)
+    axs[0, 0].plot([], [], label="line1")  # Row 0
+    axs[1, 0].plot([], [], label="line2")  # Row 1
+
+    # Place legend below Row 0 (axs[0, :]) using content from Row 1 (axs[1, :])
+    leg = fig.legend(ax=axs[1, :], ref=axs[0, :], loc="bottom")
+
+    assert leg is not None
+
+    # Should be a single legend because span is inferred from ref
+    assert not isinstance(leg, tuple)
+
+    texts = [t.get_text() for t in leg.get_texts()]
+    assert "line2" in texts
+    assert "line1" not in texts
+
+
+def test_legend_ref_argument_no_ax():
+    """Test using 'ref' where 'ax' is implied to be 'ref'."""
+    fig, axs = uplt.subplots(nrows=1, ncols=1)
+    axs[0].plot([], [], label="line1")
+
+    # ref provided, ax=None. Should behave like ax=ref.
+    leg = fig.legend(ref=axs[0], loc="bottom")
+    assert leg is not None
+
+    # Should be a single legend
+    assert not isinstance(leg, tuple)
+
+    texts = [t.get_text() for t in leg.get_texts()]
+    assert "line1" in texts