""" Drawing several variables onto one set of axes. `--overlay-style` picks how. The choice matters more than it looks, because the original answer - stack them as translucent colour and let the colours mix - has a cost the overlay module's own docstring concedes: where two layers meet, the colour belongs to neither colour bar, so the figure cannot be read back to a number exactly where it is most interesting. It also spends the base variable's colormap: to stay separable, every layer has to drop to a single-hue ramp, and the field everyone came to look at gives up `viridis` to make room. So there is more than one answer here: blend translucent colour, mixing where they overlap (the original) contour lines over a filled base, labelled with their values hatch a pattern above a threshold - presence, not amount bivariate two variables in one 2-D colour scheme, with a square key glyph the layer as sized symbols over the base Only `blend` and `bivariate` put a second variable in the colour channel at all. `contour` is the one to reach for when both fields have to stay readable: the overlay becomes geometry rather than paint, the base keeps its own colormap, and three sets of lines in three colours are still legible where three washes of alpha are mud. Everything here works on plain axes and on a cartopy GeoAxes alike - pass `transform=` for the latter - so maps, cross-sections and the polar views share one implementation. """ import contextlib import warnings import numpy as np import matplotlib.pyplot as plt from matplotlib.lines import Line2D from matplotlib.patches import Patch from matplotlib.colors import ListedColormap from .. import bivariate, overlay from ..colors import BLANK_COLOR, blank_cells, curve_color # Every style this module can draw. `panels` is not among them: it does not draw # onto shared axes at all, it makes more of them, so it is handled a level up. STYLES = ('blend', 'contour', 'hatch', 'bivariate', 'glyph') # Where the colour bars sit, in axes coordinates: the plotted variable first, # then one per overlay that needs one, a pitch apart. # # They are insets rather than `fig.colorbar(ax=ax)`. That is not a preference: # a GeoAxes with `gridlines(draw_labels=True)` and a colorbar taking its room # out of the axes collapses to nothing under `bbox_inches='tight'` - an 8x6 # figure saves as 74 pixels wide, the map gone and only the bar left. It needs # both to happen; either alone is fine. An inset takes no room from the axes, so # the question does not arise. BAR_FIRST = 1.03 BAR_PITCH = 0.13 BAR_WIDTH = 0.03 BAR_BOTTOM = 0.08 BAR_HEIGHT = 0.72 # How many contour lines a layer gets when no threshold was named. Few enough to # read the labels between them. CONTOUR_LEVELS = 6 # Hatch patterns, one per layer. Distinguishable in print, which is the whole # reason to reach for a hatch. HATCH_PATTERNS = ('//', '\\\\', 'xx', '..', '++', 'oo') @contextlib.contextmanager def quiet_wrap_warning(): """ Silence cartopy's "gouraud shading across a wrap may leave artifacts" for the polar views, and only for them. There the trigger is the *other* hemisphere. Antipodal points project to enormous coordinates in a polar stereographic, which trips cartopy's "this cell is too big to be a cell" test - restricting the field to one hemisphere silences it. Those cells are never on screen: the view is clipped to 60 degrees from the pole, so there is nothing the reader could do with the warning and nothing it would tell them about their data. On a lat/lon map the same warning is left to sound, because there it means something: cells really are overhanging the dateline, which on a grid that goes right round means the seam was built wrong. That is exactly how the duplicate cyclic column in `render_geomap` was found - suppressing it there hid a fabricated column instead of a cartopy quirk. """ with warnings.catch_warnings(): warnings.filterwarnings('ignore', message='Handling wrapped coordinates') yield def mesh_shading(interpolate, x, y, values): """ 'gouraud' when the field is to be drawn interpolated and the arrays allow it, 'auto' - one flat quad per cell - otherwise. Gouraud blends colour across each quad from its four corner values, which is what "the interpolated field" means here: the cells are an artefact of the discretisation, and on an 11.25-degree grid the mosaic is the first thing the eye reads. It needs the coordinates and the data to describe the same points. A coordinate array one longer than the data - cell edges rather than centres, a staggered pair - is exactly the case 'auto' exists to resolve, and keeps it rather than being refused. """ if not interpolate: return 'auto' shape = np.shape(values)[:2] x_shape, y_shape = np.shape(x), np.shape(y) if len(x_shape) == 1 and len(y_shape) == 1: # 1D cell centres, one per column and one per row return 'gouraud' if shape == (y_shape[0], x_shape[0]) else 'auto' # The 2D meshgrids the polar views pass return 'gouraud' if x_shape[:2] == shape and y_shape[:2] == shape else 'auto' def bar_axes(ax, index): """ Axes for the `index`-th colour bar, to the right of the plot. """ return ax.inset_axes([BAR_FIRST + BAR_PITCH * index, BAR_BOTTOM, BAR_WIDTH, BAR_HEIGHT], transform=ax.transAxes) def bar_count(composite, style): """ How many colour bars a composite will want, including the base's own. Only `blend` gives every layer a bar. A contour set is named in the legend and labelled on its own lines; a hatch means presence, which a bar cannot say; a bivariate scheme has a square key instead of two bars. Knowing this before anything is drawn is what lets the figure be made wide enough. """ extra = len(composite.layers) - 1 if composite else 0 if style == 'blend': return 1 + extra # A bivariate scheme has one image and a square key drawn inside the axes, # so it asks for no bar and no extra width - but never fewer than the one # slot every other style starts from, or the figure would come out narrower # than a plain map. return 1 def draw_base(fig, ax, values, layers, x, y, colormap, norm, label, style='blend', transform=None, interpolate=True): """ Draw the plotted variable itself, and say whether it still wants a bar. Returns (artist, wants_bar). Every style but one draws the base exactly as it would be drawn alone and then puts the layers over it. `bivariate` is the exception: it encodes both variables in the one colour channel, so there is no "base mesh plus overlays" - there is one image of the pair, and a square key instead of two bars. """ kw = _transform_kw(transform) if style == 'bivariate': pair = _bivariate_pair(layers) if pair is not None: layer, other = pair rgb = bivariate.colours(values, other) mesh = ax.pcolormesh(x, y, rgb, shading=mesh_shading( interpolate, x, y, rgb), **kw) bivariate.legend_axes(fig, ax, label, layer.label) return mesh, False underlay = paint_blank_cells(ax, x, y, values, colormap, **kw) mesh = ax.pcolormesh(x, y, values, shading=mesh_shading(interpolate, x, y, values), cmap=colormap, norm=norm, **kw) # Hung off the mesh so an animation can step the two together without the # renderers having to carry a second artist around; see `refill_blank_cells`. mesh._dispnc_blank_underlay = underlay return share_wrapped_norm(mesh), True def paint_blank_cells(ax, x, y, values, colormap, **kw): """ Lay the cells a log scale cannot place under the mesh that cannot draw them. Returns the underlay, or None when there is nothing to paint. A log scale masks every zero and negative, and a masked cell is a hole in the picture - which reads as "the field is not here" rather than "the field is zero here". This puts a grey there instead. Drawn *under* rather than set as the colormap's 'bad' colour, because cartopy needs that one transparent to wrap a quadmesh across the dateline; see `colors.log_colormap`. Only a colormap `colors.choose_style` marked has anything to paint, so an ordinary figure never gains the extra artist, and a genuine NaN stays a hole on the figures that do. Drawn per cell whatever the mesh above it does. The underlay records *which cells* the scale could not place, which is a per-cell fact; interpolating it would spread grey into neighbouring cells that do have values, which is the misreading it exists to prevent. """ if not getattr(colormap, '_dispnc_blank_note', None): return None mask = blank_cells(values) if not mask.any(): return None # NaN everywhere else, so the underlay shows through only where the mesh # above it has nothing - and stays transparent, which is what cartopy needs # in order to wrap it. return ax.pcolormesh(x, y, _blank_layer(mask), shading='auto', vmin=0, vmax=1, cmap=ListedColormap([BLANK_COLOR]), **kw) def refill_blank_cells(mesh, values): """ Step a mesh's grey underlay to the frame the mesh has just been given. A field's zeros move with it - a cap grows and retreats - so an underlay left on frame zero would grey out cells the frame on screen has values for, which is the misreading this whole mechanism exists to prevent, one frame later. """ underlay = getattr(mesh, '_dispnc_blank_underlay', None) if underlay is not None: underlay.set_array(_blank_layer(blank_cells(values))) def _blank_layer(mask): """ The underlay's own array: something to colour where the mask is set, and NaN everywhere else - which stays transparent, as cartopy needs it to be to wrap a quadmesh across the dateline. """ return np.where(mask, 1.0, np.nan) def share_wrapped_norm(mesh): """ Give cartopy's dateline collection the mesh's own norm *object*. Returns the mesh, so it can wrap a `pcolormesh` call in place. A cell that straddles 180 degrees cannot be drawn as part of a QuadMesh, so cartopy masks those cells out and redraws them with `pcolor`, hanging the result off the mesh as `_wrapped_collection_fix`. When no norm was passed - the ordinary case, since `choose_style` leaves the scaling to matplotlib - that second collection is built with a Normalize of its own, and nothing keeps the two in step afterwards. A colour bar widening a degenerate range is what makes that visible: a field that is zero everywhere autoscales to vmin == vmax, `fig.colorbar` pushes the base's limits out to +-0.1, and the wrapped cells - still on their own norm, still degenerate - are left mapping every value to the bottom of the colormap. `runoff_acc` came out uniform with a dark column at -180 and another at +180, reading as signal where the field is flat. Sharing the one object fixes that and every later rescale with it. A mesh drawn on plain axes has no such collection, and passes through untouched. Neither has an interpolated one on a lat/lon map: gouraud draws the quads between the cell centres, so nothing overhangs 180 degrees and there is nothing to wrap. This still has to be here for the per-cell figures, which do overhang by half a cell at each end. """ fix = getattr(mesh, '_wrapped_collection_fix', None) if fix is not None and fix.norm is not mesh.norm: fix.norm = mesh.norm return mesh def _bivariate_pair(layers): """ The single overlay a bivariate scheme can encode, or None having said why. Two variables have a square; three have nothing. Refusing by name beats picking two of the three and not saying which. """ if len(layers) == 1: layer, values, _ = layers[0] return layer, values names = ', '.join(f"'{layer.varname}'" for layer, _, _ in layers) print(f"Warning: --overlay-style bivariate encodes exactly two variables in " f"one colour scheme, and {len(layers) + 1} were given ({names}); " f"blending them instead.") return None def draw(fig, ax, layers, x, y, style='blend', transform=None, bar_start=1, interpolate=True): """ Draw the prepared overlay layers over a base that is already on the axes. `layers` is [(layer, values, frames)] - already sliced, oriented and put through whatever the renderer had to do to its coordinates, so that what is drawn is what was prepared. Returns [(layer, artist, frames)]; `artist` is None for the styles that cannot be refilled in place, which is how the animation knows to redraw them instead. `interpolate` reaches only `_blend`, the one style that puts a layer in the colour channel. The others are already geometry rather than a mesh, and take it only to keep one drawer signature. """ if not layers: return [] if style not in STYLES: print(f"Warning: unknown --overlay-style '{style}'; using blend.") style = 'blend' if style == 'bivariate': # Already drawn, as one image of both variables, by `draw_base` - unless # there was not exactly one layer to pair with, in which case it has # already said so and blending is what is left. Asked by counting rather # than by calling `_bivariate_pair` again, which would say it twice. if len(layers) == 1: return [] style = 'blend' drawer = _DRAWERS[style] drawn = [] handles = [] for index, (layer, values, frames) in enumerate(layers): artist, handle = drawer(fig, ax, layer, values, index, x, y, transform, bar_start, interpolate) drawn.append((layer, artist, frames)) if handle is not None: handles.append(handle) # A colour bar names its own layer; a line or a hatch does not, so those get # a legend instead. Only one of the two ever appears. if handles: ax.legend(handles=handles, loc='lower left', fontsize=8, framealpha=0.85).set_zorder(5) return drawn def _transform_kw(transform): return {'transform': transform} if transform is not None else {} def _blend(fig, ax, layer, values, index, x, y, transform, bar_start, interpolate=True): """ The original: translucent colour, its own scale, its own bar. Every layer is translucent where its own field is weak, so an overlap reads as an overlap rather than as one variable winning. That is a deliberate trade - a cell holding both CO2 and H2O comes out a colour in neither bar - and it is why the other styles exist. """ mesh = share_wrapped_norm( ax.pcolormesh(x, y, values, shading=mesh_shading(interpolate, x, y, values), cmap=layer.colormap, norm=layer.norm, alpha=overlay.layer_alpha(values, layer.norm), **_transform_kw(transform))) bar = fig.colorbar(mesh, cax=bar_axes(ax, bar_start + index)) bar.set_label(layer.label, fontsize=8) bar.ax.tick_params(labelsize=7) # The per-cell alpha would ride into the bar as stripes bar.set_alpha(1.0) return mesh, None def _contour(fig, ax, layer, values, index, x, y, transform, bar_start, interpolate=True): """ Lines over the base, labelled with their own values. Nothing mixes: the overlay is geometry, not paint, so the base keeps its full colormap and both fields stay readable everywhere, including where they cross. This is how a shaded field under contoured heights has been drawn in atmospheric science for decades, and it is the style to reach for when the question is about both variables rather than about their overlap. """ colour = curve_color(layer.requested, index) levels = _levels(values, layer) if levels is None: print(f"Warning: --overlay '{layer.varname}' is flat here; " f"there is nothing to contour.") return None, None lines = ax.contour(x, y, values, levels=levels, colors=[colour], linewidths=1.0, **_transform_kw(transform)) try: # Four significant figures: '%g' spells a surface pressure 101961.234, # which is six digits of noise crowding the line it belongs to. ax.clabel(lines, inline=True, fontsize=7, fmt='%.4g') except Exception: # Labels are a courtesy; a degenerate level set must not cost the lines pass return lines, Line2D([], [], color=colour, label=layer.label) def _hatch(fig, ax, layer, values, index, x, y, transform, bar_start, interpolate=True): """ A pattern wherever the layer is above its threshold. This encodes *presence*, not amount - which is exactly the question for "where is there any CO2 ice at all" or "where is this difference significant", and exactly the wrong tool if the reader wants a magnitude. It is also the only style that survives being photocopied. """ pattern = HATCH_PATTERNS[index % len(HATCH_PATTERNS)] threshold = _threshold(values, layer) if threshold is None: print(f"Warning: --overlay '{layer.varname}' has no finite values here; " f"nothing to hatch.") return None, None top = float(np.nanmax(values)) if not top > threshold: print(f"Warning: --overlay '{layer.varname}' never rises above " f"{threshold:g} here; nothing to hatch.") return None, None filled = ax.contourf(x, y, values, levels=[threshold, top], colors='none', hatches=[pattern], **_transform_kw(transform)) for collection in filled.collections if hasattr(filled, 'collections') else [filled]: collection.set_edgecolor(curve_color(layer.requested, index)) collection.set_linewidth(0.0) return None, Patch(facecolor='none', hatch=pattern, edgecolor=curve_color(layer.requested, index), label=f"{layer.label} > {threshold:g}") def _glyph(fig, ax, layer, values, index, x, y, transform, bar_start, interpolate=True): """ The layer as symbols, sized by its own values, over an untouched base. Good for a sparse, patchy field - an ice cap, a plume - and poor for a smooth one, where a regular lattice of dots says more about the sampling than about the data. """ colour = curve_color(layer.requested, index) step = max(1, int(np.ceil(max(np.shape(values)) / 24))) xs = np.asarray(x, dtype=float) ys = np.asarray(y, dtype=float) grid_x, grid_y = np.meshgrid(xs, ys) sub = values[::step, ::step] scale = layer.norm if layer.norm is not None else None weight = np.asarray(scale(sub) if scale is not None else sub, dtype=float) weight = np.ma.filled(np.ma.masked_invalid(weight), np.nan) finite = np.isfinite(weight) if not finite.any(): print(f"Warning: --overlay '{layer.varname}' has no finite values here; " f"nothing to draw as symbols.") return None, None sizes = 4 + 90 * np.clip(np.nan_to_num(weight), 0.0, 1.0) ax.scatter(grid_x[::step, ::step][finite], grid_y[::step, ::step][finite], s=sizes[finite], facecolors='none', edgecolors=colour, linewidths=0.7, **_transform_kw(transform)) return None, Line2D([], [], color=colour, marker='o', linestyle='none', markerfacecolor='none', label=layer.label) _DRAWERS = { 'blend': _blend, 'contour': _contour, 'hatch': _hatch, 'glyph': _glyph, } def _levels(values, layer): """ Which contour levels one layer gets. An explicitly named set - `--overlay-threshold p90,p99` - is used as given. Otherwise the range is split evenly, on the layer's own scale, so a field drawn on a log norm gets logarithmically spaced lines rather than six of them bunched at the top. """ named = overlay.resolve_thresholds(values, getattr(layer, 'thresholds', None)) if named: return named finite = np.asarray(values, dtype=float) finite = finite[np.isfinite(finite)] if finite.size < 2: return None low, high = float(finite.min()), float(finite.max()) if not high > low: return None positive = finite[finite > 0] if _is_log(layer.norm) and positive.size >= 2: return list(np.geomspace(positive.min(), positive.max(), CONTOUR_LEVELS + 2)[1:-1]) return list(np.linspace(low, high, CONTOUR_LEVELS + 2)[1:-1]) def _threshold(values, layer): """ The one level a hatch is drawn above. """ named = overlay.resolve_thresholds(values, getattr(layer, 'thresholds', None), default=('p90',)) return float(min(named)) if named else None def _is_log(norm): return norm is not None and type(norm).__name__ == 'LogNorm'