Skip to content

Commit 4ff9a38

Browse files
cvanelterenpre-commit-ci[bot]
cvanelteren
and
pre-commit-ci[bot]
authored
Internal: Refactor colorbar to decouple from axis. Introduces UltraColorbar and UltraColorbarLayout (#529)
--------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1 parent 1611e37 commit 4ff9a38

File tree

3 files changed

+1145
-293
lines changed

3 files changed

+1145
-293
lines changed

‎.github/workflows/build-ultraplot.yml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -136,9 +136,9 @@ jobs:
136136
with:
137137
path: ./ultraplot/tests/baseline # The directory to cache
138138
# Key is based on OS, Python/Matplotlib versions, and the base commit SHA
139-
key: ${{ runner.os }}-baseline-base-v4-hs${{ env.PYTHONHASHSEED }}-${{ steps.baseline-ref.outputs.base_sha }}-${{ inputs.python-version }}-${{ inputs.matplotlib-version }}
139+
key: ${{ runner.os }}-baseline-base-v5-hs${{ env.PYTHONHASHSEED }}-${{ steps.baseline-ref.outputs.base_sha }}-${{ inputs.python-version }}-${{ inputs.matplotlib-version }}
140140
restore-keys: |
141-
${{ runner.os }}-baseline-base-v4-hs${{ env.PYTHONHASHSEED }}-${{ steps.baseline-ref.outputs.base_sha }}-${{ inputs.python-version }}-${{ inputs.matplotlib-version }}-
141+
${{ runner.os }}-baseline-base-v5-hs${{ env.PYTHONHASHSEED }}-${{ steps.baseline-ref.outputs.base_sha }}-${{ inputs.python-version }}-${{ inputs.matplotlib-version }}-
142142
143143
# Conditional Baseline Generation (Only runs on cache miss)
144144
- name: Generate baseline from main

‎ultraplot/axes/base.py‎

Lines changed: 68 additions & 291 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,17 @@
4040
from .. import constructor
4141
from .. import legend as plegend
4242
from .. import ticker as pticker
43+
from ..colorbar import (
44+
UltraColorbar,
45+
_apply_inset_colorbar_layout,
46+
_determine_label_rotation,
47+
_get_axis_for,
48+
_get_colorbar_long_axis,
49+
_legacy_inset_colorbar_bounds,
50+
_reflow_inset_colorbar_frame,
51+
_register_inset_colorbar_reflow,
52+
_solve_inset_colorbar_bounds,
53+
)
4354
from ..config import rc
4455
from ..internals import (
4556
_kwargs_to_args,
@@ -1156,302 +1167,68 @@ def _add_colorbar(
11561167
center_levels=None,
11571168
**kwargs,
11581169
):
1159-
"""
1160-
The driver function for adding axes colorbars.
1161-
"""
1162-
# Parse input arguments and apply defaults
1163-
# TODO: Get the 'best' inset colorbar location using the legend algorithm
1164-
# and implement inset colorbars the same as inset legends.
1165-
grid = _not_none(
1166-
grid=grid, edges=edges, drawedges=drawedges, default=rc["colorbar.grid"]
1167-
) # noqa: E501
1168-
length = _not_none(length=length, shrink=shrink)
1169-
label = _not_none(title=title, label=label)
1170-
labelloc = _not_none(labelloc=labelloc, labellocation=labellocation)
1171-
locator = _not_none(ticks=ticks, locator=locator)
1172-
formatter = _not_none(ticklabels=ticklabels, formatter=formatter, format=format)
1173-
minorlocator = _not_none(minorticks=minorticks, minorlocator=minorlocator)
1174-
color = _not_none(c=c, color=color, default=rc["axes.edgecolor"])
1175-
linewidth = _not_none(lw=lw, linewidth=linewidth)
1176-
ticklen = units(_not_none(ticklen, rc["tick.len"]), "pt")
1177-
tickdir = _not_none(tickdir=tickdir, tickdirection=tickdirection)
1178-
tickwidth = units(_not_none(tickwidth, linewidth, rc["tick.width"]), "pt")
1179-
linewidth = units(_not_none(linewidth, default=rc["axes.linewidth"]), "pt")
1180-
ticklenratio = _not_none(ticklenratio, rc["tick.lenratio"])
1181-
tickwidthratio = _not_none(tickwidthratio, rc["tick.widthratio"])
1182-
rasterized = _not_none(rasterized, rc["colorbar.rasterized"])
1183-
center_levels = _not_none(center_levels, rc["colorbar.center_levels"])
1184-
1185-
# Build label and locator keyword argument dicts
1186-
# NOTE: This carefully handles the 'maxn' and 'maxn_minor' deprecations
1187-
kw_label = {}
1188-
locator_kw = locator_kw or {}
1189-
formatter_kw = formatter_kw or {}
1190-
minorlocator_kw = minorlocator_kw or {}
1191-
for key, value in (
1192-
("size", labelsize),
1193-
("weight", labelweight),
1194-
("color", labelcolor),
1195-
):
1196-
if value is not None:
1197-
kw_label[key] = value
1198-
kw_ticklabels = {}
1199-
for key, value in (
1200-
("size", ticklabelsize),
1201-
("weight", ticklabelweight),
1202-
("color", ticklabelcolor),
1203-
("rotation", rotation),
1204-
):
1205-
if value is not None:
1206-
kw_ticklabels[key] = value
1207-
for b, kw in enumerate((locator_kw, minorlocator_kw)):
1208-
key = "maxn_minor" if b else "maxn"
1209-
name = "minorlocator" if b else "locator"
1210-
nbins = kwargs.pop("maxn_minor" if b else "maxn", None)
1211-
if nbins is not None:
1212-
kw["nbins"] = nbins
1213-
warnings._warn_ultraplot(
1214-
f"The colorbar() keyword {key!r} was deprecated in v0.10. To "
1215-
"achieve the same effect, you can pass 'nbins' to the new default "
1216-
f"locator DiscreteLocator using {name}_kw={{'nbins': {nbins}}}. "
1217-
)
1218-
1219-
# Generate and prepare the colorbar axes
1220-
# NOTE: The inset axes function needs 'label' to know how to pad the box
1221-
# TODO: Use seperate keywords for frame properties vs. colorbar edge properties?
1222-
if loc in ("fill", "left", "right", "top", "bottom"):
1223-
length = _not_none(length, rc["colorbar.length"]) # for _add_guide_panel
1224-
kwargs.update({"align": align, "length": length})
1225-
extendsize = _not_none(extendsize, rc["colorbar.extend"])
1226-
ax = self._add_guide_panel(
1227-
loc,
1228-
align,
1229-
length=length,
1230-
width=width,
1231-
space=space,
1232-
pad=pad,
1233-
span=span,
1234-
row=row,
1235-
col=col,
1236-
rows=rows,
1237-
cols=cols,
1238-
) # noqa: E501
1239-
cax, kwargs = ax._parse_colorbar_filled(**kwargs)
1240-
else:
1241-
kwargs.update({"label": label, "length": length, "width": width})
1242-
extendsize = _not_none(extendsize, rc["colorbar.insetextend"])
1243-
cax, kwargs = self._parse_colorbar_inset(
1244-
loc=loc,
1245-
labelloc=labelloc,
1246-
labelrotation=labelrotation,
1247-
labelsize=labelsize,
1248-
pad=pad,
1249-
**kwargs,
1250-
) # noqa: E501
1251-
1252-
# Parse the colorbar mappable
1253-
# NOTE: Account for special case where auto colorbar is generated from 1D
1254-
# methods that construct an 'artist list' (i.e. colormap scatter object)
1255-
if (
1256-
np.iterable(mappable)
1257-
and len(mappable) == 1
1258-
and isinstance(mappable[0], mcm.ScalarMappable)
1259-
): # noqa: E501
1260-
mappable = mappable[0]
1261-
if not isinstance(mappable, mcm.ScalarMappable):
1262-
mappable, kwargs = cax._parse_colorbar_arg(mappable, values, **kwargs)
1263-
else:
1264-
pop = _pop_params(kwargs, cax._parse_colorbar_arg, ignore_internal=True)
1265-
if pop:
1266-
warnings._warn_ultraplot(
1267-
f"Input is already a ScalarMappable. "
1268-
f"Ignoring unused keyword arg(s): {pop}"
1269-
)
1270-
1271-
# Parse 'extendsize' and 'extendfrac' keywords
1272-
# TODO: Make this auto-adjust to the subplot size
1273-
vert = kwargs["orientation"] == "vertical"
1274-
if extendsize is not None and extendfrac is not None:
1275-
warnings._warn_ultraplot(
1276-
f"You cannot specify both an absolute extendsize={extendsize!r} "
1277-
f"and a relative extendfrac={extendfrac!r}. Ignoring 'extendfrac'."
1278-
)
1279-
extendfrac = None
1280-
if extendfrac is None:
1281-
width, height = cax._get_size_inches()
1282-
scale = height if vert else width
1283-
extendsize = units(extendsize, "em", "in")
1284-
extendfrac = extendsize / max(scale - 2 * extendsize, units(1, "em", "in"))
1285-
1286-
# Parse the tick locators and formatters
1287-
# NOTE: In presence of BoundaryNorm or similar handle ticks with special
1288-
# DiscreteLocator or else get issues (see mpl #22233).
1289-
norm = mappable.norm
1290-
formatter = _not_none(formatter, getattr(norm, "_labels", None), "auto")
1291-
formatter_kw.setdefault("tickrange", (norm.vmin, norm.vmax))
1292-
formatter = constructor.Formatter(formatter, **formatter_kw)
1293-
categorical = isinstance(formatter, mticker.FixedFormatter)
1294-
if locator is not None:
1295-
locator = constructor.Locator(locator, **locator_kw)
1296-
if minorlocator is not None: # overrides tickminor
1297-
minorlocator = constructor.Locator(minorlocator, **minorlocator_kw)
1298-
elif tickminor is None:
1299-
tickminor = False if categorical else rc["xy"[vert] + "tick.minor.visible"]
1300-
if isinstance(norm, mcolors.BoundaryNorm): # DiscreteNorm or BoundaryNorm
1301-
ticks = getattr(norm, "_ticks", norm.boundaries)
1302-
segmented = isinstance(getattr(norm, "_norm", None), pcolors.SegmentedNorm)
1303-
if locator is None:
1304-
if categorical or segmented:
1305-
locator = mticker.FixedLocator(ticks)
1306-
else:
1307-
locator = pticker.DiscreteLocator(ticks)
1308-
1309-
if tickminor and minorlocator is None:
1310-
minorlocator = pticker.DiscreteLocator(ticks, minor=True)
1311-
1312-
# Special handling for colorbar keyword arguments
1313-
# WARNING: Critical to not pass empty major locators in matplotlib < 3.5
1314-
# See this issue: https://github.com/ultraplot-dev/ultraplot/issues/301
1315-
# WARNING: ultraplot 'supports' passing one extend to a mappable function
1316-
# then overwriting by passing another 'extend' to colobar. But contour
1317-
# colorbars break when you try to change its 'extend'. Matplotlib gets
1318-
# around this by just silently ignoring 'extend' passed to colorbar() but
1319-
# we issue warning. Also note ContourSet.extend existed in matplotlib 3.0.
1320-
# WARNING: Confusingly the only default way to have auto-adjusting
1321-
# colorbar ticks is to specify no locator. Then _get_ticker_locator_formatter
1322-
# uses the default ScalarFormatter on the axis that already has a set axis.
1323-
# Otherwise it sets a default axis with locator.create_dummy_axis() in
1324-
# update_ticks() which does not track axis size. Workaround is to manually
1325-
# set the locator and formatter axis... however this messes up colorbar lengths
1326-
# in matplotlib < 3.2. So we only apply this conditionally and in earlier
1327-
# verisons recognize that DiscreteLocator will behave like FixedLocator.
1328-
axis = cax.yaxis if vert else cax.xaxis
1329-
if not isinstance(mappable, mcontour.ContourSet):
1330-
extend = _not_none(extend, "neither")
1331-
kwargs["extend"] = extend
1332-
elif extend is not None and extend != mappable.extend:
1333-
warnings._warn_ultraplot(
1334-
"Ignoring extend={extend!r}. ContourSet extend cannot be changed."
1335-
)
1336-
if (
1337-
isinstance(locator, mticker.NullLocator)
1338-
or hasattr(locator, "locs")
1339-
and len(locator.locs) == 0
1340-
):
1341-
minorlocator, tickminor = None, False # attempted fix
1342-
for ticker in (locator, formatter, minorlocator):
1343-
if version.parse(str(_version_mpl)) < version.parse("3.2"):
1344-
pass # see notes above
1345-
elif isinstance(ticker, mticker.TickHelper):
1346-
ticker.set_axis(axis)
1347-
1348-
# Create colorbar and update ticks and axis direction
1349-
# NOTE: This also adds the guides._update_ticks() monkey patch that triggers
1350-
# updates to DiscreteLocator when parent axes is drawn.
1351-
orientation = _not_none(
1352-
kwargs.pop("orientation", None), kwargs.pop("vert", None)
1353-
)
1354-
1355-
obj = cax._colorbar_fill = cax.figure.colorbar(
1170+
return UltraColorbar(self).add(
13561171
mappable,
1357-
cax=cax,
1358-
ticks=locator,
1359-
format=formatter,
1360-
drawedges=grid,
1172+
values=values,
1173+
loc=loc,
1174+
align=align,
1175+
space=space,
1176+
pad=pad,
1177+
width=width,
1178+
length=length,
1179+
span=span,
1180+
row=row,
1181+
col=col,
1182+
rows=rows,
1183+
cols=cols,
1184+
shrink=shrink,
1185+
label=label,
1186+
title=title,
1187+
reverse=reverse,
1188+
rotation=rotation,
1189+
grid=grid,
1190+
edges=edges,
1191+
drawedges=drawedges,
1192+
extend=extend,
1193+
extendsize=extendsize,
13611194
extendfrac=extendfrac,
1362-
orientation=orientation,
1363-
**kwargs,
1364-
)
1365-
outline = _not_none(outline, rc["colorbar.outline"])
1366-
obj.outline.set_visible(outline)
1367-
obj.ax.grid(False)
1368-
# obj.minorlocator = minorlocator # backwards compatibility
1369-
obj.update_ticks = guides._update_ticks.__get__(obj) # backwards compatible
1370-
if minorlocator is not None:
1371-
# Note we make use of mpl's setters and getters
1372-
current = obj.minorlocator
1373-
if current != minorlocator:
1374-
obj.minorlocator = minorlocator
1375-
obj.update_ticks()
1376-
elif tickminor:
1377-
obj.minorticks_on()
1378-
else:
1379-
obj.minorticks_off()
1380-
if getattr(norm, "descending", None):
1381-
axis.set_inverted(True)
1382-
if reverse: # potentially double reverse, although that would be weird...
1383-
axis.set_inverted(True)
1384-
1385-
# Update other colorbar settings
1386-
# WARNING: Must use the colorbar set_label to set text. Calling set_label
1387-
# on the actual axis will do nothing!
1388-
if center_levels:
1389-
# Center the ticks to the center of the colorbar
1390-
# rather than showing them on the edges
1391-
if hasattr(obj.norm, "boundaries"):
1392-
# Only apply to discrete norms
1393-
bounds = obj.norm.boundaries
1394-
centers = 0.5 * (bounds[:-1] + bounds[1:])
1395-
axis.set_ticks(centers)
1396-
ticklenratio = 0
1397-
tickwidthratio = 0
1398-
axis.set_tick_params(which="both", color=color, direction=tickdir)
1399-
axis.set_tick_params(which="major", length=ticklen, width=tickwidth)
1400-
axis.set_tick_params(
1401-
which="minor",
1402-
length=ticklen * ticklenratio,
1403-
width=tickwidth * tickwidthratio,
1404-
) # noqa: E501
1405-
1406-
# Set label and label location
1407-
long_or_short_axis = _get_axis_for(
1408-
labelloc, loc, orientation=orientation, ax=obj
1409-
)
1410-
if labelloc is None:
1411-
labelloc = long_or_short_axis.get_ticks_position()
1412-
long_or_short_axis.set_label_text(label)
1413-
long_or_short_axis.set_label_position(labelloc)
1414-
1415-
labelrotation = _not_none(labelrotation, rc["colorbar.labelrotation"])
1416-
# Note kw_label is updated in place
1417-
_determine_label_rotation(
1418-
labelrotation,
1195+
ticks=ticks,
1196+
locator=locator,
1197+
locator_kw=locator_kw,
1198+
format=format,
1199+
formatter=formatter,
1200+
ticklabels=ticklabels,
1201+
formatter_kw=formatter_kw,
1202+
minorticks=minorticks,
1203+
minorlocator=minorlocator,
1204+
minorlocator_kw=minorlocator_kw,
1205+
tickminor=tickminor,
1206+
ticklen=ticklen,
1207+
ticklenratio=ticklenratio,
1208+
tickdir=tickdir,
1209+
tickdirection=tickdirection,
1210+
tickwidth=tickwidth,
1211+
tickwidthratio=tickwidthratio,
1212+
ticklabelsize=ticklabelsize,
1213+
ticklabelweight=ticklabelweight,
1214+
ticklabelcolor=ticklabelcolor,
14191215
labelloc=labelloc,
1420-
orientation=orientation,
1421-
kw_label=kw_label,
1216+
labellocation=labellocation,
1217+
labelsize=labelsize,
1218+
labelweight=labelweight,
1219+
labelcolor=labelcolor,
1220+
c=c,
1221+
color=color,
1222+
lw=lw,
1223+
linewidth=linewidth,
1224+
edgefix=edgefix,
1225+
rasterized=rasterized,
1226+
outline=outline,
1227+
labelrotation=labelrotation,
1228+
center_levels=center_levels,
1229+
**kwargs,
14221230
)
14231231

1424-
long_or_short_axis.label.update(kw_label)
1425-
# Assume ticks are set on the long axis(!))
1426-
if hasattr(obj, "_long_axis"):
1427-
# mpl <=3.9
1428-
longaxis = obj._long_axis()
1429-
else:
1430-
# mpl >=3.10
1431-
longaxis = obj.long_axis
1432-
for label in longaxis.get_ticklabels():
1433-
label.update(kw_ticklabels)
1434-
if KIWI_AVAILABLE and getattr(cax, "_inset_colorbar_layout", None):
1435-
_reflow_inset_colorbar_frame(obj, labelloc=labelloc, ticklen=ticklen)
1436-
cax._inset_colorbar_obj = obj
1437-
cax._inset_colorbar_labelloc = labelloc
1438-
cax._inset_colorbar_ticklen = ticklen
1439-
_register_inset_colorbar_reflow(self.figure)
1440-
kw_outline = {"edgecolor": color, "linewidth": linewidth}
1441-
if obj.outline is not None:
1442-
obj.outline.update(kw_outline)
1443-
if obj.dividers is not None:
1444-
obj.dividers.update(kw_outline)
1445-
if obj.solids:
1446-
from . import PlotAxes
1447-
1448-
obj.solids.set_rasterized(rasterized)
1449-
PlotAxes._fix_patch_edges(obj.solids, edgefix=edgefix)
1450-
1451-
# Register location and return
1452-
self._register_guide("colorbar", obj, (loc, align)) # possibly replace another
1453-
return obj
1454-
14551232
def _add_legend(
14561233
self,
14571234
handles=None,

0 commit comments

Comments
 (0)