Lagrangian particles

Lagrangian particle dispersion. Tier II.

lagrangian

Tier II — Lagrangian particle dispersion (JAX).

A Markov-1 Langevin particle model: stochastic trajectories driven by mean wind plus an Ornstein–Uhlenbeck turbulent-velocity process. The bridge between the analytical Tier I (no real wind variability) and the Eulerian Tier III. See the Tier II design doc.

This implements Step 1 (the forward model) and Step 2 (model-based inference):

• turbulence — HomogeneousTurbulence (well-mixed-exact) and Hanna (1982) Monin–Obukhov-similarity profiles.
• particles — the Markov-1 integrator with the Thomson (1987) well-mixed vertical drift correction and ground / PBL reflection.
• concentration — forward steady concentration field via residence-time binning (simulate_lagrangian).
• footprint — backward source–receptor sensitivity (compute_footprint).
• inversion — closed-form Gaussian and sign-constrained lognormal Bayesian inversion of the source vector from observations, with a Matérn-3/2 spatial prior (linear_gaussian_inversion / lognormal_inversion).

The footprint emulator and the amortized predictor (Steps 3–5) are not yet implemented.

GaussianPosterior dataclass

GaussianPosterior(mean: Array, covariance: Array)

Closed-form Gaussian posterior over the source vector.

Attributes:

Name	Type	Description
mean	Array	Posterior mean emission field q* [same units as the prior mean], shape (n_grid,).
covariance	Array	Posterior covariance P* [units²], shape (n_grid, n_grid).

std property

std: Array

Posterior per-cell standard deviation sqrt(diag(P*)).

LognormalPosterior dataclass

LognormalPosterior(mean: Array, log_increment: Array, log_covariance: Array)

Linearised lognormal posterior over a non-negative source vector.

Attributes:

Name	Type	Description
mean	Array	Posterior mean emission field q* = q_a · exp(δ*) ≥ 0, shape (n_grid,).
log_increment	Array	Log-space MAP increment δ* = log(q*/q_a), shape (n_grid,).
log_covariance	Array	Posterior covariance in log-space P*_log, shape (n_grid, n_grid).

ParticleState dataclass

ParticleState(position: Array, velocity: Array)

State of an ensemble of Lagrangian particles.

Registered as a JAX pytree so it can be carried through lax.scan / vmap; position and velocity are the (array) children.

Attributes:

Name	Type	Description
position	Array	Particle positions (x, y, z) [m], shape (n, 3).
velocity	Array	Turbulent velocity perturbations [m/s], shape (n, 3).

HannaTurbulence dataclass

HannaTurbulence(u_star: float, pbl_height: float, obukhov_length: float, w_star: float = 0.0, z_floor: float = 0.1)

Height-dependent Hanna (1982) turbulence as a TurbulenceModel.

Thin adapter that binds the surface-layer scales required by :func:hanna_profiles and exposes the at(z) -> (sigma, tau) interface the integrator (:func:plumax.lagrangian.langevin_step, :func:~plumax.lagrangian.integrate_particles, :func:~plumax.lagrangian.compute_footprint) expects. Because σ_w varies with height, the well-mixed vertical drift correction ½(1+v_w²/σ_w²) ∂σ_w²/∂z is non-zero — the inhomogeneous path that :class:HomogeneousTurbulence never exercises.

Attributes:

Name	Type	Description
u_star	float	Friction velocity u_* [m/s].
pbl_height	float	Boundary-layer height h [m].
obukhov_length	float	Monin–Obukhov length L [m] (< 0 unstable).
w_star	float	Convective velocity scale w_* [m/s] (used when L < 0).
z_floor	float	Minimum height [m] at which the profiles are evaluated. The convective σ_w² term carries a (z/h)^{2/3} factor whose ∂/∂z is singular at the ground, so the integrator's well-mixed drift ½(1+v_w²/σ_w²)∂σ_w²/∂z would be inf/NaN for a ground-level source or a particle reflected exactly to z = 0. Flooring z keeps the drift finite; the default 0.1 m sits at the roughness / surface-layer scale where the similarity profiles are not valid anyway.

at

at(z: Array) -> tuple[jax.Array, jax.Array]

Return (σ, τ_L) at height(s) z from the Hanna (1982) profiles.

z is floored at :attr:z_floor before evaluation so the vertical drift stays finite at the ground (see the class docstring). σ and τ_L have a trailing axis of size 3, matching :meth:HomogeneousTurbulence.at.

Source code in src/plumax/lagrangian/turbulence.py

def at(self, z: jax.Array) -> tuple[jax.Array, jax.Array]:
    """Return ``(σ, τ_L)`` at height(s) ``z`` from the Hanna (1982) profiles.

    ``z`` is floored at :attr:`z_floor` before evaluation so the vertical
    drift stays finite at the ground (see the class docstring). ``σ`` and
    ``τ_L`` have a trailing axis of size 3, matching
    :meth:`HomogeneousTurbulence.at`.
    """
    z = jnp.maximum(jnp.asarray(z, dtype=float), self.z_floor)
    return hanna_profiles(
        z,
        u_star=self.u_star,
        pbl_height=self.pbl_height,
        obukhov_length=self.obukhov_length,
        w_star=self.w_star,
    )

HomogeneousTurbulence dataclass

HomogeneousTurbulence(sigma_u: float, sigma_v: float, sigma_w: float, tau_u: float, tau_v: float, tau_w: float)

Spatially uniform, well-mixed-exact turbulence statistics.

Attributes:

Name	Type	Description
sigma_u	float	Along-wind velocity standard deviation [m/s].
sigma_v	float	Cross-wind velocity standard deviation [m/s].
sigma_w	float	Vertical velocity standard deviation [m/s].
tau_u	float	Along-wind Lagrangian timescale [s].
tau_v	float	Cross-wind Lagrangian timescale [s].
tau_w	float	Vertical Lagrangian timescale [s].

sigma property

sigma: Array

Velocity standard deviations (σ_u, σ_v, σ_w) [m/s], shape (3,).

tau property

tau: Array

Lagrangian timescales (τ_u, τ_v, τ_w) [s], shape (3,).

isotropic classmethod

isotropic(sigma: float, tau: float) -> HomogeneousTurbulence

Build isotropic turbulence: equal σ and τ_L on all axes.

Source code in src/plumax/lagrangian/turbulence.py

@classmethod
def isotropic(cls, sigma: float, tau: float) -> HomogeneousTurbulence:
    """Build isotropic turbulence: equal ``σ`` and ``τ_L`` on all axes."""
    return cls(sigma, sigma, sigma, tau, tau, tau)

