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), Step 2 (model-based inference) and Steps 3–4 (the footprint emulator and inference on it):

• 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).
• wind_interp — GriddedWind, the les_fvm C-grid wind (from a MetField or a PrescribedWindField) sampled at each particle; accepted wherever a wind callable is.
• met_ensemble — one footprint per met-ensemble member (ensemble_footprints) and the Gaussian inversion that adds their spread as a transport-error covariance (ensemble_inversion).

• inversion — 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): vardax OptimalInterpolation over footprint_operator (the assim extra).

• emulator — FootprintEmulator, an MLP from met context and receptor to a footprint row, trained on particle-model footprints with pipekit-train (footprint_dataset, train_footprint_emulator; the emulate extra) and calibrated against them (emulator_relative_error, assert_emulator_calibrated). emulated_footprint feeds the inversions unchanged (Step 4).

The amortized predictor (Step 5) is not yet implemented.

EmulatedFootprint

Bases: Module

The emulator at fixed receptors, as the per-case map u ↦ F_emu(u).

F_emu(u)[i, s] = emulator(u, rᵢ)[s]. This is the model pipekit-train's Equinox backend fits: it maps one met context to the (obs, cell) target of a SimulationDataset pair. Gradients reach the network weights only; the standardisation, the scale s and the receptors are held fixed, as in :func:train_footprint_emulator.

Attributes:

Name	Type	Description
emulator	FootprintEmulator	The per-receptor network.
receptors	Float[Array, 'obs 3']	Receptor positions (x, y, z) [m]; z is not an input.

FootprintEmulator

Bases: Module

Per-receptor network from met context and receptor to a footprint row.

F_emu(u, r) = s · softplus(g(z; θ)), z = ([u, r_xy] − m) / σ,

with g an MLP and m, σ the mean and standard deviation of the training inputs. The softplus keeps every cell non-negative, as a residence time is, and s is the mean training footprint, so g works on order-one numbers.

Attributes:

Name	Type	Description
network	MLP	The MLP g, from feature + 2 inputs to cell outputs.
input_mean	Float[Array, ' input']	Mean m of the training inputs [u, r_xy].
input_scale	Float[Array, ' input']	Standard deviation σ of the training inputs.
footprint_scale	Float[Array, '']	Scale s [s m² kg⁻¹].

Examples:

An untrained emulator for a two-feature context on a 3-cell grid:

>>> import jax
>>> import jax.numpy as jnp
>>> import equinox as eqx
>>> network = eqx.nn.MLP(4, 3, 8, 1, key=jax.random.PRNGKey(0))
>>> emulator = FootprintEmulator(
...     network, jnp.zeros(4), jnp.ones(4), jnp.asarray(0.5)
... )
>>> row = emulator(jnp.array([3.0, 0.0]), jnp.array([0.0, 100.0]))
>>> row.shape, bool((row >= 0.0).all())
((3,), True)

FootprintSimulator

FootprintSimulator(*, wind_from_context: Callable[[Float[Array, ' feature']], GriddedWind | Callable[[Array], Array]], receptors: Float[Array, 'obs 3'], turbulence: TurbulenceModel, domain_x: tuple[float, float, int], domain_y: tuple[float, float, int], n_particles: int = 300, t_back: float = 600.0, dt: float = 1.0, pbl_height: float = 1000.0, pbl_fraction: float = 0.5, air_density: float = 1.2)

The particle model as a forward model from met context to footprint.

F(u) = compute_footprint(receptors, wind_from_context(u)), one row per receptor, cells flattened row-major over (x, y).

It satisfies pipekit_cycle.ForwardModel structurally (step and dt), so it is the forward_model of pipekit-train's SimulationDataset. The particle seed of a case is a hash of its context, so a cached case is reproducible and two cases never share their particle noise. Receptor i of a case uses seed + i, the common random numbers of :func:~plumax.lagrangian.ensemble_footprints, whose rows it reproduces for the same seed.

The particle scan is compiled once per simulator, with the context and the receptors as traced arguments and jax.vmap over the receptors, so only the first case of a simulator pays for compilation.

Parameters:

Name	Type	Description	Default
wind_from_context	Callable[[Float[Array, ' feature']], GriddedWind | Callable[[Array], Array]]	Mean wind for a context, a GriddedWind or a t -> (u, v, w) callable. It must be traceable, because it is called once with a traced context while the scan compiles: build the wind with JAX operations, as uniform_wind(u[0], u[1]), and never call float() on the context or branch on its value.	required
receptors	Float[Array, 'obs 3']	Receptor positions (x, y, z) [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
n_particles	int	Particles per receptor and case.	300
t_back	float	Backward integration time [s].	600.0
dt	float	Particle time step [s]; the forward model's dt.	1.0
pbl_height	float	Boundary-layer height [m].	1000.0
pbl_fraction	float	Fraction of the PBL that sees the surface flux.	0.5
air_density	float	Air density [kg m⁻³].	1.2

Raises:

Type	Description
ParameterError	If n_particles is below one.

Source code in src/plumax/lagrangian/_src/emulator.py

def __init__(
    self,
    *,
    wind_from_context: Callable[
        [Float[Array, " feature"]],
        GriddedWind | Callable[[jax.Array], jax.Array],
    ],
    receptors: Float[Array, "obs 3"],
    turbulence: TurbulenceModel,
    domain_x: tuple[float, float, int],
    domain_y: tuple[float, float, int],
    n_particles: int = 300,
    t_back: float = 600.0,
    dt: float = 1.0,
    pbl_height: float = 1000.0,
    pbl_fraction: float = 0.5,
    air_density: float = 1.2,
) -> None:
    if n_particles < 1:
        raise ParameterError(f"n_particles must be ≥ 1, got {n_particles}")
    self.wind_from_context = wind_from_context
    self.receptors = jnp.asarray(receptors)
    self.turbulence = turbulence
    self.domain_x = domain_x
    self.domain_y = domain_y
    self.n_particles = int(n_particles)
    self.t_back = float(t_back)
    self.dt = float(dt)
    self.pbl_height = float(pbl_height)
    self.pbl_fraction = float(pbl_fraction)
    self.air_density = float(air_density)
    self._compiled_residence = jax.jit(self._case_residence)

state_signature property

state_signature: tuple[Any, ...]

Everything that changes the footprints, for the dataset cache key.

step

step(state: Float[Array, ' feature'], dt: float) -> Float[Array, 'obs cell']

Footprint F(u) of the met context u = state.

Source code in src/plumax/lagrangian/_src/emulator.py

def step(
    self, state: Float[Array, " feature"], dt: float
) -> Float[Array, "obs cell"]:
    """Footprint F(u) of the met context u = ``state``."""
    context = np.asarray(state)
    # receptor i of a case uses seed + i, so a hash of u keeps cases apart
    seed = zlib.crc32(context.tobytes())
    receptors = np.asarray(self.receptors, dtype=float)
    keys = jnp.stack(
        [jax.random.PRNGKey(seed + index) for index in range(len(receptors))]
    )
    # the wind clock ends at the receptor, receptor_time = 0, in every case
    durations, wind_starts = _backward_steps(self.t_back, dt, 0.0)
    residence = self._compiled_residence(
        jnp.asarray(context),
        jnp.asarray(receptors, dtype=float),
        keys,
        durations,
        wind_starts,
    )
    footprint = _flux_sensitivity(
        residence,
        self.n_particles,
        self.air_density,
        self.pbl_fraction * self.pbl_height,
    )
    # F[i, s]: cells flattened row-major over (x, y)
    return einx.rearrange("obs x y -> obs (x y)", jnp.asarray(footprint))

MetContextPrior

MetContextPrior(met_sampler: Callable[[PRNGKeyArray], Float[Array, ' feature']])

Bases: Operator

The operational met distribution p(u) as pipekit-train's prior operator.

u ~ p(u | site, overpass time), one draw per integer seed. Draw u from the operational distribution at the facility, not from a uniform climatology grid: the design page notes a coarse (hour, season, stability) grid covers under 10 % of the operational probability mass.

Parameters:

Name	Type	Description	Default
met_sampler	Callable[[PRNGKeyArray], Float[Array, ' feature']]	One draw u ~ p(u) from a key.	required

Source code in src/plumax/lagrangian/_src/emulator.py

def __init__(
    self, met_sampler: Callable[[PRNGKeyArray], Float[Array, " feature"]]
) -> None:
    self.met_sampler = met_sampler

GaussianPosterior dataclass

GaussianPosterior(mean: Array, covariance: Array, prior_mean: Array | None = None, prior_covariance: Array | None = None)

Closed-form Gaussian posterior over the source vector.

It also carries the Tier V per-event payload (population.PerEventPosterior) for the event total Q = 1ᵀq, whose posterior N(1ᵀq, 1ᵀP1) and prior N(1ᵀq_a, 1ᵀB1) are exact.

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).
prior_mean	Array | None	Prior mean q_a the inversion used, shape (n_grid,) (None when built by hand).
prior_covariance	Array | None	Prior covariance B, shape (n_grid, n_grid) (None when built by hand).

