Meteorology

Gridded meteorology (MetField) on the les_fvm analysis grid, the wrfout and ERA5 loaders, and the adapter that feeds it to the Eulerian solver as wind; the shared lat/lon frame (LocalFrame) and clock (TimeAxis) every loader places data with; Monin–Obukhov surface-layer scales (MoninObukhov), a bulk-Richardson PBL height and the Hanna turbulence factory. Tier 0. The loaders and LocalFrame's transforms need pip install 'plumax[io]'.

met

Meteorological forcing for the transport tiers (Tier 0 prerequisites).

The transport models are driven by an external meteorology; this package owns the container those models consume and the loaders that populate it from numerical-weather-prediction output.

• :class:MetField — gridded wind, temperature, pressure and PBL height on a les_fvm analysis grid, on the les_fvm C-grid stagger, with a time axis.
• :func:load_wrf — read a wrfout file onto an analysis grid.
• :func:load_era5 — read ERA5 pressure- and single-level files onto an analysis grid placed by a :class:LocalFrame.
• :class:LocalFrame — the analysis grid's metric frame on lat/lon (a UTM zone shifted to an origin); built and applied outside jit.
• :class:TimeAxis — seconds since t0 ↔ UTC timestamps.
• :class:MoninObukhov — Monin–Obukhov u*, 1/L, w* (pure JAX), with the stability functions in :mod:plumax.met.surface_layer.
• :func:pbl_height_bulk_richardson / :func:hanna_turbulence_from — PBL height from the analysis profiles, and Tier II's Hanna turbulence built from the met.
• :func:to_prescribed_wind — expose a :class:MetField as the :class:~plumax.les_fvm.wind.PrescribedWindField the Eulerian solver takes, with piecewise-linear interpolation between met times.

MetField

Bases: Module

Meteorological forcing on an analysis grid with a time axis.

Attributes:

Name	Type	Description
plume_grid	PlumeGrid3D	Analysis grid every field lives on.
times	Float[Array, ' n_time']	Seconds since t0, strictly increasing, shape (n_time,).
u	Float[Array, 'n_time nz ny nx']	Eastward wind [m/s] at U-points (T-point x + dx/2), interior shape (n_time, nz, ny, nx).
v	Float[Array, 'n_time nz ny nx']	Northward wind [m/s] at V-points (T-point y + dy/2).
w	Float[Array, 'n_time nz ny nx']	Vertical wind [m/s] at T-points.
temperature	Float[Array, 'n_time nz ny nx']	Air temperature [K] at T-points.
pressure	Float[Array, 'n_time nz ny nx']	Air pressure [Pa] at T-points.
pbl_height	Float[Array, 'n_time ny nx']	Boundary-layer height above ground [m] at horizontal T-points, shape (n_time, ny, nx).
t0	str	ISO-8601 timestamp that times counts from (informational; static, so it never enters a trace).

n_time property

n_time: int

Number of met times.

wind_at

wind_at(t: Float[Array, '']) -> tuple[Float[Array, 'nz ny nx'], Float[Array, 'nz ny nx'], Float[Array, 'nz ny nx']]

Interior (u, v, w) at time t, piecewise-linear in time.

Met is typically hourly while the transport step is seconds, so the wind is interpolated linearly between the two bracketing met times (policy (b) of the Tier 0 design page). Outside [times[0], times[-1]] the nearest end value is held constant.

Source code in src/plumax/met/field.py

def wind_at(
    self, t: Float[Array, ""]
) -> tuple[
    Float[Array, "nz ny nx"],
    Float[Array, "nz ny nx"],
    Float[Array, "nz ny nx"],
]:
    """Interior ``(u, v, w)`` at time ``t``, piecewise-linear in time.

    Met is typically hourly while the transport step is seconds, so the
    wind is interpolated linearly between the two bracketing met times
    (policy (b) of the Tier 0 design page).  Outside ``[times[0],
    times[-1]]`` the nearest end value is held constant.
    """
    return (
        interpolate_in_time(self.times, self.u, t),
        interpolate_in_time(self.times, self.v, t),
        interpolate_in_time(self.times, self.w, t),
    )

LocalFrame dataclass

LocalFrame(origin_lat: float, origin_lon: float, epsg: int)

UTM metres relative to an origin, as the analysis grid's frame.

(x, y) = (E(φ, λ) − E₀, N(φ, λ) − N₀),

with E, N the easting and northing in the UTM zone epsg and (E₀, N₀) the origin's. The axes are UTM grid axes, which differ from true east and north by the meridian convergence (under a degree near the zone's central meridian, about 3° at its edges at mid-latitudes).

Parameters:

Name	Type	Description	Default
origin_lat	float	Latitude φ₀ of (x, y) = (0, 0) [degrees north].	required
origin_lon	float	Longitude λ₀ of (x, y) = (0, 0) [degrees east].	required
epsg	int	EPSG code of the projected CRS; :meth:at picks the UTM zone containing the origin.	required

Raises:

Type	Description
MissingExtraError	When a transform runs without pyproj (pip install 'plumax[io]').

Examples:

The origin maps to zero, and a round trip is exact to well under 1 mm:

>>> import numpy as np
>>> frame = LocalFrame.at(31.95, -102.15)
>>> x, y = frame.to_xy(31.95, -102.15)
>>> bool(np.allclose([x, y], 0.0, atol=1e-6))
True
>>> lat, lon = frame.to_latlon(250.0, -400.0)
>>> bool(np.allclose(frame.to_xy(lat, lon), [250.0, -400.0], atol=1e-6))
True

Arrays keep their shape:

>>> x, y = frame.to_xy(np.full((2, 3), 31.95), np.full((2, 3), -102.15))
>>> x.shape
(2, 3)

The configuration round-trips:

>>> LocalFrame.from_config(frame.get_config()) == frame
True

at classmethod

at(lat: float, lon: float) -> LocalFrame

The frame in the UTM zone containing (lat, lon), origin there.

Source code in src/plumax/met/_src/frames.py

@classmethod
def at(cls, lat: float, lon: float) -> LocalFrame:
    """The frame in the UTM zone containing ``(lat, lon)``, origin there."""
    epsg = utm_epsg(lat, lon)
    return cls(origin_lat=float(lat), origin_lon=float(lon), epsg=epsg)

to_xy

to_xy(lat: Float[ndarray, '*#points'] | float, lon: Float[ndarray, '*#points'] | float) -> tuple[Float[np.ndarray, '*shape'], Float[np.ndarray, '*shape']]

Frame metres (x, y) of points given by latitude φ and longitude λ.

Source code in src/plumax/met/_src/frames.py

def to_xy(
    self,
    lat: Float[np.ndarray, "*#points"] | float,
    lon: Float[np.ndarray, "*#points"] | float,
) -> tuple[Float[np.ndarray, "*shape"], Float[np.ndarray, "*shape"]]:
    """Frame metres ``(x, y)`` of points given by latitude φ and longitude λ."""
    forward, _, origin_easting, origin_northing = self._transformers
    lat_array, lon_array = np.broadcast_arrays(
        np.asarray(lat, dtype=float), np.asarray(lon, dtype=float)
    )
    # (E, N) = UTM(φ, λ)
    easting, northing = forward.transform(lon_array, lat_array)
    # x = E − E₀
    x = np.asarray(np.asarray(easting) - origin_easting)
    # y = N − N₀
    y = np.asarray(np.asarray(northing) - origin_northing)
    return x, y