at

at(z: Array) -> tuple[jax.Array, jax.Array]

Return (σ, τ_L) at height(s) z — constant, so z is ignored.

The signature matches :func:hanna_profiles so the integrator can treat any turbulence model uniformly. σ and τ_L broadcast to the shape of z along a trailing axis of size 3.

Source code in src/plumax/lagrangian/turbulence.py

def at(self, z: jax.Array) -> tuple[jax.Array, jax.Array]:
    """Return ``(σ, τ_L)`` at height(s) ``z`` — constant, so ``z`` is ignored.

    The signature matches :func:`hanna_profiles` so the integrator can treat
    any turbulence model uniformly. ``σ`` and ``τ_L`` broadcast to the shape
    of ``z`` along a trailing axis of size 3.
    """
    z = jnp.asarray(z)
    ones = jnp.ones((*z.shape, 3))
    return ones * self.sigma, ones * self.tau

bin_positions

bin_positions(positions: Array, x_edges: Array, y_edges: Array, z_edges: Array) -> jax.Array

Histogram particle positions (n, 3) onto a 3-D grid (counts).

Source code in src/plumax/lagrangian/concentration.py

def bin_positions(
    positions: jax.Array,
    x_edges: jax.Array,
    y_edges: jax.Array,
    z_edges: jax.Array,
) -> jax.Array:
    """Histogram particle ``positions`` ``(n, 3)`` onto a 3-D grid (counts)."""
    nx = x_edges.shape[0] - 1
    ny = y_edges.shape[0] - 1
    nz = z_edges.shape[0] - 1
    ix, gx = _cell_indices(positions[:, 0], x_edges)
    iy, gy = _cell_indices(positions[:, 1], y_edges)
    iz, gz = _cell_indices(positions[:, 2], z_edges)
    inside = gx & gy & gz
    flat = (ix * ny + iy) * nz + iz
    counts = jnp.zeros(nx * ny * nz).at[flat].add(inside.astype(float))
    return counts.reshape(nx, ny, nz)

simulate_lagrangian

simulate_lagrangian(emission_rate: float, source_location: tuple[float, float, float], turbulence: TurbulenceModel, domain_x: tuple[float, float, int], domain_y: tuple[float, float, int], domain_z: tuple[float, float, int], *, n_particles: int = 5000, t_end: float = 600.0, dt: float = 1.0, wind: Callable[[Array], Array] | None = None, wind_speed: float | None = None, wind_direction: float | None = None, pbl_height: float | None = None, background_conc: float = 0.0, seed: int = 0) -> xr.Dataset

Simulate a steady Lagrangian concentration field on a 3-D grid.

Parameters:

Name	Type	Description	Default
emission_rate	float	Continuous source strength Q [kg/s], > 0.	required
source_location	tuple[float, float, float]	(x, y, z) source coordinates [m], z ≥ 0.	required
turbulence	TurbulenceModel	Turbulence model (e.g. :class:~plumax.lagrangian.turbulence.HomogeneousTurbulence).	required
domain_x	tuple[float, float, int]	(start, stop, n_cells) along x.	required
domain_y	tuple[float, float, int]	(start, stop, n_cells) along y.	required
domain_z	tuple[float, float, int]	(start, stop, n_cells) along z.	required
n_particles	int	Number of particles in the ensemble.	5000
t_end	float	Integration horizon [s].	600.0
dt	float	Time step [s].	1.0
wind	Callable[[Array], Array] | None	Mean-wind field t -> (u, v, w). Mutually exclusive with the wind_speed / wind_direction shortcut.	None
wind_speed	float | None	Wind speed [m/s] (with wind_direction) if wind is not given.	None
wind_direction	float | None	Meteorological direction [deg from North].	None
pbl_height	float | None	Optional reflecting PBL lid [m].	None
background_conc	float	Additive background concentration [kg/m³].	0.0
seed	int	PRNG seed.	0

Returns:

Name	Type	Description
An	Dataset	class:xarray.Dataset with concentration (x, y, z) [kg/m³] and
	Dataset	column_concentration (x, y) [kg/m²].

Raises:

Type	Description
ValueError	on a non-positive emission rate, a negative source height, or an ambiguous / missing wind specification.

Source code in src/plumax/lagrangian/concentration.py