std property

std: Array

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

emission_rate property

emission_rate: Array

Posterior mean of the event total Q = 1ᵀq*.

emission_std property

emission_std: Array

Posterior standard deviation of Q, √(1ᵀP*1).

samples

samples(key: Array, n: int) -> jax.Array

n exact draws of Q from N(1ᵀq, 1ᵀP1).

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

def samples(self, key: jax.Array, n: int) -> jax.Array:
    """n exact draws of Q from N(1ᵀq*, 1ᵀP*1)."""
    return self._total_posterior().sample(key, n)

prior_logpdf

prior_logpdf(q: Array) -> jax.Array

log π(Q) of the prior on the total, N(1ᵀq_a, 1ᵀB1).

Raises:

Type	Description
PayloadError	If the posterior was built without its prior.

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

def prior_logpdf(self, q: jax.Array) -> jax.Array:
    """log π(Q) of the prior on the total, N(1ᵀq_a, 1ᵀB1).

    Raises:
        PayloadError: If the posterior was built without its prior.
    """
    if self.prior_mean is None or self.prior_covariance is None:
        raise PayloadError(
            "GaussianPosterior.prior_logpdf: this posterior carries no prior; "
            "build it with linear_gaussian_inversion."
        )
    weights = jnp.ones_like(self.prior_mean)
    prior = gaussian_total(self.prior_mean, self.prior_covariance, weights)
    return prior.log_prob(q)

LognormalPosterior dataclass

LognormalPosterior(mean: Array, log_increment: Array, log_covariance: Array, prior_mean: Array | None = None, prior_log_covariance: Array | None = None)

Linearised lognormal posterior over a non-negative source vector.

It also carries the Tier V per-event payload (population.PerEventPosterior) for the event total Q = Σᵢ q_a,i exp(δᵢ): draws are exact (δ ~ N(δ, P_log)); the prior density of Q under δ ~ N(0, B_log) is the moment-matched lognormal, exact for one cell.

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).
prior_mean	Array | None	Prior median field q_a, shape (n_grid,) (None when built by hand).
prior_log_covariance	Array | None	Log-space prior covariance B_log, shape (n_grid, n_grid) (None when built by hand).

emission_rate property

emission_rate: Array

The event total Σᵢ q*ᵢ, as population.event_from_posterior reads it.

emission_std property

emission_std: Array

Delta-method sd of the total, √(qᵀ(exp(P_log) − 1)q*).

samples

samples(key: Array, n: int) -> jax.Array

n exact draws of Q = Σᵢ qᵢ exp(εᵢ), ε ~ N(0, P_log).

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

def samples(self, key: jax.Array, n: int) -> jax.Array:
    """n exact draws of Q = Σᵢ q*ᵢ exp(εᵢ), ε ~ N(0, P*_log)."""
    zeros = jnp.zeros_like(self.mean)
    return sample_lognormal_total(key, n, zeros, self.log_covariance, self.mean)

prior_logpdf

prior_logpdf(q: Array) -> jax.Array

log π(Q) for the prior total Σᵢ q_a,i exp(δᵢ), δ ~ N(0, B_log).

Raises:

Type	Description
PayloadError	If the posterior was built without its prior.

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

def prior_logpdf(self, q: jax.Array) -> jax.Array:
    """log π(Q) for the prior total Σᵢ q_a,i exp(δᵢ), δ ~ N(0, B_log).

    Raises:
        PayloadError: If the posterior was built without its prior.
    """
    if self.prior_mean is None or self.prior_log_covariance is None:
        raise PayloadError(
            "LognormalPosterior.prior_logpdf: this posterior carries no prior; "
            "build it with lognormal_inversion."
        )
    zeros = jnp.zeros_like(self.prior_mean)
    prior = lognormal_total(zeros, self.prior_log_covariance, self.prior_mean)
    return prior.log_prob(q)

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, coriolis_parameter: float = DEFAULT_CORIOLIS_PARAMETER)

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.

To drive it from met rather than hand-typed scales, build it with :func:plumax.met.hanna_turbulence_from from a :class:plumax.met.MoninObukhov diagnosis.

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, 0 neutral; any h/|L| < 1 is neutral).
w_star	float	Convective velocity scale w_* [m/s] (used when unstable).
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.
coriolis_parameter	float	Coriolis parameter f [1/s] for the neutral profiles; defaults to FLEXPART's 1e-4 (see :func:hanna_profiles).

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,
        coriolis_parameter=self.coriolis_parameter,
    )

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

GriddedWind

Bases: Module

Mean wind on the les_fvm C-grid, trilinear in space, linear in time.

u(p, t) = Σ_corners w_corner u_corner over the 16 samples of u that bracket (t, z, y, x), and likewise for v and w on their own points.

Attributes:

Name	Type	Description
x	Float[Array, ' x']	T-point x-coordinates x_i [m], uniformly spaced.
y	Float[Array, ' y']	T-point y-coordinates y_j [m], uniformly spaced.
z	Float[Array, ' z']	T-point z-coordinates z_k [m], uniformly spaced.
times	Float[Array, ' time']	Sample times [s], strictly increasing.
u	Float[Array, 'time z y x']	Eastward wind [m/s] at U-points (x_i + Δx/2).
v	Float[Array, 'time z y x']	Northward wind [m/s] at V-points (y_j + Δy/2).
w	Float[Array, 'time z y x']	Vertical wind [m/s] at T-points.

Examples:

A uniform wind is reproduced exactly, anywhere and at any time:

>>> import jax.numpy as jnp
>>> axis = jnp.arange(4.0) * 10.0 + 5.0
>>> ones = jnp.ones((2, 4, 4, 4))
>>> wind = GriddedWind(
...     axis, axis, axis, jnp.array([0.0, 60.0]), 3.0 * ones, -ones, 0 * ones
... )
>>> wind(jnp.array([[12.0, 7.0, 31.0], [-50.0, 80.0, 0.0]]), 30.0).tolist()
[[3.0, -1.0, 0.0], [3.0, -1.0, 0.0]]

u is read on U-points, half a cell east of the T-points, so a field linear in x comes back as x itself at the T-points:

>>> u = jnp.broadcast_to(axis + 5.0, (2, 4, 4, 4))
>>> sheared = GriddedWind(axis, axis, axis, jnp.array([0.0, 60.0]), u, ones, ones)
>>> sheared(jnp.array([[15.0, 15.0, 15.0]]), 0.0)[:, 0].tolist()
[15.0]

Piecewise-linear in time between the stored times:

>>> ramp = jnp.stack([ones[0], 3.0 * ones[0]])
>>> gusty = GriddedWind(axis, axis, axis, jnp.array([0.0, 60.0]), ramp, ones, ones)
>>> gusty(jnp.array([[15.0, 15.0, 15.0]]), 15.0)[:, 0].tolist()
[1.5]

from_grid classmethod

from_grid(plume_grid: PlumeGrid3D, times: Float[Array, ' time'], u: Float[Array, 'time z y x'], v: Float[Array, 'time z y x'], w: Float[Array, 'time z y x']) -> GriddedWind