to_latlon

to_latlon(x: Float[ndarray, '*#points'] | float, y: Float[ndarray, '*#points'] | float) -> tuple[Float[np.ndarray, '*shape'], Float[np.ndarray, '*shape']]

Latitude φ and longitude λ of points given in frame metres.

Source code in src/plumax/met/_src/frames.py

def to_latlon(
    self,
    x: Float[np.ndarray, "*#points"] | float,
    y: Float[np.ndarray, "*#points"] | float,
) -> tuple[Float[np.ndarray, "*shape"], Float[np.ndarray, "*shape"]]:
    """Latitude φ and longitude λ of points given in frame metres."""
    _, inverse, origin_easting, origin_northing = self._transformers
    x_array, y_array = np.broadcast_arrays(
        np.asarray(x, dtype=float), np.asarray(y, dtype=float)
    )
    # E = x + E₀
    easting = x_array + origin_easting
    # N = y + N₀
    northing = y_array + origin_northing
    # (φ, λ) = UTM⁻¹(E, N)
    lon, lat = inverse.transform(easting, northing)
    return np.asarray(lat, dtype=float), np.asarray(lon, dtype=float)

east_north_axes

east_north_axes(lat: Float[ndarray, '*#points'] | float, lon: Float[ndarray, '*#points'] | float) -> tuple[Float[np.ndarray, '*shape 2'], Float[np.ndarray, '*shape 2']]

Unit vectors of true east ê and true north n̂, in frame components.

ê = ∂(x, y)/∂λ / ‖∂(x, y)/∂λ‖, n̂ = ∂(x, y)/∂φ / ‖∂(x, y)/∂φ‖.

A wind (u_e, v_n) has frame components u_e ê + v_n n̂. The derivatives are forward differences of :meth:to_xy over :data:AXIS_STEP_DEGREES, so the vectors carry the meridian convergence and nothing is assumed about the projection.

Parameters:

Name	Type	Description	Default
lat	Float[ndarray, '*#points'] | float	Latitudes φ [degrees north].	required
lon	Float[ndarray, '*#points'] | float	Longitudes λ [degrees east], broadcastable with lat.	required

Returns:

Type	Description
Float[ndarray, '*shape 2']	(east, north), each holding the (x, y) components of a unit
Float[ndarray, '*shape 2']	vector in its last axis.

Examples:

Near a zone's central meridian the frame barely rotates; 2.85° east of it, at 32° N, grid north leans by about sin φ · Δλ ≈ 1.5°:

>>> import numpy as np
>>> frame = LocalFrame.at(31.95, -102.15)
>>> east, north = frame.east_north_axes(31.95, -102.15)
>>> east.shape
(2,)
>>> lean = np.degrees(np.arctan2(north[0], north[1]))
>>> bool(1.0 < abs(lean) < 2.0)
True

Source code in src/plumax/met/_src/frames.py

def east_north_axes(
    self,
    lat: Float[np.ndarray, "*#points"] | float,
    lon: Float[np.ndarray, "*#points"] | float,
) -> tuple[Float[np.ndarray, "*shape 2"], Float[np.ndarray, "*shape 2"]]:
    """Unit vectors of true east ê and true north n̂, in frame components.

    ê = ∂(x, y)/∂λ / ‖∂(x, y)/∂λ‖,  n̂ = ∂(x, y)/∂φ / ‖∂(x, y)/∂φ‖.

    A wind ``(u_e, v_n)`` has frame components ``u_e ê + v_n n̂``. The
    derivatives are forward differences of :meth:`to_xy` over
    :data:`AXIS_STEP_DEGREES`, so the vectors carry the meridian convergence
    and nothing is assumed about the projection.

    Args:
        lat: Latitudes φ [degrees north].
        lon: Longitudes λ [degrees east], broadcastable with ``lat``.

    Returns:
        ``(east, north)``, each holding the ``(x, y)`` components of a unit
        vector in its last axis.

    Examples:
        Near a zone's central meridian the frame barely rotates; 2.85° east
        of it, at 32° N, grid north leans by about sin φ · Δλ ≈ 1.5°:

        >>> import numpy as np
        >>> frame = LocalFrame.at(31.95, -102.15)
        >>> east, north = frame.east_north_axes(31.95, -102.15)
        >>> east.shape
        (2,)
        >>> lean = np.degrees(np.arctan2(north[0], north[1]))
        >>> bool(1.0 < abs(lean) < 2.0)
        True
    """
    lat_array, lon_array = np.broadcast_arrays(
        np.asarray(lat, dtype=float), np.asarray(lon, dtype=float)
    )
    x0, y0 = self.to_xy(lat_array, lon_array)
    x_east, y_east = self.to_xy(lat_array, lon_array + AXIS_STEP_DEGREES)
    x_north, y_north = self.to_xy(lat_array + AXIS_STEP_DEGREES, lon_array)
    # ∂(x, y)/∂λ ≈ ((x, y)(φ, λ + δ) − (x, y)(φ, λ)) / δ
    east = einx.id("..., ... -> ... (1 + 1)", x_east - x0, y_east - y0)
    # ∂(x, y)/∂φ ≈ ((x, y)(φ + δ, λ) − (x, y)(φ, λ)) / δ
    north = einx.id("..., ... -> ... (1 + 1)", x_north - x0, y_north - y0)
    # ‖e‖ = √(Σ_c e_c²)
    east_norm = np.sqrt(einx.sum("... [c]", east**2))
    north_norm = np.sqrt(einx.sum("... [c]", north**2))
    # ê = e / ‖e‖
    east_unit = einx.divide("... c, ... -> ... c", east, east_norm)
    # n̂ = n / ‖n‖
    north_unit = einx.divide("... c, ... -> ... c", north, north_norm)
    return east_unit, north_unit

grid_indices

grid_indices(lat: Float[ndarray, '*#points'] | float, lon: Float[ndarray, '*#points'] | float, grid: PlumeGrid3D) -> tuple[Int[np.ndarray, '*shape'], Int[np.ndarray, '*shape']]

Indices (i, j) of the analysis T-cells containing lat/lon points.

i = ⌊(x − x₀ + Δx/2) / Δx⌋, j = ⌊(y − y₀ + Δy/2) / Δy⌋,

with x₀ = grid.x[0]; points outside the grid get -1.

Parameters:

Name	Type	Description	Default
lat	Float[ndarray, '*#points'] | float	Latitudes φ [degrees north].	required
lon	Float[ndarray, '*#points'] | float	Longitudes λ [degrees east], broadcastable with lat.	required
grid	PlumeGrid3D	Analysis grid, in this frame's metres.	required

Returns:

