""" The path from a variable name to a figure. select -> remap -> resolve axes -> infer -> orient -> render Each step is a function in its own module; this file only sequences them. The pre-refactor code did all of it inside one 144-line function that dispatched on the number of dimensions. """ from dataclasses import replace import numpy as np from .colors import choose_style from .io import materialize from .coords import Axis, resolve_variable_axes from .derive import (DeriveError, apply_anomaly, apply_difference, apply_reduce, check_reduce_conflicts, parse_reduce) from .inference import SubsurfaceShell, TooManyDimensions, infer from . import overlay as overlay_mod from .interactive import prompt_choice from .render import RENDERERS, PlotContext from .selection import apply_selection, select_lazy from .stats import area_weights, print_summary from .unstructured import build_grid, to_grid from . import vertical # Plot kinds that can step through a dimension. A curve, a profile or a section # has no spare axis to spend on time. ANIMATABLE = frozenset({'geomap', 'globe'}) # Plot kinds that can carry more than one variable in one frame. The rest build # the composite, hand it to a renderer that has never heard of it, and draw one # variable - which is indistinguishable, from the terminal, from a layer that # was refused. `--animate` has said which kinds it applies to since it was # written; this is the same courtesy. # # Kinds join this set as their renderers learn to read `ctx.composite`. Adding # one here before that is what the gate exists to prevent. OVERLAYABLE = frozenset({'geomap', 'section', 'globe', 'line', 'timeseries', 'profile'}) def _slice_and_reduce(da, extra_indices, reductions): """ The lazy selection and reduction every field on the plan goes through. Returns (array, remaining dims, note, dimensions reduced away). Reductions run after the selection, because `-e '{"Time": "avg"}' --reduce max:lon` is the maximum of the mean and not the mean of the maxima; `check_reduce_conflicts` has already made sure the two cannot name the same dimension. Nothing is read until the last line. """ lazy = select_lazy(da, dict(extra_indices or {})) note, reduced = '', [] if reductions: lazy, note, reduced = apply_reduce(lazy, reductions) return materialize(lazy), list(lazy.dims), note, reduced def _axes_for(ds, varname, remaining, unstructured): """ Resolve the axes of a variable and keep those that survived the selection, in data order. """ resolved = resolve_variable_axes(ds, varname, unstructured) by_dim = {a.dim: a for a in resolved} return [by_dim[d] for d in remaining if d in by_dim] def _remap_unstructured(ds, data, axes, grid, info): """ Turn a field on the flattened physics grid into a latitude/longitude field. Only done when the physics dimension is the only one left, i.e. when the result is actually a map. With other dimensions still present the physics dimension stays an index axis and the field is drawn as a cross-section, which is what the original code did too. """ data = to_grid(data, grid) lat_axis = Axis(role='Y', dim='lat', size=grid.lats.size, coord=info['lat'], values=grid.lats, units='degrees_north', long_name='Latitude', ascending=True, source='unstructured', confidence=95) lon_axis = Axis(role='X', dim='lon', size=grid.lons.size, coord=info['lon'], values=grid.lons, units='degrees_east', long_name='Longitude', ascending=True, is_cyclic=grid.seam_source is not None, source='unstructured', confidence=95) return data, [lat_axis, lon_axis] def _orient(data, axes, plan): """ Transpose a field into (y, x) order, or (z, y, x) for the 3D globe, so renderers never deal with dimension order themselves. """ if plan.x is None or plan.y is None: return data dims = [a.dim for a in axes] if plan.z is not None and data.ndim == 3: return np.moveaxis( data, [dims.index(plan.z.dim), dims.index(plan.y.dim), dims.index(plan.x.dim)], [0, 1, 2]) if data.ndim != 2: return data return np.moveaxis(data, [dims.index(plan.y.dim), dims.index(plan.x.dim)], [0, 1]) def plot_variable(src, report, varname, colormap='jet', output_path=None, extra_indices=None, interactive=True, show_polar=False, show_3d=False, x_dim=None, show_topo=True, plot_kind=None, remap=True, norm=None, explain=False, anomaly=None, reduce=None, stats=False, stats_only=False, vmin=None, vmax=None, globe_options=None, animate=None, overlay=None, figsize=None, dpi=None, title_override=None, overlay_style='blend', overlay_threshold=None, overlay_scale='own', diff=None, interpolate=True): """ Plot one variable. Returns a process exit status. """ ds = src.ds da = ds[varname] # Character/string variables (file descriptors and the like) are not plottable if da.dtype.kind not in 'iufc': print(f"Error: '{varname}' holds {da.dtype} data, which is not numeric; nothing to plot.") return 1 empty = [d for d, n in zip(da.dims, da.shape) if n == 0] if empty: print(f"Error: '{varname}' has dimension(s) of size 0 ({', '.join(empty)}); " f"there is nothing to plot.") return 1 extra_indices = dict(extra_indices or {}) signed = False # The anomaly reference is taken before slicing, so that "--anomaly time" # with "-e '{\"Time\": 0}'" means the first step's departure from the mean # of the whole run rather than a departure from itself. try: reductions = parse_reduce(reduce, list(da.dims)) check_reduce_conflicts(reductions, extra_indices) if anomaly: full_axes = resolve_variable_axes(ds, varname, report.unstructured) da, description = apply_anomaly(da, anomaly, full_axes) signed = True print(f"Anomaly: {description}") if diff: if diff not in ds: raise DeriveError(f"--diff: '{diff}' is not in the file") da, description = apply_difference(da, ds[diff], varname, diff) signed = True print(f"Difference: {description}") except DeriveError as err: print(f"Error: {err}") return 2 try: data, remaining, note, reduced = _slice_and_reduce(da, extra_indices, reductions) except Exception as err: print(f"Error: Cannot read data for '{varname}': {err}") return 1 axes = _axes_for(ds, varname, remaining, report.unstructured) if note: print(f"Reduced: {note}") if stats or stats_only: weights, weight_note = area_weights(ds, report, axes, data.shape, reduced=reduced) print_summary(varname, data, weights, weight_note, da.attrs.get('units')) if stats_only: return 0 # Rebuild a map from the unstructured physics grid when that is what is left if remap and len(axes) == 1 and axes[0].role == 'U' and report.unstructured: grid = build_grid(ds, report.unstructured) if grid is not None: data, axes = _remap_unstructured(ds, data, axes, grid, report.unstructured) # An animated dimension is held back from the plan rather than sliced away, # so what is inferred is the shape of a single frame frames = frame_axis = None if animate: data, axes, frames, frame_axis, note = _split_frames(data, axes, animate) if note: print(f"Animate: {note}") try: plan = infer(axes, x_dim=x_dim, plot_kind=plot_kind) except SubsurfaceShell as err: print(f"Error: '{err.axis.dim}' measures depth below the surface, so it " f"cannot be drawn as an atmospheric shell. Use --plot-kind section " f"for a cross-section, or select a level with --extra-indices.") return 1 except TooManyDimensions as err: names = [a.dim for a in err.axes] print(f"Error: {len(names)} dimensions remain ({', '.join(names)}); " f"restrict more of them via --extra-indices to get a plot.") return 1 # Maps and the globe step through a dimension; a curve or a section does # not, and would silently drop every frame but the first if frames is not None and plan.kind not in ANIMATABLE: print(f"Animate: a {plan.kind} plot is not animated " f"(only {', '.join(sorted(ANIMATABLE))}); showing one step") frames = frame_axis = None if interactive and plan.kind in ('section', 'geomap') and x_dim is None: chosen = _prompt_for_x(plan, axes) if chosen is not None and chosen != plan.x.dim: plan = infer(axes, x_dim=chosen, plot_kind=plot_kind) data = _orient(data, axes, plan) units = da.attrs.get('units') long_name = da.attrs.get('long_name') # Colours are chosen from the data unless the user named a colormap. For an # animation they come from every frame at once: a scale re-derived per frame # would shift under the data and make the whole globe flicker. if plan.kind in ('geomap', 'section', 'globe'): reference = frames if frames is not None else data colormap, norm, why = choose_style(reference, varname, cmap=colormap, norm=norm, vmin=vmin, vmax=vmax, signed=signed) if explain: # A Colormap object when the log scale had cells to paint, a name # otherwise; either way the line reads as a name. print(f"Colour: {getattr(colormap, 'name', colormap)} - {why}") base_layer = overlay_mod.Layer(varname=varname, data=data, colormap=colormap, norm=norm, units=units, long_name=long_name, frames=frames) composite = None if overlay: # `panels` needs nothing of the renderer beyond drawing into the axes it # is handed, so it works for every plot kind - which is why the refusal # below can offer an alternative instead of just saying no. if plan.kind in OVERLAYABLE or overlay_style == 'panels': composite = _build_composite( ds, report, overlay, base_layer, plan, extra_indices, reductions, # Only an animation that actually happened has anything to say # about a layer holding still: `--animate` over a single step is # already refused above, and reporting the layers of a figure # that is not moving would be noise about a problem that is not # there. animate if frames is not None else None, thresholds=overlay_threshold) else: print(f"Overlay: a {plan.kind} plot does not stack variables " f"(only {', '.join(sorted(OVERLAYABLE))}); use " f"--overlay-style panels for one panel each. " f"Showing '{varname}' alone.") ctx = PlotContext( varname=varname, data=data, plan=plan, label=(long_name or varname) + (f" ({units})" if units else ""), units=units, long_name=long_name, colormap=colormap, norm=norm, output_path=output_path, show_topo=show_topo, show_polar=show_polar, show_3d=show_3d, interactive=interactive, subtitle=_scalar_coord_subtitle(da, extra_indices), globe=_globe_options(globe_options, ds, report, extra_indices, reductions, data, plan), frames=frames, frame_dim=animate if frames is not None else None, frame_axis=frame_axis, fps=getattr(globe_options, 'fps', 12) if globe_options else 12, composite=composite, figsize=figsize, dpi=dpi, title_override=title_override, interpolate=interpolate, overlay_style=overlay_style, overlay_scale=overlay_scale, ) if composite and overlay_style == 'panels': return _render_panels(ctx, composite) return RENDERERS[plan.kind](ctx) def _render_panels(ctx, composite): """ One panel per variable, drawn by the same renderer the plan already chose. The renderers are not told about this. Each is handed a context whose data and labels are one layer's and whose axes come from the grid, so a map is still drawn by `render_geomap` and a curve by `render_timeseries` - the only difference is where it lands. Nothing is saved until the last panel, or each would write over the one before it. """ from .render.panels import layout, shared_style # A scalar has no figure to put in a panel; each layer simply prints itself, # which is what one panel each already means for a single number. grid = None if ctx.plan.kind != 'scalar': grid = layout(len(composite.layers), figsize=ctx.figsize) # Before anything is drawn: the last panel saves and closes the figure, # and a title set after that lands on nothing. # The slice is the same in every panel, so it is said once up here # rather than three times across three colliding panel titles. shared = ctx.title_override or overlay_mod.name_list(composite) if ctx.subtitle and not ctx.title_override: shared = f"{shared} - {ctx.subtitle}" grid.fig.suptitle(shared, fontweight='bold') shared_norm = None if ctx.overlay_scale == 'shared': shared_norm, note = shared_style(composite.layers, ctx.varname) if note: print(f"Panels: {note}.") status = 0 for index, layer in enumerate(composite.layers): last = index == len(composite.layers) - 1 panel_ctx = replace( ctx, varname=layer.varname, data=layer.data, label=layer.label, units=layer.units, long_name=layer.long_name, colormap=layer.colormap, norm=shared_norm if shared_norm is not None else layer.norm, frames=layer.frames, composite=None, # each panel is one variable again panel=grid, # A panel is named by its variable and nothing else. The renderers # each decorate a title in their own way - "over time", "profile", # "(soildepth vs lat)" - and three of those plus a shared slice is # more text than plot. title_override=layer.long_name or layer.varname, subtitle='', # Only the last panel writes the figure; the others would each save # a half-drawn copy over the one before. output_path=ctx.output_path if last else None, interactive=ctx.interactive and last, show_polar=False, # secondary views of a panel grid are not a thing show_3d=False, ) status = RENDERERS[ctx.plan.kind](panel_ctx) or status return status def _build_composite(ds, report, overlay, base, plan, extra_indices, reductions, animate, thresholds=None): """ The `--overlay` variables prepared onto the same plan as the base one, or None when none were asked for or none survived. Every layer goes through `_prepare_layer`, which repeats the selection, reduction and orientation the base variable has just been through - so an overlay cannot end up sliced differently from the thing it is drawn over. A layer that cannot be laid over the plan is refused by name rather than quietly dropped or, worse, broadcast into place. """ specs = overlay_mod.parse_specs(overlay) if not specs: return None layers = [base] assigned = overlay_mod.assign_colormaps(specs, taken=[base.colormap]) # The raw text after the colon travels with the layer alongside the colormap # it resolved to: a blended layer needs the ramp, a contour or a curve needs # the same word read as one colour, and only the caller knows which. for (name, requested), (_, colormap) in zip(specs, assigned): if name not in ds: print(f"Warning: --overlay '{name}' is not in the file; skipping it.") continue layer = _prepare_layer(ds, report, name, base, plan, extra_indices, reductions, animate, colormap) if layer is not None: layer.requested = requested layer.thresholds = thresholds layers.append(layer) if len(layers) < 2: return None return overlay_mod.Composite(layers=layers) def _prepare_layer(ds, report, varname, base, plan, extra_indices, reductions, animate, colormap): """ One overlay variable, sliced and oriented exactly like the base one. Returns None, having said why, when the variable cannot be laid over the plan: the wrong grid, a dimension the base one does not have, or nothing numeric to draw. """ da = ds[varname] if da.dtype.kind not in 'iufc': print(f"Warning: --overlay '{varname}' holds {da.dtype} data; skipping it.") return None try: data, remaining, _, _ = _slice_and_reduce(da, extra_indices, reductions) axes = _axes_for(ds, varname, remaining, report.unstructured) frames = None if animate: if any(a.dim == animate for a in axes): data, axes, frames, _, _ = _split_frames(data, axes, animate) else: # Said out loud, because a layer that silently holds still under # a field that is moving looks exactly like a layer that is not # changing over the run print(f"Overlay: '{varname}' has no '{animate}'; " f"it stays fixed across the frames.") # The same dimensions, not merely the same number of points, and asked # before orienting rather than after: these files are 33x33 and 33x32, # so a layer on the wrong pair of axes - or on the right pair # transposed - can match the base's shape by coincidence. `_orient` # would then either raise something incidental about a missing index or, # for a plan with a free axis, quietly draw the layer sideways. Naming # the dimensions is what turns the shape check below into a guarantee. wanted = {a.dim for a in (plan.z, plan.y, plan.x) if a is not None} found = {a.dim for a in axes} if wanted and found != wanted: print(f"Warning: --overlay '{varname}' is on " f"{', '.join(sorted(found)) or 'no dimensions'} where the " f"{plan.kind} is on {', '.join(sorted(wanted))}; skipping it.") return None shaped = _orient(data, axes, plan) except Exception as err: print(f"Warning: --overlay '{varname}' cannot be laid over the " f"{plan.kind} ({err}); skipping it.") return None # The layers are drawn on 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. if np.shape(shaped) != np.shape(base.data): print(f"Warning: --overlay '{varname}' is {np.shape(shaped)} where " f"'{base.varname}' is {np.shape(base.data)}; skipping it.") return None _, norm, _ = choose_style(frames if frames is not None else shaped, varname, cmap=colormap) return overlay_mod.Layer(varname=varname, data=shaped, colormap=colormap, norm=norm, units=da.attrs.get('units'), long_name=da.attrs.get('long_name'), frames=frames) def _split_frames(data, axes, animate): """ (first frame, remaining axes, every frame with that axis first, that axis, note). The animated dimension is moved to the front and taken out of the axis list, so `infer` sees only what one frame looks like. Returns the data untouched when the request cannot be honoured. Which renderers can then use the frames is not decided here: the plan does not exist yet at this point, so the check belongs after `infer`. """ names = [a.dim for a in axes] if animate not in names: return data, axes, None, None, ( f"'{animate}' is not among the remaining dimensions " f"({', '.join(names) or 'none'}); ignored") index = names.index(animate) if data.shape[index] < 2: return data, axes, None, None, \ f"'{animate}' has a single step; nothing to animate" frames = np.moveaxis(data, index, 0) return (frames[0], [a for a in axes if a.dim != animate], frames, axes[index], f"{frames.shape[0]} frames over '{animate}'") def _globe_options(options, ds, report, extra_indices, reductions, data, plan): """ Resolve what the globe needs from the file: `--globe-relief`, and the true altitudes of an atmospheric shell. Everything is read here and handed on as arrays, so the renderer never needs dataset access. """ if options is None: return None options.altitudes = _shell_altitudes(ds, report, extra_indices, reductions, plan) mode = options.relief_mode or 'topo' if mode in ('topo', 'none') or plan.kind != 'geomap': return options if mode == 'data': options.relief, options.relief_label = data, 'the plotted field' return options if mode not in ds: print(f"Warning: --globe-relief '{mode}' is not in the file; using topography.") options.relief_mode = 'topo' return options try: options.relief = _field_on_plan(ds, report, mode, extra_indices, reductions, plan) options.relief_label = mode except Exception as err: print(f"Warning: could not use '{mode}' as relief ({err}); using topography.") options.relief_mode = 'topo' return options def _field_on_plan(ds, report, varname, extra_indices, reductions, plan): """ Another variable of the file, sliced, reduced and oriented like the plotted one - the same three steps `_prepare_layer` puts an --overlay through. Only what applies is used: a `Time` pinned or averaged on the plotted variable is pinned or averaged here too, while a dimension this variable does not have is simply not its business. """ da = ds[varname] usable = {d: i for d, i in (extra_indices or {}).items() if d in da.dims} values, remaining, _, _ = _slice_and_reduce(da, usable, reductions) axes = _axes_for(ds, varname, remaining, report.unstructured) return _orient(values, axes, plan) def _shell_altitudes(ds, report, extra_indices, reductions, plan): """ Where the levels of an atmospheric shell really are, or None. The vertical axis of an LMDZ file is a pseudo-altitude: one number per level for the whole planet, when the levels the model integrates follow the ground at the bottom and flatten into isobars at the top. `vertical.air_altitudes` rebuilds them from the hybrid coefficients and the surface pressure; all it needs from here is that pressure on the plan's own grid, which is the same slicing every other field goes through. The coefficients are looked for before the pressure is read: `ps` over a 669-step year is 50 MB, and reading it to find out there is no hybrid coordinate to use it with would be 50 MB spent on nothing. """ if plan.kind != 'globe' or plan.z is None: return None if vertical.mid_coefficients(ds, np.size(plan.z.values)) is None: return None name = vertical.surface_pressure_name(ds) if name is None: return None try: return vertical.air_altitudes( ds, plan, _field_on_plan(ds, report, name, extra_indices, reductions, plan)) except Exception as err: print(f"Warning: could not place the levels from the hybrid coordinate " f"({err}); using '{plan.z.dim}' as it is.") return None def plot_vector_field(src, report, u_name, v_name, colormap='auto', output_path=None, extra_indices=None, show_topo=True, style='quiver', density=None, background='magnitude', norm=None, vmin=None, vmax=None, interactive=False, show_polar=False, show_3d=False, globe_options=None, figsize=None, dpi=None, title_override=None, interpolate=True): """ Plot a vector field from two named components. Returns an exit status. """ from .render.vectors import VectorError, check_pair, render_vectors ds = src.ds for name in (u_name, v_name): if name not in src: print(f"Error: variable '{name}' not in file.") return 1 u_da, v_da = ds[u_name], ds[v_name] try: check_pair(u_da, v_da) except VectorError as err: print(f"Error: {err}") return 1 extra_indices = dict(extra_indices or {}) autoselected = {d: 0 for d, n in zip(u_da.dims, u_da.shape) if n == 1 and d not in extra_indices} extra_indices.update(autoselected) try: u, remaining = apply_selection(u_da, extra_indices) v, _ = apply_selection(v_da, extra_indices) except Exception as err: print(f"Error: Cannot read vector components: {err}") return 1 axes = _axes_for(ds, u_name, remaining, report.unstructured) if len(axes) != 2 or {a.role for a in axes} != {'X', 'Y'}: print(f"Error: a vector plot needs exactly a longitude and a latitude axis; " f"what remains is {', '.join(f'{a.dim}[{a.role}]' for a in axes) or 'nothing'}. " f"Restrict the other dimensions with --extra-indices.") return 1 try: plan = infer(axes) except TooManyDimensions: print("Error: too many dimensions remain for a vector plot.") return 1 u = _orient(u, axes, plan) v = _orient(v, axes, plan) field = None if background not in ('none', 'magnitude') and background in src: field, _ = apply_selection(ds[background], extra_indices) field = _orient(field, axes, plan) reference = field if field is not None else np.hypot(u, v) colormap, norm, _ = choose_style(reference, background if field is not None else u_name, cmap=colormap, norm=norm, vmin=vmin, vmax=vmax) units = u_da.attrs.get('units') ctx = PlotContext( varname=f"{u_name},{v_name}", data=field, plan=plan, label=(u_da.attrs.get('long_name') or u_name) + (f" ({units})" if units else ""), units=units, long_name=f"{u_name} / {v_name}", colormap=colormap, norm=norm, output_path=output_path, show_topo=show_topo, show_polar=show_polar, show_3d=show_3d, interactive=interactive, globe=globe_options, figsize=figsize, dpi=dpi, title_override=title_override, interpolate=interpolate, ) try: return render_vectors(ctx, u, v, style=style, density=density, background=background) except VectorError as err: print(f"Error: {err}") return 1 def _prompt_for_x(plan, axes): """ Preserve the interactive "which dimension on X" question for 2D plots. The answer is a *dimension* name, but a file usually advertises a longer coordinate name for the same axis - 'lon' the dimension, 'longitude' the coordinate - so both are offered as aliases and completion is restricted to the dimensions themselves. """ names = [a.dim for a in axes] aliases = {} for axis in axes: for other in (axis.coord, axis.long_name): if other: aliases[str(other)] = axis.dim return prompt_choice("Which dimension on X?", names, aliases=aliases, default=plan.x.dim if plan.x is not None else None) def _scalar_coord_subtitle(da, extra_indices): """ Describe the slice that was taken, e.g. "Time = 3 year", so a map says which step it shows. Only the dimensions the user pinned are reported. """ if not extra_indices: return '' bits = [] for dim, sel in extra_indices.items(): if sel == 'avg': bits.append(f"mean over {dim}") elif isinstance(sel, int) and dim in da.coords and da.coords[dim].ndim == 1: try: value = float(da.coords[dim].values[sel]) except (IndexError, TypeError, ValueError): continue units = str(da.coords[dim].attrs.get('units', '')) # "days since 0001-01-01 00:00:00" is an epoch, not a label; the # period alone is what belongs next to the number in a title. units = units.split(' since ')[0].strip() bits.append(f"{dim} = {value:g}{' ' + units if units else ''}") return ', '.join(bits)