Wind on a les_fvm analysis grid, interior arrays on the C-grid.

Parameters:

Name	Type	Description	Default
plume_grid	PlumeGrid3D	The analysis grid; its T-point axes place the samples.	required
times	Float[Array, ' time']	Sample times [s], strictly increasing.	required
u	Float[Array, 'time z y x']	Eastward wind [m/s] at U-points.	required
v	Float[Array, 'time z y x']	Northward wind [m/s] at V-points.	required
w	Float[Array, 'time z y x']	Vertical wind [m/s] at T-points.	required

Returns:

Type	Description
GriddedWind	The gridded wind.

Source code in src/plumax/lagrangian/_src/wind_interp.py

@classmethod
def from_grid(
    cls,
    plume_grid: PlumeGrid3D,
    times: Float[Array, " time"],
    u: Float[Array, "time z y x"],
    v: Float[Array, "time z y x"],
    w: Float[Array, "time z y x"],
) -> GriddedWind:
    """Wind on a ``les_fvm`` analysis grid, interior arrays on the C-grid.

    Args:
        plume_grid: The analysis grid; its T-point axes place the samples.
        times: Sample times [s], strictly increasing.
        u: Eastward wind [m/s] at U-points.
        v: Northward wind [m/s] at V-points.
        w: Vertical wind [m/s] at T-points.

    Returns:
        The gridded wind.
    """
    return cls(plume_grid.x, plume_grid.y, plume_grid.z, times, u, v, w)

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] | GriddedWind | 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] | GriddedWind | None	Mean-wind field t -> (u, v, w), or a :class:~plumax.lagrangian.GriddedWind sampled at each particle. 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] | GriddedWind | 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)``, or a
            :class:`~plumax.lagrangian.GriddedWind` sampled at each particle.
            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

assert_emulator_calibrated

assert_emulator_calibrated(emulator: FootprintEmulator, reference_footprints: Float[Array, 'case obs cell'], *, met_context: Float[Array, 'case feature'], receptors: Float[Array, 'obs 3'], key: PRNGKeyArray, threshold: float = 0.05, n_probes: int = 10) -> None

Check the emulated source–receptor maps against the particle model's.

‖F_emu v − F v‖ / ‖F v‖ < threshold for n_probes random source fields v, through vardax.assert_adjoint_calibrated. The maps are q ↦ (F_emu,ᵢ q)ᵢ and q ↦ (Fᵢ q)ᵢ stacked over the held-out cases i, so this is the probed form of the Frobenius test over the same set, and it checks the sensitivity the inversion sees.

Parameters:

Name	Type	Description	Default
emulator	FootprintEmulator	The trained emulator.	required
reference_footprints	Float[Array, 'case obs cell']	Particle-model footprints Fᵢ at met_context.	required
met_context	Float[Array, 'case feature']	Held-out met contexts uᵢ.	required
receptors	Float[Array, 'obs 3']	Receptor positions (x, y, z) [m].	required
key	PRNGKeyArray	PRNG key for the probe vectors.	required
threshold	float	Largest relative error allowed.	0.05
n_probes	int	Number of probe vectors.	10

Raises:

Type	Description
AssertionError	If a probe exceeds threshold.
ImportError	Without the assim extra.

Source code in src/plumax/lagrangian/_src/emulator.py

def assert_emulator_calibrated(
    emulator: FootprintEmulator,
    reference_footprints: Float[Array, "case obs cell"],
    *,
    met_context: Float[Array, "case feature"],
    receptors: Float[Array, "obs 3"],
    key: PRNGKeyArray,
    threshold: float = 0.05,
    n_probes: int = 10,
) -> None:
    """Check the emulated source–receptor maps against the particle model's.

    ‖F_emu v − F v‖ / ‖F v‖ < threshold for ``n_probes`` random source
    fields v, through ``vardax.assert_adjoint_calibrated``. The maps are
    q ↦ (F_emu,ᵢ q)ᵢ and q ↦ (Fᵢ q)ᵢ stacked over the held-out cases i, so
    this is the probed form of the Frobenius test over the same set, and it
    checks the sensitivity the inversion sees.

    Args:
        emulator: The trained emulator.
        reference_footprints: Particle-model footprints Fᵢ at ``met_context``.
        met_context: Held-out met contexts uᵢ.
        receptors: Receptor positions (x, y, z) [m].
        key: PRNG key for the probe vectors.
        threshold: Largest relative error allowed.
        n_probes: Number of probe vectors.

    Raises:
        AssertionError: If a probe exceeds ``threshold``.
        ImportError: Without the ``assim`` extra.
    """
    vdx = require_vardax("assert_emulator_calibrated")
    emulated = jax.vmap(emulated_footprint, in_axes=(None, 0, None))(
        emulator, met_context, receptors
    )
    source = jnp.ones(emulated.shape[-1], dtype=emulated.dtype)

    def emulated_map(q: Float[Array, " cell"]) -> Float[Array, "case obs"]:
        # F_emu,ᵢ q for every case i
        return einx.dot("case obs cell, cell -> case obs", emulated, q)

    def reference_map(q: Float[Array, " cell"]) -> Float[Array, "case obs"]:
        # Fᵢ q for every case i
        return einx.dot("case obs cell, cell -> case obs", reference_footprints, q)

    vdx.assert_adjoint_calibrated(
        emulated_map,
        reference_map,
        source,
        key=key,
        threshold=threshold,
        n_probes=n_probes,
    )

emulated_footprint

emulated_footprint(emulator: FootprintEmulator, met_context: Float[Array, ' feature'], receptors: Float[Array, 'obs 3']) -> Float[Array, 'obs cell']

The emulated footprint matrix F_emu of one met context.

F_emu[i, s] = emulator(u, rᵢ)[s], in the layout of :func:~plumax.lagrangian.ensemble_footprints and of the inversions.

Parameters:

Name	Type	Description	Default
emulator	FootprintEmulator	The trained emulator.	required
met_context	Float[Array, ' feature']	Met context u.	required
receptors	Float[Array, 'obs 3']	Receptor positions (x, y, z) [m]; z is not an input.	required

Returns:

Type	Description
Float[Array, 'obs cell']	The footprint rows [s m² kg⁻¹].

Examples:

Two receptors on a 3-cell grid:

>>> import jax
>>> import jax.numpy as jnp
>>> import equinox as eqx
>>> network = eqx.nn.MLP(4, 3, 8, 1, key=jax.random.PRNGKey(0))
>>> emulator = FootprintEmulator(
...     network, jnp.zeros(4), jnp.ones(4), jnp.asarray(0.5)
... )
>>> receptors = jnp.array([[0.0, 0.0, 10.0], [0.0, 50.0, 10.0]])
>>> emulated_footprint(emulator, jnp.array([3.0, 0.0]), receptors).shape
(2, 3)

Source code in src/plumax/lagrangian/_src/emulator.py

def emulated_footprint(
    emulator: FootprintEmulator,
    met_context: Float[Array, " feature"],
    receptors: Float[Array, "obs 3"],
) -> Float[Array, "obs cell"]:
    """The emulated footprint matrix F_emu of one met context.

    F_emu[i, s] = emulator(u, rᵢ)[s], in the layout of
    :func:`~plumax.lagrangian.ensemble_footprints` and of the inversions.

    Args:
        emulator: The trained emulator.
        met_context: Met context u.
        receptors: Receptor positions (x, y, z) [m]; z is not an input.

    Returns:
        The footprint rows [s m² kg⁻¹].

    Examples:
        Two receptors on a 3-cell grid:

        >>> import jax
        >>> import jax.numpy as jnp
        >>> import equinox as eqx
        >>> network = eqx.nn.MLP(4, 3, 8, 1, key=jax.random.PRNGKey(0))
        >>> emulator = FootprintEmulator(
        ...     network, jnp.zeros(4), jnp.ones(4), jnp.asarray(0.5)
        ... )
        >>> receptors = jnp.array([[0.0, 0.0, 10.0], [0.0, 50.0, 10.0]])
        >>> emulated_footprint(emulator, jnp.array([3.0, 0.0]), receptors).shape
        (2, 3)
    """
    receptor_xy = receptors[:, :2]
    return jax.vmap(emulator, in_axes=(None, 0))(met_context, receptor_xy)