Type	Description
tuple[Int[ndarray, '*shape'], Int[ndarray, '*shape']]	(i, j): x and y cell indices, -1 outside.

Examples:

>>> import jax.numpy as jnp
>>> import numpy as np
>>> from plumax.les_fvm.grid import make_grid
>>> frame = LocalFrame.at(31.95, -102.15)
>>> grid = make_grid(
...     (0.0, 400.0, 4), (0.0, 300.0, 4), (0.0, 50.0, 4), dtype=jnp.float64
... )
>>> x, y = np.array([60.0, -10.0]), np.array([140.0, 50.0])
>>> lat, lon = frame.to_latlon(x, y)
>>> i, j = frame.grid_indices(lat, lon, grid)
>>> i.tolist(), j.tolist()
([0, -1], [1, -1])

Source code in src/plumax/met/_src/frames.py

def grid_indices(
    self,
    lat: Float[np.ndarray, "*#points"] | float,
    lon: Float[np.ndarray, "*#points"] | float,
    grid: PlumeGrid3D,
) -> tuple[Int[np.ndarray, "*shape"], Int[np.ndarray, "*shape"]]:
    """Indices ``(i, j)`` of the analysis T-cells containing lat/lon points.

    i = ⌊(x − x₀ + Δx/2) / Δx⌋,  j = ⌊(y − y₀ + Δy/2) / Δy⌋,

    with ``x₀ = grid.x[0]``; points outside the grid get ``-1``.

    Args:
        lat: Latitudes φ [degrees north].
        lon: Longitudes λ [degrees east], broadcastable with ``lat``.
        grid: Analysis grid, in this frame's metres.

    Returns:
        ``(i, j)``: x and y cell indices, ``-1`` outside.

    Examples:
        >>> import jax.numpy as jnp
        >>> import numpy as np
        >>> from plumax.les_fvm.grid import make_grid
        >>> frame = LocalFrame.at(31.95, -102.15)
        >>> grid = make_grid(
        ...     (0.0, 400.0, 4), (0.0, 300.0, 4), (0.0, 50.0, 4), dtype=jnp.float64
        ... )
        >>> x, y = np.array([60.0, -10.0]), np.array([140.0, 50.0])
        >>> lat, lon = frame.to_latlon(x, y)
        >>> i, j = frame.grid_indices(lat, lon, grid)
        >>> i.tolist(), j.tolist()
        ([0, -1], [1, -1])
    """
    x, y = self.to_xy(lat, lon)
    grid_x = np.asarray(grid.x, dtype=float)
    grid_y = np.asarray(grid.y, dtype=float)
    # i = ⌊(x − x₀ + Δx/2) / Δx⌋
    i = np.floor((x - grid_x[0] + 0.5 * grid.dx) / grid.dx).astype(int)
    # j = ⌊(y − y₀ + Δy/2) / Δy⌋
    j = np.floor((y - grid_y[0] + 0.5 * grid.dy) / grid.dy).astype(int)
    outside_x = (i < 0) | (i >= grid_x.size)
    outside_y = (j < 0) | (j >= grid_y.size)
    outside = outside_x | outside_y
    return np.where(outside, -1, i), np.where(outside, -1, j)

get_config

get_config() -> dict[str, Any]

The frame's fields, for from_config.

Source code in src/plumax/met/_src/frames.py

def get_config(self) -> dict[str, Any]:
    """The frame's fields, for ``from_config``."""
    return dataclasses.asdict(self)

from_config classmethod

from_config(config: dict[str, Any]) -> LocalFrame

The frame described by get_config's output.

Source code in src/plumax/met/_src/frames.py

@classmethod
def from_config(cls, config: dict[str, Any]) -> LocalFrame:
    """The frame described by ``get_config``'s output."""
    return cls(**config)

MoninObukhov

Bases: Module

Diagnose u, 1/L and w from a MetField's lowest level (dry air).

u = κ U / (ln(z/z₀) − ψ_m(z/L) + ψ_m(z₀/L)), 1/L = −κ g H / (ρ c_p T u³), w* = (g H h / (ρ c_p T))^{1/3} for H > 0, else 0.

The usual forms carry g w'θ'/θ. H/(ρ c_p) is the actual-temperature flux w'T' and w'θ' = (θ/T) w'T', so the potential-temperature factor cancels and T remains.

Algorithm

