Add ref argument to fig.legend for decoupled placement and support 1D slicing in SubplotGrid

Author: cvanelteren

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,51 @@ 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)
+                # Handle single axis
+                else:
+                    handles, labels = content_ax.get_legend_handles_labels()
+
+            # 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))
+            ):
                 try:
-                    ax_single = next(iter(ax))
+                    ax_single = next(iter(loc_ax))
                 except (TypeError, StopIteration):
-                    ax_single = ax
+                    ax_single = loc_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,41 @@ 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
+
+    legs = leg if isinstance(leg, tuple) else (leg,)
+
+    for l in legs:
+        texts = [t.get_text() for t in l.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
+
+    legs = leg if isinstance(leg, tuple) else (leg,)
+
+    for l in legs:
+        texts = [t.get_text() for t in l.get_texts()]
+        assert "line1" in texts