emulator_relative_error

emulator_relative_error(emulated: Float[Array, '*case obs cell'], reference: Float[Array, '*case obs cell']) -> Float[Array, '']

Relative Frobenius error ‖F_emu − F‖_F / ‖F‖_F over every case.

Parameters:

Name	Type	Description	Default
emulated	Float[Array, '*case obs cell']	Emulated footprints F_emu.	required
reference	Float[Array, '*case obs cell']	Particle-model footprints F on the same met and receptors.	required

Returns:

Type	Description
Float[Array, '']	The relative error; the design page's hard test asks for < 0.05.

Examples:

A uniform 2 % overestimate is a 2 % error:

>>> import jax.numpy as jnp
>>> reference = jnp.ones((2, 3, 4))
>>> round(float(emulator_relative_error(1.02 * reference, reference)), 6)
0.02

Source code in src/plumax/lagrangian/_src/emulator.py

def emulator_relative_error(
    emulated: Float[Array, "*case obs cell"],
    reference: Float[Array, "*case obs cell"],
) -> Float[Array, ""]:
    """Relative Frobenius error ‖F_emu − F‖_F / ‖F‖_F over every case.

    Args:
        emulated: Emulated footprints F_emu.
        reference: Particle-model footprints F on the same met and receptors.

    Returns:
        The relative error; the design page's hard test asks for < 0.05.

    Examples:
        A uniform 2 % overestimate is a 2 % error:

        >>> import jax.numpy as jnp
        >>> reference = jnp.ones((2, 3, 4))
        >>> round(float(emulator_relative_error(1.02 * reference, reference)), 6)
        0.02
    """
    # ‖F_emu − F‖²_F
    squared_error = einx.sum("[...]", (emulated - reference) ** 2)
    # ‖F‖²_F
    squared_norm = einx.sum("[...]", reference**2)
    return jnp.sqrt(squared_error / squared_norm)

footprint_dataset

footprint_dataset(simulator: FootprintSimulator, met_sampler: Callable[[PRNGKeyArray], Float[Array, ' feature']], *, n_cases: int, cache_dir: str, seed: int = 0, split: str = 'train') -> Any

Simulated (met context, footprint) pairs, cached once on disk.

uᵢ ~ p(u), Fᵢ = simulator.step(uᵢ), as pipekit-train's SimulationDataset wrapped in its CachedDataset: the particle model runs once per case on first use, and every later epoch reads the zarr cache under cache_dir, keyed on the simulator, the sampler, the seed and the split.

The cache is held in memory after its first read, because training reads it case by case and a zarr read costs about 1 ms per array.

SimulationDataset partitions the n_cases seeds 80 / 10 / 10 into the "train", "val" and "test" splits, so a split holds a fraction of n_cases and the splits never share a case.

Parameters:

Name	Type	Description	Default
simulator	FootprintSimulator	The particle model.	required
met_sampler	Callable[[PRNGKeyArray], Float[Array, ' feature']]	One draw u ~ p(u) from a key.	required
n_cases	int	Number of met contexts N over all three splits.	required
cache_dir	str	Directory (or fsspec URI) of the cache.	required
seed	int	Seed of the first case.	0
split	str	"train", "val" or "test".	'train'

Returns:

Type	Description
Any	The cached pipekit-train dataset.

Raises:

Type	Description
ParameterError	If n_cases is below one.
ImportError	Without the emulate extra.

Source code in src/plumax/lagrangian/_src/emulator.py

def footprint_dataset(
    simulator: FootprintSimulator,
    met_sampler: Callable[[PRNGKeyArray], Float[Array, " feature"]],
    *,
    n_cases: int,
    cache_dir: str,
    seed: int = 0,
    split: str = "train",
) -> Any:
    """Simulated ``(met context, footprint)`` pairs, cached once on disk.

    uᵢ ~ p(u), Fᵢ = simulator.step(uᵢ), as pipekit-train's
    ``SimulationDataset`` wrapped in its ``CachedDataset``: the particle
    model runs once per case on first use, and every later epoch reads the
    zarr cache under ``cache_dir``, keyed on the simulator, the sampler, the
    seed and the split.

    The cache is held in memory after its first read, because training reads
    it case by case and a zarr read costs about 1 ms per array.

    ``SimulationDataset`` partitions the ``n_cases`` seeds 80 / 10 / 10 into
    the ``"train"``, ``"val"`` and ``"test"`` splits, so a split holds a
    fraction of ``n_cases`` and the splits never share a case.

    Args:
        simulator: The particle model.
        met_sampler: One draw u ~ p(u) from a key.
        n_cases: Number of met contexts N over all three splits.
        cache_dir: Directory (or fsspec URI) of the cache.
        seed: Seed of the first case.
        split: ``"train"``, ``"val"`` or ``"test"``.

    Returns:
        The cached pipekit-train dataset.

    Raises:
        ParameterError: If ``n_cases`` is below one.
        ImportError: Without the ``emulate`` extra.
    """
    if n_cases < 1:
        raise ParameterError(f"n_cases must be ≥ 1, got {n_cases}")
    pt = _require_pipekit_train("footprint_dataset")
    source = pt.SimulationDataset(
        forward_model=simulator,
        prior=MetContextPrior(met_sampler),
        n_samples=n_cases,
        seed=seed,
        split=split,
    )
    return pt.CachedDataset(source, cache_dir=cache_dir, in_memory=True)

init_footprint_emulator

init_footprint_emulator(dataset: Any, receptors: Float[Array, 'obs 3'], *, width: int = 32, depth: int = 2, key: PRNGKeyArray) -> FootprintEmulator

An untrained emulator, standardised to the training set.

Parameters:

Name	Type	Description	Default
dataset	Any	The training (u, F) pairs, e.g. from :func:footprint_dataset; fixes the input and output sizes, the input standardisation and the footprint scale s.	required
receptors	Float[Array, 'obs 3']	Receptor positions (x, y, z) [m] of the footprints.	required
width	int	Hidden width of the MLP.	32
depth	int	Number of hidden layers.	2
key	PRNGKeyArray	PRNG key for the MLP weights.	required

Returns:

Type	Description
FootprintEmulator	The emulator, ready for :func:train_footprint_emulator.

Raises:

Type	Description
ParameterError	If width or depth is below one.

Source code in src/plumax/lagrangian/_src/emulator.py

def init_footprint_emulator(
    dataset: Any,
    receptors: Float[Array, "obs 3"],
    *,
    width: int = 32,
    depth: int = 2,
    key: PRNGKeyArray,
) -> FootprintEmulator:
    """An untrained emulator, standardised to the training set.

    Args:
        dataset: The training ``(u, F)`` pairs, e.g. from
            :func:`footprint_dataset`; fixes the input and output sizes, the
            input standardisation and the footprint scale s.
        receptors: Receptor positions (x, y, z) [m] of the footprints.
        width: Hidden width of the MLP.
        depth: Number of hidden layers.
        key: PRNG key for the MLP weights.

    Returns:
        The emulator, ready for :func:`train_footprint_emulator`.

    Raises:
        ParameterError: If ``width`` or ``depth`` is below one.
    """
    if width < 1 or depth < 1:
        raise ParameterError(f"width and depth must be ≥ 1, got {width}, {depth}")
    contexts, footprints = _stack_dataset(dataset)
    inputs = _case_inputs(contexts, jnp.asarray(receptors, dtype=contexts.dtype))
    input_mean = einx.mean("[case obs] input", inputs)
    input_std = einx.std("[case obs] input", inputs)
    # constant inputs (one case, one receptor) keep a unit scale
    input_scale = jnp.where(input_std > 0.0, input_std, 1.0)
    footprint_scale = jnp.mean(footprints)
    network = eqx.nn.MLP(
        inputs.shape[-1],
        footprints.shape[-1],
        width,
        depth,
        activation=jax.nn.tanh,
        key=key,
    )
    return FootprintEmulator(network, input_mean, input_scale, footprint_scale)

