Add bbox anchors for inset colorbars

docs/colorbars_legends.py

@@ -54,6 +54,9 @@
 #   (e.g., ``loc='upper right'`` or the shorthand ``loc='ur'``). Inset
 #   colorbars have optional background "frames" that can be configured
 #   with various :func:`~ultraplot.axes.Axes.colorbar` keywords.
+#   They also accept ``bbox_to_anchor`` with the same two- or four-value
+#   anchor semantics as inset legends when the default edge-aware placement
+#   should be explicitly overridden.
 
 # :func:`~ultraplot.axes.Axes.colorbar` and :meth:`~ultraplot.axes.Axes.legend` also both accept
 # `space` and `pad` keywords. `space` controls the absolute separation of the
@@ -92,6 +95,7 @@
 ax.colorbar(m, loc="r")
 ax.colorbar(m, loc="t")  # title is automatically adjusted
 ax.colorbar(m, loc="ll", label="colorbar label")  # inset colorbar demonstration
+ax.colorbar(m, loc="ur", bbox_to_anchor=(0.92, 0.92))
 
 # Legends
 ax = fig.subplot(122, title="Axes legends", titlepad="0em")

ultraplot/axes/base.py

@@ -43,6 +43,7 @@
 from .. import ticker as pticker
 from ..colorbar import (
     UltraColorbar,
+    _anchor_inset_colorbar_bounds,
     _apply_inset_colorbar_layout,
     _determine_label_rotation,
     _get_axis_for,
@@ -2041,6 +2042,7 @@ def _parse_colorbar_filled(
     def _parse_colorbar_inset(
         self,
         loc=None,
+        bbox_to_anchor=None,
         width=None,
         length=None,
         shrink=None,
@@ -2117,6 +2119,9 @@ def _parse_colorbar_inset(
                 tick_fontsize=tick_fontsize,
                 label_fontsize=label_fontsize,
             )
+        bounds_inset, bounds_frame = _anchor_inset_colorbar_bounds(
+            bounds_inset, bounds_frame, loc, bbox_to_anchor
+        )
 
         # Create axes and frame
         ax = self._add_colorbar_child_axes(bounds_inset)
@@ -2137,6 +2142,7 @@ def _parse_colorbar_inset(
             "length_raw": length_raw,
             "width_raw": width_raw,
             "pad_raw": pad_raw,
+            "bbox_to_anchor": bbox_to_anchor,
         }
         ax._inset_colorbar_frame = frame_artist
 
@@ -3455,9 +3461,13 @@ def draw(self, renderer=None, *args, **kwargs):
             self.indicate_inset_zoom()
         self._apply_align_text(renderer)
         needs_inset_reflow = bool(getattr(self, "_inset_colorbar_needs_reflow", False))
+        has_inset_colorbar = bool(
+            getattr(self, "_inset_colorbar_layout", None)
+            and getattr(self, "_inset_colorbar_obj", None)
+        )
         has_inset_frame = bool(
             getattr(self, "_inset_colorbar_frame", None) is not None
-            and getattr(self, "_inset_colorbar_obj", None)
+            and has_inset_colorbar
         )
         super().draw(renderer, *args, **kwargs)
         if has_inset_frame:
@@ -3467,7 +3477,7 @@ def draw(self, renderer=None, *args, **kwargs):
                     labelloc=getattr(self, "_inset_colorbar_labelloc", None),
                     renderer=renderer,
                 )
-        if has_inset_frame and needs_inset_reflow:
+        if has_inset_colorbar and needs_inset_reflow:
             _reflow_inset_colorbar_frame(
                 self._inset_colorbar_obj,
                 labelloc=getattr(self, "_inset_colorbar_labelloc", None),
@@ -3706,6 +3716,11 @@ def colorbar(self, mappable, values=None, loc=None, location=None, **kwargs):
                 Strings are interpreted by `~ultraplot.utils.units`.
             %(axes.colorbar_space)s
                 Has no visible effect if `length` is ``1``.
+            bbox_to_anchor : 2-tuple, 4-tuple, or `matplotlib.transforms.Bbox`, optional
+                For inset colorbars, anchor the full colorbar footprint using the
+                same semantics as `~matplotlib.axes.Axes.legend`. The colorbar
+                `loc` selects the corresponding anchor corner. Outer colorbar
+                placement is unchanged.
             Other parameters
             ----------------
             %(axes.colorbar_kwargs)s
@@ -5423,9 +5438,15 @@ def _reflow_inset_colorbar_frame(
         bounds = solver.solve()
     except Exception:
         return
+    bounds_inset, bounds_frame = _anchor_inset_colorbar_bounds(
+        list(bounds["inset"]),
+        list(bounds["frame"]),
+        loc,
+        layout.get("bbox_to_anchor"),
+    )
     _apply_inset_colorbar_layout(
         cax,
-        bounds_inset=list(bounds["inset"]),
-        bounds_frame=list(bounds["frame"]),
+        bounds_inset=bounds_inset,
+        bounds_frame=bounds_frame,
         frame=frame,
     )

ultraplot/colorbar.py

@@ -163,6 +163,7 @@ def add(
         # NOTE: The inset axes function needs 'label' to know how to pad the box
         # TODO: Use seperate keywords for frame properties vs. colorbar edge properties?
         frame = _not_none(frame=frame, frameon=frameon)
+        bbox_to_anchor = kwargs.pop("bbox_to_anchor", None)
         inset_side = loc in ("left", "right", "top", "bottom") and getattr(
             ax, "_inset_parent", None
         )
@@ -197,7 +198,14 @@ def add(
                     **kwargs,
                 )
             else:
-                kwargs.update({"label": label, "length": length, "width": width})
+                kwargs.update(
+                    {
+                        "bbox_to_anchor": bbox_to_anchor,
+                        "label": label,
+                        "length": length,
+                        "width": width,
+                    }
+                )
                 extendsize = _not_none(extendsize, rc["colorbar.insetextend"])
                 cax, kwargs = ax._parse_colorbar_inset(
                     loc=loc,
@@ -366,9 +374,7 @@ def add(
             cax._inset_colorbar_obj = obj
             cax._inset_colorbar_labelloc = labelloc
             cax._inset_colorbar_ticklen = ticklen
-            has_frame = getattr(cax, "_inset_colorbar_frame", None) is not None
-            if has_frame:
-                _register_inset_colorbar_reflow(ax.figure)
+            _register_inset_colorbar_reflow(ax.figure)
         kw_outline = {"edgecolor": color, "linewidth": linewidth}
         if obj.outline is not None:
             obj.outline.update(kw_outline)
@@ -806,6 +812,52 @@ def _solve_inset_colorbar_bounds(
     return list(layout["inset"]), list(layout["frame"])
 
 
+def _anchor_inset_colorbar_bounds(
+    bounds_inset: list[float],
+    bounds_frame: list[float],
+    loc: str,
+    bbox_to_anchor,
+) -> Tuple[list[float], list[float]]:
+    """Align an inset colorbar footprint to a legend-style anchor box."""
+    if bbox_to_anchor is None:
+        return bounds_inset, bounds_frame
+    if isinstance(bbox_to_anchor, mtransforms.BboxBase):
+        bbox = bbox_to_anchor
+    else:
+        try:
+            values = tuple(bbox_to_anchor)
+        except TypeError as exc:
+            raise ValueError(
+                "bbox_to_anchor must be a 2- or 4-tuple, or a matplotlib Bbox."
+            ) from exc
+        if len(values) == 2:
+            bbox = mtransforms.Bbox.from_bounds(*values, 0, 0)
+        elif len(values) == 4:
+            bbox = mtransforms.Bbox.from_bounds(*values)
+        else:
+            raise ValueError(
+                "bbox_to_anchor must be a 2- or 4-tuple, or a matplotlib Bbox."
+            )
+
+    x, y, width, height = bounds_frame
+    if loc == "upper left":
+        source, target = (x, y + height), (bbox.x0, bbox.y1)
+    elif loc == "lower left":
+        source, target = (x, y), (bbox.x0, bbox.y0)
+    elif loc == "lower right":
+        source, target = (x + width, y), (bbox.x1, bbox.y0)
+    else:  # ``best`` resolves to upper right in the inset layout.
+        source, target = (x + width, y + height), (bbox.x1, bbox.y1)
+    dx, dy = target[0] - source[0], target[1] - source[1]
+    inset = list(bounds_inset)
+    frame = list(bounds_frame)
+    inset[0] += dx
+    inset[1] += dy
+    frame[0] += dx
+    frame[1] += dy
+    return inset, frame
+
+
 def _legacy_inset_colorbar_bounds(
     *,
     axes: maxes.Axes,
@@ -1089,9 +1141,15 @@ def _reflow_inset_colorbar_frame(
         bounds = solver.solve()
     except Exception:
         return
+    bounds_inset, bounds_frame = _anchor_inset_colorbar_bounds(
+        list(bounds["inset"]),
+        list(bounds["frame"]),
+        loc,
+        layout.get("bbox_to_anchor"),
+    )
     _apply_inset_colorbar_layout(
         cax,
-        bounds_inset=list(bounds["inset"]),
-        bounds_frame=list(bounds["frame"]),
+        bounds_inset=bounds_inset,
+        bounds_frame=bounds_frame,
         frame=frame,
     )

ultraplot/tests/test_colorbar.py

@@ -5,6 +5,7 @@
 
 import numpy as np
 import pytest
+from matplotlib.transforms import Bbox
 
 import ultraplot as uplt
 
@@ -79,6 +80,49 @@ def test_inset_colorbar_frame_alias_still_controls_frame(rng, kwargs):
     assert cb.ax._inset_colorbar_frame is None
 
 
+@pytest.mark.parametrize(
+    "loc, bbox_to_anchor, corners",
+    [
+        ("ur", (0.8, 0.75), ("x1", "y1", 0.8, 0.75)),
+        ("ll", (0.2, 0.15, 0.5, 0.4), ("x0", "y0", 0.2, 0.15)),
+    ],
+)
+def test_inset_colorbar_bbox_to_anchor(rng, loc, bbox_to_anchor, corners):
+    fig, ax = uplt.subplots()
+    mappable = ax.pcolormesh(rng.random((8, 8)))
+    colorbar = ax.colorbar(
+        mappable,
+        loc=loc,
+        label="A label that must be included",
+        bbox_to_anchor=bbox_to_anchor,
+    )
+    fig.canvas.draw()
+    frame = Bbox.from_bounds(*colorbar.ax._inset_colorbar_bounds["frame"])
+    xattr, yattr, xanchor, yanchor = corners
+    assert getattr(frame, xattr) == pytest.approx(xanchor)
+    assert getattr(frame, yattr) == pytest.approx(yanchor)
+    fig.set_size_inches(7, 4.5)
+    fig.canvas.draw()
+    frame = Bbox.from_bounds(*colorbar.ax._inset_colorbar_bounds["frame"])
+    assert getattr(frame, xattr) == pytest.approx(xanchor)
+    assert getattr(frame, yattr) == pytest.approx(yanchor)
+
+
+def test_unanchored_inset_colorbar_label_stays_inside_axes(rng):
+    fig, ax = uplt.subplots()
+    mappable = ax.pcolormesh(rng.random((8, 8)))
+    colorbar = ax.colorbar(
+        mappable, loc="ur", label="Inset colorbar label", frame=False
+    )
+    fig.canvas.draw()
+    fig.set_size_inches(7, 4.5)
+    fig.canvas.draw()
+    renderer = fig.canvas.get_renderer()
+    bbox = colorbar.ax.get_tightbbox(renderer).transformed(ax.transAxes.inverted())
+    assert bbox.x1 <= 1
+    assert bbox.y1 <= 1
+
+
 def test_colorbar_side_locations_work_on_inset_axes(rng):
     fig, ax = uplt.subplots()
     ix = ax.inset_axes([0.55, 0.55, 0.35, 0.35], zoom=False)