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

Author: cvanelteren

ultraplot/figure.py

@@ -2594,6 +2594,8 @@ def colorbar(
         """
         # Backwards compatibility
         ax = kwargs.pop("ax", None)
+        ref = kwargs.pop("ref", None)
+        loc_ax = ref if ref is not None else ax
         cax = kwargs.pop("cax", None)
         if isinstance(values, maxes.Axes):
             cax = _not_none(cax_positional=values, cax=cax)
@@ -2613,20 +2615,96 @@ def colorbar(
             with context._state_context(cax, _internal_call=True):  # do not wrap pcolor
                 cb = super().colorbar(mappable, cax=cax, **kwargs)
         # Axes panel colorbar
-        elif ax is not None:
+        elif loc_ax is not None:
             # Check if span parameters are provided
             has_span = _not_none(span, row, col, rows, cols) is not None
 
+            # 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, "colorbar", default=rc["colorbar.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
             # Otherwise, pass the array as-is for normal colorbar behavior
-            if has_span and np.iterable(ax) and not isinstance(ax, (str, maxes.Axes)):
-                try:
-                    ax_single = next(iter(ax))
+            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 colorbar side
+                loc_trans = _translate_loc(loc, "colorbar", default=rc["colorbar.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
 
-                except (TypeError, StopIteration):
-                    ax_single = ax
+                # Fallback to first axis
+                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
 
             # Pass span parameters through to axes colorbar
             cb = ax_single.colorbar(

ultraplot/tests/test_legend.py

@@ -531,11 +531,6 @@ def test_legend_explicit_handles_labels_override_auto_collection():
     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)