train_footprint_emulator

train_footprint_emulator(emulator: FootprintEmulator, dataset: Any, receptors: Float[Array, 'obs 3'], *, val_dataset: Any | None = None, max_steps: int = 2000, batch_size: int = 16, learning_rate: float = 0.01, patience: int = 10, eval_every_n_steps: int = 100, seed: int = 0) -> FootprintEmulator

Fit the network weights θ with pipekit-train's Equinox backend.

θ* = argmin_θ 𝓛(θ), 𝓛(θ) = mean over cases, receptors and cells of (F_emu(uᵢ, r; θ) − Fᵢ)²,

by Adam on mini-batches of cases. The standardisation and the scale s stay fixed.

Algorithm

1. Wrap the emulator at receptors as an EmulatedFootprint, the model u ↦ F_emu(u) that matches the dataset's pairs.
2. Run a pipekit-train TrainingLoop on the "equinox" backend with its MSE loss for at most max_steps steps.
3. With val_dataset, evaluate every eval_every_n_steps steps and stop through pipekit-train's EarlyStopping once the validation MSE has not improved for patience evaluations.

Parameters:

Name	Type	Description	Default
emulator	FootprintEmulator	Initial emulator, e.g. from :func:init_footprint_emulator.	required
dataset	Any	The training pairs, e.g. from :func:footprint_dataset.	required
receptors	Float[Array, 'obs 3']	Receptor positions (x, y, z) [m] of the footprints.	required
val_dataset	Any | None	Optional held-out pairs for early stopping.	None
max_steps	int	Maximum number of optimiser steps.	2000
batch_size	int	Cases per step; at most the number of training cases.	16
learning_rate	float	Adam learning rate.	0.01
patience	int	Evaluations without improvement before stopping.	10
eval_every_n_steps	int	Validation cadence in steps.	100
seed	int	Seed of the batch shuffle.	0

Returns:

Type	Description
FootprintEmulator	The emulator with the fitted weights.

Raises:

Type	Description
ImportError	Without the emulate extra.

Source code in src/plumax/lagrangian/_src/emulator.py

def train_footprint_emulator(
    emulator: FootprintEmulator,
    dataset: Any,
    receptors: Float[Array, "obs 3"],
    *,
    val_dataset: Any | None = None,
    max_steps: int = 2000,
    batch_size: int = 16,
    learning_rate: float = 1e-2,
    patience: int = 10,
    eval_every_n_steps: int = 100,
    seed: int = 0,
) -> FootprintEmulator:
    """Fit the network weights θ with pipekit-train's Equinox backend.

    θ* = argmin_θ 𝓛(θ),  𝓛(θ) = mean over cases, receptors and cells of
    (F_emu(uᵢ, r; θ) − Fᵢ)²,

    by Adam on mini-batches of cases. The standardisation and the scale s
    stay fixed.

    Algorithm:
        1. Wrap the emulator at ``receptors`` as an ``EmulatedFootprint``,
           the model u ↦ F_emu(u) that matches the dataset's pairs.
        2. Run a pipekit-train ``TrainingLoop`` on the ``"equinox"`` backend
           with its ``MSE`` loss for at most ``max_steps`` steps.
        3. With ``val_dataset``, evaluate every ``eval_every_n_steps`` steps
           and stop through pipekit-train's ``EarlyStopping`` once the
           validation MSE has not improved for ``patience`` evaluations.

    Args:
        emulator: Initial emulator, e.g. from :func:`init_footprint_emulator`.
        dataset: The training pairs, e.g. from :func:`footprint_dataset`.
        receptors: Receptor positions (x, y, z) [m] of the footprints.
        val_dataset: Optional held-out pairs for early stopping.
        max_steps: Maximum number of optimiser steps.
        batch_size: Cases per step; at most the number of training cases.
        learning_rate: Adam learning rate.
        patience: Evaluations without improvement before stopping.
        eval_every_n_steps: Validation cadence in steps.
        seed: Seed of the batch shuffle.

    Returns:
        The emulator with the fitted weights.

    Raises:
        ImportError: Without the ``emulate`` extra.
    """
    pt = _require_pipekit_train("train_footprint_emulator")
    from pipekit_train.adapters.equinox import EquinoxModelOp

    model = EmulatedFootprint(emulator, jnp.asarray(receptors))
    callbacks = ()
    if val_dataset is not None:
        callbacks = (pt.EarlyStopping(metric="mse", patience=patience),)
    # the loss is in the footprint's units, so Adam's ε must sit below s²
    epsilon = 1e-8 * float(emulator.footprint_scale) ** 2
    loop = pt.TrainingLoop(
        model_op=EquinoxModelOp(model),
        dataset=dataset,
        val_dataset=val_dataset,
        loss=pt.MSE(),
        optimizer_config={"name": "adam", "lr": learning_rate, "eps": epsilon},
        max_steps=max_steps,
        batch_size=batch_size,
        backend="equinox",
        callbacks=callbacks,
        eval_every_n_steps=eval_every_n_steps,
        log_every_n_steps=max_steps,
        seed=seed,
    )
    trained_op, _ = loop.run()
    return trained_op.module.emulator

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] | GriddedWind, 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] | GriddedWind	Forward mean-wind field t -> (u, v, w) or a :class:~plumax.lagrangian.GriddedWind sampled at each particle; 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] | GriddedWind,
    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)`` or a
            :class:`~plumax.lagrangian.GriddedWind` sampled at each particle;
            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, y_edges = _surface_edges(domain_x, domain_y)
    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
    durations, wind_starts = _backward_steps(t_back, dt, receptor_time)
    residence = _surface_residence(
        wind,
        jnp.asarray(receptor_location, dtype=float),
        jax.random.PRNGKey(seed),
        turbulence,
        jnp.asarray(x_edges),
        jnp.asarray(y_edges),
        durations,
        wind_starts,
        n_particles=n_particles,
        pbl_height=pbl_height,
        mix_height=mix_height,
    )
    footprint = _flux_sensitivity(residence, n_particles, air_density, mix_height)
    return footprint, x_c, y_c

footprint_operator

footprint_operator(footprint: Array) -> vardax.CallableObs

The Tier II source–receptor map 𝒢(q) = F q as a vardax operator.

vardax analyses carry the control as a (T, N) state; for a source inversion T = 1. The operator maps that (1, n_grid) row to the (1, n_obs) predicted observations, and vardax linearises it by autodiff.

Parameters:

Name	Type	Description	Default
footprint	Array	Source–receptor matrix F, shape (n_obs, n_grid).	required

Returns:

Type	Description
CallableObs	A vardax.CallableObs (1, n_grid) → (1, n_obs).

Raises:

Type	Description
ImportError	Without the assim extra.

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

def footprint_operator(footprint: jax.Array) -> vardax.CallableObs:
    """The Tier II source–receptor map ``𝒢(q) = F q`` as a vardax operator.

    ``vardax`` analyses carry the control as a ``(T, N)`` state; for a source
    inversion ``T = 1``. The operator maps that ``(1, n_grid)`` row to the
    ``(1, n_obs)`` predicted observations, and vardax linearises it by
    autodiff.

    Args:
        footprint: Source–receptor matrix ``F``, shape ``(n_obs, n_grid)``.

    Returns:
        A ``vardax.CallableObs`` ``(1, n_grid) → (1, n_obs)``.

    Raises:
        ImportError: Without the ``assim`` extra.
    """
    vdx = require_vardax("footprint_operator")
    f = jnp.asarray(footprint, dtype=float)
    return vdx.CallableObs(fn=lambda q: einx.dot("t n, m n -> t m", q, f))

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.

Raises:

Type	Description
ImportError	Without the assim extra (pip install 'plumax[assim]'): the update is a vardax.OptimalInterpolation analysis over :func:footprint_operator.

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``.

    Raises:
        ImportError: Without the ``assim`` extra (``pip install
            'plumax[assim]'``): the update is a ``vardax.OptimalInterpolation``
            analysis over :func:`footprint_operator`.
    """
    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, caller="linear_gaussian_inversion"
    )
    return GaussianPosterior(
        mean=qa + increment,
        covariance=posterior_cov,
        prior_mean=qa,
        prior_covariance=b,
    )

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.

Raises:

Type	Description
ImportError	Without the assim extra: the log-space update is a vardax.OptimalInterpolation analysis over δ ↦ F̃ δ.

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``.

    Raises:
        ImportError: Without the ``assim`` extra: the log-space update is a
            ``vardax.OptimalInterpolation`` analysis over ``δ ↦ F̃ δ``.
    """
    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, caller="lognormal_inversion"
    )
    mean = qa * jnp.exp(log_increment)
    return LognormalPosterior(
        mean=mean,
        log_increment=log_increment,
        log_covariance=log_cov,
        prior_mean=qa,
        prior_log_covariance=b_log,
    )

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_

ensemble_footprints

ensemble_footprints(winds: Sequence[GriddedWind | Callable[[Array], Array]], receptors: Float[Array, 'obs 3'], turbulence: TurbulenceModel, domain_x: tuple[float, float, int], domain_y: tuple[float, float, int], *, 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) -> Float[Array, 'member obs cell']

Footprint matrix F_e of every ensemble member.

F_e[i, s] = compute_footprint(receptor i, wind e)[s], cells flattened row-major over (x, y).

Every member reuses the same particle draws for a receptor (common random numbers: receptor i uses seed seed + i in every member). The spread between members is then the effect of the winds alone; independent draws would add particle-sampling noise that :func:ensemble_inversion would count as transport error.

Parameters:

Name	Type	Description	Default
winds	Sequence[GriddedWind | Callable[[Array], Array]]	One mean wind per member: a :class:~plumax.lagrangian.wind_interp.GriddedWind (e.g. gridded_wind_from_met of each member's MetField) or a uniform t -> (u, v, w) callable.	required
receptors	Float[Array, 'obs 3']	Receptor positions (x, y, z) [m].	required
turbulence	TurbulenceModel	Turbulence model, shared by the members.	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
n_particles	int	Particles per receptor and member.	5000
t_back	float	Backward integration time [s].	600.0
dt	float	Time step [s].	1.0
pbl_height	float	Boundary-layer height [m].	1000.0
pbl_fraction	float	Fraction of the PBL that sees the surface flux.	0.5
air_density	float	Air density [kg m⁻³].	1.2
receptor_time	float	Observation time [s] on the winds' clock.	0.0
seed	int	Base PRNG seed.	0

Returns:

Type	Description
Float[Array, 'member obs cell']	The member footprints, units s m² kg⁻¹ (see compute_footprint).

Raises:

Type	Description
ValueError	If winds is empty.

Source code in src/plumax/lagrangian/_src/met_ensemble.py

def ensemble_footprints(
    winds: Sequence[GriddedWind | Callable[[jax.Array], jax.Array]],
    receptors: Float[Array, "obs 3"],
    turbulence: TurbulenceModel,
    domain_x: tuple[float, float, int],
    domain_y: tuple[float, float, int],
    *,
    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,
) -> Float[Array, "member obs cell"]:
    """Footprint matrix F_e of every ensemble member.

    F_e[i, s] = compute_footprint(receptor i, wind e)[s], cells flattened
    row-major over ``(x, y)``.

    Every member reuses the same particle draws for a receptor (common random
    numbers: receptor ``i`` uses seed ``seed + i`` in every member). The spread
    between members is then the effect of the winds alone; independent draws
    would add particle-sampling noise that :func:`ensemble_inversion` would count
    as transport error.

    Args:
        winds: One mean wind per member: a
            :class:`~plumax.lagrangian.wind_interp.GriddedWind` (e.g.
            ``gridded_wind_from_met`` of each member's ``MetField``) or a
            uniform ``t -> (u, v, w)`` callable.
        receptors: Receptor positions (x, y, z) [m].
        turbulence: Turbulence model, shared by the members.
        domain_x: ``(start, stop, n_cells)`` along x for the surface grid.
        domain_y: ``(start, stop, n_cells)`` along y for the surface grid.
        n_particles: Particles per receptor and member.
        t_back: Backward integration time [s].
        dt: Time step [s].
        pbl_height: Boundary-layer height [m].
        pbl_fraction: Fraction of the PBL that sees the surface flux.
        air_density: Air density [kg m⁻³].
        receptor_time: Observation time [s] on the winds' clock.
        seed: Base PRNG seed.

    Returns:
        The member footprints, units s m² kg⁻¹ (see ``compute_footprint``).

    Raises:
        ValueError: If ``winds`` is empty.
    """
    if len(winds) == 0:
        raise ValueError("ensemble_footprints: need at least one member wind.")
    receptor_points = np.asarray(receptors, dtype=float)
    members = []
    for wind in winds:
        rows = []
        for index, receptor in enumerate(receptor_points):
            footprint, _, _ = compute_footprint(
                tuple(receptor),
                turbulence,
                domain_x,
                domain_y,
                wind=wind,
                n_particles=n_particles,
                t_back=t_back,
                dt=dt,
                pbl_height=pbl_height,
                pbl_fraction=pbl_fraction,
                air_density=air_density,
                receptor_time=receptor_time,
                seed=seed + index,
            )
            rows.append(einx.rearrange("x y -> (x y)", jnp.asarray(footprint)))
        members.append(jnp.stack(rows))
    return jnp.stack(members)

ensemble_inversion

ensemble_inversion(footprints: Float[Array, 'member obs cell'], observation: Float[Array, ' obs'], *, prior_mean: Float[Array, ' cell'], prior_cov: Float[Array, 'cell cell'], obs_variance: Float[Array, ' obs']) -> GaussianPosterior

Gaussian inversion that marginalises the footprint over a met ensemble.

q = q_a + B F̄ᵀ (F̄ B F̄ᵀ + R + R_transport)⁻¹ (y − F̄ q_a), P = B − B F̄ᵀ (F̄ B F̄ᵀ + R + R_transport)⁻¹ F̄ B,

with F̄ and R_transport from :func:transport_covariance. One member gives :func:~plumax.lagrangian.linear_gaussian_inversion on its footprint.

Parameters:

Name	Type	Description	Default
footprints	Float[Array, 'member obs cell']	Member footprints F_e (e.g. :func:ensemble_footprints).	required
observation	Float[Array, ' obs']	Observed values y.	required
prior_mean	Float[Array, ' cell']	Prior mean q_a.	required
prior_cov	Float[Array, 'cell cell']	Prior covariance B.	required
obs_variance	Float[Array, ' obs']	Diagonal of the retrieval + representation covariance R.	required

Returns:

Name	Type	Description
The	GaussianPosterior	class:~plumax.lagrangian.GaussianPosterior over q.

Raises:

Type	Description
ImportError	Without the assim extra (pip install 'plumax[assim]'), as linear_gaussian_inversion.

Source code in src/plumax/lagrangian/_src/met_ensemble.py

def ensemble_inversion(
    footprints: Float[Array, "member obs cell"],
    observation: Float[Array, " obs"],
    *,
    prior_mean: Float[Array, " cell"],
    prior_cov: Float[Array, "cell cell"],
    obs_variance: Float[Array, " obs"],
) -> GaussianPosterior:
    """Gaussian inversion that marginalises the footprint over a met ensemble.

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

    with F̄ and R_transport from :func:`transport_covariance`. One member gives
    :func:`~plumax.lagrangian.linear_gaussian_inversion` on its footprint.

    Args:
        footprints: Member footprints F_e (e.g. :func:`ensemble_footprints`).
        observation: Observed values y.
        prior_mean: Prior mean q_a.
        prior_cov: Prior covariance B.
        obs_variance: Diagonal of the retrieval + representation covariance R.

    Returns:
        The :class:`~plumax.lagrangian.GaussianPosterior` over q.

    Raises:
        ImportError: Without the ``assim`` extra (``pip install
            'plumax[assim]'``), as ``linear_gaussian_inversion``.
    """
    member_footprints = jnp.asarray(footprints, dtype=float)
    if member_footprints.ndim != 3:
        raise ValueError(
            "ensemble_inversion: `footprints` must be 3-D (member, obs, cell)."
        )
    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(member_footprints[0], y, qa, b, r)

    mean_footprint, transport = transport_covariance(member_footprints, qa)
    # R_total = R + R_transport
    total_covariance = jnp.diag(r) + transport
    # y − F̄ q_a
    innovation = y - einx.dot("obs cell, cell -> obs", mean_footprint, qa)
    increment, posterior_cov = _blue_update(
        mean_footprint,
        innovation,
        b,
        total_covariance,
        caller="ensemble_inversion",
    )
    return GaussianPosterior(
        mean=qa + increment,
        covariance=posterior_cov,
        prior_mean=qa,
        prior_covariance=b,
    )

transport_covariance

transport_covariance(footprints: Float[Array, 'member obs cell'], prior_mean: Float[Array, ' cell']) -> tuple[Float[Array, 'obs cell'], Float[Array, 'obs obs']]

Mean footprint F̄ and linearised transport-error covariance at q_a.

R_transport = Σ_e d_e d_eᵀ / (E − 1), d_e = (F_e − F̄) q_a, computed by gaussx.ensemble_covariance(..., bessel=True) on F_e q_a,

and zero for a single member.

Parameters:

Name	Type	Description	Default
footprints	Float[Array, 'member obs cell']	Member footprints F_e.	required
prior_mean	Float[Array, ' cell']	Prior mean q_a, the linearisation point.	required

Returns:

Type	Description
tuple[Float[Array, 'obs cell'], Float[Array, 'obs obs']]	(F̄, R_transport).

Examples:

Two members that disagree on one receptor:

>>> import jax.numpy as jnp
>>> footprints = jnp.array([[[1.0, 0.0], [0.0, 1.0]], [[3.0, 0.0], [0.0, 1.0]]])
>>> mean_footprint, covariance = transport_covariance(
...     footprints, jnp.array([1.0, 1.0])
... )
>>> mean_footprint.tolist(), covariance.tolist()
([[2.0, 0.0], [0.0, 1.0]], [[2.0, 0.0], [0.0, 0.0]])

Identical members carry no transport error:

>>> same = jnp.stack([footprints[0], footprints[0]])
>>> bool((transport_covariance(same, jnp.ones(2))[1] == 0.0).all())
True

Source code in src/plumax/lagrangian/_src/met_ensemble.py

def transport_covariance(
    footprints: Float[Array, "member obs cell"],
    prior_mean: Float[Array, " cell"],
) -> tuple[Float[Array, "obs cell"], Float[Array, "obs obs"]]:
    """Mean footprint F̄ and linearised transport-error covariance at q_a.

    R_transport = Σ_e d_e d_eᵀ / (E − 1),   d_e = (F_e − F̄) q_a,
    computed by ``gaussx.ensemble_covariance(..., bessel=True)`` on F_e q_a,

    and zero for a single member.

    Args:
        footprints: Member footprints F_e.
        prior_mean: Prior mean q_a, the linearisation point.

    Returns:
        ``(F̄, R_transport)``.

    Examples:
        Two members that disagree on one receptor:

        >>> import jax.numpy as jnp
        >>> footprints = jnp.array([[[1.0, 0.0], [0.0, 1.0]], [[3.0, 0.0], [0.0, 1.0]]])
        >>> mean_footprint, covariance = transport_covariance(
        ...     footprints, jnp.array([1.0, 1.0])
        ... )
        >>> mean_footprint.tolist(), covariance.tolist()
        ([[2.0, 0.0], [0.0, 1.0]], [[2.0, 0.0], [0.0, 0.0]])

        Identical members carry no transport error:

        >>> same = jnp.stack([footprints[0], footprints[0]])
        >>> bool((transport_covariance(same, jnp.ones(2))[1] == 0.0).all())
        True
    """
    n_members = footprints.shape[0]
    # F̄ = F_0 + mean_e(F_e − F_0): identical members give F̄ = F_0 exactly
    offsets = footprints - footprints[0]
    mean_footprint = footprints[0] + einx.mean("member obs cell -> obs cell", offsets)
    # r_e = F_e q_a
    response = einx.dot("member obs cell, cell -> member obs", footprints, prior_mean)
    if n_members == 1:
        # gaussx divides by E − 1, so a lone member is an explicit zero
        n_obs = response.shape[1]
        return mean_footprint, jnp.zeros((n_obs, n_obs), dtype=response.dtype)
    # R_transport = Σ_e (r_e − r̄)(r_e − r̄)ᵀ / (E − 1), as a low-rank operator
    operator = gaussx.ensemble_covariance(response, bessel=True)
    covariance = operator.as_matrix()
    return mean_footprint, covariance

integrate_particles

integrate_particles(state: ParticleState, wind: Callable[[Array], Array] | GriddedWind, 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] | GriddedWind	Mean-wind field: t -> (u, v, w) [m/s], spatially uniform, or a :class:~plumax.lagrangian.GriddedWind sampled at each particle.	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] | GriddedWind,
    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,
            or a :class:`~plumax.lagrangian.GriddedWind` sampled at each
            particle.
        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,
            mean_wind(wind, carry.position, 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,) for a uniform wind or (n, 3) for one per particle.	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,)`` for
            a uniform wind or ``(n, 3)`` for one per particle.
        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 + 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, coriolis_parameter: float = DEFAULT_CORIOLIS_PARAMETER) -> tuple[jax.Array, jax.Array]

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