1. Take the lowest T-level (height z = Δz/2 above the grid's surface): wind speed from the neighbouring U- and V-points, T, p and ρ.
2. Start neutral (1/L = 0) and iterate u ← log law, 1/L ← −b/u³ n_iterations times, with u* floored at :data:U_STAR_FLOOR where it divides.
3. Evaluate w* where H > 0.

Calm and neutral gives u = 0, 1/L = 0, w = 0. Very stable cases (z/L ≳ 1) converge slowly; check the residual there.

Attributes:

Name	Type	Description
n_iterations	int	Fixed-point iterations of the u*–L system (≥ 1).

Call arguments: the MetField; z0, the roughness length z₀ [m]; and sensible_heat_flux, H [W m⁻²] positive upward. z0 and H are scalars or broadcast against (time, y, x).

Raises:

Type	Description
PlumaxError	If n_iterations is below 1, or (under jit too) if z0 is not positive.

Examples:

Neutral air over 5 m s⁻¹ wind: the log law, and 1/L = 0:

>>> import jax.numpy as jnp
>>> from plumax.les_fvm.grid import make_grid
>>> from plumax.met.field import MetField
>>> grid = make_grid(
...     (0.0, 400.0, 4), (0.0, 400.0, 4), (0.0, 200.0, 4), dtype=jnp.float64
... )
>>> shape = (1, 4, 4, 4)
>>> field = MetField(
...     plume_grid=grid, times=jnp.zeros(1), u=jnp.full(shape, 5.0),
...     v=jnp.zeros(shape), w=jnp.zeros(shape),
...     temperature=jnp.full(shape, 300.0), pressure=jnp.full(shape, 1.0e5),
...     pbl_height=jnp.full((1, 4, 4), 800.0),
... )
>>> layer = MoninObukhov()(field, z0=0.1, sensible_heat_flux=0.0)
>>> layer.u_star.shape
(1, 4, 4)
>>> bool(jnp.allclose(layer.u_star, 0.4 * 5.0 / jnp.log(25.0 / 0.1)))
True

Daytime heating makes the column unstable and convective:

>>> sunny = MoninObukhov()(field, z0=0.1, sensible_heat_flux=200.0)
>>> bool((sunny.inverse_obukhov_length < 0).all() & (sunny.w_star > 0).all())
True

The configuration round-trips:

>>> MoninObukhov.from_config(MoninObukhov(n_iterations=8).get_config())
MoninObukhov(n_iterations=8)

References

Monin & Obukhov (1954); Businger et al. (1971), JAS 28, 181–189.

with_n_iterations

with_n_iterations(n_iterations: int) -> MoninObukhov

A copy with another iteration count (re-runs the field check).

Source code in src/plumax/met/_src/surface_layer.py

def with_n_iterations(self, n_iterations: int) -> MoninObukhov:
    """A copy with another iteration count (re-runs the field check)."""
    return dataclasses.replace(self, n_iterations=n_iterations)

get_config

get_config() -> dict[str, Any]

The hyperparameters, for from_config.

Source code in src/plumax/met/_src/surface_layer.py

def get_config(self) -> dict[str, Any]:
    """The hyperparameters, for ``from_config``."""
    return {"n_iterations": self.n_iterations}

from_config classmethod

from_config(config: dict[str, Any]) -> MoninObukhov

The diagnosis described by get_config's output.

Source code in src/plumax/met/_src/surface_layer.py

@classmethod
def from_config(cls, config: dict[str, Any]) -> MoninObukhov:
    """The diagnosis described by ``get_config``'s output."""
    return cls(**config)

SurfaceLayer

Bases: Module

Surface-layer scales on the analysis grid's horizontal T-points.

Attributes:

Name	Type	Description
u_star	Float[Array, 'time y x']	Friction velocity u* [m s⁻¹].
inverse_obukhov_length	Float[Array, 'time y x']	1/L [m⁻¹]; zero is neutral, negative unstable. Stored as 1/L because it is finite at neutral.
w_star	Float[Array, 'time y x']	Convective velocity w* [m s⁻¹]; zero unless H > 0.

obukhov_length property

obukhov_length: Float[Array, 'time y x']

L [m]; ±inf at neutral.

TimeAxis dataclass

TimeAxis(t0: datetime64, seconds: Float[ndarray, ' time'])

Seconds since t0 for a set of UTC sample times.

s_k = (τ_k − t₀) / 1 s,

with τ_k the sample timestamps and t₀ the reference the model counts from (by default the first sample).

Parameters:

Name	Type	Description	Default
t0	datetime64	Reference time t₀, datetime64 (naive, UTC).	required
seconds	Float[ndarray, ' time']	Sample times s_k as seconds since t0.	required

Examples:

Hourly samples, and back from seconds to timestamps:

>>> axis = TimeAxis.from_datetimes(
...     ["2024-06-01T00:00", "2024-06-01T01:00", "2024-06-01T03:30"]
... )
>>> axis.seconds.tolist()
[0.0, 3600.0, 12600.0]
>>> str(axis.to_datetime(5400.0))
'2024-06-01T01:30:00.000'

An explicit origin, and the stamp MetField.t0 stores:

>>> late = TimeAxis.from_datetimes(["2024-06-01T01:00"], t0="2024-06-01")
>>> late.seconds.tolist(), late.t0_iso
([3600.0], '2024-06-01_00:00:00')

Unsorted timestamps are rejected:

>>> TimeAxis.from_datetimes(["2024-06-02", "2024-06-01"])
...
Traceback (most recent call last):
...
TimeAxisError: TimeAxis timestamps must be strictly increasing

stamps property

stamps: Shaped[ndarray, ' time']

The sample times τ_k as datetime64[ms].

t0_iso property

t0_iso: str

t0 as YYYY-MM-DD_HH:MM:SS[.fff], the stamp MetField.t0 stores.

Milliseconds are kept when t0 has them, so t0_iso plus seconds reproduces every sample time exactly.

from_datetimes classmethod

from_datetimes(stamps: object, *, t0: object | None = None) -> TimeAxis

Build the axis from timestamps τ_k (anything datetime64 accepts).

Raises:

Type	Description
TimeAxisError	If stamps is empty or not strictly increasing.

Source code in src/plumax/met/_src/time.py

@classmethod
def from_datetimes(cls, stamps: object, *, t0: object | None = None) -> TimeAxis:
    """Build the axis from timestamps τ_k (anything ``datetime64`` accepts).

    Raises:
        TimeAxisError: If ``stamps`` is empty or not strictly increasing.
    """
    values = np.atleast_1d(_as_datetime64(stamps))
    if values.size == 0:
        raise TimeAxisError("TimeAxis needs at least one timestamp")
    # Δτ_k = τ_{k+1} − τ_k > 0
    steps = np.diff(values)
    if np.any(steps <= np.timedelta64(0, "ms")):
        raise TimeAxisError("TimeAxis timestamps must be strictly increasing")
    reference = values[0] if t0 is None else _as_datetime64(t0)
    # s_k = (τ_k − t₀) / 1 s
    seconds = (values - reference) / np.timedelta64(1, "s")
    return cls(t0=reference.astype(TIMESTAMP_UNIT)[()], seconds=seconds)

to_seconds

to_seconds(when: object) -> Float[np.ndarray, ' ...']

Seconds since t0 of timestamps: s = (τ − t₀) / 1 s.

Source code in src/plumax/met/_src/time.py

def to_seconds(self, when: object) -> Float[np.ndarray, " ..."]:
    """Seconds since ``t0`` of timestamps: s = (τ − t₀) / 1 s."""
    # s = (τ − t₀) / 1 s
    seconds = (_as_datetime64(when) - self.t0) / np.timedelta64(1, "s")
    return np.asarray(seconds)

to_datetime

to_datetime(seconds: object) -> Shaped[np.ndarray, ' ...']

Timestamps τ = t₀ + s, rounded to the millisecond.

Source code in src/plumax/met/_src/time.py

def to_datetime(self, seconds: object) -> Shaped[np.ndarray, " ..."]:
    """Timestamps τ = t₀ + s, rounded to the millisecond."""
    # n_ms = round(1000 s)
    milliseconds = np.round(np.asarray(seconds, dtype=float) * 1e3)
    offsets = milliseconds.astype(np.int64).astype("timedelta64[ms]")
    # τ = t₀ + n_ms · 1 ms
    return np.asarray(self.t0 + offsets)

load_era5

load_era5(pressure_levels: str | Path, single_levels: str | Path, *, plume_grid: PlumeGrid3D, frame: LocalFrame, times: slice | None = None) -> MetField

Load ERA5 pressure- and single-level files onto an analysis grid.

w = −ω / (ρ g), ρ = p / (R_d T), z_k = hypsometric height of level k.

Algorithm

1. Read both files, check they share times and a lat/lon grid, and order the levels surface first.
2. Integrate each level's height above ground from sp.
3. Convert ω to w. Unwrap a regional subset's longitudes so they increase across 0°/360° or ±180°; close the seam of a global grid.
4. Interpolate every column to the analysis heights, then sample the analysis T-, U- and V-points bilinearly and project the wind onto the frame's axes.

Parameters:

Name	Type	Description	Default
pressure_levels	str | Path	NetCDF with u, v, w (ω), t and optionally q on pressure levels in hPa.	required
single_levels	str | Path	NetCDF with sp and blh at the same times and on the same grid.	required
plume_grid	PlumeGrid3D	Analysis grid, in frame metres.	required
frame	LocalFrame	Frame the analysis grid is laid out in.	required
times	slice | None	Optional slice over the shared time axis.	None

Returns:

Name	Type	Description
A	MetField	class:MetField on plume_grid.

Raises:

Type	Description
MissingExtraError	If the io extra (xrtoolz) is not installed.
MissingVariableError	If a required variable is missing.
GridMismatchError	If the two files' times or lat/lon grids differ, or the longitudes are unevenly spaced.
TimeAxisError	If times selects nothing.
DomainError	If an analysis point falls outside the ERA5 domain.

Source code in src/plumax/met/_src/era5.py

def load_era5(
    pressure_levels: str | Path,
    single_levels: str | Path,
    *,
    plume_grid: PlumeGrid3D,
    frame: LocalFrame,
    times: slice | None = None,
) -> MetField:
    """Load ERA5 pressure- and single-level files onto an analysis grid.

    w = −ω / (ρ g),  ρ = p / (R_d T),  z_k = hypsometric height of level k.

    Algorithm:
        1. Read both files, check they share times and a lat/lon grid, and
           order the levels surface first.
        2. Integrate each level's height above ground from ``sp``.
        3. Convert ω to w. Unwrap a regional subset's longitudes so they
           increase across 0°/360° or ±180°; close the seam of a global grid.
        4. Interpolate every column to the analysis heights, then sample the
           analysis T-, U- and V-points bilinearly and project the wind onto
           the frame's axes.

    Args:
        pressure_levels: NetCDF with ``u``, ``v``, ``w`` (ω), ``t`` and
            optionally ``q`` on pressure levels in hPa.
        single_levels: NetCDF with ``sp`` and ``blh`` at the same times and on
            the same grid.
        plume_grid: Analysis grid, in ``frame`` metres.
        frame: Frame the analysis grid is laid out in.
        times: Optional ``slice`` over the shared time axis.

    Returns:
        A :class:`MetField` on ``plume_grid``.

    Raises:
        MissingExtraError: If the ``io`` extra (xrtoolz) is not installed.
        MissingVariableError: If a required variable is missing.
        GridMismatchError: If the two files' times or lat/lon grids differ,
            or the longitudes are unevenly spaced.
        TimeAxisError: If ``times`` selects nothing.
        DomainError: If an analysis point falls outside the ERA5 domain.
    """
    hypsometric_height = _require_hypsometric_height()

    # --- read and check ---
    with (
        xr.open_dataset(pressure_levels) as pressure_raw,
        xr.open_dataset(single_levels) as single_raw,
    ):
        levels = _normalise(pressure_raw)
        surface = _normalise(single_raw)
        _require(levels, ("u", "v", "w", "t"), pressure_levels)
        _require(surface, ("sp", "blh"), single_levels)
        if not np.array_equal(levels["time"].values, surface["time"].values):
            raise GridMismatchError("load_era5: the two files cover different times")
        same_lat = np.array_equal(levels["lat"].values, surface["lat"].values)
        same_lon = np.array_equal(levels["lon"].values, surface["lon"].values)
        if not (same_lat and same_lon):
            raise GridMismatchError(
                "load_era5: the two files are on different lat/lon grids; request "
                "both with the same `area` and `grid` so `sp` and `blh` line up "
                "with the pressure levels"
            )
        if times is not None:
            levels = levels.isel(time=times)
            surface = surface.isel(time=times)
        if levels.sizes["time"] == 0:
            raise TimeAxisError("load_era5: `times` selects no time steps")
        # surface first, so heights increase along `level`
        levels = levels.sortby("level", ascending=False)

        # --- heights ---
        # p_k = 100 · level_k  (hPa → Pa)
        pressure = (100.0 * levels["level"]).broadcast_like(levels["t"])
        columns = levels.assign(pressure=pressure, sp=surface["sp"])
        # z_k: hypsometric integration up from p_s
        height = hypsometric_height(
            columns,
            pressure="pressure",
            temperature="t",
            specific_humidity="q" if "q" in levels else None,
            surface_pressure="sp",
            level="level",
            name="z_agl",
        )["z_agl"]

        dims = ("time", "level", "lat", "lon")
        u_east = levels["u"].transpose(*dims).values
        v_north = levels["v"].transpose(*dims).values
        omega = levels["w"].transpose(*dims).values
        temperature = levels["t"].transpose(*dims).values
        pressure_pa = pressure.transpose(*dims).values
        height_agl = height.transpose(*dims).values
        pbl_height = surface["blh"].transpose("time", "lat", "lon").values
        lat = levels["lat"].values
        lon, lon_order = _unwrap_longitudes(levels["lon"].values)
        clock = TimeAxis.from_datetimes(levels["time"].values)

    # --- longitudes in file order → increasing from the subset's west edge ---
    reorder = lambda f: einx.get_at("... [lon], k -> ... k", f, lon_order)
    fields = (u_east, v_north, omega, temperature, pressure_pa, height_agl)
    u_east, v_north, omega, temperature, pressure_pa, height_agl = map(reorder, fields)
    pbl_height = reorder(pbl_height)

    # --- vertical velocity ---
    # ρ = p / (R_d T)
    density = pressure_pa / (R_DRY * temperature)
    # w = −ω / (ρ g)
    w = -omega / (density * GRAVITY)

    # --- longitude seam ---
    if _is_periodic(lon):
        # append the first column after the last, so points between it and 360°
        # interpolate across the seam
        fields = (u_east, v_north, temperature, pressure_pa, w, height_agl)
        u_east, v_north, temperature, pressure_pa, w, height_agl = map(
            _close_seam, fields
        )
        pbl_height = _close_seam(pbl_height)

    # --- onto the analysis grid ---
    grid_z = np.asarray(plume_grid.z, dtype=np.float64)
    # z_surface = z₀ − Δz/2
    grid_surface = grid_z[0] - 0.5 * plume_grid.dz
    # h_k = z_k − z_surface
    target_heights = grid_z - grid_surface
    to_index = _regular_index_map(lat, lon, frame)
    points = stagger_points(plume_grid, frame, to_index, loader="load_era5")
    at_t = points["t"]
    on_levels = lambda f: _interpolate_vertical(height_agl, f, target_heights)
    u_frame, v_frame = frame_wind(on_levels(u_east), on_levels(v_north), points)
    w_t = bilinear_at(on_levels(w), at_t.i, at_t.j)
    temperature_t = bilinear_at(on_levels(temperature), at_t.i, at_t.j)
    pressure_t = bilinear_at(on_levels(pressure_pa), at_t.i, at_t.j)
    pbl_height_t = bilinear_at(pbl_height, at_t.i, at_t.j)
    dtype = plume_grid.x.dtype
    as_jax = lambda a: jnp.asarray(a, dtype=dtype)
    return MetField(
        plume_grid=plume_grid,
        times=as_jax(clock.seconds),
        u=as_jax(u_frame),
        v=as_jax(v_frame),
        w=as_jax(w_t),
        temperature=as_jax(temperature_t),
        pressure=as_jax(pressure_t),
        pbl_height=as_jax(pbl_height_t),
        t0=clock.t0_iso,
    )

to_prescribed_wind

to_prescribed_wind(field: MetField) -> PrescribedWindField

Expose a :class:MetField as the wind object the Eulerian solver takes.

The returned :class:~plumax.les_fvm.wind.PrescribedWindField samples :meth:MetField.wind_at — so simulate_eulerian_dispersion( wind_field=to_prescribed_wind(field), ...) runs the FV transport on loaded meteorology with no other change. The caller's domain_* triples must describe the same grid field.plume_grid was built on.

Source code in src/plumax/met/field.py

def to_prescribed_wind(field: MetField) -> PrescribedWindField:
    """Expose a :class:`MetField` as the wind object the Eulerian solver takes.

    The returned :class:`~plumax.les_fvm.wind.PrescribedWindField` samples
    :meth:`MetField.wind_at` — so ``simulate_eulerian_dispersion(
    wind_field=to_prescribed_wind(field), ...)`` runs the FV transport on
    loaded meteorology with no other change.  The caller's ``domain_*``
    triples must describe the same grid ``field.plume_grid`` was built on.
    """

    def interior_fn(t):
        return field.wind_at(t)

    return PrescribedWindField(plume_grid=field.plume_grid, interior_fn=interior_fn)

hanna_turbulence_from

hanna_turbulence_from(field: MetField, layer: SurfaceLayer, *, it: int, iy: int, ix: int, coriolis_parameter: float = DEFAULT_CORIOLIS_PARAMETER) -> HannaTurbulence

Hanna (1982) turbulence at one analysis column and met time.

L = 1 / (1/L), passed as 0 at neutral.

HannaTurbulence treats L = 0 as neutral (as well as any h/|L| < 1), so a neutral column (1/L = 0) is passed as L = 0 rather than ±inf. MetField carries no latitude, so the Coriolis parameter f is an argument with FLEXPART's default.

Parameters:

Name	Type	Description	Default
field	MetField	Met on the analysis grid (for the PBL height h).	required
layer	SurfaceLayer	Surface-layer scales from :class:MoninObukhov.	required
it	int	Met time index.	required
iy	int	Analysis y index.	required
ix	int	Analysis x index.	required
coriolis_parameter	float	Coriolis parameter f = 2Ω sin φ [1/s] for the neutral profiles; the default 1e-4 s⁻¹ (latitude ≈ 43°) is the value FLEXPART's hanna.f90 hard-codes.	DEFAULT_CORIOLIS_PARAMETER

Returns:

Type	Description
HannaTurbulence	The turbulence model, with plain-float parameters.

Examples:

A neutral column maps to Hanna's L = 0, an unstable one keeps its sign:

>>> import jax.numpy as jnp
>>> from plumax.les_fvm.grid import make_grid
>>> from plumax.met._src.surface_layer import MoninObukhov
>>> from plumax.met.field import MetField
>>> grid = make_grid(
...     (0.0, 400.0, 4), (0.0, 400.0, 4), (0.0, 200.0, 4), dtype=jnp.float64
... )
>>> shape = (1, 4, 4, 4)
>>> field = MetField(
...     plume_grid=grid, times=jnp.zeros(1), u=jnp.full(shape, 5.0),
...     v=jnp.zeros(shape), w=jnp.zeros(shape),
...     temperature=jnp.full(shape, 300.0), pressure=jnp.full(shape, 1.0e5),
...     pbl_height=jnp.full((1, 4, 4), 800.0),
... )
>>> neutral = MoninObukhov()(field, z0=0.1, sensible_heat_flux=0.0)
>>> turbulence = hanna_turbulence_from(field, neutral, it=0, iy=1, ix=0)
>>> turbulence.obukhov_length, turbulence.pbl_height
(0.0, 800.0)
>>> sunny = MoninObukhov()(field, z0=0.1, sensible_heat_flux=200.0)
>>> hanna_turbulence_from(field, sunny, it=0, iy=1, ix=0).obukhov_length < 0
True

Source code in src/plumax/met/_src/pbl.py

def hanna_turbulence_from(
    field: MetField,
    layer: SurfaceLayer,
    *,
    it: int,
    iy: int,
    ix: int,
    coriolis_parameter: float = DEFAULT_CORIOLIS_PARAMETER,
) -> HannaTurbulence:
    """Hanna (1982) turbulence at one analysis column and met time.

    L = 1 / (1/L), passed as 0 at neutral.

    ``HannaTurbulence`` treats L = 0 as neutral (as well as any h/|L| < 1),
    so a neutral column (1/L = 0) is passed as L = 0 rather than ±inf.
    ``MetField`` carries no latitude, so the Coriolis parameter f is an
    argument with FLEXPART's default.

    Args:
        field: Met on the analysis grid (for the PBL height h).
        layer: Surface-layer scales from :class:`MoninObukhov`.
        it: Met time index.
        iy: Analysis y index.
        ix: Analysis x index.
        coriolis_parameter: Coriolis parameter f = 2Ω sin φ [1/s] for the
            neutral profiles; the default 1e-4 s⁻¹ (latitude ≈ 43°) is the
            value FLEXPART's ``hanna.f90`` hard-codes.

    Returns:
        The turbulence model, with plain-float parameters.

    Examples:
        A neutral column maps to Hanna's L = 0, an unstable one keeps its sign:

        >>> import jax.numpy as jnp
        >>> from plumax.les_fvm.grid import make_grid
        >>> from plumax.met._src.surface_layer import MoninObukhov
        >>> from plumax.met.field import MetField
        >>> grid = make_grid(
        ...     (0.0, 400.0, 4), (0.0, 400.0, 4), (0.0, 200.0, 4), dtype=jnp.float64
        ... )
        >>> shape = (1, 4, 4, 4)
        >>> field = MetField(
        ...     plume_grid=grid, times=jnp.zeros(1), u=jnp.full(shape, 5.0),
        ...     v=jnp.zeros(shape), w=jnp.zeros(shape),
        ...     temperature=jnp.full(shape, 300.0), pressure=jnp.full(shape, 1.0e5),
        ...     pbl_height=jnp.full((1, 4, 4), 800.0),
        ... )
        >>> neutral = MoninObukhov()(field, z0=0.1, sensible_heat_flux=0.0)
        >>> turbulence = hanna_turbulence_from(field, neutral, it=0, iy=1, ix=0)
        >>> turbulence.obukhov_length, turbulence.pbl_height
        (0.0, 800.0)
        >>> sunny = MoninObukhov()(field, z0=0.1, sensible_heat_flux=200.0)
        >>> hanna_turbulence_from(field, sunny, it=0, iy=1, ix=0).obukhov_length < 0
        True
    """
    inverse_length = float(layer.inverse_obukhov_length[it, iy, ix])
    # L = 1 / (1/L), or Hanna's neutral sentinel 0
    obukhov_length = 0.0 if inverse_length == 0.0 else 1.0 / inverse_length
    return HannaTurbulence(
        u_star=float(layer.u_star[it, iy, ix]),
        pbl_height=float(field.pbl_height[it, iy, ix]),
        obukhov_length=obukhov_length,
        w_star=float(layer.w_star[it, iy, ix]),
        coriolis_parameter=coriolis_parameter,
    )

pbl_height_bulk_richardson

pbl_height_bulk_richardson(field: MetField, *, ri_crit: float = 0.25) -> Float[np.ndarray, 'time y x']

PBL height h from the bulk Richardson number of the analysis profiles.

Ri_b(z) = g (θ(z) − θ_s)(z − z_s) / (θ_s |𝐮(z) − 𝐮_s|²), h: first Ri_b = Ri_c,

with θ = T (p_ref/p)^{R_d/c_p} (a MetField has no humidity, so θ_v = θ) and the surface values taken from the lowest level.

Algorithm

1. Build θ, the T-point wind and the heights above the grid's surface.
2. Hand them to :func:xrtoolz.atm.pbl_height_bulk_richardson, which finds the first crossing and interpolates between levels.

Parameters:

Name	Type	Description	Default
field	MetField	Met on the analysis grid.	required
ri_crit	float	Critical bulk Richardson number Ri_c.	0.25

Returns:

Type	Description
Float[ndarray, 'time y x']	h [m]; NaN where no level reaches Ri_c inside the grid.

Raises:

Type	Description
MissingExtraError	If the io extra (xrtoolz) is not installed.

Examples:

A neutral layer under an inversion at 400 m, with shear:

>>> import jax.numpy as jnp
>>> from plumax.les_fvm.grid import make_grid
>>> from plumax.met.field import MetField
>>> grid = make_grid(
...     (0.0, 400.0, 4), (0.0, 400.0, 4), (0.0, 1000.0, 20), dtype=jnp.float64
... )
>>> z = grid.z - (grid.z[0] - 0.5 * grid.dz)
>>> theta = 300.0 + 0.01 * jnp.maximum(z - 400.0, 0.0)
>>> column = lambda f: jnp.broadcast_to(f[None, :, None, None], (1, 20, 4, 4))
>>> field = MetField(
...     plume_grid=grid, times=jnp.zeros(1), u=column(2.0 + 0.005 * z),
...     v=jnp.zeros((1, 20, 4, 4)), w=jnp.zeros((1, 20, 4, 4)),
...     temperature=column(theta), pressure=jnp.full((1, 20, 4, 4), 1.0e5),
...     pbl_height=jnp.zeros((1, 4, 4)),
... )
>>> height = pbl_height_bulk_richardson(field)
>>> height.shape
(1, 4, 4)
>>> bool(((height > 350.0) & (height < 1000.0)).all())
True

Source code in src/plumax/met/_src/pbl.py

def pbl_height_bulk_richardson(
    field: MetField, *, ri_crit: float = 0.25
) -> Float[np.ndarray, "time y x"]:
    """PBL height h from the bulk Richardson number of the analysis profiles.

    Ri_b(z) = g (θ(z) − θ_s)(z − z_s) / (θ_s |𝐮(z) − 𝐮_s|²),  h: first Ri_b = Ri_c,

    with θ = T (p_ref/p)^{R_d/c_p} (a ``MetField`` has no humidity, so θ_v = θ)
    and the surface values taken from the lowest level.

    Algorithm:
        1. Build θ, the T-point wind and the heights above the grid's surface.
        2. Hand them to :func:`xrtoolz.atm.pbl_height_bulk_richardson`, which
           finds the first crossing and interpolates between levels.

    Args:
        field: Met on the analysis grid.
        ri_crit: Critical bulk Richardson number Ri_c.

    Returns:
        h [m]; NaN where no level reaches Ri_c inside the grid.

    Raises:
        MissingExtraError: If the ``io`` extra (xrtoolz) is not installed.

    Examples:
        A neutral layer under an inversion at 400 m, with shear:

        >>> import jax.numpy as jnp
        >>> from plumax.les_fvm.grid import make_grid
        >>> from plumax.met.field import MetField
        >>> grid = make_grid(
        ...     (0.0, 400.0, 4), (0.0, 400.0, 4), (0.0, 1000.0, 20), dtype=jnp.float64
        ... )
        >>> z = grid.z - (grid.z[0] - 0.5 * grid.dz)
        >>> theta = 300.0 + 0.01 * jnp.maximum(z - 400.0, 0.0)
        >>> column = lambda f: jnp.broadcast_to(f[None, :, None, None], (1, 20, 4, 4))
        >>> field = MetField(
        ...     plume_grid=grid, times=jnp.zeros(1), u=column(2.0 + 0.005 * z),
        ...     v=jnp.zeros((1, 20, 4, 4)), w=jnp.zeros((1, 20, 4, 4)),
        ...     temperature=column(theta), pressure=jnp.full((1, 20, 4, 4), 1.0e5),
        ...     pbl_height=jnp.zeros((1, 4, 4)),
        ... )
        >>> height = pbl_height_bulk_richardson(field)
        >>> height.shape
        (1, 4, 4)
        >>> bool(((height > 350.0) & (height < 1000.0)).all())
        True
    """
    try:
        import xrtoolz.atm
    except ImportError as err:
        raise MissingExtraError(
            "pbl_height_bulk_richardson runs through xrtoolz; install it with "
            "`pip install 'plumax[io]'`"
        ) from err
    grid = field.plume_grid
    grid_z = np.asarray(grid.z, dtype=np.float64)
    # z_surface = z₀ − Δz/2
    grid_surface = grid_z[0] - 0.5 * grid.dz
    # h_k = z_k − z_surface
    heights = grid_z - grid_surface
    u = np.asarray(field.u, dtype=np.float64)
    v = np.asarray(field.v, dtype=np.float64)
    # each U-point's western neighbour, the edge column repeating itself
    u_west = einx.id(
        "t z y first, t z y rest -> t z y (first + rest)", u[..., :1], u[..., :-1]
    )
    # each V-point's southern neighbour
    v_south = einx.id(
        "t z first x, t z rest x -> t z (first + rest) x",
        v[..., :1, :],
        v[..., :-1, :],
    )
    # u_T = (u_{i−1/2} + u_{i+1/2}) / 2
    u_centre = 0.5 * (u + u_west)
    # v_T = (v_{j−1/2} + v_{j+1/2}) / 2
    v_centre = 0.5 * (v + v_south)
    temperature = np.asarray(field.temperature, dtype=np.float64)
    pressure = np.asarray(field.pressure, dtype=np.float64)
    # θ = T (p_ref / p)^{R_d / c_p}
    theta = temperature * (P_REFERENCE / pressure) ** (R_DRY / CP_DRY)
    dims = ("time", "level", "y", "x")
    profiles = xr.Dataset(
        {"theta_v": (dims, theta), "u": (dims, u_centre), "v": (dims, v_centre)},
        coords={"z_agl": ("level", heights)},
    )
    result = xrtoolz.atm.pbl_height_bulk_richardson(
        profiles, theta_v="theta_v", ri_crit=ri_crit, height="z_agl"
    )
    return result["pbl_height"].transpose("time", "y", "x").values

load_wrf

load_wrf(path: str | Path, *, plume_grid: PlumeGrid3D, origin: tuple[float, float] | LocalFrame = (0.0, 0.0), times: slice | None = None) -> MetField

Load a wrfout file onto an analysis grid.

Parameters:

Name	Type	Description	Default
path	str | Path	NetCDF wrfout file.	required
plume_grid	PlumeGrid3D	Analysis grid (from :func:plumax.les_fvm.make_grid).	required
origin	tuple[float, float] | LocalFrame	WRF-frame metres of the analysis grid's (x=0, y=0), or a :class:~plumax.met.LocalFrame the grid is laid out in; see the module docstring.	(0.0, 0.0)
times	slice | None	Optional slice over the file's Time axis.	None

Returns:

Name	Type	Description
A	MetField	class:MetField on plume_grid.

Raises:

Type	Description
MissingExtraError	If the io extra (xrtoolz) is not installed.
TimeAxisError	If times selects nothing.
MissingVariableError	With a LocalFrame, if the file lacks XLAT / XLONG or the COSALPHA / SINALPHA its earth-relative wind needs.
DomainError	With a LocalFrame, if an analysis point falls outside the WRF domain.
ValueError	With metre offsets, if an analysis point falls outside the WRF mass-point domain (the loader never extrapolates horizontally); or if the file has no usable time axis (Times, or XTIME with the SIMULATION_START_DATE attribute).

Source code in src/plumax/met/wrf.py

def load_wrf(
    path: str | Path,
    *,
    plume_grid: PlumeGrid3D,
    origin: tuple[float, float] | LocalFrame = (0.0, 0.0),
    times: slice | None = None,
) -> MetField:
    """Load a ``wrfout`` file onto an analysis grid.

    Args:
        path: NetCDF ``wrfout`` file.
        plume_grid: Analysis grid (from :func:`plumax.les_fvm.make_grid`).
        origin: WRF-frame metres of the analysis grid's ``(x=0, y=0)``,
            or a :class:`~plumax.met.LocalFrame` the grid is laid
            out in; see the module docstring.
        times: Optional ``slice`` over the file's ``Time`` axis.

    Returns:
        A :class:`MetField` on ``plume_grid``.

    Raises:
        MissingExtraError: If the ``io`` extra (xrtoolz) is not installed.
        TimeAxisError: If ``times`` selects nothing.
        MissingVariableError: With a ``LocalFrame``, if the file lacks
            ``XLAT`` / ``XLONG`` or the ``COSALPHA`` / ``SINALPHA`` its
            earth-relative wind needs.
        DomainError: With a ``LocalFrame``, if an analysis point falls outside
            the WRF domain.
        ValueError: With metre offsets, if an analysis point falls outside
            the WRF mass-point domain (the loader never extrapolates
            horizontally); or if the file has no usable time axis (``Times``,
            or ``XTIME`` with the ``SIMULATION_START_DATE`` attribute).
    """
    in_frame = isinstance(origin, LocalFrame)
    open_wrfout = _require_open_wrfout()

    # --- read ---
    rotation = ["COSALPHA", "SINALPHA"]
    with xr.open_dataset(path, decode_times=False) as raw:
        has_rotation = all(name in raw for name in rotation)
        selected_steps = range(raw.sizes["Time"])[times or slice(None)]
    if len(selected_steps) == 0:
        raise TimeAxisError("load_wrf: `times` selects no time steps")
    # xrtoolz destaggers to mass points, rebuilds height above ground and
    # temperature, and decodes the time axis.
    with open_wrfout(
        path, times=times, variables=rotation if has_rotation else None
    ) as ds:
        dx_wrf = float(ds.attrs["DX"])
        dy_wrf = float(ds.attrs["DY"])
        if in_frame:
            u_mass, v_mass = _earth_relative_wind(ds)
            lat_mass, lon_mass = _latlon(ds)
        else:
            u_mass, v_mass = _grid_relative_wind(ds)
        w_mass = ds["w"].values
        temperature = ds["temperature"].values
        pressure = ds["pressure"].values
        height_agl = np.broadcast_to(ds["z_agl"].values, u_mass.shape)
        pbl_height = ds["pbl_height"].values
        clock = TimeAxis.from_datetimes(ds["time"].values)

    # --- vertical targets ---
    # The analysis grid's lower boundary is its surface (``make_grid``
    # semantics), and WRF heights are above ground.
    grid_z = np.asarray(plume_grid.z, dtype=np.float64)
    # z_surface = z₀ − Δz/2
    grid_surface = grid_z[0] - 0.5 * plume_grid.dz
    # h_k = z_k − z_surface
    target_heights = grid_z - grid_surface
    dtype = plume_grid.x.dtype
    as_jax = lambda a: jnp.asarray(a, dtype=dtype)

    if in_frame:
        # --- horizontal, by lat/lon ---
        to_index = _wrf_index_map(lat_mass, lon_mass, origin)
        points = stagger_points(plume_grid, origin, to_index, loader="load_wrf")
        at_t = points["t"]
        on_levels = lambda f: _interpolate_vertical(height_agl, f, target_heights)
        u_frame, v_frame = frame_wind(on_levels(u_mass), on_levels(v_mass), points)
        w_t = bilinear_at(on_levels(w_mass), at_t.i, at_t.j)
        temperature_t = bilinear_at(on_levels(temperature), at_t.i, at_t.j)
        pressure_t = bilinear_at(on_levels(pressure), at_t.i, at_t.j)
        pbl_height_t = bilinear_at(pbl_height, at_t.i, at_t.j)
        return MetField(
            plume_grid=plume_grid,
            times=as_jax(clock.seconds),
            u=as_jax(u_frame),
            v=as_jax(v_frame),
            w=as_jax(w_t),
            temperature=as_jax(temperature_t),
            pressure=as_jax(pressure_t),
            pbl_height=as_jax(pbl_height_t),
            t0=clock.t0_iso,
        )

    # --- horizontal, by WRF metres ---
    n_y_w, n_x_w = u_mass.shape[-2:]
    x_t = np.asarray(plume_grid.x, dtype=np.float64) + origin[0]
    y_t = np.asarray(plume_grid.y, dtype=np.float64) + origin[1]
    x_u = x_t + 0.5 * plume_grid.dx
    y_v = y_t + 0.5 * plume_grid.dy
    _check_inside(x_u, y_v, dx_w=dx_wrf, dy_w=dy_wrf, n_x_w=n_x_w, n_y_w=n_y_w)
    _check_inside(x_t, y_t, dx_w=dx_wrf, dy_w=dy_wrf, n_x_w=n_x_w, n_y_w=n_y_w)

    def onto(field: np.ndarray, x: np.ndarray, y: np.ndarray) -> np.ndarray:
        column = _interpolate_vertical(height_agl, field, target_heights)
        return _bilinear(column, x / dx_wrf, y / dy_wrf)

    return MetField(
        plume_grid=plume_grid,
        times=as_jax(clock.seconds),
        u=as_jax(onto(u_mass, x_u, y_t)),
        v=as_jax(onto(v_mass, x_t, y_v)),
        w=as_jax(onto(w_mass, x_t, y_t)),
        temperature=as_jax(onto(temperature, x_t, y_t)),
        pressure=as_jax(onto(pressure, x_t, y_t)),
        pbl_height=as_jax(_bilinear(pbl_height, x_t / dx_wrf, y_t / dy_wrf)),
        t0=clock.t0_iso,
    )