""" Where a model level actually is: the hybrid sigma-pressure coordinate. An LMDZ vertical axis is not a height. `altitude` in a diagfi or a start is a *pseudo*-altitude, one number per level for the whole planet, while the levels the model really integrates are surfaces of p = A + B*ps: they follow the terrain at the bottom, where B is 1 and the level is the ground, and flatten into isobars at the top, where A carries the pressure and B is nothing. Drawing them all at one altitude puts the atmosphere in the wrong place over every mountain and every basin on the planet. So this module turns the file's own `ap`/`bp` (or `aps`/`bps`) and its surface pressure into a height per column, which the globe then lifts above the relief. Everything here is plain numpy over an xarray dataset - no vedo, no rendering - so a shell's geometry can be tested without a render window. The scale height is fitted from the file rather than assumed. A run of the generic model on a thick atmosphere has neither Mars' pressures nor Mars' scale height, but it does carry `altitude` and the coefficients that produced it, and the two together say what H the model used. """ from dataclasses import dataclass import numpy as np from . import paths from .conventions import (CONTROL_GRAVITY, CONTROL_NAME, HYBRID_INTERFACE_NAMES, HYBRID_MID_NAMES, SURFACE_PRESSURE_NAMES) from .coords import find_coord_var # Metres per unit, for a vertical coordinate that measures a length. LENGTH_UNITS = { 'm': 1.0, 'meter': 1.0, 'meters': 1.0, 'metre': 1.0, 'metres': 1.0, 'km': 1000.0, 'kilometer': 1000.0, 'kilometers': 1000.0, 'kilometre': 1000.0, 'kilometres': 1000.0, } def length_scale(units): """ Metres per unit of a vertical coordinate, or None when it is not a length. """ return LENGTH_UNITS.get((units or '').strip().lower()) @dataclass(frozen=True) class Altitudes: """ Where each model level sits, per column. `above_ground` is what the globe wants: metres above the *local* surface, since `Globe.to_cartesian(altitude=...)` adds it on top of the relief. The terrain-following comes out of that for free - over high ground `ps` is low, so `H*ln(ps/p)` is smaller by about the height of the ground, and a level whose pressure is fixed ends up at the same absolute altitude everywhere. `interfaces` is the pressure at the layer boundaries, which is the only thing a column integral can be weighed with, and None when the file gives no way to work them out. """ above_ground: np.ndarray # (level, lat, lon) metres above the ground interfaces: np.ndarray # (level+1, lat, lon) Pa, or None gravity: float # m/s2 note: str = '' def thickness(self): """ The mass of each layer per unit area, |dp|/g in kg/m2, or None when the file gives no interlayer coefficients to take a thickness from. """ if self.interfaces is None: return None return np.abs(np.diff(self.interfaces, axis=0)) / self.gravity def surface_pressure_name(ds): """ The name of the surface pressure field, or None. Case-insensitive, through the same lookup the coordinate resolution uses: the difference between `ps` and `PS` is not a difference worth falling back to pseudo-altitudes over. """ return find_coord_var(ds, SURFACE_PRESSURE_NAMES) def gravity(ds): """ Surface gravity in m/s2, from the file when it says so. LMDZ writes its own physical constants into `controle`, so a run of the generic model on another planet is weighed with that planet's gravity rather than with Mars'. """ if CONTROL_NAME in ds.variables: values = np.asarray(ds[CONTROL_NAME].values, dtype=float).ravel() if values.size > CONTROL_GRAVITY: g = float(values[CONTROL_GRAVITY]) if np.isfinite(g) and g > 0: return g return float(paths.planet_gravity) def _pair(ds, candidates, size): """ The first (A, B) pair of coefficients present in the file with `size` entries each, as float arrays, or None. The names are matched case-insensitively, and in the order they are listed rather than the file's, so a file carrying both LMDZ's and CF's spellings gets the one this tool knows the layout of. """ for names in candidates: found = [find_coord_var(ds, (name,)) for name in names] if not all(found): continue a, b = (np.asarray(ds[name].values, dtype=float).ravel() for name in found) if a.size == b.size == size: return a, b return None def mid_coefficients(ds, nlev): """ (A, B) on the mid-layers, which is where a level's own pressure comes from. Derived from the interlayer pair when a file carries only that one: the mid-layer value LMDZ writes is the arithmetic mean of its two interfaces. """ mids = _pair(ds, HYBRID_MID_NAMES, nlev) if mids is not None: return mids edges = _pair(ds, HYBRID_INTERFACE_NAMES, nlev + 1) if edges is None: return None a, b = edges return 0.5 * (a[:-1] + a[1:]), 0.5 * (b[:-1] + b[1:]) def interface_coefficients(ds, nlev): """ (A, B) on the interlayers, which is where a layer thickness comes from. Rebuilt from the mid-layers when the file has only those: each inner interface is the mean of the two levels around it, and the two ends are extrapolated from the nearest pair. The top interface is then pinned to zero pressure, which is what the model itself does - `ap[-1]` and `bp[-1]` are both 0 in every LMDZ file - so a rebuilt column still holds the whole mass of the atmosphere. """ edges = _pair(ds, HYBRID_INTERFACE_NAMES, nlev + 1) if edges is not None: return edges mids = _pair(ds, HYBRID_MID_NAMES, nlev) if mids is None or nlev < 2: return None a, b = mids def rebuild(values): inner = 0.5 * (values[:-1] + values[1:]) first = values[0] + (values[0] - inner[0]) last = values[-1] + (values[-1] - inner[-1]) return np.concatenate([[first], inner, [last]]) edge_a, edge_b = rebuild(a), rebuild(b) edge_a[-1], edge_b[-1] = 0.0, 0.0 return edge_a, edge_b def fit_scale_height(levels, units, reference_pressure): """ (H in metres, note): the scale height the file's own vertical implies. LMDZ's `altitude` is -H*ln(p/p_ref) over a reference column, so H is the slope of the levels against -ln(p) and a straight line fit recovers it to the metre. That is worth more than any constant: this same tool is pointed at the generic model, where the atmosphere is another planet's and 10.3 km would be wrong by a factor of several. Falls back on `paths.scale_height` when the vertical is not a length - pressure, sigma, a bare index - because then there is nothing to fit. """ scale = length_scale(units) metres = np.asarray(levels, dtype=float) * scale if scale else None pressure = np.asarray(reference_pressure, dtype=float) usable = (metres is not None and metres.size > 1 and pressure.size == metres.size and np.all(np.isfinite(metres)) and np.all(pressure > 0)) if not usable: return float(paths.scale_height), ( f"no metric vertical to fit a scale height to; using " f"{paths.scale_height / 1000:g} km") slope, _ = np.polyfit(-np.log(pressure), metres, 1) if not np.isfinite(slope) or slope <= 0: return float(paths.scale_height), ( f"the vertical does not follow log-pressure; using a scale height " f"of {paths.scale_height / 1000:g} km") return float(slope), '' def air_altitudes(ds, plan, surface_pressure): """ An `Altitudes` for the plan's vertical axis, or None when the file cannot give one. `surface_pressure` is the 2D field on the plan's own grid, already sliced the way the plotted variable was - the caller has that machinery, and this module has no business repeating it. None means "there is nothing here to improve on": no hybrid coefficients, no surface pressure, or a vertical axis the coefficients do not describe (a field on the interlayers, say). The renderer then keeps the file's own vertical, exactly as before. """ if plan.z is None or surface_pressure is None: return None levels = np.asarray(plan.z.values, dtype=float) nlev = levels.size mids = mid_coefficients(ds, nlev) if mids is None: return None ps = np.asarray(surface_pressure, dtype=float) # The same grid as the field, or the altitudes would not line up with it. # A file can hold a surface pressure on another grid entirely - a staggered # one, a different resolution - and broadcasting that into place would put # the levels somewhere plausible and wrong. grid = (np.size(plan.y.values), np.size(plan.x.values)) if ps.shape != grid or not np.any(np.isfinite(ps)) or np.nanmax(ps) <= 0: return None # A gap in the surface pressure would put a whole column of the shell at # nowhere in particular; the planetary mean is the least surprising stand-in ps = np.where(np.isfinite(ps) & (ps > 0), ps, np.nanmean(ps[np.isfinite(ps)])) a, b = mids pressure = a[:, None, None] + b[:, None, None] * ps[None, :, :] height, note = fit_scale_height(levels, plan.z.units, a + b * float(np.mean(ps))) # Height above the local ground, which is what the globe adds to the relief. # Clipped at zero: the lowest level sits a few metres up, but rounding in a # rebuilt coefficient could otherwise push it just below the terrain. above_ground = np.maximum(height * np.log(ps[None, :, :] / pressure), 0.0) edges = interface_coefficients(ds, nlev) interfaces = None if edges is not None: edge_a, edge_b = edges interfaces = edge_a[:, None, None] + edge_b[:, None, None] * ps[None, :, :] detail = (f"altitudes from the hybrid coordinate and surface pressure " f"({np.nanmin(ps) / 100:.0f}-{np.nanmax(ps) / 100:.0f} hPa), " f"scale height {height / 1000:.1f} km") return Altitudes(above_ground=above_ground, interfaces=interfaces, gravity=gravity(ds), note='; '.join(part for part in (detail, note) if part))