def simulate_lagrangian(
    emission_rate: float,
    source_location: tuple[float, float, float],
    turbulence: TurbulenceModel,
    domain_x: tuple[float, float, int],
    domain_y: tuple[float, float, int],
    domain_z: tuple[float, float, int],
    *,
    n_particles: int = 5000,
    t_end: float = 600.0,
    dt: float = 1.0,
    wind: Callable[[jax.Array], jax.Array] | None = None,
    wind_speed: float | None = None,
    wind_direction: float | None = None,
    pbl_height: float | None = None,
    background_conc: float = 0.0,
    seed: int = 0,
) -> xr.Dataset:
    """Simulate a steady Lagrangian concentration field on a 3-D grid.

    Args:
        emission_rate: Continuous source strength ``Q`` [kg/s], ``> 0``.
        source_location: ``(x, y, z)`` source coordinates [m], ``z ≥ 0``.
        turbulence: Turbulence model (e.g.
            :class:`~plumax.lagrangian.turbulence.HomogeneousTurbulence`).
        domain_x: ``(start, stop, n_cells)`` along x.
        domain_y: ``(start, stop, n_cells)`` along y.
        domain_z: ``(start, stop, n_cells)`` along z.
        n_particles: Number of particles in the ensemble.
        t_end: Integration horizon [s].
        dt: Time step [s].
        wind: Mean-wind field ``t -> (u, v, w)``. Mutually exclusive with the
            ``wind_speed`` / ``wind_direction`` shortcut.
        wind_speed: Wind speed [m/s] (with ``wind_direction``) if ``wind`` is
            not given.
        wind_direction: Meteorological direction [deg from North].
        pbl_height: Optional reflecting PBL lid [m].
        background_conc: Additive background concentration [kg/m³].
        seed: PRNG seed.

    Returns:
        An :class:`xarray.Dataset` with ``concentration`` (x, y, z) [kg/m³] and
        ``column_concentration`` (x, y) [kg/m²].

    Raises:
        ValueError: on a non-positive emission rate, a negative source height,
            or an ambiguous / missing wind specification.
    """
    import xarray as xr

    if not (emission_rate > 0.0):
        raise ValueError("simulate_lagrangian: `emission_rate` must be > 0")
    if source_location[2] < 0.0:
        raise ValueError("simulate_lagrangian: source height must be ≥ 0")

    if wind is None:
        if wind_speed is None or wind_direction is None:
            raise ValueError(
                "simulate_lagrangian: provide either `wind` or both "
                "`wind_speed` and `wind_direction`."
            )
        wind = wind_from_speed_direction(wind_speed, wind_direction)
    elif wind_speed is not None or wind_direction is not None:
        raise ValueError(
            "simulate_lagrangian: pass `wind` or the speed/direction pair, not both."
        )

    x_edges = np.linspace(domain_x[0], domain_x[1], int(domain_x[2]) + 1)
    y_edges = np.linspace(domain_y[0], domain_y[1], int(domain_y[2]) + 1)
    z_edges = np.linspace(domain_z[0], domain_z[1], int(domain_z[2]) + 1)
    x_c = 0.5 * (x_edges[:-1] + x_edges[1:])
    y_c = 0.5 * (y_edges[:-1] + y_edges[1:])
    z_c = 0.5 * (z_edges[:-1] + z_edges[1:])
    cell_volume = float(
        (x_edges[1] - x_edges[0])
        * (y_edges[1] - y_edges[0])
        * (z_edges[1] - z_edges[0])
    )

    key = jax.random.PRNGKey(seed)
    key, vkey = jax.random.split(key)

    src = jnp.asarray(source_location, dtype=float)
    sigma0, _ = turbulence.at(src[2])
    vel0 = jax.random.normal(vkey, (n_particles, 3)) * sigma0
    pos0 = jnp.broadcast_to(src, (n_particles, 3))
    state = ParticleState(position=pos0, velocity=vel0)

    xe = jnp.asarray(x_edges)
    ye = jnp.asarray(y_edges)
    ze = jnp.asarray(z_edges)
    residence = _accumulate_residence(
        state, wind, turbulence, t_end, dt, key, pbl_height, xe, ye, ze
    )

    conc = emission_rate * np.asarray(residence) / (n_particles * cell_volume)
    conc = conc + background_conc
    dz = float(z_edges[1] - z_edges[0])
    column = conc.sum(axis=2) * dz

    ds = xr.Dataset(
        data_vars={
            "concentration": (["x", "y", "z"], conc),
            "column_concentration": (["x", "y"], column),
        },
        coords={"x": x_c, "y": y_c, "z": z_c},
        attrs={
            "title": "Lagrangian particle dispersion (steady field, JAX)",
            "model": "markov-1-langevin",
            "emission_rate": emission_rate,
            "emission_rate_units": "kg/s",
            "n_particles": n_particles,
            "t_end": t_end,
            "dt": dt,
            "background_concentration": background_conc,
        },
    )
    ds["concentration"].attrs = {"long_name": "Mass concentration", "units": "kg/m^3"}
    ds["column_concentration"].attrs = {
        "long_name": "Column-integrated concentration",
        "units": "kg/m^2",
    }
    return ds

compute_footprint

compute_footprint(receptor_location: tuple[float, float, float], turbulence: TurbulenceModel, domain_x: tuple[float, float, int], domain_y: tuple[float, float, int], *, wind: Callable[[Array], Array], n_particles: int = 5000, t_back: float = 600.0, dt: float = 1.0, pbl_height: float = 1000.0, pbl_fraction: float = 0.5, air_density: float = 1.2, receptor_time: float = 0.0, seed: int = 0) -> tuple[np.ndarray, np.ndarray, np.ndarray]

Compute a single receptor's surface footprint by backward integration.

Parameters:

Name	Type	Description	Default
receptor_location	tuple[float, float, float]	(x, y, z) receptor coordinates [m].	required
turbulence	TurbulenceModel	Turbulence model.	required
domain_x	tuple[float, float, int]	(start, stop, n_cells) along x for the surface grid.	required
domain_y	tuple[float, float, int]	(start, stop, n_cells) along y for the surface grid.	required
wind	Callable[[Array], Array]	Forward mean-wind field t -> (u, v, w); integration reverses it internally and samples it on the receptor clock (see receptor_time).	required
n_particles	int	Ensemble size.	5000
t_back	float	Backward integration time [s].	600.0
dt	float	Time step [s].	1.0
pbl_height	float	Boundary-layer height h [m] (reflecting lid).	1000.0
pbl_fraction	float	Fraction f_pbl of h defining the surface layer in which surface flux is "seen".	0.5
air_density	float	Air density ρ_air [kg/m³].	1.2
receptor_time	float	Physical time T of the receptor observation [s]. Backward steps undo the forward intervals in reverse order (the partial final interval first) and sample −wind at each interval's physical start time, so a time-varying wind is evaluated at the correct time and the run is the exact discrete reverse of a forward one. The default 0.0 is exact for a stationary wind (its argument is ignored) and simply reverses it.	0.0
seed	int	PRNG seed.	0

Returns:

Type	Description
ndarray	(footprint, x_centers, y_centers) where footprint has shape
ndarray	(nx, ny) and units s·m²·kg⁻¹ (flux sensitivity; see the module
ndarray	docstring).

Source code in src/plumax/lagrangian/footprint.py

