""" Several variables on one 1-D plot. A curve has no colour channel to share and no per-layer colour bar to say which scale a mark belongs to, so the whole question is which curves may share an axis - and what happens to the ones that may not. """ import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt import numpy as np import pytest from dispnc.coords import Axis from dispnc.inference import PlotPlan from dispnc.overlay import Composite, Layer, axis_groups, units_of from dispnc.render.context import PlotContext from dispnc.render.curves import render_profile, render_timeseries TIME = Axis(role='T', dim='Time', size=6, values=np.arange(6.0), units='year', long_name='Year of run') DEPTH = Axis(role='Z', dim='soildepth', size=6, values=np.linspace(0.1, 3.0, 6), units='m', long_name='Soil depth', positive='down') def layer(name, units, seed=0): rng = np.random.default_rng(seed) return Layer(varname=name, data=rng.random(6) * 100, colormap='viridis', units=units, long_name=name) def context(kind, layers, **kwargs): plan = (PlotPlan(kind='timeseries', x=TIME) if kind == 'timeseries' else PlotPlan(kind='profile', y=DEPTH, invert_y=True)) base = layers[0] return PlotContext(varname=base.varname, data=base.data, plan=plan, label=f"{base.long_name} ({base.units})", units=base.units, long_name=base.long_name, interactive=False, composite=Composite(layers=layers) if len(layers) > 1 else None, **kwargs) # --- the decision table ------------------------------------------------------- @pytest.mark.parametrize('units, expected', [ (['K', 'K', 'K'], (3, 0, 0)), (['K', 'Pa'], (1, 1, 0)), (['K', 'Pa', 'W/m2'], (1, 1, 1)), (['K', 'K', 'Pa'], (2, 1, 0)), ([None, None], (2, 0, 0)), (['K', None], (1, 1, 0)), (['K'], (1, 0, 0)), ]) def test_which_curves_may_share_an_axis(units, expected): """ One unit is one axis; two get a twin; past that there is no third side to the figure. A variable with no units is its own group - "unknown" cannot be claimed to match "K". """ layers = [layer(f'v{i}', u) for i, u in enumerate(units)] primary, secondary, refused = axis_groups(layers) assert (len(primary), len(secondary), len(refused)) == expected def test_a_group_is_named_by_the_unit_it_shares(): assert units_of([layer('a', 'K'), layer('b', 'K')]) == 'K' assert units_of([layer('a', None)]) == '' # --- what actually lands on the axes ----------------------------------------- def test_several_variables_share_one_plot_with_a_legend(): ctx = context('timeseries', [layer('tsurf', 'K', 0), layer('tsoil', 'K', 1)]) fig = _draw(ctx, render_timeseries) try: ax = fig.axes[0] assert len(ax.lines) == 2 legend = ax.get_legend() assert legend is not None assert len(legend.get_texts()) == 2 # one shared axis, named by the unit rather than by either variable assert ax.get_ylabel() == '[K]' assert len(fig.axes) == 1, 'same units must not create a twin' finally: plt.close(fig) def test_a_second_unit_gets_the_opposite_axis(tmp_path): ctx = context('timeseries', [layer('tsurf', 'K', 0), layer('ps', 'Pa', 1)]) fig = _draw(ctx, render_timeseries) try: assert len(fig.axes) == 2, 'a second unit needs its own scale' left, right = fig.axes assert left.get_ylabel() == 'tsurf [K]' assert right.get_ylabel() == 'ps [Pa]' labels = [t.get_text() for t in left.get_legend().get_texts()] assert labels == ['tsurf [K]', 'ps [Pa] [right axis]'] finally: plt.close(fig) def test_a_profile_twins_on_x_because_its_shared_axis_is_the_vertical(): """ The value goes on X and depth up Y, so a second unit is a second *x* scale. """ ctx = context('profile', [layer('tsoil', 'K', 0), layer('rho', 'kg/m3', 1)]) fig = _draw(ctx, render_profile) try: assert len(fig.axes) == 2 bottom, top = fig.axes assert bottom.get_xlabel() == 'tsoil [K]' assert top.get_xlabel() == 'rho [kg/m3]' assert bottom.get_ylabel() == 'Soil depth (m)' labels = [t.get_text() for t in bottom.get_legend().get_texts()] assert labels[1].endswith('[top axis]') finally: plt.close(fig) def test_a_third_unit_is_refused_by_name_and_points_at_panels(capsys): ctx = context('timeseries', [layer('tsurf', 'K', 0), layer('ps', 'Pa', 1), layer('rad', 'W/m2', 2)]) fig = _draw(ctx, render_timeseries) try: said = capsys.readouterr().out assert "'rad' is in W/m2" in said assert 'two scales at most' in said assert '--overlay-style panels' in said # the base always survives, and so does the one that did fit assert len(fig.axes[0].lines) == 1 and len(fig.axes[1].lines) == 1 finally: plt.close(fig) def test_a_single_variable_plot_gains_no_legend_and_keeps_its_label(): """ The regression guard: a legend on every existing curve, or a relabelled axis, would move every one of them for nothing. """ ctx = context('timeseries', [layer('tsurf', 'K', 0)]) fig = _draw(ctx, render_timeseries) try: ax = fig.axes[0] assert ax.get_legend() is None assert ax.get_ylabel() == 'tsurf (K)', 'the long-standing label style' finally: plt.close(fig) def test_every_curve_gets_its_own_colour(): ctx = context('timeseries', [layer('a', 'K', 0), layer('b', 'K', 1), layer('c', 'K', 2)]) fig = _draw(ctx, render_timeseries) try: colours = {line.get_color() for line in fig.axes[0].lines} assert len(colours) == 3 finally: plt.close(fig) def test_a_named_colour_is_used_for_that_curve(): second = layer('co2_ice', 'K', 1) second.requested = 'red' ctx = context('timeseries', [layer('tsurf', 'K', 0), second]) fig = _draw(ctx, render_timeseries) try: assert fig.axes[0].lines[1].get_color() == '#ff0000' finally: plt.close(fig) def _draw(ctx, renderer): """ Run a renderer without letting it save or show, and hand back its figure. """ import dispnc.figure as figure_mod saved = figure_mod.plt.show figure_mod.plt.show = lambda *a, **k: None before = set(plt.get_fignums()) try: renderer(ctx) finally: figure_mod.plt.show = saved new = set(plt.get_fignums()) - before return plt.figure(sorted(new)[-1])