Returns height-dependent velocity standard deviations and Lagrangian timescales for three stability regimes, chosen as in FLEXPART's hanna.f90:

• Neutral (h/|L| < 1, or L = 0): Stohl et al. (2005) eqs. 26–28, σ_u = 2.0 u_ exp(−3fz/u_), σ_v = σ_w = 1.3 u_ exp(−2fz/u_), τ_L = 0.5 z / σ_w / (1 + 15fz/u_*) on all three axes.
• Unstable / convective (h/|L| ≥ 1 and L < 0): convective scaling with w_*.
• Stable (h/|L| ≥ 1 and L > 0): eqs. 29–33, σ_u = 2.0 u_ (1 − z/h), σ_v = σ_w = 1.3 u_ (1 − z/h), τ_L,u = 0.15 (h/σ_u)(z/h)^0.5, τ_L,v = 0.07 (h/σ_v)(z/h)^0.5, τ_L,w = 0.1 (h/σ_w)(z/h)^0.8.

The stable τ_L,w exponent is 0.8, as in Hanna (1982) and hanna.f90; Stohl et al. (2005) print 0.5 in eq. 33. L = 0 is the neutral sentinel (1/L = 0) that :func:plumax.met.hanna_turbulence_from passes.

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, 0 neutral).	required
w_star	float	Convective velocity scale [m/s] (used when unstable).	0.0
coriolis_parameter	float	Coriolis parameter f = 2Ω sin φ [1/s], used by the neutral profiles; only |f| enters. The default 1e-4 s⁻¹ (latitude ≈ 43°) is the value FLEXPART's hanna.f90 hard-codes; pass 2 · 7.292e-5 · sin(latitude) for a specific site.	DEFAULT_CORIOLIS_PARAMETER

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).

References

