""" Drawing several variables in one picture. The layering rules are pure numpy, so none of this needs a figure. """ import numpy as np import pytest from matplotlib.colors import LogNorm, Normalize from dispnc.colors import LAYER_COLORMAPS from dispnc.overlay import (Composite, Layer, assign_colormaps, describe, layer_alpha, parse_specs) def test_specs_take_an_optional_colormap_after_a_colon(): """ A colon, not a comma: a comma already means "two components of one vector" in --vector, and anyone who knows that would read a comma wrongly here. """ assert parse_specs(['h2o_ice']) == [('h2o_ice', None)] assert parse_specs(['h2o_ice:Blues']) == [('h2o_ice', 'Blues')] assert parse_specs(['a:Blues', 'b']) == [('a', 'Blues'), ('b', None)] # Repeated flags and one comma-separated flag mean the same thing assert parse_specs(['a,b']) == parse_specs(['a', 'b']) assert parse_specs(None) == [] and parse_specs([' ']) == [] def test_every_layer_gets_a_colour_no_other_layer_is_using(): """ Two layers in one colour is the single thing a composite cannot survive, so the rota has to skip what is already spoken for - including the plotted variable's own map, chosen before any of this ran. """ assigned = assign_colormaps(parse_specs(['a', 'b', 'c']), taken=['viridis']) colormaps = [c for _, c in assigned] assert len(set(colormaps)) == 3 assert 'viridis' not in colormaps assert all(c in LAYER_COLORMAPS for c in colormaps) def test_a_named_colormap_is_kept_and_still_blocks_the_rota(): assigned = assign_colormaps(parse_specs(['a:Reds', 'b']), taken=[]) colormaps = [c for _, c in assigned] assert colormaps[0] == 'Reds' assert colormaps[1] != 'Reds' def test_layers_are_invisible_where_their_own_field_is_absent(): """ A layer is drawn over the ones below it, so a layer that renders its own floor as a faint wash does not merely look wrong - it hides everything beneath it, everywhere. """ values = np.array([[0.0, 0.5, 1.0], [np.nan, 0.01, 0.9]]) alpha = layer_alpha(values, Normalize(vmin=0.0, vmax=1.0)) assert alpha[0, 0] == 0.0, "the bottom of the scale is nothing" assert alpha[1, 0] == 0.0, "and so is missing data" assert alpha[1, 1] == 0.0, "and so is a value below the floor" assert alpha[0, 2] > alpha[0, 1] > 0.0 assert alpha.max() <= 1.0 def test_layer_alpha_follows_a_log_scale_where_the_colour_does(): """ Transparency and colour have to say the same thing, or a field spanning decades is drawn in full colour and no opacity at all. """ values = np.array([[1e-9, 1e-5, 1e-1]]) alpha = layer_alpha(values, LogNorm(vmin=1e-9, vmax=1e-1)) assert alpha[0, 0] == 0.0 assert 0.0 < alpha[0, 1] < alpha[0, 2] def test_a_composite_of_one_is_not_a_composite(): """ The renderers ask `if ctx.composite`, so a lone layer has to be falsy or every ordinary plot takes the overlay path. """ lone = Composite([Layer('a', np.zeros((2, 2)), 'viridis')]) pair = Composite([Layer('a', np.zeros((2, 2)), 'viridis'), Layer('b', np.zeros((2, 2)), 'Reds')]) assert not lone assert pair assert describe(pair) == 'b over a' def test_a_layer_labels_itself_with_its_units(): layer = Layer('co2_ice', np.zeros((2, 2)), 'Reds', units='kg/m2', long_name='CO2 ice') assert layer.label == 'CO2 ice [kg/m2]' assert Layer('x', np.zeros((2, 2)), 'Reds').label == 'x' def test_a_layer_on_the_wrong_grid_is_refused_by_name(capsys, tmp_path): """ The layers share one set of axes, so they have to be one shape. Broadcasting a mismatch would put the data in the wrong place rather than failing, which is the sort of wrong that survives review. """ from dispnc.conventions import normalize from dispnc.io import open_source from dispnc.pipeline import plot_variable src = open_source('sample_files/start.nc') try: _, report = normalize(src.ds) # `vcov` is on the staggered rlatv grid: 32 latitudes where the others # have 33, so it cannot be laid over them plot_variable(src, report, 'h2o_ice', interactive=False, output_path=str(tmp_path / 'map.png'), extra_indices={'Time': 0, 'altitude': 0}, plot_kind='geomap', overlay=['vcov']) finally: src.close() said = capsys.readouterr().out assert 'vcov' in said, 'the layer that was dropped has to be named' assert 'skipping' in said # And the figure is still drawn, without it, rather than the whole command # failing over one layer assert 'Saved' in said def test_a_layer_is_reordered_with_the_longitudes_it_is_drawn_against(): """ start.nc's `rlonu` reaches 185.6 degrees, so wrapping it into [-180, 180] moves a column to the front and the base field's columns move with it. The layers used to be drawn raw against the reordered axis, so two *identical* fields came out a column apart - a mis-registration over the whole map that nothing announced. """ from dispnc.render.geomap import _normalize_longitudes, _prepare_layers lons = np.linspace(-174.375, 185.625, 33) field = np.tile(np.arange(33, dtype=float), (5, 1)) _, base = _normalize_longitudes(lons, field) assert not np.array_equal(base, field), 'this axis has to actually reorder' layer = Layer(varname='cu', data=field.copy(), colormap='Reds') composite = Composite(layers=[Layer('ucov', field.copy(), 'Blues'), layer]) (_, values, _), = _prepare_layers(composite, lons, cyclic=False) assert np.array_equal(values, base) def test_an_animated_layer_has_its_frames_reordered_too(): """ The frames are what a movie steps through, so a layer whose values were reordered and whose frames were not would register correctly on the first frame and wrongly on every other one. """ from dispnc.render.geomap import _normalize_longitudes, _prepare_layers lons = np.linspace(-174.375, 185.625, 33) frames = np.arange(4 * 5 * 33, dtype=float).reshape(4, 5, 33) _, expected = _normalize_longitudes(lons, frames) layer = Layer(varname='co2_ice', data=frames[0], colormap='Reds', frames=frames) composite = Composite(layers=[Layer('ps', frames[0], 'Blues'), layer]) (_, _, prepared), = _prepare_layers(composite, lons, cyclic=False) assert prepared is not None, 'the frames have to survive preparation' assert np.array_equal(prepared, expected) def test_a_layer_with_no_frames_keeps_the_values_it_was_prepared_with(): """ A mask or a topography carries no time, and must hold still rather than vanish when the field under it starts moving. """ from dispnc.render.geomap import _prepare_layers lons = np.linspace(-180, 168.75, 32) field = np.ones((5, 32)) composite = Composite(layers=[Layer('ps', field, 'Blues'), Layer('phisinit', field, 'Reds')]) (_, values, frames), = _prepare_layers(composite, lons, cyclic=False) assert frames is None assert np.array_equal(values, field) def test_every_layer_steps_with_the_field_under_it(monkeypatch): """ The regression: `_animate` used to refill the base mesh alone, so a year of surface pressure ran underneath an ice field frozen on step zero. The transparency has to follow, too - it is derived from the values, so a stale alpha shows this frame's field through the last one's holes. """ import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt import cartopy.crs as ccrs from dispnc.render import geomap, layers as layer_draw lons = np.linspace(-180, 168.75, 32) lats = np.linspace(-87.5, 87.5, 12) rng = np.random.default_rng(0) base_frames, layer_frames = rng.random((4, 12, 32)), rng.random((4, 12, 32)) layer = Layer('co2_ice', layer_frames[0], 'Reds', norm=Normalize(0, 1), frames=layer_frames) composite = Composite(layers=[Layer('ps', base_frames[0], 'Blues'), layer]) # A real PlotContext, not a stand-in: a hand-rolled stub drifts from the # dataclass the moment a renderer starts reading a new field off it. from dispnc.render.context import PlotContext ctx = PlotContext(varname='ps', data=base_frames[0], plan=None, frame_dim='time_counter', interactive=False) stepped = {} def fake_drive(fig, ctx, on_frame, output_path, label): on_frame(3) return False monkeypatch.setattr(geomap, 'drive', fake_drive) proj = ccrs.PlateCarree() fig, ax = plt.subplots(subplot_kw=dict(projection=proj)) try: prepared = geomap._prepare_layers(composite, lons, cyclic=False) drawn = layer_draw.draw(fig, ax, prepared, lons, lats, transform=proj) mesh = ax.pcolormesh(lons, lats, base_frames[0], shading='auto', transform=proj) geomap._animate(fig, ax, mesh, ctx, base_frames, lons, lats, drawn) _, layer_mesh, _ = drawn[0] stepped['layer'] = np.asarray(layer_mesh.get_array()).reshape(12, 32) stepped['alpha'] = np.asarray(layer_mesh.get_alpha()).reshape(12, 32) finally: plt.close(fig) assert np.allclose(stepped['layer'], layer_frames[3]), \ 'the overlay stayed on frame 0 while the base moved' assert np.allclose(stepped['alpha'], layer_alpha(layer_frames[3], Normalize(0, 1))) @pytest.mark.parametrize('kind', ['scalar']) def test_a_plot_that_cannot_stack_variables_says_so(kind, capsys, pem_ds): """ The composite used to be built in full - every layer sliced, reduced, shape-checked and given a colormap - and then handed to a renderer that had never heard of it. One variable was drawn, and nothing said why. From the terminal that is indistinguishable from a layer that was refused. """ from dispnc.conventions import normalize from dispnc.pipeline import OVERLAYABLE, plot_variable assert kind not in OVERLAYABLE, 'this test is about the kinds left out' ds, report = normalize(pem_ds) class Src: def __init__(self, ds): self.ds = ds def __contains__(self, name): return name in self.ds.variables plot_variable(Src(ds), report, 'tsurf', interactive=False, plot_kind=kind, extra_indices={'Time': 0, 'latitude': 0, 'longitude': 0}, overlay=['tsurf']) said = capsys.readouterr().out assert 'does not stack variables' in said assert "'tsurf' alone" in said def test_the_rota_says_so_rather_than_quietly_repeating_a_colour(capsys): """ Past the end of the rota two layers share a colour, which is the one outcome a composite cannot survive. There is no fixing it - seven translucent fields on one map cannot be read - so it is reported instead. """ specs = [(f'v{i}', None) for i in range(len(LAYER_COLORMAPS) + 2)] assigned = assign_colormaps(specs) said = capsys.readouterr().out assert 'reuses the' in said and 'colours' in said assert len(assigned) == len(specs)