def compute_footprint(
    receptor_location: tuple[float, float, float],
    turbulence: TurbulenceModel,
    domain_x: tuple[float, float, int],
    domain_y: tuple[float, float, int],
    *,
    wind: Callable[[jax.Array], jax.Array],
    n_particles: int = 5000,
    t_back: float = 600.0,
    dt: float = 1.0,
    pbl_height: float = 1000.0,
    pbl_fraction: float = 0.5,
    air_density: float = 1.2,
    receptor_time: float = 0.0,
    seed: int = 0,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Compute a single receptor's surface footprint by backward integration.

    Args:
        receptor_location: ``(x, y, z)`` receptor coordinates [m].
        turbulence: Turbulence model.
        domain_x: ``(start, stop, n_cells)`` along x for the surface grid.
        domain_y: ``(start, stop, n_cells)`` along y for the surface grid.
        wind: Forward mean-wind field ``t -> (u, v, w)``; integration reverses
            it internally and samples it on the receptor clock (see
            ``receptor_time``).
        n_particles: Ensemble size.
        t_back: Backward integration time [s].
        dt: Time step [s].
        pbl_height: Boundary-layer height ``h`` [m] (reflecting lid).
        pbl_fraction: Fraction ``f_pbl`` of ``h`` defining the surface layer in
            which surface flux is "seen".
        air_density: Air density ``ρ_air`` [kg/m³].
        receptor_time: Physical time ``T`` of the receptor observation [s].
            Backward steps undo the forward intervals in reverse order (the
            partial final interval first) and sample ``−wind`` at each interval's
            physical start time, so a time-varying ``wind`` is evaluated at the
            correct time and the run is the exact discrete reverse of a forward
            one. The default ``0.0`` is exact for a stationary ``wind`` (its
            argument is ignored) and simply reverses it.
        seed: PRNG seed.

    Returns:
        ``(footprint, x_centers, y_centers)`` where ``footprint`` has shape
        ``(nx, ny)`` and units ``s·m²·kg⁻¹`` (flux sensitivity; see the module
        docstring).
    """
    x_edges = np.linspace(domain_x[0], domain_x[1], int(domain_x[2]) + 1)
    y_edges = np.linspace(domain_y[0], domain_y[1], int(domain_y[2]) + 1)
    x_c = 0.5 * (x_edges[:-1] + x_edges[1:])
    y_c = 0.5 * (y_edges[:-1] + y_edges[1:])
    mix_height = pbl_fraction * pbl_height

    key = jax.random.PRNGKey(seed)
    key, vkey = jax.random.split(key)
    rec = jnp.asarray(receptor_location, dtype=float)
    sigma0, _ = turbulence.at(rec[2])
    vel0 = jax.random.normal(vkey, (n_particles, 3)) * sigma0
    pos0 = jnp.broadcast_to(rec, (n_particles, 3))
    state = ParticleState(position=pos0, velocity=vel0)

    xe, ye = jnp.asarray(x_edges), jnp.asarray(y_edges)
    n_steps = n_steps_for_horizon(t_back, dt)
    keys = jax.random.split(key, n_steps)
    # Per-step durations summing to exactly ``t_back``: the final forward step is
    # shortened to the remainder when ``t_back`` is not a multiple of ``dt``, so
    # the footprint integrates over the requested horizon rather than
    # ``n_steps * dt`` (which would over-count surface residence).
    #
    # Backward integration undoes the forward intervals in reverse order, so the
    # first backward step consumes the (possibly partial) *final* forward
    # interval — hence the reversal. ``wind_starts[j]`` is the physical start
    # time of the interval step ``j`` undoes; sampling ``−wind(wind_starts[j])``
    # (step-start, matching ``integrate_particles``) makes the backward run the
    # exact discrete reverse of the forward one for any horizon.
    dts = step_durations(t_back, dt, n_steps)[::-1]
    wind_starts = receptor_time - jnp.cumsum(dts)
    nx, ny = len(x_c), len(y_c)

    def body(carry, inputs):
        st, hist = carry
        t_start, k, dt_i = inputs
        st = langevin_step(
            st,
            -jnp.asarray(wind(t_start)),
            turbulence,
            dt_i,
            k,
            pbl_height=pbl_height,
        )
        below = st.position[:, 2] < mix_height
        hist = hist + _bin_surface(st.position, xe, ye, below) * dt_i
        return (st, hist), None

    (_, residence), _ = jax.lax.scan(
        body, (state, jnp.zeros((nx, ny))), (wind_starts, keys, dts)
    )

    # Flux-sensitivity normalisation: residence-time per particle divided by the
    # surface-layer air mass column density ρ·h (no per-cell-area division), so
    # ``F`` has units s·m²·kg⁻¹ and pairs with a surface flux in kg·m⁻²·s⁻¹.
    footprint = np.asarray(residence) / (n_particles * air_density * mix_height)
    return footprint, x_c, y_c

linear_gaussian_inversion

linear_gaussian_inversion(footprint: Array, observation: Array, *, prior_mean: Array, prior_cov: Array, obs_variance: Array) -> GaussianPosterior

Closed-form Gaussian–Gaussian (BLUE) inversion of the source vector.

For a Gaussian prior q ~ N(q_a, B) and Gaussian likelihood y = F q + ε, ε ~ N(0, R) (R diagonal), returns the exact posterior N(q*, P*) with::

q* = q_a + B Fᵀ (F B Fᵀ + R)⁻¹ (y − F q_a)
P* = B − B Fᵀ (F B Fᵀ + R)⁻¹ F B

Parameters:

Name	Type	Description	Default
footprint	Array	Source–receptor matrix F, shape (n_obs, n_grid).	required
observation	Array	Observed values y, shape (n_obs,).	required
prior_mean	Array	Prior mean q_a, shape (n_grid,).	required
prior_cov	Array	Prior covariance B, shape (n_grid, n_grid) (e.g. from :func:matern32_covariance).	required
obs_variance	Array	Diagonal of R, shape (n_obs,) (e.g. from :func:observation_covariance).	required

Returns:

Name	Type	Description
The	GaussianPosterior	class:GaussianPosterior over q.

Source code in src/plumax/lagrangian/inversion.py

def linear_gaussian_inversion(
    footprint: jax.Array,
    observation: jax.Array,
    *,
    prior_mean: jax.Array,
    prior_cov: jax.Array,
    obs_variance: jax.Array,
) -> GaussianPosterior:
    """Closed-form Gaussian–Gaussian (BLUE) inversion of the source vector.

    For a Gaussian prior ``q ~ N(q_a, B)`` and Gaussian likelihood
    ``y = F q + ε``, ``ε ~ N(0, R)`` (``R`` diagonal), returns the exact
    posterior ``N(q*, P*)`` with::

        q* = q_a + B Fᵀ (F B Fᵀ + R)⁻¹ (y − F q_a)
        P* = B − B Fᵀ (F B Fᵀ + R)⁻¹ F B

    Args:
        footprint: Source–receptor matrix ``F``, shape ``(n_obs, n_grid)``.
        observation: Observed values ``y``, shape ``(n_obs,)``.
        prior_mean: Prior mean ``q_a``, shape ``(n_grid,)``.
        prior_cov: Prior covariance ``B``, shape ``(n_grid, n_grid)`` (e.g. from
            :func:`matern32_covariance`).
        obs_variance: Diagonal of ``R``, shape ``(n_obs,)`` (e.g. from
            :func:`observation_covariance`).

    Returns:
        The :class:`GaussianPosterior` over ``q``.
    """
    f = jnp.asarray(footprint, dtype=float)
    y = jnp.asarray(observation, dtype=float)
    qa = jnp.asarray(prior_mean, dtype=float)
    b = jnp.asarray(prior_cov, dtype=float)
    r = jnp.asarray(obs_variance, dtype=float)
    _check_inversion_shapes(f, y, qa, b, r)

    innovation = y - f @ qa
    increment, posterior_cov = _blue_update(f, innovation, b, r)
    return GaussianPosterior(mean=qa + increment, covariance=posterior_cov)

lognormal_inversion

lognormal_inversion(footprint: Array, observation: Array, *, prior_mean: Array, prior_log_cov: Array, obs_variance: Array) -> LognormalPosterior

Sign-constrained inversion via a lognormal prior linearised in log-space.

With log q ~ N(log q_a, B_log) and the forward linearised about the prior mean, the tangent-linear Jacobian is F̃ = F · diag(q_a) and the log-space increment δ = log(q / q_a) follows the same BLUE update::

δ* = B_log F̃ᵀ (F̃ B_log F̃ᵀ + R)⁻¹ (y − F q_a)
q* = q_a · exp(δ*)          (≥ 0 by construction)

This is the design note's recommended v1 prior: it enforces non-negative emissions while keeping the update closed-form (no projected-gradient or NNLS inner loop). The innovation y − F q_a is the observation minus the forward evaluated at the prior mean (δ = 0 ⇒ q = q_a).

Parameters:

Name	Type	Description	Default
footprint	Array	Source–receptor matrix F, shape (n_obs, n_grid).	required
observation	Array	Observed values y, shape (n_obs,).	required
prior_mean	Array	Prior median emission field q_a > 0, shape (n_grid,).	required
prior_log_cov	Array	Log-space prior covariance B_log, shape (n_grid, n_grid).	required
obs_variance	Array	Diagonal of R, shape (n_obs,).	required

Returns:

Name	Type	Description
The	LognormalPosterior	class:LognormalPosterior over q.

Source code in src/plumax/lagrangian/inversion.py

def lognormal_inversion(
    footprint: jax.Array,
    observation: jax.Array,
    *,
    prior_mean: jax.Array,
    prior_log_cov: jax.Array,
    obs_variance: jax.Array,
) -> LognormalPosterior:
    """Sign-constrained inversion via a lognormal prior linearised in log-space.

    With ``log q ~ N(log q_a, B_log)`` and the forward linearised about the
    prior mean, the tangent-linear Jacobian is ``F̃ = F · diag(q_a)`` and the
    log-space increment ``δ = log(q / q_a)`` follows the same BLUE update::

        δ* = B_log F̃ᵀ (F̃ B_log F̃ᵀ + R)⁻¹ (y − F q_a)
        q* = q_a · exp(δ*)          (≥ 0 by construction)

    This is the design note's recommended v1 prior: it enforces non-negative
    emissions while keeping the update closed-form (no projected-gradient or
    NNLS inner loop). The innovation ``y − F q_a`` is the observation minus the
    forward evaluated at the prior mean (``δ = 0 ⇒ q = q_a``).

    Args:
        footprint: Source–receptor matrix ``F``, shape ``(n_obs, n_grid)``.
        observation: Observed values ``y``, shape ``(n_obs,)``.
        prior_mean: Prior median emission field ``q_a > 0``, shape
            ``(n_grid,)``.
        prior_log_cov: Log-space prior covariance ``B_log``, shape
            ``(n_grid, n_grid)``.
        obs_variance: Diagonal of ``R``, shape ``(n_obs,)``.

    Returns:
        The :class:`LognormalPosterior` over ``q``.
    """
    f = jnp.asarray(footprint, dtype=float)
    y = jnp.asarray(observation, dtype=float)
    qa = jnp.asarray(prior_mean, dtype=float)
    b_log = jnp.asarray(prior_log_cov, dtype=float)
    r = jnp.asarray(obs_variance, dtype=float)
    _check_inversion_shapes(f, y, qa, b_log, r)
    # Eager positivity check; skipped under jit/vmap over `prior_mean` (where a
    # host-bool conversion would raise). The lognormal form q = q_a·exp(δ)
    # implicitly assumes q_a > 0; a non-positive concrete q_a is a usage error.
    if not _is_traced(prior_mean) and bool(jnp.any(qa <= 0.0)):
        raise ValueError("lognormal_inversion: `prior_mean` (q_a) must be > 0.")

    jac_tilde = f * qa[None, :]  # F̃ = F diag(q_a), (n_obs, n_grid)
    innovation = y - f @ qa
    log_increment, log_cov = _blue_update(jac_tilde, innovation, b_log, r)
    mean = qa * jnp.exp(log_increment)
    return LognormalPosterior(
        mean=mean, log_increment=log_increment, log_covariance=log_cov
    )

matern32_covariance

matern32_covariance(coordinates: Array, *, variance: float, length_scale: float) -> jax.Array

Build a Matérn-3/2 prior covariance over grid-cell coordinates.

The Matérn-3/2 kernel

k(r) = σ² (1 + √3 r / ℓ) exp(−√3 r / ℓ),   r = ‖xᵢ − xⱼ‖

gives a once-differentiable random field — smoother than the exponential (AR(1) / Matérn-1/2) kernel but far less rigid than a squared-exponential. It is the design note's recommended spatial prior for regional emission inversions: a real correlation length ℓ regularises the posterior so unresolved structure spreads over neighbouring cells instead of collapsing into single-cell spikes.

Parameters:

Name	Type	Description	Default
coordinates	Array	Cell-centre coordinates [m], shape (n_grid, n_dim) (1-D coordinates of shape (n_grid,) are accepted and promoted).	required
variance	float	Marginal prior variance σ² (per cell), > 0.	required
length_scale	float	Correlation length ℓ [m], > 0.	required

Returns:

Type	Description
Array	The (n_grid, n_grid) covariance matrix, symmetric positive-definite.

Source code in src/plumax/lagrangian/inversion.py

def matern32_covariance(
    coordinates: jax.Array,
    *,
    variance: float,
    length_scale: float,
) -> jax.Array:
    """Build a Matérn-3/2 prior covariance over grid-cell coordinates.

    The Matérn-3/2 kernel

        k(r) = σ² (1 + √3 r / ℓ) exp(−√3 r / ℓ),   r = ‖xᵢ − xⱼ‖

    gives a once-differentiable random field — smoother than the exponential
    (AR(1) / Matérn-1/2) kernel but far less rigid than a squared-exponential.
    It is the design note's recommended spatial prior for regional emission
    inversions: a real correlation length ``ℓ`` regularises the posterior so
    unresolved structure spreads over neighbouring cells instead of collapsing
    into single-cell spikes.

    Args:
        coordinates: Cell-centre coordinates [m], shape ``(n_grid, n_dim)``
            (1-D coordinates of shape ``(n_grid,)`` are accepted and promoted).
        variance: Marginal prior variance ``σ²`` (per cell), ``> 0``.
        length_scale: Correlation length ``ℓ`` [m], ``> 0``.

    Returns:
        The ``(n_grid, n_grid)`` covariance matrix, symmetric positive-definite.
    """
    # Validate eagerly for concrete inputs; skip under jit/grad/vmap (the doc
    # advertises gradient-based tuning of variance / length_scale, where these
    # would be tracers and a Python ``<=`` test would raise).
    if not _is_traced(variance, length_scale):
        if variance <= 0.0:
            raise ValueError("matern32_covariance: `variance` must be > 0.")
        if length_scale <= 0.0:
            raise ValueError("matern32_covariance: `length_scale` must be > 0.")
    coords = jnp.asarray(coordinates, dtype=float)
    if coords.ndim == 1:
        coords = coords[:, None]
    diff = coords[:, None, :] - coords[None, :, :]
    r = jnp.sqrt(jnp.sum(diff**2, axis=-1) + 1e-300)
    scaled = jnp.sqrt(3.0) * r / length_scale
    return variance * (1.0 + scaled) * jnp.exp(-scaled)

observation_covariance

observation_covariance(retrieval_variance: Array, *, representation_variance: float | Array = 0.0) -> jax.Array

Assemble the diagonal observation-error covariance R.

R = R_retr + R_repr: the per-pixel L2 retrieval error plus a representation error that accounts for model-vs-observation footprint mismatch. The design note is explicit that omitting R_repr (using the retrieval error alone) overweights the observations and yields overconfident posteriors, so the representation term is a first-class argument here.

Parameters:

Name	Type	Description	Default
retrieval_variance	Array	Per-observation retrieval variance σ²_retr, shape (n_obs,). Must be > 0.	required
representation_variance	float | Array	Representation variance σ²_repr added to every observation — a scalar or a (n_obs,) array. Default 0.	0.0

Returns:

Type	Description
Array	The diagonal of R, shape (n_obs,).

Source code in src/plumax/lagrangian/inversion.py

def observation_covariance(
    retrieval_variance: jax.Array,
    *,
    representation_variance: float | jax.Array = 0.0,
) -> jax.Array:
    """Assemble the diagonal observation-error covariance ``R``.

    ``R = R_retr + R_repr``: the per-pixel L2 **retrieval** error plus a
    **representation** error that accounts for model-vs-observation footprint
    mismatch. The design note is explicit that omitting ``R_repr`` (using the
    retrieval error alone) overweights the observations and yields overconfident
    posteriors, so the representation term is a first-class argument here.

    Args:
        retrieval_variance: Per-observation retrieval variance ``σ²_retr``,
            shape ``(n_obs,)``. Must be ``> 0``.
        representation_variance: Representation variance ``σ²_repr`` added to
            every observation — a scalar or a ``(n_obs,)`` array. Default ``0``.

    Returns:
        The diagonal of ``R``, shape ``(n_obs,)``.
    """
    retr = jnp.asarray(retrieval_variance, dtype=float)
    if jnp.ndim(retr) != 1:
        raise ValueError("observation_covariance: `retrieval_variance` must be 1-D.")
    repr_ = jnp.asarray(representation_variance, dtype=float)
    return retr + repr_

integrate_particles

integrate_particles(state: ParticleState, wind: Callable[[Array], Array], turbulence: TurbulenceModel, *, t0: float, t1: float, dt: float, key: Array, pbl_height: float | None = None, save_trajectory: bool = False) -> tuple[ParticleState, jax.Array | None]

Integrate the ensemble from t0 to t1 with fixed step dt.

Parameters:

Name	Type	Description	Default
state	ParticleState	Initial particle state.	required
wind	Callable[[Array], Array]	Mean-wind field t -> (u, v, w) [m/s], spatially uniform.	required
turbulence	TurbulenceModel	Turbulence model.	required
t0	float	Start time [s].	required
t1	float	End time [s].	required
dt	float	Time step [s].	required
key	Array	PRNG key.	required
pbl_height	float | None	Optional PBL-top reflecting lid [m].	None
save_trajectory	bool	If True, also return the stacked positions at every step, shape (n_steps + 1, n, 3).	False

Returns:

Type	Description
ParticleState	(final_state, trajectory) where trajectory is None unless
Array | None	save_trajectory is set.

Source code in src/plumax/lagrangian/particles.py

def integrate_particles(
    state: ParticleState,
    wind: Callable[[jax.Array], jax.Array],
    turbulence: TurbulenceModel,
    *,
    t0: float,
    t1: float,
    dt: float,
    key: jax.Array,
    pbl_height: float | None = None,
    save_trajectory: bool = False,
) -> tuple[ParticleState, jax.Array | None]:
    """Integrate the ensemble from ``t0`` to ``t1`` with fixed step ``dt``.

    Args:
        state: Initial particle state.
        wind: Mean-wind field ``t -> (u, v, w)`` [m/s], spatially uniform.
        turbulence: Turbulence model.
        t0: Start time [s].
        t1: End time [s].
        dt: Time step [s].
        key: PRNG key.
        pbl_height: Optional PBL-top reflecting lid [m].
        save_trajectory: If ``True``, also return the stacked positions at every
            step, shape ``(n_steps + 1, n, 3)``.

    Returns:
        ``(final_state, trajectory)`` where ``trajectory`` is ``None`` unless
        ``save_trajectory`` is set.
    """
    n_steps = n_steps_for_horizon(t1 - t0, dt)
    keys = jax.random.split(key, n_steps)
    # Step start times and per-step durations. When ``t1 - t0`` is not an exact
    # multiple of ``dt`` the final step is shortened to the remainder so the
    # ensemble is advanced to exactly ``t1`` (not ``t0 + n_steps * dt``).
    times = t0 + dt * jnp.arange(n_steps)
    dts = step_durations(t1 - t0, dt, n_steps)

    def body(
        carry: ParticleState, inputs: tuple[jax.Array, jax.Array, jax.Array]
    ) -> tuple[ParticleState, jax.Array]:
        t, k, dt_i = inputs
        nxt = langevin_step(
            carry, jnp.asarray(wind(t)), turbulence, dt_i, k, pbl_height=pbl_height
        )
        return nxt, nxt.position

    final, positions = jax.lax.scan(body, state, (times, keys, dts))
    if not save_trajectory:
        return final, None
    trajectory = jnp.concatenate([state.position[None], positions], axis=0)
    return final, trajectory

langevin_step

langevin_step(state: ParticleState, wind: Array, turbulence: TurbulenceModel, dt: float, key: Array, *, pbl_height: float | None = None) -> ParticleState

Advance the ensemble one step dt with the Markov-1 update.

Parameters:

Name	Type	Description	Default
state	ParticleState	Current particle state.	required
wind	Array	Mean wind (u, v, w) [m/s] at this step, shape (3,).	required
turbulence	TurbulenceModel	Turbulence model supplying σ(z) and τ_L(z).	required
dt	float	Time step [s].	required
key	Array	PRNG key for the Wiener increment.	required
pbl_height	float | None	If given, reflect particles off the PBL top [m].	None

Returns:

Type	Description
ParticleState	The advanced :class:ParticleState.

Source code in src/plumax/lagrangian/particles.py

def langevin_step(
    state: ParticleState,
    wind: jax.Array,
    turbulence: TurbulenceModel,
    dt: float,
    key: jax.Array,
    *,
    pbl_height: float | None = None,
) -> ParticleState:
    """Advance the ensemble one step ``dt`` with the Markov-1 update.

    Args:
        state: Current particle state.
        wind: Mean wind ``(u, v, w)`` [m/s] at this step, shape ``(3,)``.
        turbulence: Turbulence model supplying ``σ(z)`` and ``τ_L(z)``.
        dt: Time step [s].
        key: PRNG key for the Wiener increment.
        pbl_height: If given, reflect particles off the PBL top [m].

    Returns:
        The advanced :class:`ParticleState`.
    """
    pos, vel = state.position, state.velocity
    z = pos[:, 2]

    sigma, tau = turbulence.at(z)  # each (n, 3)
    sigma = jnp.maximum(sigma, 1e-12)

    decay = jnp.exp(-dt / tau)  # (n, 3)
    ou_std = sigma * jnp.sqrt(jnp.clip(1.0 - decay**2, 0.0, None))
    noise = jax.random.normal(key, vel.shape)
    vel_ou = decay * vel + ou_std * noise

    # Well-mixed vertical drift correction: ½(1 + v_w²/σ_w²) ∂σ_w²/∂z · dt.
    dsigw2_dz = jax.vmap(jax.grad(lambda zz: _sigma_w_sq(turbulence, zz)))(z)
    vw = vel_ou[:, 2]
    drift_w = 0.5 * (1.0 + vw**2 / sigma[:, 2] ** 2) * dsigw2_dz * dt
    vel_new = vel_ou.at[:, 2].add(drift_w)

    pos_new = pos + (wind[None, :] + vel_new) * dt

    # Ground reflection at z = 0.
    zc = pos_new[:, 2]
    vc = vel_new[:, 2]
    below = zc < 0.0
    zc = jnp.where(below, -zc, zc)
    vc = jnp.where(below, -vc, vc)
    # PBL-top reflection at z = h.
    if pbl_height is not None:
        above = zc > pbl_height
        zc = jnp.where(above, 2.0 * pbl_height - zc, zc)
        vc = jnp.where(above, -vc, vc)
    pos_new = pos_new.at[:, 2].set(zc)
    vel_new = vel_new.at[:, 2].set(vc)

    return ParticleState(position=pos_new, velocity=vel_new)

n_steps_for_horizon

n_steps_for_horizon(horizon: float, dt: float, *, rtol: float = 1e-09) -> int

Number of fixed-dt steps covering horizon (ceil, float-robust).

Plain ceil(horizon / dt) adds a spurious extra step when the quotient is mathematically an integer but floating-point rounds it just above (e.g. 2.1 / 0.15 → 14.000…02). Snapping the quotient to the nearest integer when it is within rtol avoids that degenerate final zero-length step. Returns 0 for a non-positive horizon.

Source code in src/plumax/lagrangian/particles.py

def n_steps_for_horizon(horizon: float, dt: float, *, rtol: float = 1e-9) -> int:
    """Number of fixed-``dt`` steps covering ``horizon`` (ceil, float-robust).

    Plain ``ceil(horizon / dt)`` adds a spurious extra step when the quotient is
    mathematically an integer but floating-point rounds it just above (e.g.
    ``2.1 / 0.15 → 14.000…02``). Snapping the quotient to the nearest integer
    when it is within ``rtol`` avoids that degenerate final zero-length step.
    Returns 0 for a non-positive horizon.
    """
    if horizon <= 0.0:
        return 0
    quotient = horizon / dt
    nearest = round(quotient)
    if abs(quotient - nearest) <= rtol * max(1.0, abs(nearest)):
        return max(int(nearest), 1)
    return int(np.ceil(quotient))

step_durations

step_durations(horizon: float, dt: float, n_steps: int) -> jax.Array

Per-step durations summing to horizon; the final step takes the rest.

All steps are dt except the last, which is the remainder horizon - (n_steps - 1) * dt (equal to dt when horizon is an exact multiple). Returns an empty array when n_steps == 0. Pair with :func:n_steps_for_horizon so the final duration is always > 0 (no degenerate zero-length step from float round-up). Used by the integrator / residence / footprint paths so a non-divisible horizon advances over exactly horizon rather than n_steps * dt.

Source code in src/plumax/lagrangian/particles.py

def step_durations(horizon: float, dt: float, n_steps: int) -> jax.Array:
    """Per-step durations summing to ``horizon``; the final step takes the rest.

    All steps are ``dt`` except the last, which is the remainder
    ``horizon - (n_steps - 1) * dt`` (equal to ``dt`` when ``horizon`` is an
    exact multiple). Returns an empty array when ``n_steps == 0``. Pair with
    :func:`n_steps_for_horizon` so the final duration is always ``> 0`` (no
    degenerate zero-length step from float round-up). Used by the integrator /
    residence / footprint paths so a non-divisible horizon advances over exactly
    ``horizon`` rather than ``n_steps * dt``.
    """
    if n_steps == 0:
        return jnp.zeros((0,))
    dts = np.full(n_steps, dt, dtype=float)
    dts[-1] = horizon - dt * (n_steps - 1)
    return jnp.asarray(dts)

uniform_wind

uniform_wind(u: float, v: float, w: float = 0.0) -> Callable[[jax.Array], jax.Array]

Return a constant mean-wind field t -> (u, v, w).

Source code in src/plumax/lagrangian/particles.py

def uniform_wind(
    u: float, v: float, w: float = 0.0
) -> Callable[[jax.Array], jax.Array]:
    """Return a constant mean-wind field ``t -> (u, v, w)``."""
    vec = jnp.array([u, v, w])

    def field(t: jax.Array) -> jax.Array:
        del t
        return vec

    return field

wind_from_speed_direction

wind_from_speed_direction(speed: float, direction_deg: float, w: float = 0.0) -> Callable[[jax.Array], jax.Array]

Constant wind from speed + meteorological direction (degrees from North).

direction_deg = 270 is a wind from the west (flowing east), matching the convention used by :func:plumax.gauss_plume.simulate_plume.

Source code in src/plumax/lagrangian/particles.py

def wind_from_speed_direction(
    speed: float, direction_deg: float, w: float = 0.0
) -> Callable[[jax.Array], jax.Array]:
    """Constant wind from speed + meteorological direction (degrees *from* North).

    ``direction_deg = 270`` is a wind from the west (flowing east), matching the
    convention used by :func:`plumax.gauss_plume.simulate_plume`.
    """
    theta = jnp.deg2rad(direction_deg)
    u = -speed * jnp.sin(theta)
    v = -speed * jnp.cos(theta)
    return uniform_wind(float(u), float(v), w)

hanna_profiles

hanna_profiles(z: Array, *, u_star: float, pbl_height: float, obukhov_length: float, w_star: float = 0.0) -> tuple[jax.Array, jax.Array]

Hanna (1982) Monin–Obukhov-similarity turbulence profiles.

Returns height-dependent velocity standard deviations and Lagrangian timescales for the three stability regimes selected by the sign of the Obukhov length L:

• Unstable / convective (L < 0): convective scaling with w_*.
• Stable (L > 0): surface-layer scaling decaying with z / h.
• Neutral (|L| large): σ ∝ u_* with an exponential PBL decay.

Parameters:

Name	Type	Description	Default
z	Array	Height(s) above ground [m]; any shape.	required
u_star	float	Friction velocity [m/s].	required
pbl_height	float	Boundary-layer height h [m].	required
obukhov_length	float	Monin–Obukhov length L [m] (<0 unstable).	required
w_star	float	Convective velocity scale [m/s] (used when L < 0).	0.0

Returns:

Type	Description
Array	(sigma, tau) arrays, each of shape z.shape + (3,): velocity
Array	standard deviations [m/s] and Lagrangian timescales [s] for
tuple[Array, Array]	(u, v, w).

Source code in src/plumax/lagrangian/turbulence.py

def hanna_profiles(
    z: jax.Array,
    *,
    u_star: float,
    pbl_height: float,
    obukhov_length: float,
    w_star: float = 0.0,
) -> tuple[jax.Array, jax.Array]:
    """Hanna (1982) Monin–Obukhov-similarity turbulence profiles.

    Returns height-dependent velocity standard deviations and Lagrangian
    timescales for the three stability regimes selected by the sign of the
    Obukhov length ``L``:

    - **Unstable / convective** (``L < 0``): convective scaling with ``w_*``.
    - **Stable** (``L > 0``): surface-layer scaling decaying with ``z / h``.
    - **Neutral** (``|L|`` large): ``σ ∝ u_*`` with an exponential PBL decay.

    Args:
        z: Height(s) above ground [m]; any shape.
        u_star: Friction velocity [m/s].
        pbl_height: Boundary-layer height ``h`` [m].
        obukhov_length: Monin–Obukhov length ``L`` [m] (``<0`` unstable).
        w_star: Convective velocity scale [m/s] (used when ``L < 0``).

    Returns:
        ``(sigma, tau)`` arrays, each of shape ``z.shape + (3,)``: velocity
        standard deviations [m/s] and Lagrangian timescales [s] for
        ``(u, v, w)``.
    """
    z = jnp.asarray(z, dtype=float)
    h = pbl_height
    zh = jnp.clip(z / h, 0.0, 1.0)
    us = u_star
    ws = w_star

    # --- Unstable (convective), L < 0 ---
    sig_u2_conv = (12.0 + 0.5 * h / jnp.maximum(-obukhov_length, 1e-6)) ** (
        2.0 / 3.0
    ) * us**2
    sig_w2_conv = 1.74 * us**2 * (1.0 - 0.8 * zh) ** 2 + (
        0.35 * ws**2 * (zh ** (2.0 / 3.0)) * (1.0 - 0.8 * zh) ** 2
    )
    sig_w_conv = jnp.sqrt(jnp.maximum(sig_w2_conv, 0.0))
    sig_uv_conv = jnp.sqrt(jnp.maximum(sig_u2_conv, 0.0)) * jnp.ones_like(z)
    tau_w_conv = 0.15 * h / jnp.maximum(sig_w_conv, _SIGMA_FLOOR)
    tau_uv_conv = 0.15 * h / jnp.maximum(sig_uv_conv, _SIGMA_FLOOR)

    # --- Stable, L > 0 ---
    decay = jnp.maximum(1.0 - zh, 1e-3)
    sig_uv_stab = 2.0 * us * decay
    sig_w_stab = 1.3 * us * decay
    tau_uv_stab = 0.07 * h / jnp.maximum(us, _SIGMA_FLOOR) * jnp.sqrt(zh + 1e-3)
    tau_w_stab = 0.1 * h / jnp.maximum(us, _SIGMA_FLOOR) * (zh + 1e-3) ** 0.8

    # --- Neutral ---
    exp_decay = jnp.exp(-2.0 * zh)
    sig_uv_neut = 2.0 * us * exp_decay
    sig_w_neut = 1.3 * us * exp_decay
    f_c = 1e-4  # Coriolis parameter [1/s], mid-latitude default
    tau_neut = (
        0.5
        * z
        / jnp.maximum(sig_w_neut, _SIGMA_FLOOR)
        / (1.0 + 15.0 * f_c * z / jnp.maximum(us, _SIGMA_FLOOR))
    )
    tau_neut = jnp.maximum(tau_neut, 1.0)

    unstable = obukhov_length < 0.0
    stable = obukhov_length > 0.0

    def pick(conv: jax.Array, stab: jax.Array, neut: jax.Array) -> jax.Array:
        return jnp.where(unstable, conv, jnp.where(stable, stab, neut))

    sigma_u = pick(sig_uv_conv, sig_uv_stab, sig_uv_neut)
    sigma_v = pick(0.8 * sig_uv_conv, sig_uv_stab, sig_uv_neut)
    sigma_w = pick(sig_w_conv, sig_w_stab, sig_w_neut)
    tau_u = pick(tau_uv_conv, tau_uv_stab, tau_neut)
    tau_v = pick(tau_uv_conv, tau_uv_stab, tau_neut)
    tau_w = pick(tau_w_conv, tau_w_stab, tau_neut)

    sigma = jnp.stack([sigma_u, sigma_v, sigma_w], axis=-1)
    tau = jnp.stack([tau_u, tau_v, tau_w], axis=-1)
    tau = jnp.maximum(tau, 1.0)
    return sigma, tau