Hanna, S. R. (1982). Applications in air pollution modeling. Stohl, A. et al. (2005). Technical note: The Lagrangian particle dispersion model FLEXPART version 6.2. Atmos. Chem. Phys. 5, 2461–2474, eqs. 20–33.

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,
    coriolis_parameter: float = DEFAULT_CORIOLIS_PARAMETER,
) -> tuple[jax.Array, jax.Array]:
    """Hanna (1982) Monin–Obukhov-similarity turbulence profiles.

    Returns height-dependent velocity standard deviations and Lagrangian
    timescales for three stability regimes, chosen as in FLEXPART's
    ``hanna.f90``:

    - **Neutral** (``h/|L| < 1``, or ``L = 0``): Stohl et al. (2005) eqs. 26–28,
      σ_u = 2.0 u_* exp(−3fz/u_*), σ_v = σ_w = 1.3 u_* exp(−2fz/u_*),
      τ_L = 0.5 z / σ_w / (1 + 15fz/u_*) on all three axes.
    - **Unstable / convective** (``h/|L| ≥ 1`` and ``L < 0``): convective
      scaling with ``w_*``.
    - **Stable** (``h/|L| ≥ 1`` and ``L > 0``): eqs. 29–33,
      σ_u = 2.0 u_* (1 − z/h), σ_v = σ_w = 1.3 u_* (1 − z/h),
      τ_L,u = 0.15 (h/σ_u)(z/h)^0.5, τ_L,v = 0.07 (h/σ_v)(z/h)^0.5,
      τ_L,w = 0.1 (h/σ_w)(z/h)^0.8.

    The stable τ_L,w exponent is 0.8, as in Hanna (1982) and ``hanna.f90``;
    Stohl et al. (2005) print 0.5 in eq. 33. ``L = 0`` is the neutral sentinel
    (``1/L = 0``) that :func:`plumax.met.hanna_turbulence_from` passes.

    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,
            ``0`` neutral).
        w_star: Convective velocity scale [m/s] (used when unstable).
        coriolis_parameter: Coriolis parameter f = 2Ω sin φ [1/s], used by the
            neutral profiles; only ``|f|`` enters. The default ``1e-4`` s⁻¹
            (latitude ≈ 43°) is the value FLEXPART's ``hanna.f90`` hard-codes;
            pass ``2 · 7.292e-5 · sin(latitude)`` for a specific site.

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

    References:
        Hanna, S. R. (1982). Applications in air pollution modeling.
        Stohl, A. et al. (2005). Technical note: The Lagrangian particle
        dispersion model FLEXPART version 6.2. *Atmos. Chem. Phys.* 5,
        2461–2474, eqs. 20–33.
    """
    z = jnp.asarray(z, dtype=float)
    h = pbl_height
    zh = jnp.clip(z / h, 0.0, 1.0)
    us = u_star
    ws = w_star
    us_safe = jnp.maximum(us, _SIGMA_FLOOR)
    f_c = jnp.abs(coriolis_parameter)

    # --- 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: Stohl et al. (2005) eqs. 29–33 ---
    decay = jnp.maximum(1.0 - zh, 1e-3)
    # σ_u = 2.0 u_* (1 − z/h)
    sig_u_stab = 2.0 * us * decay
    # σ_v = σ_w = 1.3 u_* (1 − z/h)
    sig_vw_stab = 1.3 * us * decay
    # τ_L,u = 0.15 (h/σ_u)(z/h)^0.5
    tau_u_stab = (
        0.15
        * h
        / jnp.maximum(sig_u_stab, _SIGMA_FLOOR)
        * jnp.sqrt(jnp.maximum(zh, 1e-3))
    )
    # τ_L,v = 0.07 (h/σ_v)(z/h)^0.5
    tau_v_stab = (
        0.07
        * h
        / jnp.maximum(sig_vw_stab, _SIGMA_FLOOR)
        * jnp.sqrt(jnp.maximum(zh, 1e-3))
    )
    # τ_L,w = 0.1 (h/σ_w)(z/h)^0.8
    tau_w_stab = (
        0.1 * h / jnp.maximum(sig_vw_stab, _SIGMA_FLOOR) * jnp.maximum(zh, 1e-3) ** 0.8
    )

    # --- Neutral: Stohl et al. (2005) eqs. 26–28 ---
    # σ_u = 2.0 u_* exp(−3fz/u_*)
    sig_u_neut = 2.0 * us * jnp.exp(-3.0 * f_c * z / us_safe)
    # σ_v = σ_w = 1.3 u_* exp(−2fz/u_*)
    sig_vw_neut = 1.3 * us * jnp.exp(-2.0 * f_c * z / us_safe)
    # τ_L = 0.5 z / σ_w / (1 + 15fz/u_*)
    tau_neut = (
        0.5
        * z
        / jnp.maximum(sig_vw_neut, _SIGMA_FLOOR)
        / (1.0 + 15.0 * f_c * z / us_safe)
    )
    tau_neut = jnp.maximum(tau_neut, 1.0)

    # --- regime: neutral when h/|L| < 1 (FLEXPART hanna.f90); L = 0 is neutral ---
    is_neutral_sentinel = obukhov_length == 0.0
    safe_length = jnp.where(is_neutral_sentinel, 1.0, obukhov_length)
    # h/|L|, with 1/L = 0 at the sentinel
    stability_ratio = jnp.where(is_neutral_sentinel, 0.0, h / jnp.abs(safe_length))
    neutral = stability_ratio < 1.0
    unstable = jnp.logical_and(~neutral, obukhov_length < 0.0)
    stable = jnp.logical_and(~neutral, 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_u_stab, sig_u_neut)
    sigma_v = pick(0.8 * sig_uv_conv, sig_vw_stab, sig_vw_neut)
    sigma_w = pick(sig_w_conv, sig_vw_stab, sig_vw_neut)
    tau_u = pick(tau_uv_conv, tau_u_stab, tau_neut)
    tau_v = pick(tau_uv_conv, tau_v_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

gridded_wind_from_met

gridded_wind_from_met(field: MetField) -> GriddedWind

The wind of a :class:~plumax.met.MetField, for particles.

MetField already stores u, v, w on the C-grid with a time axis, so this is a relabelling; no value changes.

Parameters:

Name	Type	Description	Default
field	MetField	Meteorology on an analysis grid.	required

Returns:

Type	Description
GriddedWind	The gridded wind on field.plume_grid at field.times.

Source code in src/plumax/lagrangian/_src/wind_interp.py

def gridded_wind_from_met(field: MetField) -> GriddedWind:
    """The wind of a :class:`~plumax.met.MetField`, for particles.

    ``MetField`` already stores ``u``, ``v``, ``w`` on the C-grid with a time
    axis, so this is a relabelling; no value changes.

    Args:
        field: Meteorology on an analysis grid.

    Returns:
        The gridded wind on ``field.plume_grid`` at ``field.times``.
    """
    return GriddedWind.from_grid(
        field.plume_grid, field.times, field.u, field.v, field.w
    )

gridded_wind_from_prescribed

gridded_wind_from_prescribed(wind_field: PrescribedWindField, times: Float[Array, ' time']) -> GriddedWind

Sample a les_fvm :class:~plumax.les_fvm.wind.PrescribedWindField.

The field is evaluated at each of times; between them the particles see the piecewise-linear interpolant, so times should resolve the field's own time variation.

Parameters:

Name	Type	Description	Default
wind_field	PrescribedWindField	The Eulerian solver's wind.	required
times	Float[Array, ' time']	Times [s] to sample it at, strictly increasing.	required

Returns:

Type	Description
GriddedWind	The gridded wind on wind_field.plume_grid.

Source code in src/plumax/lagrangian/_src/wind_interp.py

def gridded_wind_from_prescribed(
    wind_field: PrescribedWindField, times: Float[Array, " time"]
) -> GriddedWind:
    """Sample a ``les_fvm`` :class:`~plumax.les_fvm.wind.PrescribedWindField`.

    The field is evaluated at each of ``times``; between them the particles see
    the piecewise-linear interpolant, so ``times`` should resolve the field's
    own time variation.

    Args:
        wind_field: The Eulerian solver's wind.
        times: Times [s] to sample it at, strictly increasing.

    Returns:
        The gridded wind on ``wind_field.plume_grid``.
    """
    times = jnp.asarray(times)
    samples = [wind_field(t) for t in times]
    # interior points: drop the one-cell ghost ring the solver pads with
    u, v, w = (
        jnp.stack([sample[c][1:-1, 1:-1, 1:-1] for sample in samples]) for c in range(3)
    )
    return GriddedWind.from_grid(wind_field.plume_grid, times, u, v, w)