""" The atmosphere above the globe: clouds, dust and winds as a 3D shell. The surface modes put the field on a `vedo.StructuredGrid` whose points sit on the sphere, so the spherical grid is used directly and never resampled. That grid supports contouring, which is what turns a cloud field into a body with a surface. The `cloud` mode cannot work that way - ray casting needs a uniform box - so it resamples into cartesian voxels in `cloud.py`, which is the one place that happens. Four ways of showing the same volume, because no single one answers every question about it: cloud the field ray-cast as a translucent volume - what does it look like? iso one or more isosurfaces - where does the field exceed this? column the vertical integral on a single shell - where is there any of it? layers one translucent surface per level - how is it stacked? """ import contextlib from dataclasses import dataclass, field import numpy as np from ...overlay import resolve_thresholds from matplotlib.colors import LogNorm, Normalize from ...colors import SHELL_RANGE, truncated from ...coords import close_longitude_seam from ...vertical import length_scale from .geometry import (NO_DATA_TRANSPARENT, add_bar, apply_colormap, finite_range) try: import vtkmodules.all as vtk except ImportError: # pragma: no cover try: import vtk except ImportError: vtk = None try: from vedo import Arrows, StructuredGrid vedo_available = True except ImportError: vedo_available = False # Where the shell tops out when the vertical coordinate is not a length, as a # multiple of the exaggerated relief. Pressure and sigma levels carry no metric # altitude, so the levels are simply spread over this depth. DEFAULT_SHELL_TOP_KM = 50.0 # How many arrows to keep along each axis of a wind field. WIND_TARGET = 24 # Most surfaces --shell-mode layers will stack. The stack is cut open, so this # is about how many altitudes stay tellable apart rather than about how many the # eye can see through. MAX_LAYERS = 8 # Relaxation passes over a marching-cubes surface. Enough to lose the cell # facets of a coarse model grid, few enough to keep the body's own shape. SMOOTH_ITERATIONS = 20 # A shell's colour bar, as (bottom-left, top-right) in viewport coordinates. # The right-hand side, opposite the surface's slot: the two used to be the same # rectangle, so a terrain globe under a stack of layers drew the elevation bar # and the field bar on top of each other, one legible line of text out of two. SHELL_BAR_POS = ((0.90, 0.25), (0.94, 0.75)) # numpy 2 renamed trapz; the environment pins numpy 2.5 but the tool is also run # against older interpreters on shared machines. _trapezoid = getattr(np, 'trapezoid', None) or getattr(np, 'trapz') @dataclass class Shell: """ What `build_shell` produced: the actors to draw, a one-line note for the terminal, and the caption lines that describe the view on screen. """ actors: list = field(default_factory=list) note: str = '' caption: list = field(default_factory=list) depth_peeling: bool = False cut: bool = False # slice it open, so the view gets inside slicer: object = None # (origin, normal) -> a cross-section Mesh @contextlib.contextmanager def _quiet_vtk(): """ Silence VTK's warning stream. A latitude/longitude sphere has degenerate cells at the poles, where every longitude collapses onto one point, so contouring floods the terminal with 'Unable to factor linear system'. The degeneracy is inherent to the grid, not a defect in the data, and the warnings drown out real messages. """ if vtk is None: yield return was_on = vtk.vtkObject.GetGlobalWarningDisplay() vtk.vtkObject.GlobalWarningDisplayOff() try: yield finally: if was_on: vtk.vtkObject.GlobalWarningDisplayOn() def altitudes_in_metres(levels, units, top_km=None): """ (altitudes in metres, note). A vertical coordinate that measures a length is converted directly. Anything else - pressure in Pa, a sigma or hybrid coordinate, a bare level index - has no metric altitude, so the levels are spread evenly over `top_km` in their own order and the caller is told the vertical is indexed. """ levels = np.asarray(levels, dtype=float) scale = length_scale(units) if scale is not None: return levels * scale, None top = (top_km if top_km is not None else DEFAULT_SHELL_TOP_KM) * 1000.0 count = levels.size if count < 2: return np.full(count, top), f"'{units or 'level'}' is not a length; placed at {top / 1000:g} km" # Always ascending: the caller reorders the data instead, so that a # descending coordinate such as pressure still stacks bottom-to-top return np.linspace(0.0, top, count), ( f"'{units or 'level'}' is not a length; the {count} levels are " f"spread evenly to {top / 1000:g} km") def altitude_grid(altitudes, lats, lons): """ (altitude, latitude, longitude) broadcast to (level, lat, lon) each. The altitude is one number per level - a shell of uniform height - or one per column, which is what the hybrid coordinate gives and what makes the lower levels follow the terrain. """ altitudes = np.asarray(altitudes, dtype=float) shape = (altitudes.shape[0], np.size(lats), np.size(lons)) if altitudes.ndim == 1: altitudes = altitudes[:, None, None] lat3d = np.asarray(lats, dtype=float)[None, :, None] lon3d = np.asarray(lons, dtype=float)[None, None, :] return (np.broadcast_to(altitudes, shape), np.broadcast_to(lat3d, shape), np.broadcast_to(lon3d, shape)) def build_structured_grid(globe, altitudes, lats, lons, cube): """ A StructuredGrid of the atmosphere, with `cube` attached as point scalars. `cube` is (level, lat, lon). VTK orders structured points with the first dimension varying fastest, which for a grid built in (level, lat, lon) order means Fortran order - attaching the scalars in C order instead would silently scramble the field. A grid that goes right round the planet has its seam closed first. Nothing else here is periodic - a StructuredGrid is a box, and marching cubes has no cells outside it - so the last eleven degrees before the date line held no surface at all, and every cloud crossing it came out as two bodies with a flat face each. """ # The altitudes are wrapped against the *incoming* axis, so the field and # the geometry are closed by the same decision on the same longitudes source_lons = lons lons, cube = close_longitude_seam(source_lons, cube) if np.ndim(altitudes) > 1: altitudes = close_longitude_seam(source_lons, altitudes)[1] lev3d, lat3d, lon3d = altitude_grid(altitudes, lats, lons) points = globe.to_cartesian(lat3d, lon3d, altitude=lev3d) x, y, z = (points[:, 0].reshape(lev3d.shape), points[:, 1].reshape(lev3d.shape), points[:, 2].reshape(lev3d.shape)) grid = StructuredGrid([x, y, z]) grid.pointdata['Scalars'] = np.asarray(cube, dtype=float).ravel(order='F') return grid def drop_duplicate_longitudes(lons, cube): """ Remove a repeated longitude column. A field whose longitudes span -180 to +180 inclusive carries the same meridian twice. Left in place it makes a zero-width cell, which contouring reports as a degenerate cell for every level. """ lons = np.asarray(lons, dtype=float) keep = np.ones(lons.size, dtype=bool) if lons.size > 1 and np.isclose(abs(lons[-1] - lons[0]), 360.0): keep[-1] = False return lons[keep], cube[..., keep] def threshold_values(cube, spec): """ Turn a `--shell-threshold` into a sorted list of values. A bare number is used as-is; 'pNN' is that percentile of the finite values, which is what makes one setting work across fields whose magnitudes differ by decades. Several may be given, comma-separated, to nest one isosurface inside another: 'p90,p99,p99.9' draws the whole cloud, its body and its core at once, which reads as a concentration where a single surface cannot. """ return resolve_thresholds(cube, spec, default=('p99',)) def _field_norm(cube, norm): """ A scale for the field, falling back on its own finite range. A constant-valued isosurface has to be handed explicit limits: left to rescale itself, vedo would map its single value onto the whole colormap and every surface would come out the same colour. """ if norm is not None: return norm low, high = finite_range(cube) return Normalize(vmin=low, vmax=high) def _threshold_norm(values, norm): """ The scale nested isosurfaces are coloured on: the thresholds' own range. The field's full range is the wrong one here. `h2o_ice` runs from 1e-40 to 1e-5, so its p90, p99 and p99.9 all sit in the top two decades of a thirty-five decade scale and come out the same yellow - three surfaces the colour cannot tell apart, which is the whole reason for drawing three. Stretching the scale over the thresholds keeps the units and the numbers honest while making the surfaces distinguishable, which is what they are for. """ if len(values) < 2: return norm low, high = float(min(values)), float(max(values)) if isinstance(norm, LogNorm) and low > 0: return LogNorm(vmin=low, vmax=high) return Normalize(vmin=low, vmax=high) @dataclass class Vertical: """ Where the levels are, resolved once for every mode to share. `field` is what the meshes are built on: metres above the ground, either one number per level or - when the file's hybrid coordinate could be read - one per column, so the shell drapes the terrain at the bottom and flattens into an isobar at the top. `profile` is the one-dimensional reading of the same thing, for the ray-cast cloud and its cross-section, which resample into a uniform box and so can only have one vertical axis for the whole planet. """ field: np.ndarray # (level,) or (level, lat, lon), metres profile: np.ndarray # (level,), metres mass: np.ndarray = None # (level, lat, lon) kg/m2 in each layer note: str = '' descending: bool = False # the coordinate ran top-down and was turned over @property def terrain_following(self): """ Whether each level is a map of heights rather than a single one. """ return np.ndim(self.field) > 1 def at(self, index): """ The altitude of one level: a number, or a map of one when the levels follow the terrain. """ return self.field[index] def _resolve_vertical(plan, cube_lons, air, top_km): """ The `Vertical` a shell is built on. `air` is the hybrid-coordinate altitudes when the file could give them, and is what this prefers: it is the only reading of the vertical that knows where a level actually is over a given column. Without it the levels are the coordinate's own values, converted when they are a length and spread evenly when they are not, exactly as before. """ levels = np.asarray(plan.z.values, dtype=float) descending = bool(levels.size > 1 and levels[0] > levels[-1]) if air is not None: _, field = drop_duplicate_longitudes(cube_lons, air.above_ground) mass = air.thickness() if mass is not None: _, mass = drop_duplicate_longitudes(cube_lons, mass) if descending: field = field[::-1] mass = None if mass is None else mass[::-1] # Averaged after the flip and after the seam column has gone, so the # profile is the mean of exactly the columns that are drawn profile = field.reshape(field.shape[0], -1).mean(axis=1) return Vertical(field=field, profile=profile, mass=mass, note=air.note, descending=descending) altitudes, note = altitudes_in_metres(levels, plan.z.units, top_km) # A length coordinate keeps its own values, so a descending one has to be # turned over with the data. The evenly spread fallback is built ascending # already and must not be. if descending and length_scale(plan.z.units) is not None: altitudes = altitudes[::-1] return Vertical(field=altitudes, profile=altitudes, note=note or '', descending=descending) def build_shell(globe, plan, cube, colormap, varname, units=None, norm=None, mode='cloud', threshold=None, level=None, opacity=0.35, top_km=None, color='value', lit=False, bar_pos=None, resolution=None, cut=None, altitudes=None): """ A `Shell` describing the atmosphere over the globe. `cloud` ray-casts the volume itself, which is the only one of these that shows an atmosphere without drawing an edge it does not have. `iso` contours it at one or more thresholds, which is what shows a cloud or a dust plume as a body. `column` integrates it over the vertical onto a single translucent shell. `layers` draws one surface per level, for vertical context. `level` draws a single level as a solid shell. `altitudes` is a `vertical.Altitudes` when the file carries a hybrid coordinate and a surface pressure; None falls back on the vertical axis' own values. """ if not vedo_available: return Shell(note="vedo missing") lats = np.asarray(plan.y.values, dtype=float) lons = np.asarray(plan.x.values, dtype=float) air = _resolve_vertical(plan, lons, altitudes, top_km) lons, cube = drop_duplicate_longitudes(lons, np.asarray(cube, dtype=float)) note = air.note # A descending vertical (pressure) has to be ordered bottom-to-top before # the levels can be stacked into a grid; `_resolve_vertical` has already # turned the altitudes over the same way if air.descending: cube = cube[::-1] title = varname + (f' [{units}]' if units else '') span = (f"levels {air.profile.min() / 1000:.0f}-{air.profile.max() / 1000:.0f} km" f"{' above the ground' if air.terrain_following else ''}, " f"relief x{globe.exaggeration:g}, altitude x{globe.air_exaggeration:g}") if level is not None: shell = _one_level(globe, air, lats, lons, cube, level, colormap, norm, title, lit) shell.note = note or '' shell.caption = [f"single level {level}", span] return shell if mode == 'layers': shell = _layers(globe, air, lats, lons, cube, colormap, norm, title, opacity, bar_pos) elif mode == 'column': shell = _column(globe, air, lats, lons, cube, colormap, norm, varname, units, opacity, bar_pos) else: shell = None if mode == 'cloud': shell = _cloud(globe, air, lats, lons, cube, colormap, norm, title, opacity, lit, bar_pos, resolution) if shell is None: # Contouring needs neither scipy nor a volume mapper, so it is # what is left when the cloud cannot be built at all from . import cloud as cloud_module note = _join(note, "vedo Volume or scipy missing; contouring instead" if not cloud_module.available() else "nothing to ray-cast; contouring instead") elif air.terrain_following: note = _join(note, "the ray-cast volume uses the mean profile") if shell is None: shell = _isosurfaces(globe, air, lats, lons, cube, colormap, norm, title, threshold, opacity, color, bar_pos) shell.cut = shell.cut if cut is None else bool(cut) if shell.cut and shell.actors: shell.slicer = _slicer(globe, air, lats, lons, cube, colormap, norm, resolution) if shell.slicer is None: shell.cut = False note = _join(note, "no cross-section available; not cutting") shell.note = _join(note, shell.note) shell.caption.append(span) return shell def _slicer(globe, air, lats, lons, cube, colormap, norm, resolution): """ The cross-section maker a cutaway is drawn with, or None when there is none. """ from . import cloud as cloud_module return cloud_module.build_slicer(globe, air.profile, lats, lons, cube, colormap, _field_norm(cube, norm), resolution=resolution) def _join(note, detail): """ Join the vertical-coordinate note to a mode's own note, dropping the empties. """ return '; '.join(part for part in (note, detail) if part) def _cloud(globe, air, lats, lons, cube, colormap, norm, title, opacity, lit, bar_pos, resolution): """ The field ray-cast as one translucent cloud, or None to fall back on `iso`. The scale is resolved here rather than in `cloud`, so the cloud is coloured and faded on exactly the scale the flat map used, and falls back on the field's own range the same way an isosurface does. Ray casting needs a uniform box, and a box has one vertical axis, so this is the mean profile rather than the altitude per column. The sampler already subtracts the relief at every voxel, so the cloud still sits on the terrain; what it loses is the way a pressure level rises over a basin. """ from . import cloud as cloud_module return cloud_module.build_cloud(globe, air.profile, lats, lons, cube, colormap, _field_norm(cube, norm), title, Shell, opacity=opacity, lit=lit, bar_pos=bar_pos, resolution=resolution) def _isosurfaces(globe, air, lats, lons, cube, colormap, norm, title, threshold, opacity, color, bar_pos): """ One nested isosurface per threshold, smoothed and shaded. Colouring by `value` is what lets the shell share the field's own scale: a surface is one value everywhere, so it is drawn in that value's colour and the bar spans the whole field. `height` keeps the older reading, where the colour separates ground fog from high cloud instead. Either way the colour says almost nothing - one flat tone over a whole body - so the shape has to say the rest, which is why these are the one thing on the globe that is always lit. No `lit`: this mode does not take the globe's answer. """ values = threshold_values(cube, threshold) if not values: return Shell(note="the field is empty") # A grid one point thick has no cells, and VTK's marching cubes does not # survive being handed one - it takes the process down with it if len(air.profile) < 2: return Shell(note="a single level has no volume to contour; " "try --shell-level 0, or --shell-mode column") grid = build_structured_grid(globe, air.field, lats, lons, cube) scale = _threshold_norm(values, _field_norm(cube, norm)) alphas = _nested_alphas(len(values), opacity) actors, drawn, empty = [], [], [] for value, alpha in zip(values, alphas): with _quiet_vtk(): surface = grid.isosurface(value=value) if surface.npoints == 0: empty.append(value) continue _shape_isosurface(surface) _paint_isosurface(surface, value, colormap, scale, color, globe) # Always shaded, unlike everything else on the globe. The globe's own # "lighting off" default exists because shading multiplies the colormap # by the angle to the light, so a value reads differently depending on # where it sits; an isosurface is one value everywhere, so it has no # scale to distort - and shape is the only thing it has to say. surface.lighting('glossy').opacity(alpha) actors.append(surface) drawn.append(value) if not actors: return Shell(note=f"nothing reaches {', '.join(f'{v:.3g}' for v in values)}") # The innermost surface carries the bar: it is the one drawn last and least # transparent, so its colour is the one the eye is already reading. if color == 'height': add_bar(actors[-1], "shell height [km]", bar_pos or SHELL_BAR_POS) else: add_bar(actors[-1], title, bar_pos) heights = _height_km(globe, actors[-1].coordinates) detail = (f"isosurface{'s' if len(drawn) > 1 else ''} at " f"{', '.join(f'{v:.3g}' for v in drawn)}, " f"up to {heights.max():.1f} km") if empty: detail += f" (nothing reaches {', '.join(f'{v:.3g}' for v in empty)})" caption = [f"iso at {', '.join(f'{v:.3g}' for v in drawn)}, coloured by {color}"] return Shell(actors=actors, note=detail, caption=caption, depth_peeling=len(actors) > 1) def _shape_isosurface(surface, niter=SMOOTH_ITERATIONS): """ Take the staircase off a marching-cubes surface and give it normals. Contouring a 33x33x26 model grid produces a surface whose facets are the cells themselves, and lighting a facetted surface only draws attention to the grid. The windowed-sinc filter relaxes the vertices without shrinking the body the way a Laplacian smoother would, so the surface stays where its threshold put it. `boundary=True` keeps the rim of an open surface - one that runs into the ground or off the top of the model - from being pulled inwards and leaving a gap. Normals come after, or they would describe the shape before smoothing. """ with _quiet_vtk(): surface.smooth(niter=niter, boundary=True) surface.compute_normals() return surface def _paint_isosurface(surface, value, colormap, scale, color, globe): """ Colour one isosurface, either by its own value or by height. Drawn in the bright part of the colormap rather than the whole of it. With several thresholds the scale is stretched across them, which puts the *outermost* envelope at the bottom of the range - and the outermost is the one in view, so the surface that matters most came out in the colour that shows least, a near-black purple over a globe of much the same tone. The same narrowing goes on the terrain underneath, from the other end. """ if color == 'height': surface.cmap(truncated('plasma', *SHELL_RANGE), _height_km(globe, surface.coordinates)) return # A constant scalar field plus explicit limits: the mesh comes out in the # single colour its threshold has on the field's scale, while the bar built # from it still spans the whole field. apply_colormap(surface, truncated(colormap, *SHELL_RANGE), np.full(surface.npoints, float(value)), scale, '') def _nested_alphas(count, opacity): """ Opacity per nested surface, rising inward. A single surface is drawn at the opacity that was asked for. It used to be scaled up like the innermost of a nest, which at the default came out at 0.875 - solid enough to hide the planet behind the very thing the surface is meant to sit in. Shading now carries the shape, so the surface no longer has to be opaque to be read. With several, the outer envelope has to stay faint enough to see the core through it, which is the whole point of nesting them. """ if count < 2: return [min(1.0, opacity)] return list(np.linspace(opacity, min(1.0, opacity * 2.5), count)) def _column(globe, air, lats, lons, cube, colormap, norm, varname, units, opacity, bar_pos): """ The vertical integral of the field, on one shell over the whole planet. This is the view that answers "where is there any of this at all", which neither an isosurface nor a stack of layers does: a thin dense cloud and a deep faint haze both show up, in proportion to how much is actually there. The integral is over mass wherever the file allows it - the layer thicknesses come from the hybrid coordinate and the surface pressure, so what is drawn is what a column of atmosphere actually holds, in kg/m2 for a mixing ratio. The geometric integral it falls back on answers a subtly different and much less useful question: weighted by metres rather than by mass, it is dominated by the deep upper layers, where a tracer is thin and almost horizontally uniform, and the map comes out flat. It is drawn high enough to clear the highest ground on the planet, and no higher: at the model top - the old choice, and the only safe one when relief and air shared a factor - the shell is so far out that it reads as a lid around the planet rather than as a map of it. """ finite = np.nan_to_num(np.asarray(cube, dtype=float), nan=0.0) if air.profile.size < 2: amount = finite[0] column_units = units or '' detail = "single level; column is that level's value" elif air.mass is not None: # sum(q * dp/g): the mass of the field in the column, which for a mixing # ratio in kg/kg comes out in kg/m2 amount = np.sum(finite * air.mass, axis=0) column_units = _mass_column_units(units) detail = f"mass-weighted column over {air.profile.size} levels" else: amount = _trapezoid(finite, x=air.profile, axis=0) column_units = f"{units}.m" if units else 'm' detail = (f"column integral over {air.profile.size} levels, by height " f"rather than by mass") # Cells the field never covered stay gaps rather than reading as zero amount = np.where(np.any(np.isfinite(cube), axis=0), amount, np.nan) if not np.any(np.isfinite(amount)): return Shell(note="the column integral is empty") amount, lats, lons, smoothed = _on_render_grid(amount, lats, lons) low, high = finite_range(amount) scale = _column_scale(norm, amount, low, high) altitude = _clearing_altitude(globe, air) shell = _level_mesh(globe, altitude, lats, lons, on_datum=True) flat = _wrap_scalars(amount) # Transparent where the field never reached, evenly translucent everywhere # else. Alpha used to follow the value as well as the colour, which sounds # right and reads badly: a low cell came out as a hole onto the grey terrain # instead of as the dark end of the colour bar, so half the map was a # different picture. A map should be read through its colours; what is left # for transparency to say is "there is no data here". # # Twice the opacity rather than two and a half times it: this shell wraps # the whole planet, and at 0.875 it is simply a lid. --shell-opacity is the # dial. solid = min(1.0, opacity * 2.0) alpha = np.where(np.isfinite(flat), solid, 0.0) title = f"column {varname}" + (f" [{column_units}]" if column_units else '') bar_title = apply_colormap(shell, colormap, flat, scale, title, alpha=alpha, no_data=NO_DATA_TRANSPARENT) add_bar(shell, bar_title, bar_pos, use_alpha=False) shell.lighting('off') shell_km = altitude / 1000.0 detail = (f"{detail}, {low:.3g} to {high:.3g} {column_units}, " f"on a shell {shell_km:.0f} km up") if smoothed: detail += "; interpolated to the 1 degree grid" scale_note = 'log' if isinstance(scale, LogNorm) else 'linear' return Shell(actors=[shell], note=detail, caption=[f"column integral on a {scale_note} scale, " f"units {column_units or 'none'}, " f"drawn at {shell_km:.0f} km"], depth_peeling=True) def _mass_column_units(units): """ The units of a mass-weighted column: kg/m2 of whatever the field measures. A mixing ratio is kilograms per kilogram, so the kilograms cancel and the column is simply kg/m2 - which is what everyone calls it, and writing it 'kg/kg.kg/m2' would only make it harder to read. """ if not units: return 'kg/m2' cleaned = str(units).strip().lower().replace(' ', '') if cleaned in ('kg/kg', 'kgkg-1', 'kg.kg-1', '1', ''): return 'kg/m2' return f"{units}.kg/m2" # The number of decades a log scale is allowed to span before its floor is # lifted. `h2o_ice` integrates to 2e-35 in the driest cell and 4e-2 in the # wettest, and a scale over all thirty-four of those decades puts every value # that matters within a per cent of the top: the globe comes back one flat # yellow, which is what it did. MAX_COLUMN_DECADES = 6 def _column_scale(norm, amount, low, high): """ The scale the column integral is drawn on. The integral is not the field - its units are not the field's - so the field's own limits mean nothing here and the range has to be the integral's own. What does carry over is the *shape* of the scale: a field the tool drew on a log scale is one that spans decades, and its integral spans them too. Drawn linearly, as it used to be, `h2o_ice` put everything but its single brightest cell at alpha zero and the globe came back empty. A log scale then gets a floor under it. Its bottom is a numerical zero - what is left of 1e-40 after an integration - and a bottom nothing is ever near is not a scale, it is a single colour. The floor is the higher of the bulk's own lower quartile and six decades down from the top, so the bar covers the values that are actually on the globe. """ # A field with no range at all cannot be scaled by it if high <= low: return Normalize(vmin=low, vmax=low + 1.0) if isinstance(norm, LogNorm) and low > 0: return LogNorm(vmin=_log_floor(amount, low, high), vmax=high) return Normalize(vmin=low, vmax=high) def _log_floor(amount, low, high, decades=MAX_COLUMN_DECADES): """ Where a log scale over `amount` should start. """ if high / low <= 10.0 ** decades: return low values = np.asarray(amount, dtype=float) positive = values[np.isfinite(values) & (values > 0)] quartile = float(np.percentile(positive, 25)) if positive.size else low return float(max(low, min(quartile, high / 10.0 ** decades))) def _clearing_altitude(globe, air, margin=1.1): """ How high the column shell goes: over the highest ground, and no higher. The shell is a map, so it wants to be near the surface it describes; what stops it is that a globe drawing 21 km of Olympus Mons at ten times its size has terrain standing 182 km off the datum, and a shell below that comes out in pieces with a quarter of the planet growing through it. Since the air is stretched harder than the relief, that is bought cheaply - at x30 against x10, under seven kilometres of air clears every summit on Mars. The relief asked for is the one on screen, not MOLA: with `--globe-relief none` there is nothing to clear and the map can lie on the sphere, and with a globe shaped by a field the shell has to clear that instead. """ return float(np.clip(globe.clearance_altitude(margin), air.profile.min(), air.profile.max())) def _on_render_grid(values, lats, lons): """ (values, lats, lons, whether it was interpolated) on the 1 degree grid. A model grid is coarse - 33x33 on the sample files - and a shell meshed straight from it shows its own cells, which against the globe's own 1 degree surface reads as moire rather than as data. Interpolating onto the same grid the globe is meshed on removes that, and costs one pass over 64800 points. Falls back on the original grid when scipy is missing, since the shell is still correct there, only blockier. """ from . import regrid, render_grid, scipy_available if not scipy_available: return values, lats, lons, False grid_lats, grid_lons, lat_grid, lon_grid = render_grid() resampled = regrid(lats, lons, values, lat_grid, lon_grid) if not np.any(np.isfinite(resampled)): return values, lats, lons, False return resampled, grid_lats, grid_lons, True def _height_km(globe, points): """ Kilometres of air under each point. Height above the ground rather than above the datum: a shell carries the relief's exaggeration under it and its own above, so the two have to be taken off separately, and what an atmosphere is measured in is its height over the terrain anyway. """ return np.atleast_1d(globe.height_above_ground(points)) / 1000.0 def _one_level(globe, air, lats, lons, cube, level, colormap, norm, title, lit): """ A single vertical level as one solid shell. """ index = int(np.clip(level, 0, len(air.profile) - 1)) shell = _level_mesh(globe, air.at(index), lats, lons) apply_colormap(shell, colormap, _wrap_scalars(cube[index]), norm, title, no_data=NO_DATA_TRANSPARENT) add_bar(shell, title) shell.lighting('default' if lit else 'off') return Shell(actors=[shell]) def _layer_indices(altitudes, limit=MAX_LAYERS): """ Which levels to draw, spread evenly over altitude rather than over index. Stepping by index bunches every drawn layer near the ground on a stretched vertical - a hybrid coordinate puts half its levels in the lowest few kilometres - so the stack ends up describing the model's resolution instead of the atmosphere. """ count = len(altitudes) if count <= limit: return list(range(count)) wanted = np.linspace(altitudes[0], altitudes[-1], limit) return sorted(set(int(np.abs(altitudes - target).argmin()) for target in wanted)) def _layers(globe, air, lats, lons, cube, colormap, norm, title, opacity, bar_pos): """ One surface per level, cut open so the stack can be looked into. The cut is what makes this mode mean anything. Concentric closed shells cannot be seen into from outside however transparent they are made - the outermost covers every other, and VTK cannot even order them, since sorting translucent actors by centroid puts every shell here at the same distance. Slicing the stack turns it into what it was always trying to be: a stack, seen edge-on, one band per altitude. Alpha still follows the value, so a level with nothing in it does not lay a blank lid over the ones below, and the levels are still thinned to MAX_LAYERS so that neighbouring bands stay tellable apart. """ scale = norm or _field_norm(cube, None) indices = _layer_indices(air.profile) actors = [] for index in indices: values = cube[index] if not np.any(np.isfinite(values)): continue shell = _level_mesh(globe, air.at(index), lats, lons) flat = _wrap_scalars(values) # Alpha follows the same scale as the colour, so a cell that reads as # "nothing here" is also invisible alpha = np.ma.filled(np.ma.masked_invalid(scale(flat)), 0.0) alpha = np.clip(np.asarray(alpha, dtype=float), 0.0, 1.0) * opacity apply_colormap(shell, colormap, flat, norm, title, alpha=alpha, no_data=NO_DATA_TRANSPARENT) shell.lighting('off') actors.append(shell) if actors: # use_alpha=False: the layers carry a per-vertex alpha ramp, which the # bar would otherwise render as stripes instead of a colour scale add_bar(actors[-1], title, bar_pos or SHELL_BAR_POS, use_alpha=False) detail = (f"{len(actors)} of {len(air.profile)} levels, " f"{air.profile[indices[0]] / 1000:.0f} to " f"{air.profile[indices[-1]] / 1000:.0f} km") return Shell(actors=actors, note=detail if actors else "every level is empty", caption=[f"{len(actors)} layers, cut open"], depth_peeling=True, cut=True) def _level_mesh(globe, altitude, lats, lons, on_datum=False): """ A closed lat/lon surface at one altitude above the relief. `altitude` is a single height for the whole surface, or one per point of the lat/lon grid - which is what a hybrid level is, riding up over a basin and down over a shield volcano. `on_datum` measures that altitude from the datum instead, giving a sphere. A model level belongs over the ground it sits on; a map of a column integral does not, and drawing it that way gave the integral a ripple that came from the topography rather than from the field. The longitude seam is wrapped the same way the surface globe wraps it, so the shell has no slit down one meridian. """ from vedo import Mesh lon2d, lat2d = np.meshgrid(lons, lats) points = globe.to_cartesian(lat2d, lon2d, altitude=altitude, elevation=0.0 if on_datum else None) nlat, nlon = lat2d.shape i = np.arange(nlat - 1)[:, None] j = np.arange(nlon)[None, :] p0 = i * nlon + j p1 = i * nlon + (j + 1) % nlon p2 = p0 + nlon p3 = p1 + nlon faces = np.concatenate([ np.stack(np.broadcast_arrays(p0, p2, p1), axis=-1).reshape(-1, 3), np.stack(np.broadcast_arrays(p1, p2, p3), axis=-1).reshape(-1, 3), ]) return Mesh([points, faces]) def _wrap_scalars(values): return np.asarray(values, dtype=float).ravel() def wind_arrows(globe, lats, lons, u, v, altitude=0.0, scale=None, target=WIND_TARGET): """ Wind vectors as 3D arrows tangent to the sphere. (u, v) are eastward and northward components, so each is projected onto the local tangent basis: east is d/d(longitude) and north is -d/d(colatitude), both normalized. """ if not vedo_available: return None lats = np.asarray(lats, dtype=float) lons = np.asarray(lons, dtype=float) u = np.asarray(u, dtype=float) v = np.asarray(v, dtype=float) step_y = max(1, int(np.ceil(lats.size / float(target)))) step_x = max(1, int(np.ceil(lons.size / float(target)))) lats, lons = lats[::step_y], lons[::step_x] u, v = u[::step_y, ::step_x], v[::step_y, ::step_x] lon2d, lat2d = np.meshgrid(lons, lats) phi = np.deg2rad(90.0 - lat2d).ravel() theta = np.deg2rad(lon2d).ravel() east = np.column_stack([-np.sin(theta), np.cos(theta), np.zeros_like(theta)]) north = np.column_stack([-np.cos(phi) * np.cos(theta), -np.cos(phi) * np.sin(theta), np.sin(phi)]) start = globe.to_cartesian(lat2d, lon2d, altitude=altitude) vectors = east * u.ravel()[:, None] + north * v.ravel()[:, None] speed = np.linalg.norm(vectors, axis=1) fastest = np.nanmax(speed) if np.any(np.isfinite(speed)) else 0.0 if fastest <= 0: return None if scale is None: # The longest arrow spans about a tenth of the radius scale = globe.radius * 0.10 / fastest end = start + vectors * scale good = np.isfinite(end).all(axis=1) & (speed > 0) if not np.any(good): return None return Arrows(start[good], end[good], c='white', alpha=0.85)