Internal: Refactor colorbar to decouple from axis. Introduces UltraColorbar and UltraColorbarLayout

ultraplot/axes/base.py

@@ -40,6 +40,17 @@
 from .. import constructor
 from .. import legend as plegend
 from .. import ticker as pticker
+from ..colorbar import (
+    UltraColorbar,
+    _apply_inset_colorbar_layout,
+    _determine_label_rotation,
+    _get_axis_for,
+    _get_colorbar_long_axis,
+    _legacy_inset_colorbar_bounds,
+    _reflow_inset_colorbar_frame,
+    _register_inset_colorbar_reflow,
+    _solve_inset_colorbar_bounds,
+)
 from ..config import rc
 from ..internals import (
     _kwargs_to_args,
@@ -1156,302 +1167,68 @@ def _add_colorbar(
         center_levels=None,
         **kwargs,
     ):
-        """
-        The driver function for adding axes colorbars.
-        """
-        # Parse input arguments and apply defaults
-        # TODO: Get the 'best' inset colorbar location using the legend algorithm
-        # and implement inset colorbars the same as inset legends.
-        grid = _not_none(
-            grid=grid, edges=edges, drawedges=drawedges, default=rc["colorbar.grid"]
-        )  # noqa: E501
-        length = _not_none(length=length, shrink=shrink)
-        label = _not_none(title=title, label=label)
-        labelloc = _not_none(labelloc=labelloc, labellocation=labellocation)
-        locator = _not_none(ticks=ticks, locator=locator)
-        formatter = _not_none(ticklabels=ticklabels, formatter=formatter, format=format)
-        minorlocator = _not_none(minorticks=minorticks, minorlocator=minorlocator)
-        color = _not_none(c=c, color=color, default=rc["axes.edgecolor"])
-        linewidth = _not_none(lw=lw, linewidth=linewidth)
-        ticklen = units(_not_none(ticklen, rc["tick.len"]), "pt")
-        tickdir = _not_none(tickdir=tickdir, tickdirection=tickdirection)
-        tickwidth = units(_not_none(tickwidth, linewidth, rc["tick.width"]), "pt")
-        linewidth = units(_not_none(linewidth, default=rc["axes.linewidth"]), "pt")
-        ticklenratio = _not_none(ticklenratio, rc["tick.lenratio"])
-        tickwidthratio = _not_none(tickwidthratio, rc["tick.widthratio"])
-        rasterized = _not_none(rasterized, rc["colorbar.rasterized"])
-        center_levels = _not_none(center_levels, rc["colorbar.center_levels"])
-
-        # Build label and locator keyword argument dicts
-        # NOTE: This carefully handles the 'maxn' and 'maxn_minor' deprecations
-        kw_label = {}
-        locator_kw = locator_kw or {}
-        formatter_kw = formatter_kw or {}
-        minorlocator_kw = minorlocator_kw or {}
-        for key, value in (
-            ("size", labelsize),
-            ("weight", labelweight),
-            ("color", labelcolor),
-        ):
-            if value is not None:
-                kw_label[key] = value
-        kw_ticklabels = {}
-        for key, value in (
-            ("size", ticklabelsize),
-            ("weight", ticklabelweight),
-            ("color", ticklabelcolor),
-            ("rotation", rotation),
-        ):
-            if value is not None:
-                kw_ticklabels[key] = value
-        for b, kw in enumerate((locator_kw, minorlocator_kw)):
-            key = "maxn_minor" if b else "maxn"
-            name = "minorlocator" if b else "locator"
-            nbins = kwargs.pop("maxn_minor" if b else "maxn", None)
-            if nbins is not None:
-                kw["nbins"] = nbins
-                warnings._warn_ultraplot(
-                    f"The colorbar() keyword {key!r} was deprecated in v0.10. To "
-                    "achieve the same effect, you can pass 'nbins' to the new default "
-                    f"locator DiscreteLocator using {name}_kw={{'nbins': {nbins}}}. "
-                )
-
-        # Generate and prepare the colorbar axes
-        # 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?
-        if loc in ("fill", "left", "right", "top", "bottom"):
-            length = _not_none(length, rc["colorbar.length"])  # for _add_guide_panel
-            kwargs.update({"align": align, "length": length})
-            extendsize = _not_none(extendsize, rc["colorbar.extend"])
-            ax = self._add_guide_panel(
-                loc,
-                align,
-                length=length,
-                width=width,
-                space=space,
-                pad=pad,
-                span=span,
-                row=row,
-                col=col,
-                rows=rows,
-                cols=cols,
-            )  # noqa: E501
-            cax, kwargs = ax._parse_colorbar_filled(**kwargs)
-        else:
-            kwargs.update({"label": label, "length": length, "width": width})
-            extendsize = _not_none(extendsize, rc["colorbar.insetextend"])
-            cax, kwargs = self._parse_colorbar_inset(
-                loc=loc,
-                labelloc=labelloc,
-                labelrotation=labelrotation,
-                labelsize=labelsize,
-                pad=pad,
-                **kwargs,
-            )  # noqa: E501
-
-        # Parse the colorbar mappable
-        # NOTE: Account for special case where auto colorbar is generated from 1D
-        # methods that construct an 'artist list' (i.e. colormap scatter object)
-        if (
-            np.iterable(mappable)
-            and len(mappable) == 1
-            and isinstance(mappable[0], mcm.ScalarMappable)
-        ):  # noqa: E501
-            mappable = mappable[0]
-        if not isinstance(mappable, mcm.ScalarMappable):
-            mappable, kwargs = cax._parse_colorbar_arg(mappable, values, **kwargs)
-        else:
-            pop = _pop_params(kwargs, cax._parse_colorbar_arg, ignore_internal=True)
-            if pop:
-                warnings._warn_ultraplot(
-                    f"Input is already a ScalarMappable. "
-                    f"Ignoring unused keyword arg(s): {pop}"
-                )
-
-        # Parse 'extendsize' and 'extendfrac' keywords
-        # TODO: Make this auto-adjust to the subplot size
-        vert = kwargs["orientation"] == "vertical"
-        if extendsize is not None and extendfrac is not None:
-            warnings._warn_ultraplot(
-                f"You cannot specify both an absolute extendsize={extendsize!r} "
-                f"and a relative extendfrac={extendfrac!r}. Ignoring 'extendfrac'."
-            )
-            extendfrac = None
-        if extendfrac is None:
-            width, height = cax._get_size_inches()
-            scale = height if vert else width
-            extendsize = units(extendsize, "em", "in")
-            extendfrac = extendsize / max(scale - 2 * extendsize, units(1, "em", "in"))
-
-        # Parse the tick locators and formatters
-        # NOTE: In presence of BoundaryNorm or similar handle ticks with special
-        # DiscreteLocator or else get issues (see mpl #22233).
-        norm = mappable.norm
-        formatter = _not_none(formatter, getattr(norm, "_labels", None), "auto")
-        formatter_kw.setdefault("tickrange", (norm.vmin, norm.vmax))
-        formatter = constructor.Formatter(formatter, **formatter_kw)
-        categorical = isinstance(formatter, mticker.FixedFormatter)
-        if locator is not None:
-            locator = constructor.Locator(locator, **locator_kw)
-        if minorlocator is not None:  # overrides tickminor
-            minorlocator = constructor.Locator(minorlocator, **minorlocator_kw)
-        elif tickminor is None:
-            tickminor = False if categorical else rc["xy"[vert] + "tick.minor.visible"]
-        if isinstance(norm, mcolors.BoundaryNorm):  # DiscreteNorm or BoundaryNorm
-            ticks = getattr(norm, "_ticks", norm.boundaries)
-            segmented = isinstance(getattr(norm, "_norm", None), pcolors.SegmentedNorm)
-            if locator is None:
-                if categorical or segmented:
-                    locator = mticker.FixedLocator(ticks)
-                else:
-                    locator = pticker.DiscreteLocator(ticks)
-
-            if tickminor and minorlocator is None:
-                minorlocator = pticker.DiscreteLocator(ticks, minor=True)
-
-        # Special handling for colorbar keyword arguments
-        # WARNING: Critical to not pass empty major locators in matplotlib < 3.5
-        # See this issue: https://github.com/ultraplot-dev/ultraplot/issues/301
-        # WARNING: ultraplot 'supports' passing one extend to a mappable function
-        # then overwriting by passing another 'extend' to colobar. But contour
-        # colorbars break when you try to change its 'extend'. Matplotlib gets
-        # around this by just silently ignoring 'extend' passed to colorbar() but
-        # we issue warning. Also note ContourSet.extend existed in matplotlib 3.0.
-        # WARNING: Confusingly the only default way to have auto-adjusting
-        # colorbar ticks is to specify no locator. Then _get_ticker_locator_formatter
-        # uses the default ScalarFormatter on the axis that already has a set axis.
-        # Otherwise it sets a default axis with locator.create_dummy_axis() in
-        # update_ticks() which does not track axis size. Workaround is to manually
-        # set the locator and formatter axis... however this messes up colorbar lengths
-        # in matplotlib < 3.2. So we only apply this conditionally and in earlier
-        # verisons recognize that DiscreteLocator will behave like FixedLocator.
-        axis = cax.yaxis if vert else cax.xaxis
-        if not isinstance(mappable, mcontour.ContourSet):
-            extend = _not_none(extend, "neither")
-            kwargs["extend"] = extend
-        elif extend is not None and extend != mappable.extend:
-            warnings._warn_ultraplot(
-                "Ignoring extend={extend!r}. ContourSet extend cannot be changed."
-            )
-        if (
-            isinstance(locator, mticker.NullLocator)
-            or hasattr(locator, "locs")
-            and len(locator.locs) == 0
-        ):
-            minorlocator, tickminor = None, False  # attempted fix
-        for ticker in (locator, formatter, minorlocator):
-            if version.parse(str(_version_mpl)) < version.parse("3.2"):
-                pass  # see notes above
-            elif isinstance(ticker, mticker.TickHelper):
-                ticker.set_axis(axis)
-
-        # Create colorbar and update ticks and axis direction
-        # NOTE: This also adds the guides._update_ticks() monkey patch that triggers
-        # updates to DiscreteLocator when parent axes is drawn.
-        orientation = _not_none(
-            kwargs.pop("orientation", None), kwargs.pop("vert", None)
-        )
-
-        obj = cax._colorbar_fill = cax.figure.colorbar(
+        return UltraColorbar(self).add(
             mappable,
-            cax=cax,
-            ticks=locator,
-            format=formatter,
-            drawedges=grid,
+            values=values,
+            loc=loc,
+            align=align,
+            space=space,
+            pad=pad,
+            width=width,
+            length=length,
+            span=span,
+            row=row,
+            col=col,
+            rows=rows,
+            cols=cols,
+            shrink=shrink,
+            label=label,
+            title=title,
+            reverse=reverse,
+            rotation=rotation,
+            grid=grid,
+            edges=edges,
+            drawedges=drawedges,
+            extend=extend,
+            extendsize=extendsize,
             extendfrac=extendfrac,
-            orientation=orientation,
-            **kwargs,
-        )
-        outline = _not_none(outline, rc["colorbar.outline"])
-        obj.outline.set_visible(outline)
-        obj.ax.grid(False)
-        # obj.minorlocator = minorlocator  # backwards compatibility
-        obj.update_ticks = guides._update_ticks.__get__(obj)  # backwards compatible
-        if minorlocator is not None:
-            # Note we make use of mpl's setters and getters
-            current = obj.minorlocator
-            if current != minorlocator:
-                obj.minorlocator = minorlocator
-            obj.update_ticks()
-        elif tickminor:
-            obj.minorticks_on()
-        else:
-            obj.minorticks_off()
-        if getattr(norm, "descending", None):
-            axis.set_inverted(True)
-        if reverse:  # potentially double reverse, although that would be weird...
-            axis.set_inverted(True)
-
-        # Update other colorbar settings
-        # WARNING: Must use the colorbar set_label to set text. Calling set_label
-        # on the actual axis will do nothing!
-        if center_levels:
-            # Center the ticks to the center of the colorbar
-            # rather than showing them on  the edges
-            if hasattr(obj.norm, "boundaries"):
-                # Only apply to discrete norms
-                bounds = obj.norm.boundaries
-                centers = 0.5 * (bounds[:-1] + bounds[1:])
-                axis.set_ticks(centers)
-                ticklenratio = 0
-                tickwidthratio = 0
-        axis.set_tick_params(which="both", color=color, direction=tickdir)
-        axis.set_tick_params(which="major", length=ticklen, width=tickwidth)
-        axis.set_tick_params(
-            which="minor",
-            length=ticklen * ticklenratio,
-            width=tickwidth * tickwidthratio,
-        )  # noqa: E501
-
-        # Set label and label location
-        long_or_short_axis = _get_axis_for(
-            labelloc, loc, orientation=orientation, ax=obj
-        )
-        if labelloc is None:
-            labelloc = long_or_short_axis.get_ticks_position()
-        long_or_short_axis.set_label_text(label)
-        long_or_short_axis.set_label_position(labelloc)
-
-        labelrotation = _not_none(labelrotation, rc["colorbar.labelrotation"])
-        # Note kw_label is updated in place
-        _determine_label_rotation(
-            labelrotation,
+            ticks=ticks,
+            locator=locator,
+            locator_kw=locator_kw,
+            format=format,
+            formatter=formatter,
+            ticklabels=ticklabels,
+            formatter_kw=formatter_kw,
+            minorticks=minorticks,
+            minorlocator=minorlocator,
+            minorlocator_kw=minorlocator_kw,
+            tickminor=tickminor,
+            ticklen=ticklen,
+            ticklenratio=ticklenratio,
+            tickdir=tickdir,
+            tickdirection=tickdirection,
+            tickwidth=tickwidth,
+            tickwidthratio=tickwidthratio,
+            ticklabelsize=ticklabelsize,
+            ticklabelweight=ticklabelweight,
+            ticklabelcolor=ticklabelcolor,
             labelloc=labelloc,
-            orientation=orientation,
-            kw_label=kw_label,
+            labellocation=labellocation,
+            labelsize=labelsize,
+            labelweight=labelweight,
+            labelcolor=labelcolor,
+            c=c,
+            color=color,
+            lw=lw,
+            linewidth=linewidth,
+            edgefix=edgefix,
+            rasterized=rasterized,
+            outline=outline,
+            labelrotation=labelrotation,
+            center_levels=center_levels,
+            **kwargs,
         )
 
-        long_or_short_axis.label.update(kw_label)
-        # Assume ticks are set on the long axis(!))
-        if hasattr(obj, "_long_axis"):
-            # mpl <=3.9
-            longaxis = obj._long_axis()
-        else:
-            # mpl >=3.10
-            longaxis = obj.long_axis
-        for label in longaxis.get_ticklabels():
-            label.update(kw_ticklabels)
-        if KIWI_AVAILABLE and getattr(cax, "_inset_colorbar_layout", None):
-            _reflow_inset_colorbar_frame(obj, labelloc=labelloc, ticklen=ticklen)
-            cax._inset_colorbar_obj = obj
-            cax._inset_colorbar_labelloc = labelloc
-            cax._inset_colorbar_ticklen = ticklen
-            _register_inset_colorbar_reflow(self.figure)
-        kw_outline = {"edgecolor": color, "linewidth": linewidth}
-        if obj.outline is not None:
-            obj.outline.update(kw_outline)
-        if obj.dividers is not None:
-            obj.dividers.update(kw_outline)
-        if obj.solids:
-            from . import PlotAxes
-
-            obj.solids.set_rasterized(rasterized)
-            PlotAxes._fix_patch_edges(obj.solids, edgefix=edgefix)
-
-        # Register location and return
-        self._register_guide("colorbar", obj, (loc, align))  # possibly replace another
-        return obj
-
     def _add_legend(
         self,
         handles=None,