Eulerian finite volume

Eulerian 3-D advection–diffusion on an Arakawa C-grid via finitevolX. Tier III.

les_fvm

Eulerian 3-D advection-diffusion of a methane tracer (Level-2 fidelity).

This sub-package implements a finite-volume solver for the passive-tracer transport equation

∂C/∂t + ∇·(u C) = ∇·(K ∇C) + S(x, t)

on an Arakawa C-grid (via finitevolX), with a prescribed wind field u = (u, v, w) and K-theory (possibly anisotropic) eddy diffusivity K = diag(K_h, K_h, K_z). No turbulence is resolved — this model sits between the analytic Gaussian puff (L1) and a full resolved-flow LES (L3).

Submodules

• grid : CartesianGrid3D wrapper + coordinate helpers.
• wind : prescribed 3-D wind fields (uniform, sheared, time-varying).
• source : Gaussian point-source and per-cell surface-flux terms.
• advection : horizontal WENO advection (finitevolX) + vertical upwind.
• diffusion : K-theory eddy diffusivity + anisotropic Laplacian.
• boundary : BC application for 3-D fields (horizontal vmap + vertical).
• dynamics : diffrax-compatible RHS for the scalar transport equation.
• simulate : high-level xarray-returning runner.
• _vertical_ops : private — vertical ∂_z / flux / Laplacian helpers.

HorizontalBC

Bases: Module

Apply a :class:BoundaryConditionSet to every z-slice of a 3-D field.

advection_variant

advection_variant() -> HorizontalBC

Return a copy whose Dirichlet faces use a cell-value ghost.

Open-mode advection reads the ghost as an exterior cell value; a face-Dirichlet ghost (2*value - interior) would make an inflow carry the wrong state. Outflow / periodic / Neumann faces are unchanged, so this is a no-op unless a Dirichlet face is present.

Source code in src/plumax/les_fvm/boundary.py

def advection_variant(self) -> HorizontalBC:
    """Return a copy whose Dirichlet faces use a cell-value ghost.

    Open-mode advection reads the ghost as an exterior *cell* value; a
    face-Dirichlet ghost (``2*value - interior``) would make an inflow
    carry the wrong state.  Outflow / periodic / Neumann faces are
    unchanged, so this is a no-op unless a Dirichlet face is present.
    """
    bc_set = self.bc_set
    return HorizontalBC(
        bc_set=BoundaryConditionSet(
            south=_to_cell_dirichlet(bc_set.south),
            north=_to_cell_dirichlet(bc_set.north),
            west=_to_cell_dirichlet(bc_set.west),
            east=_to_cell_dirichlet(bc_set.east),
            mask=bc_set.mask,
        )
    )

VerticalBC

Bases: Module

Top/bottom ghost-slice update for a 3-D field.

Parameters

bottom_kind : {"dirichlet", "neumann", "outflow", "periodic"} Ground-boundary behaviour. bottom_value : float, default 0.0 For dirichlet: the boundary value. For neumann: the coordinate gradient ∂C/∂z along +z at the face — a positive value means C increases with height. The ghost cell is set so the finite-difference ∂C/∂z across the face equals value. Ignored for outflow / periodic. top_kind, top_value : same, with the Neumann value again the coordinate gradient ∂C/∂z along +z (same convention on both faces; see the sign note in :func:_apply_vertical_face).

advection_variant

advection_variant() -> VerticalBC

Return a copy whose Dirichlet faces use a cell-value ghost.

The vertical advection term also upwinds the ghost slice, so a Dirichlet vertical inflow needs the same cell-value convention as the horizontal faces. A no-op unless a vertical Dirichlet face is set.

Source code in src/plumax/les_fvm/boundary.py

def advection_variant(self) -> VerticalBC:
    """Return a copy whose Dirichlet faces use a cell-value ghost.

    The vertical advection term also upwinds the ghost slice, so a
    Dirichlet vertical inflow needs the same cell-value convention as the
    horizontal faces.  A no-op unless a vertical Dirichlet face is set.
    """
    return VerticalBC(
        bottom_kind=self.bottom_kind,
        top_kind=self.top_kind,
        bottom_value=self.bottom_value,
        top_value=self.top_value,
        dirichlet_cell=True,
    )

EddyDiffusivity

Bases: Module

Scalar or field-valued (K_h, K_z) eddy diffusivity.

Attributes

horizontal : float or Float[Array, "nz ny nx"] Horizontal eddy diffusivity [m²/s]. A scalar is broadcast over the grid. A field is expected on interior T-points. vertical : float or Float[Array, "nz ny nx"] Vertical eddy diffusivity [m²/s]. Same broadcasting rules.

as_arrays

as_arrays() -> tuple[float | Float[Array, 'Nz Ny Nx'], float | Float[Array, 'Nz Ny Nx']]

Return (K_h, K_z) in forms consumable by the tendency calls.

Source code in src/plumax/les_fvm/diffusion.py

def as_arrays(
    self,
) -> tuple[
    float | Float[Array, "Nz Ny Nx"],
    float | Float[Array, "Nz Ny Nx"],
]:
    """Return ``(K_h, K_z)`` in forms consumable by the tendency calls."""
    return _broadcast_or_scalar(self.horizontal), _broadcast_or_scalar(
        self.vertical
    )

EulerianDispersionRHS

Bases: Module

Full tracer-transport tendency for diffrax.

Parameters

plume_grid : PlumeGrid3D wind_field : PrescribedWindField Time-queryable prescribed wind. eddy_diffusivity : EddyDiffusivity (K_h, K_z) eddy diffusivity. source : GaussianSource or SurfaceFluxSource Methane source term. horizontal_bc : HorizontalBC vertical_bc : VerticalBC advection_scheme : str, default "weno5" Horizontal reconstruction scheme for :func:advection.advection_tendency.

BackgroundAnalysis dataclass

BackgroundAnalysis(initial_concentration: Array, bc_scale: Array)

The MAP of the :class:BackgroundControl part of a 4D-Var control.

Attributes:

Name	Type	Description
initial_concentration	Array	c₀ = c_b + δc₀ on the interior [kg m⁻³], shape (nz, ny, nx).
bc_scale	Array	Per-face factors α*, ordered (west, east, south, north).

BackgroundControl dataclass

BackgroundControl(c_b: Array, B_c: AbstractLinearOperator, bc_scale_prior: tuple[float, float] = (1.0, 0.1))

Initial concentration and per-face inflow scaling in the control.

The flat control becomes (S, δc₀, α), with priors independent of the source's,

c₀ = c_b + δc₀,   δc₀ ~ N(0, B_c),   α_f ~ N(α̅, σ_α²) for each face f,

and every Dirichlet lateral boundary value c_f of the forward becomes α_f c_f (:func:~plumax.les_fvm.boundary.scale_dirichlet_faces). The factors are ordered (west, east, south, north). This is the per-face Gaussian of the design page; a per-face Matérn along the boundary is its open v2.

Attributes:

Name	Type	Description
c_b	Array	Background initial concentration c_b on the interior cells [kg m⁻³], shape (nz, ny, nx).
B_c	AbstractLinearOperator	Covariance of δc₀, a lineax operator on the flattened interior.
bc_scale_prior	tuple[float, float]	Mean α̅ and standard deviation σ_α of each factor.

Raises:

Type	Description
ValueError	If c_b is not a finite 3-D array, B_c has the wrong size, or bc_scale_prior has no positive standard deviation.

Examples:

A background on a 2 × 3 × 4 interior adds 24 + 4 entries to the control:

>>> import jax
>>> import jax.numpy as jnp
>>> import lineax as lx
>>> covariance = lx.DiagonalLinearOperator(jnp.ones(24))
>>> background = BackgroundControl(jnp.zeros((2, 3, 4)), covariance)
>>> background.size
28

size property

size: int

Length of the control's (δc₀, α) part.

split

split(flat: Array) -> tuple[jax.Array, jax.Array, jax.Array]

Split a flat control into (S, δc₀ as a field, α).

Source code in src/plumax/les_fvm/fourdvar.py

def split(self, flat: jax.Array) -> tuple[jax.Array, jax.Array, jax.Array]:
    """Split a flat control into (S, δc₀ as a field, α)."""
    n_source = flat.shape[-1] - self.size
    increment = flat[..., n_source:-_N_FACES]
    increment_field = increment.reshape(
        (*increment.shape[:-1], *np.shape(self.c_b))
    )
    return flat[..., :n_source], increment_field, flat[..., -_N_FACES:]

flat_prior_mean

flat_prior_mean() -> jax.Array

Prior mean of the (δc₀, α) part, (0, α̅).

Source code in src/plumax/les_fvm/fourdvar.py

def flat_prior_mean(self) -> jax.Array:
    """Prior mean of the (δc₀, α) part, (0, α̅)."""
    increment_mean = jnp.zeros(np.size(self.c_b), dtype=float)
    scale_mean = jnp.full(_N_FACES, self.bc_scale_prior[0], dtype=float)
    return jnp.concatenate([increment_mean, scale_mean])

ColumnObservationOperator dataclass

ColumnObservationOperator(dz: float, receptor_index: Array | None = None)

Sample the column-integrated field at fixed receptor cells.

H_t in the 4D-Var cost: vertically integrate the 3-D concentration to a column (∑_z c · dz), then gather the values at a set of receptor pixels. With receptor_index = None every interior column is an observation. A thin wrapper over the shared :class:~plumax.gauss_plume.observation.ColumnObservation (identity averaging kernel, h_k = dz), whose module docstring states the units.

Attributes:

Name	Type	Description
dz	float	Vertical cell size [m] for the column integral.
receptor_index	Array | None	Optional (n_obs, 2) int array of (iy, ix) interior-cell indices to sample. None flattens the full (ny, nx) column grid.

EulerianForward4DVar dataclass

EulerianForward4DVar(rhs_factory: Callable[..., EulerianDispersionRHS], observation_op: ColumnObservationOperator, save_times: Array, initial_concentration: Array, t0: float, t1: float, solver_kwargs: dict = dict(), plume_grid: PlumeGrid3D | None = None, control: SpatialControl | None = None, background: BackgroundControl | None = None)

Differentiable map from an emission time series to column observations.

Bundles the fixed transport configuration (grid, wind, diffusivity, BCs, source geometry, save times) so that :meth:predict is a pure function of the control vector q — the time-resolved emission rate at the save times, or the flat per-cell flux of a :class:SpatialControl — and is therefore jax.grad / jax.linearize-able for 4D-Var.

Build with :func:build_forward rather than directly.

Attributes:

Name	Type	Description
rhs_factory	Callable[..., EulerianDispersionRHS]	(q_series, bc_scale=None) -> EulerianDispersionRHS closing over the static transport configuration; the emission rate (linearly interpolated between save times) and the optional per-face boundary scaling are its only free inputs.
observation_op	ColumnObservationOperator	The column H_t.
save_times	Array	Observation/integration times [s], shape (n_t,).
initial_concentration	Array	Padded initial field (Nz, Ny, Nx).
solver_kwargs	dict	diffrax solve settings (solver, dt0, tolerances).
plume_grid	PlumeGrid3D | None	The analysis grid; its T-cells place a spatial prior.
control	SpatialControl | None	The spatial control the flat control vector is laid out by; None for the emission time series.
background	BackgroundControl | None	The initial-concentration and boundary-scaling part of the flat control; None pins both.

with_background

with_background(background: BackgroundControl) -> EulerianForward4DVar

This forward over the flat control (S, δc₀, α) of background.

Source code in src/plumax/les_fvm/fourdvar.py

def with_background(self, background: BackgroundControl) -> EulerianForward4DVar:
    """This forward over the flat control (S, δc₀, α) of ``background``."""
    return dataclasses.replace(self, background=background)

concentration_history

concentration_history(q_series: Array) -> jax.Array

Solve the FV transport for q_series; return interior fields.

Returns the (n_t, nz, ny, nx) interior concentration history (ghost rings stripped). Differentiable in q_series via the diffrax solve.

Source code in src/plumax/les_fvm/fourdvar.py

def concentration_history(self, q_series: jax.Array) -> jax.Array:
    """Solve the FV transport for ``q_series``; return interior fields.

    Returns the ``(n_t, nz, ny, nx)`` interior concentration history (ghost
    rings stripped). Differentiable in ``q_series`` via the diffrax solve.
    """
    import diffrax

    flat = jnp.asarray(q_series)
    if self.background is None:
        rhs = self.rhs_factory(flat)
        initial_concentration = self.initial_concentration
    else:
        source, increment, bc_scale = self.background.split(flat)
        rhs = self.rhs_factory(source, bc_scale)
        # c₀ = c_b + δc₀ on the interior; the first RHS call fills the ghosts
        interior = jnp.asarray(self.background.c_b, flat.dtype) + increment
        initial_concentration = self.initial_concentration.at[1:-1, 1:-1, 1:-1].set(
            interior
        )
    kw = {
        "solver": diffrax.Heun(),
        "dt0": 1.0,
        "stepsize_controller": diffrax.ConstantStepSize(),
        "max_steps": 100_000,
        **self.solver_kwargs,
    }
    # diffrax's StepTo controller prescribes its own step points and requires
    # dt0=None; the default dt0=1.0 above would otherwise make the solve
    # fail, so drop it whenever a StepTo controller is in effect.
    if isinstance(kw["stepsize_controller"], diffrax.StepTo):
        kw["dt0"] = None
    sol = diffrax.diffeqsolve(
        diffrax.ODETerm(rhs),
        kw["solver"],
        t0=self.t0,
        t1=self.t1,
        dt0=kw["dt0"],
        y0=initial_concentration,
        saveat=diffrax.SaveAt(ts=self.save_times),
        stepsize_controller=kw["stepsize_controller"],
        max_steps=kw["max_steps"],
        adjoint=diffrax.RecursiveCheckpointAdjoint(),
    )
    return sol.ys[:, 1:-1, 1:-1, 1:-1]

predict

predict(q_series: Array) -> jax.Array

Column observations stacked over time, shape (n_t, n_obs).

Source code in src/plumax/les_fvm/fourdvar.py

def predict(self, q_series: jax.Array) -> jax.Array:
    """Column observations stacked over time, shape ``(n_t, n_obs)``."""
    history = self.concentration_history(q_series)
    return jax.vmap(self.observation_op)(history)

FourDVarProblem dataclass

FourDVarProblem(forward: EulerianForward4DVar, prior_mean: Array, prior_chol: Array, observations: Array, obs_variance: Array, control: SpatialControl | None = None, prior_sqrt_op: AbstractLinearOperator | None = None)

A strong-constraint 4D-Var problem over a flat source control.

The control is a scalar emission time series or the per-cell flux of a :class:SpatialControl, optionally stacked with a :class:BackgroundControl; n_state below is its length. Holds the differentiable forward, the prior (S_b, B = L Lᵀ), and the observations (y_t, R_t). The cost is evaluated in whitened space χ = L⁻¹ (S − S_b) so the prior term is the identity ½‖χ‖² and the Hessian is well-conditioned (design §control-transform). observations has shape (n_t, n_obs) and obs_variance broadcasts to it.

Build with :func:build_problem.

n_source property

n_source: int

Length of the source part of the flat control.

prior_sqrt_operator

prior_sqrt_operator() -> lx.AbstractLinearOperator

The prior square root L as an operator, structured where it can be.

Source code in src/plumax/les_fvm/fourdvar.py

def prior_sqrt_operator(self) -> lx.AbstractLinearOperator:
    """The prior square root L as an operator, structured where it can be."""
    if self.prior_sqrt_op is not None:
        return self.prior_sqrt_op
    return lx.MatrixLinearOperator(self.prior_chol)

source_from_whitened

source_from_whitened(chi: Array) -> jax.Array

Map whitened control χ back to the flat control, S_b + L χ.

Source code in src/plumax/les_fvm/fourdvar.py

def source_from_whitened(self, chi: jax.Array) -> jax.Array:
    """Map whitened control ``χ`` back to the flat control, ``S_b + L χ``."""
    return self.prior_mean + self.prior_chol @ chi

cost

cost(chi: Array) -> jax.Array

The 4D-Var cost J at whitened control χ (scalar).

A diagnostic: the solve minimises the same cost inside vardax.ThreeDVar, not through this method.

Source code in src/plumax/les_fvm/fourdvar.py

def cost(self, chi: jax.Array) -> jax.Array:
    """The 4D-Var cost ``J`` at whitened control ``χ`` (scalar).

    A diagnostic: the solve minimises the same cost inside
    ``vardax.ThreeDVar``, not through this method.
    """
    q = self.source_from_whitened(chi)
    predicted = self.forward.predict(q)  # (n_t, n_obs)
    residual = predicted - self.observations
    obs_term = 0.5 * jnp.sum(residual**2 / self.obs_variance)
    prior_term = 0.5 * jnp.sum(chi**2)
    return obs_term + prior_term

value_and_grad

value_and_grad(chi: Array) -> tuple[jax.Array, jax.Array]

(J, ∂J/∂χ) — the gradient is the exact adjoint via reverse-mode AD.

Source code in src/plumax/les_fvm/fourdvar.py

def value_and_grad(self, chi: jax.Array) -> tuple[jax.Array, jax.Array]:
    """``(J, ∂J/∂χ)`` — the gradient is the exact adjoint via reverse-mode AD."""
    return jax.value_and_grad(self.cost)(chi)

source_receptor_operator

source_receptor_operator() -> vardax.CallableObs

The whitened source–receptor map as a vardax observation operator.

𝒢(S) = vec(forward.predict(S)) / σ on a (1, n_t) source row, so the observation-error covariance is the identity. The diffrax solve defines only a VJP, whose while_loop JAX cannot transpose, so the operator linearises by forming the Jacobian with one VJP per observation (jac="bwd", materialise=True), so the cost scales with the number of observations rather than the length of the control.

Returns:

Type	Description
CallableObs	A vardax.CallableObs (1, n_state) → (1, n_t·n_obs).

Raises:

Type	Description
ImportError	Without the assim extra.

Source code in src/plumax/les_fvm/fourdvar.py

def source_receptor_operator(self) -> vardax.CallableObs:
    """The whitened source–receptor map as a vardax observation operator.

    ``𝒢(S) = vec(forward.predict(S)) / σ`` on a ``(1, n_t)`` source row,
    so the observation-error covariance is the identity. The diffrax
    solve defines only a VJP, whose ``while_loop`` JAX cannot transpose,
    so the operator linearises by forming the Jacobian with one VJP per
    observation (``jac="bwd", materialise=True``), so the cost scales with
    the number of observations rather than the length of the control.

    Returns:
        A ``vardax.CallableObs`` ``(1, n_state) → (1, n_t·n_obs)``.

    Raises:
        ImportError: Without the ``assim`` extra.
    """
    vdx = require_vardax("FourDVarProblem.source_receptor_operator")
    sigma = jnp.sqrt(self.obs_variance)  # (n_t, n_obs)
    predict = self.forward.predict

    def whitened_prediction(source: jax.Array) -> jax.Array:
        # 𝒢(S) / σ, flattened over (time, receptor)
        return einx.rearrange("t o -> 1 (t o)", predict(source[0]) / sigma)

    return vdx.CallableObs(fn=whitened_prediction, jac="bwd", materialise=True)

incremental_analysis

incremental_analysis(*, outer_iters: int = 5, inner_iters: int = 50, rtol: float = 1e-06, atol: float = 1e-08) -> tuple[vardax.IncrementalFourDVar, vardax.Batch1D]

The Gauss–Newton analysis of this problem, and its one-window batch.

vardax.IncrementalFourDVar with an identity forward over a one-time window, because the control is a source rather than an initial state, over the whitened source–receptor map with R = I. Each outer iteration forms the Jacobian with one VJP per observation (jac="bwd") and solves I + (J̃ L)ᵀ (J̃ L) in the control-variable transform by CG.

Parameters:

Name	Type	Description	Default
outer_iters	int	Gauss–Newton outer iterations.	5
inner_iters	int	CG iterations per outer iteration.	50
rtol	float	Relative tolerance of the CG solves.	1e-06
atol	float	Absolute tolerance of the CG solves.	1e-08

Returns:

Type	Description
IncrementalFourDVar	(model, batch): model(batch) is the MAP flat control, shape
Batch1D	(1, n_state).

Raises:

Type	Description
ImportError	Without the assim extra.

Source code in src/plumax/les_fvm/fourdvar.py

def incremental_analysis(
    self,
    *,
    outer_iters: int = 5,
    inner_iters: int = 50,
    rtol: float = 1e-6,
    atol: float = 1e-8,
) -> tuple[vardax.IncrementalFourDVar, vardax.Batch1D]:
    """The Gauss–Newton analysis of this problem, and its one-window batch.

    ``vardax.IncrementalFourDVar`` with an identity forward over a
    one-time window, because the control is a source rather than an initial
    state, over the whitened source–receptor map with ``R = I``. Each outer
    iteration forms the Jacobian with one VJP per observation
    (``jac="bwd"``) and solves I + (J̃ L)ᵀ (J̃ L) in the control-variable
    transform by CG.

    Args:
        outer_iters: Gauss–Newton outer iterations.
        inner_iters: CG iterations per outer iteration.
        rtol: Relative tolerance of the CG solves.
        atol: Absolute tolerance of the CG solves.

    Returns:
        ``(model, batch)``: ``model(batch)`` is the MAP flat control, shape
        ``(1, n_state)``.

    Raises:
        ImportError: Without the ``assim`` extra.
    """
    vdx = require_vardax("FourDVarProblem.incremental_analysis")
    sigma = jnp.sqrt(self.obs_variance)
    predict = self.forward.predict

    def whitened_prediction(state: jax.Array) -> jax.Array:
        # 𝒢(q) / σ, flattened over (time, receptor)
        return einx.rearrange("t o -> (t o)", predict(state) / sigma)

    observations = self.observations / sigma
    n_flat = observations.size
    dtype = self.prior_mean.dtype
    prior_sqrt = self.prior_sqrt_operator()
    model = vdx.IncrementalFourDVar(
        forward=_StaticForward(),
        obs_op=whitened_prediction,
        prior_mean=self.prior_mean,
        prior_cov_op=prior_sqrt @ prior_sqrt.T,
        obs_cov_op=lx.IdentityLinearOperator(
            jax.ShapeDtypeStruct((n_flat,), dtype)
        ),
        prior_sqrt_op=prior_sqrt,
        config=vdx.IncrementalConfig(
            n_outer=outer_iters,
            n_inner=inner_iters,
            cg_rtol=rtol,
            cg_atol=atol,
            covariance_tol=atol,
            jac="bwd",
        ),
    )
    batch = vdx.Batch1D(
        input=einx.rearrange("t o -> 1 1 (t o)", observations),
        mask=jnp.ones((1, 1, n_flat), dtype=dtype),
    )
    return model, batch

analysis

analysis(*, minimiser: AbstractMinimiser | None = None, max_steps: int = 100, initial_whitened: Array | None = None) -> tuple[vardax.ThreeDVar, vardax.Batch1D]

The 3D-Var analysis of this problem, and its one-window batch.

vardax.ThreeDVar over :meth:source_receptor_operator, with the whitened control S = S_b + L χ, ½‖χ‖², as a vardax.LinearParameterization and R = I on the whitened observations y / σ.

Parameters:

Name	Type	Description	Default
minimiser	AbstractMinimiser | None	An optimistix minimiser. Default optx.LBFGS(rtol=1e-6, atol=1e-8).	None
max_steps	int	Iteration cap.	100
initial_whitened	Array | None	First guess χ₀. Default 0 (the prior mean).	None

Returns:

Type	Description
ThreeDVar	(model, batch): model(batch) is the MAP source, shape
Batch1D	(1, 1, n_state), model.analyse(batch) adds the minimiser's
tuple[ThreeDVar, Batch1D]	diagnostics, and model.posterior(batch) is its Laplace
tuple[ThreeDVar, Batch1D]	posterior.

Raises:

Type	Description
ImportError	Without the assim extra.

Source code in src/plumax/les_fvm/fourdvar.py

def analysis(
    self,
    *,
    minimiser: optx.AbstractMinimiser | None = None,
    max_steps: int = 100,
    initial_whitened: jax.Array | None = None,
) -> tuple[vardax.ThreeDVar, vardax.Batch1D]:
    """The 3D-Var analysis of this problem, and its one-window batch.

    ``vardax.ThreeDVar`` over :meth:`source_receptor_operator`, with the
    whitened control ``S = S_b + L χ``, ``½‖χ‖²``, as a
    ``vardax.LinearParameterization`` and ``R = I`` on the whitened
    observations ``y / σ``.

    Args:
        minimiser: An optimistix minimiser. Default
            ``optx.LBFGS(rtol=1e-6, atol=1e-8)``.
        max_steps: Iteration cap.
        initial_whitened: First guess ``χ₀``. Default ``0`` (the prior
            mean).

    Returns:
        ``(model, batch)``: ``model(batch)`` is the MAP source, shape
        ``(1, 1, n_state)``, ``model.analyse(batch)`` adds the minimiser's
        diagnostics, and ``model.posterior(batch)`` is its Laplace
        posterior.

    Raises:
        ImportError: Without the ``assim`` extra.
    """
    import optimistix as optx

    vdx = require_vardax("FourDVarProblem.analysis")
    n_state = self.prior_mean.shape[0]
    dtype = self.prior_mean.dtype
    basis = vdx.linear_basis(self.prior_chol, variance=1.0)
    if initial_whitened is None:
        control = vdx.LinearParameterization(
            basis, background=self.prior_mean[None]
        )
    else:
        control = _warm_start_control()(
            basis,
            background=self.prior_mean[None],
            start=jnp.asarray(initial_whitened),
        )
    observations = self.observations / jnp.sqrt(self.obs_variance)
    n_flat = observations.size

    def row_operator(op: lx.AbstractLinearOperator, n: int):
        return lx.FunctionLinearOperator(
            lambda x: op.mv(x[0])[None],
            jax.ShapeDtypeStruct((1, n), dtype),
            frozenset({lx.symmetric_tag, lx.positive_semidefinite_tag}),
        )

    psd = frozenset({lx.symmetric_tag, lx.positive_semidefinite_tag})
    model = vdx.ThreeDVar(
        obs_op=self.source_receptor_operator(),
        prior_mean=self.prior_mean[None],
        # unused by the solve and the posterior (the control's prior is
        # ½‖χ‖²); B = L Lᵀ for consistency
        prior_cov_op=row_operator(
            lx.MatrixLinearOperator(self.prior_chol @ self.prior_chol.T, psd),
            n_state,
        ),
        obs_cov_op=lx.IdentityLinearOperator(
            jax.ShapeDtypeStruct((1, n_flat), dtype)
        ),
        minimiser=vdx.OptimistixMinimiser(
            minimiser or optx.LBFGS(rtol=1e-6, atol=1e-8), max_steps=max_steps
        ),
        parameterization=control,
        cov_solver=lx.AutoLinearSolver(well_posed=True),
    )
    batch = vdx.Batch1D(
        input=einx.rearrange("t o -> 1 1 (t o)", observations),
        mask=jnp.ones((1, 1, n_flat), dtype=dtype),
    )
    return model, batch

FourDVarResult dataclass

FourDVarResult(source: Array, whitened: Array, cost: float, n_steps: int, posterior: PosteriorCovariance | None = None, converged: bool = True, prior_mean: Array | None = None, prior_covariance: Array | None = None, control: SpatialControl | None = None, background: BackgroundAnalysis | None = None)

Outcome of a 4D-Var minimisation.

Attributes:

Name	Type	Description
source	Array	MAP emission time series S* = S_b + L χ*, shape (n_t,).
whitened	Array	Optimal whitened control χ*, of the whole flat control.
cost	float	Final cost value J(χ*).
n_steps	int	Optimiser iterations consumed (-1 if not reported).
posterior	PosteriorCovariance | None	Optional :class:PosteriorCovariance around the MAP, attached when :func:solve_4dvar is called with compute_posterior=True.
converged	bool	optimistix's strict success flag (result == successful), as vardax.MinimiseResult.converged reports it. Appended after posterior to preserve the positional constructor contract. L-BFGS on these stiff problems frequently exhausts max_steps while still returning a good, finite MAP, so converged is False with a finite source is common and not by itself alarming. The genuine failure to watch for is a non-finite source (a diverged solve), which :func:solve_4dvar warns about loudly.
prior_mean	Array | None	Source prior mean S_b the solve used, shape (n_t,) (None when built by hand).
prior_covariance	Array | None	Source prior covariance B = L Lᵀ, shape (n_t, n_t) (None when built by hand).
control	SpatialControl | None	The :class:SpatialControl of a per-cell flux solve, whose source is a flat vector of length control.size, mapped by :attr:source_field. None for the emission time series.
background	BackgroundAnalysis | None	MAP of the initial concentration and the boundary scaling, when the problem has a :class:BackgroundControl.

For an emission time series, the result also carries the Tier V per-event payload (population.PerEventPosterior) for the event emission rate, the time mean Q = wᵀS with w = 1/n_t: its Laplace posterior N(wᵀS*, wᵀP_S w) and its prior N(wᵀS_b, wᵀBw) are exact images of the Gaussian source laws. That time mean means nothing over cells, so a per-cell result raises PayloadError there.

source_field property

source_field: Array

MAP flux coefficient fields [kg m⁻² s⁻¹], shape (K, ny, nx).

Raises:

Type	Description
PayloadError	For an emission time series.

emission_rate property

emission_rate: Array

Time-mean MAP emission rate Q = wᵀS.

emission_std property

emission_std: Array

Laplace posterior sd of Q, √(wᵀP_S w); NaN without the posterior.

samples

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

n draws of Q from the Laplace posterior N(wᵀS*, wᵀP_S w).

The same law as averaging :func:laplace_sample source draws over time, without re-solving for the covariance.

Raises:

Type	Description
PayloadError	Without the posterior (compute_posterior=False), or for a per-cell flux control.

Source code in src/plumax/les_fvm/fourdvar.py

def samples(self, key: jax.Array, n: int) -> jax.Array:
    """n draws of Q from the Laplace posterior N(wᵀS*, wᵀP_S w).

    The same law as averaging :func:`laplace_sample` source draws over
    time, without re-solving for the covariance.

    Raises:
        PayloadError: Without the posterior (``compute_posterior=False``), or
            for a per-cell flux control.
    """
    self._require_time_series()
    if self.posterior is None:
        raise PayloadError(
            "FourDVarResult.samples needs the Laplace posterior; solve with "
            "compute_posterior=True."
        )
    law = self._mean_law(self.source, self.posterior.source_covariance)
    return law.sample(key, n)

prior_logpdf

prior_logpdf(q: Array) -> jax.Array

log π(Q) of the time-mean source prior, N(wᵀS_b, wᵀBw).

Raises:

Type	Description
PayloadError	If the result was built without its prior, or for a per-cell flux control.

Source code in src/plumax/les_fvm/fourdvar.py

def prior_logpdf(self, q: jax.Array) -> jax.Array:
    """log π(Q) of the time-mean source prior, N(wᵀS_b, wᵀBw).

    Raises:
        PayloadError: If the result was built without its prior, or for a
            per-cell flux control.
    """
    self._require_time_series()
    if self.prior_mean is None or self.prior_covariance is None:
        raise PayloadError(
            "FourDVarResult.prior_logpdf: this result carries no prior; "
            "build it with solve_4dvar."
        )
    return self._mean_law(self.prior_mean, self.prior_covariance).log_prob(q)

PosteriorCovariance dataclass

PosteriorCovariance(whitened_covariance: Array, source_covariance: Array, control: SpatialControl | None = None, control_covariance: Array | None = None)

Laplace (Gauss-Newton) posterior covariance of a 4D-Var solution.

A Gaussian approximation N(S*, P_S) to the posterior, built from the Gauss-Newton Hessian of the whitened-space cost at the MAP χ* (design §posterior-covariance). The whitened-space Hessian is

H_GN = I + J̃ᵀ R⁻¹ J̃,   J̃ = ∂(predicted_obs)/∂χ,

so the whitened posterior covariance is P_χ = H_GN⁻¹ and — because S = S_b + L χ — the source-space covariance is P_S = L P_χ Lᵀ. With a :class:BackgroundControl, P_S is the source block of the covariance of the whole flat control.

Attributes:

Name	Type	Description
whitened_covariance	Array	P_χ, shape (n_state, n_state).
source_covariance	Array	P_S, shape (n_source, n_source).
control	SpatialControl | None	The :class:SpatialControl of a per-cell flux problem.
control_covariance	Array | None	Covariance of the whole flat control, shape (n_state, n_state), when it has a :class:BackgroundControl.

source_std property

source_std: Array

Marginal source-space posterior standard deviations √diag(P_S).

bc_scale_std property

bc_scale_std: Array

Posterior sd of the per-face factors, ordered (west, east, south, north).

Raises:

Type	Description
PayloadError	Without a :class:BackgroundControl.

std_field property

std_field: Array

Marginal posterior sd as fields [kg m⁻² s⁻¹], shape (K, ny, nx).

Raises:

Type	Description
PayloadError	For an emission time series.

SpatialControl dataclass

SpatialControl(mask: Array, prior_mean: Array, length_scale_m: float, prior_std: Array, time_basis: Array | None = None, separable: bool = True)

Per-cell surface flux control with a Matérn-3/2 prior.

q(s, t) = Σ_k φ_k(t) q_k(s) on the active cells and 0 elsewhere, with independent coefficient fields sharing the spatial covariance,

q_k(s) ~ N(a_k m(s), B_s),   B_s = Σ_σ ρ Σ_σ,   Φ a = 1.

The coefficients a make the prior mean of the flux m(s) at every save time, so the time basis must span the constant. The flux at the save times then has covariance Φ Φᵀ ⊗ B_s, so a dense temporal prior B_t = L_t L_tᵀ is the basis Φ = L_t.

The correlation ρ depends on separable, and the two are different kernels. With d = s − s′ and h(r) = (1 + √3 r / ℓ) exp(−√3 r / ℓ):

• separable=True is the product kernel ρ = h(|d_y|) h(|d_x|), stored as gaussx.Kronecker(ρ_y, ρ_x) when every cell is active.
• separable=False is the isotropic kernel ρ = h(‖d‖), the dense :func:~plumax.lagrangian.inversion.matern32_covariance.

The prior is Gaussian, so the MAP flux is not constrained to be positive.

Attributes:

Name	Type	Description
mask	Array	Active T-cells, shape (ny, nx). Flat control vectors list them row-major over (y, x), one block per basis function.
prior_mean	Array	Prior mean flux m [kg m⁻² s⁻¹], shape (ny, nx).
length_scale_m	float	Correlation length ℓ [m].
prior_std	Array	Prior standard deviation σ [kg m⁻² s⁻¹], shape (ny, nx).
time_basis	Array | None	Basis Φ at the save times, shape (n_t, K), linearly interpolated between them. None makes the flux constant in time.
separable	bool	Use the product kernel rather than the isotropic one.

Raises:

Type	Description
ValueError	If the shapes disagree, no cell is active, length_scale_m or an active prior_std is not positive, an input is not finite, or the basis does not span the constant.

Examples:

Two active cells of a 2 × 2 grid, constant in time:

>>> import jax.numpy as jnp
>>> control = SpatialControl(
...     mask=jnp.array([[True, False], [False, True]]),
...     prior_mean=jnp.ones((2, 2)),
...     length_scale_m=100.0,
...     prior_std=0.5 * jnp.ones((2, 2)),
... )
>>> control.size
2
>>> control.to_field(jnp.array([2.0, 3.0])).tolist()
[[[2.0, 0.0], [0.0, 3.0]]]

n_active property

n_active: int

Number of active cells.

n_basis property

n_basis: int

Number K of basis functions.

size property

size: int

Length of the flat control vector, n_active · K.

flat_prior_mean

flat_prior_mean() -> jax.Array

Prior mean μ_q of the flat control, shape (size,).

Source code in src/plumax/les_fvm/fourdvar.py

def flat_prior_mean(self) -> jax.Array:
    """Prior mean μ_q of the flat control, shape ``(size,)``."""
    mean = jnp.asarray(self.prior_mean, dtype=float)[np.asarray(self.mask, bool)]
    # μ_q = a ⊗ m
    return jnp.kron(self._mean_coefficients(), mean)

to_field

to_field(flat: Array) -> jax.Array

Map a flat control vector onto the grid, with 0 on inactive cells.

Parameters:

Name	Type	Description	Default
flat	Array	Flat control values, shape (..., size).	required

Returns:

Type	Description
Array	The coefficient fields, shape (..., K, ny, nx).

Source code in src/plumax/les_fvm/fourdvar.py

def to_field(self, flat: jax.Array) -> jax.Array:
    """Map a flat control vector onto the grid, with 0 on inactive cells.

    Args:
        flat: Flat control values, shape ``(..., size)``.

    Returns:
        The coefficient fields, shape ``(..., K, ny, nx)``.
    """
    ny, nx = np.shape(self.mask)
    rows, columns = np.nonzero(np.asarray(self.mask, dtype=bool))
    per_cell = einx.rearrange(
        "... (basis cell) -> ... basis cell", flat, basis=self.n_basis
    )
    field_shape = (*per_cell.shape[:-1], ny, nx)
    return (
        jnp.zeros(field_shape, dtype=flat.dtype)
        .at[..., rows, columns]
        .set(per_cell)
    )

flux_function

flux_function(flat: Array, save_times: Array) -> Callable[[jax.Array], jax.Array]

The flux t -> q(·, t) of a flat control, on the (ny, nx) grid.

Source code in src/plumax/les_fvm/fourdvar.py

def flux_function(
    self, flat: jax.Array, save_times: jax.Array
) -> Callable[[jax.Array], jax.Array]:
    """The flux t -> q(·, t) of a flat control, on the ``(ny, nx)`` grid."""
    fields = self.to_field(flat)
    if self.time_basis is None:
        constant_flux = fields[0]

        def flux(t: jax.Array) -> jax.Array:
            del t
            return constant_flux

        return flux

    basis = einx.rearrange(
        "time basis -> basis time", jnp.asarray(self.time_basis, dtype=flat.dtype)
    )

    def flux(t: jax.Array) -> jax.Array:
        # φ_k(t), linear between the save times
        basis_at_t = jax.vmap(lambda phi: jnp.interp(t, save_times, phi))(basis)
        # q(s, t) = Σ_k φ_k(t) q_k(s)
        return einx.dot("basis, basis y x -> y x", basis_at_t, fields)

    return flux

prior_covariance

prior_covariance(y: Array, x: Array) -> lx.AbstractLinearOperator

Prior covariance of the flat control, I_K ⊗ Σ_σ ρ Σ_σ.

Parameters:

Name	Type	Description	Default
y	Array	T-cell centres along y [m], shape (ny,).	required
x	Array	T-cell centres along x [m], shape (nx,).	required

Returns:

Type	Description
AbstractLinearOperator	The covariance, with ρ a gaussx.Kronecker(ρ_y, ρ_x) when the
AbstractLinearOperator	kernel is separable and every cell is active, dense otherwise.

Source code in src/plumax/les_fvm/fourdvar.py

def prior_covariance(self, y: jax.Array, x: jax.Array) -> lx.AbstractLinearOperator:
    """Prior covariance of the flat control, I_K ⊗ Σ_σ ρ Σ_σ.

    Args:
        y: T-cell centres along y [m], shape ``(ny,)``.
        x: T-cell centres along x [m], shape ``(nx,)``.

    Returns:
        The covariance, with ρ a ``gaussx.Kronecker(ρ_y, ρ_x)`` when the
        kernel is separable and every cell is active, dense otherwise.
    """
    if self._is_kronecker():
        rho_y, rho_x = self._correlation_factors(y, x)
        rho = gx.Kronecker(
            lx.MatrixLinearOperator(rho_y, _PSD_TAGS),
            lx.MatrixLinearOperator(rho_x, _PSD_TAGS),
        )
    else:
        rho = lx.MatrixLinearOperator(self._active_correlation(y, x), _PSD_TAGS)
    # B = (I_K ⊗ Σ_σ) (I_K ⊗ ρ) (I_K ⊗ Σ_σ)
    covariance = self._scale() @ self._over_basis(rho) @ self._scale()
    return lx.TaggedLinearOperator(covariance, lx.positive_semidefinite_tag)

prior_sqrt

prior_sqrt(y: Array, x: Array) -> lx.AbstractLinearOperator

Lower-triangular square root S of :meth:prior_covariance, S Sᵀ = B.

Factorises each block with gaussx.safe_cholesky, whose adaptive jitter survives a smooth Matérn on a fine grid; a Kronecker block needs only its two 1-D factors.

Parameters:

Name	Type	Description	Default
y	Array	T-cell centres along y [m], shape (ny,).	required
x	Array	T-cell centres along x [m], shape (nx,).	required

Returns:

Type	Description
AbstractLinearOperator	S = (I_K ⊗ Σ_σ) (I_K ⊗ chol ρ).

Source code in src/plumax/les_fvm/fourdvar.py

def prior_sqrt(self, y: jax.Array, x: jax.Array) -> lx.AbstractLinearOperator:
    """Lower-triangular square root S of :meth:`prior_covariance`, S Sᵀ = B.

    Factorises each block with ``gaussx.safe_cholesky``, whose adaptive
    jitter survives a smooth Matérn on a fine grid; a Kronecker block needs
    only its two 1-D factors.

    Args:
        y: T-cell centres along y [m], shape ``(ny,)``.
        x: T-cell centres along x [m], shape ``(nx,)``.

    Returns:
        S = (I_K ⊗ Σ_σ) (I_K ⊗ chol ρ).
    """
    if self._is_kronecker():
        rho_y, rho_x = self._correlation_factors(y, x)
        root = gx.Kronecker(
            lx.MatrixLinearOperator(_cholesky(rho_y)),
            lx.MatrixLinearOperator(_cholesky(rho_x)),
        )
    else:
        rho = self._active_correlation(y, x)
        root = lx.MatrixLinearOperator(_cholesky(rho))
    # S = (I_K ⊗ Σ_σ) (I_K ⊗ chol ρ)
    return self._scale() @ self._over_basis(root)

TemporalControl dataclass

TemporalControl(prior_mean: Array, prior_covariance: Array)

Prior of the scalar emission time series at a known source.

The control of :func:build_problem when the forward has a source location.

Attributes:

Name	Type	Description
prior_mean	Array	Source prior mean S_b, shape (n_t,).
prior_covariance	Array	Source prior covariance B, shape (n_t, n_t).

PlumeGrid3D

Bases: Module

Arakawa C-grid + cached interior coordinate arrays.

Attributes

grid : CartesianGrid3D Underlying finitevolX grid. All field operators (advection, diffusion, interpolation, difference) take this as input. x : Float[Array, "nx_interior"] Interior T-point x-coordinates [m]. y : Float[Array, "ny_interior"] Interior T-point y-coordinates [m]. z : Float[Array, "nz_interior"] Interior T-point z-coordinates [m].

Notes

x, y, z are the interior cell-centre coordinates — they drop the one-cell ghost rings on either side of each axis. The sizes therefore match the interior slice shapes used by the advection and diffusion operators, not the raw field-array shapes.

shape property

shape: tuple[int, int, int]

Full field shape (Nz, Ny, Nx) including ghost rings.

interior_shape property

interior_shape: tuple[int, int, int]

Interior shape (nz, ny, nx) excluding ghost rings.

GaussianSource

Bases: Module

Gaussian point source with constant or time-varying emission rate.

Attributes

emission_fn : Callable[[Float[Array, '']], Float[Array, '']] Pure function t -> q(t) returning the instantaneous emission rate [kg/s]. For a constant rate, pass a closure that discards t. density : Float[Array, 'nz ny nx'] Interior-shaped unit-mass spatial density [1/m³] centred at the release point. Integrates to 1/Δx/Δy/Δz in discrete grid units (so that q · density * Δx Δy Δz sums to q).

SurfaceFluxSource

Bases: Module

Per-cell surface flux injected into the lowest interior layer.

∂C/∂t ⊇ q(s, t) / Δz in the lowest layer and 0 above it, so the tendency integrated over the cell volume Δx Δy Δz is the flux through the cell's base area, q Δx Δy [kg/s], whatever the layer thickness.

Attributes:

Name	Type	Description
flux_fn	Callable[[Float[Array, '']], Float[Array, 'ny nx']]	Pure function t -> q(·, t) returning the surface flux [kg m⁻² s⁻¹] on the interior T-cells, shape (ny, nx).
layer_thickness	float	Thickness Δz of the lowest layer [m].
n_layers	int	Number of interior layers.

Examples:

The tendency sits in the lowest layer, scaled by the thickness:

>>> import jax.numpy as jnp
>>> source = SurfaceFluxSource(lambda t: 2.0 * jnp.ones((3, 4)), 4.0, 2)
>>> tendency = source(jnp.asarray(0.0))
>>> tendency.shape
(4, 5, 6)
>>> float(tendency[1, 1, 1]), float(tendency[2, 1, 1])
(0.5, 0.0)

PrescribedWindField

Bases: Module

Wind velocity evaluator on the full [Nz, Ny, Nx] field layout.

The interior field values are populated by _interior_fn; ghost rings are padded with the interior boundary values so downstream advection operators see a well-defined field everywhere.

Parameters

plume_grid : PlumeGrid3D Grid the field lives on. interior_fn : Callable[[float], tuple[Float[Array, 'nz ny nx'], ...]] Pure function t -> (u_int, v_int, w_int) returning interior velocity components at the T/U/V stagger points (the les_fvm convention). See :func:uniform_wind_field for the simplest example.

advection_tendency

advection_tendency(concentration: Float[Array, 'Nz Ny Nx'], u: Float[Array, 'Nz Ny Nx'], v: Float[Array, 'Nz Ny Nx'], w: Float[Array, 'Nz Ny Nx'], plume_grid: PlumeGrid3D, method: str = 'weno5', periodic: tuple[bool, bool] = (False, False)) -> Float[Array, 'Nz Ny Nx']

Flux-form advective tendency -∇·(u C) at interior T-points.

Parameters

concentration : Float[Array, "Nz Ny Nx"] Tracer at T-points. u : Float[Array, "Nz Ny Nx"] x-velocity at U-points (east faces). v : Float[Array, "Nz Ny Nx"] y-velocity at V-points (north faces). w : Float[Array, "Nz Ny Nx"] Vertical velocity collocated at T-points (see grid.py). plume_grid : PlumeGrid3D Grid the fields live on. method : str, default "weno5" Horizontal reconstruction scheme. Passed through to :class:finitevolx.Advection3D. periodic : tuple[bool, bool], default (False, False) (x_periodic, y_periodic). On a periodic axis the normal-velocity ghost face is wrapped (rather than left edge-padded) so the two representations of the periodic seam share one velocity and the open-wall fluxes cancel — otherwise a spatially-varying normal wind breaks exact mass conservation at the seam. See :func:~plumax.les_fvm.boundary.periodic_axes.

Returns

Float[Array, "Nz Ny Nx"] Advective tendency at T-points, zero on every ghost face.

Notes

The horizontal operator runs in finitevolX wall="open" mode, so the lateral domain-wall face fluxes are assembled from the ghost ring (first-order upwind against the wall, high-order method in the interior). The caller must fill the ghost cells with the desired lateral BC (Dirichlet / outflow / periodic) before calling — see :func:~plumax.les_fvm.boundary.apply_boundary_conditions, which :class:~plumax.les_fvm.dynamics.EulerianDispersionRHS runs every step. This is what makes horizontal transport honor bc_x / bc_y.

Source code in src/plumax/les_fvm/advection.py

@eqx.filter_jit
def advection_tendency(
    concentration: Float[Array, "Nz Ny Nx"],
    u: Float[Array, "Nz Ny Nx"],
    v: Float[Array, "Nz Ny Nx"],
    w: Float[Array, "Nz Ny Nx"],
    plume_grid: PlumeGrid3D,
    method: str = "weno5",
    periodic: tuple[bool, bool] = (False, False),
) -> Float[Array, "Nz Ny Nx"]:
    """Flux-form advective tendency ``-∇·(u C)`` at interior T-points.

    Parameters
    ----------
    concentration : Float[Array, "Nz Ny Nx"]
        Tracer at T-points.
    u : Float[Array, "Nz Ny Nx"]
        x-velocity at U-points (east faces).
    v : Float[Array, "Nz Ny Nx"]
        y-velocity at V-points (north faces).
    w : Float[Array, "Nz Ny Nx"]
        Vertical velocity collocated at T-points (see ``grid.py``).
    plume_grid : PlumeGrid3D
        Grid the fields live on.
    method : str, default ``"weno5"``
        Horizontal reconstruction scheme.  Passed through to
        :class:`finitevolx.Advection3D`.
    periodic : tuple[bool, bool], default ``(False, False)``
        ``(x_periodic, y_periodic)``.  On a periodic axis the normal-velocity
        ghost face is *wrapped* (rather than left edge-padded) so the two
        representations of the periodic seam share one velocity and the
        open-wall fluxes cancel — otherwise a spatially-varying normal wind
        breaks exact mass conservation at the seam.  See
        :func:`~plumax.les_fvm.boundary.periodic_axes`.

    Returns
    -------
    Float[Array, "Nz Ny Nx"]
        Advective tendency at T-points, zero on every ghost face.

    Notes
    -----
    The horizontal operator runs in finitevolX ``wall="open"`` mode, so the
    lateral domain-wall face fluxes are assembled from the ghost ring
    (first-order upwind against the wall, high-order ``method`` in the
    interior).  The caller must fill the ghost cells with the desired lateral
    BC (Dirichlet / outflow / periodic) *before* calling — see
    :func:`~plumax.les_fvm.boundary.apply_boundary_conditions`, which
    :class:`~plumax.les_fvm.dynamics.EulerianDispersionRHS` runs every step.
    This is what makes horizontal transport honor ``bc_x`` / ``bc_y``.
    """
    x_periodic, y_periodic = periodic
    if x_periodic:
        # Wrap the x-normal velocity seam: west wall face (u[..., 0]) and east
        # wall face (u[..., -2]) are the same periodic seam and must match.
        u = u.at[:, :, 0].set(u[:, :, -2]).at[:, :, -1].set(u[:, :, 1])
    if y_periodic:
        v = v.at[:, 0, :].set(v[:, -2, :]).at[:, -1, :].set(v[:, 1, :])
    horizontal_op = Advection3D(grid=plume_grid.grid)
    horizontal = horizontal_op(concentration, u, v, method=method, wall="open")
    vertical = vertical_advection_tendency(concentration, w, plume_grid.dz)
    return horizontal + vertical

apply_boundary_conditions

apply_boundary_conditions(field: Float[Array, 'Nz Ny Nx'], horizontal_bc: HorizontalBC, vertical_bc: VerticalBC, plume_grid: PlumeGrid3D) -> Float[Array, 'Nz Ny Nx']

Apply horizontal then vertical BCs to a 3-D tracer field.

Source code in src/plumax/les_fvm/boundary.py

def apply_boundary_conditions(
    field: Float[Array, "Nz Ny Nx"],
    horizontal_bc: HorizontalBC,
    vertical_bc: VerticalBC,
    plume_grid: PlumeGrid3D,
) -> Float[Array, "Nz Ny Nx"]:
    """Apply horizontal then vertical BCs to a 3-D tracer field."""
    out = horizontal_bc(field, dx=plume_grid.dx, dy=plume_grid.dy)
    out = vertical_bc(out, dz=plume_grid.dz)
    return out

build_default_concentration_bc

build_default_concentration_bc(bc_x: str | tuple[str, str] | tuple[tuple[str, float], tuple[str, float]] = ('dirichlet', 'outflow'), bc_y: str | tuple[str, str] | tuple[tuple[str, float], tuple[str, float]] = 'periodic', bc_z: str | tuple[str, str] | tuple[tuple[str, float], tuple[str, float]] = ('neumann', 'neumann')) -> tuple[HorizontalBC, VerticalBC]

Build (HorizontalBC, VerticalBC) from user-facing BC specs.

Each of bc_x, bc_y, bc_z can be:

• "periodic" — periodic on both faces of that axis.
• (west, east) (for bc_x) / (south, north) (bc_y) / (bottom, top) (bc_z) where each entry is either a BC-kind string ("dirichlet", "neumann", "outflow") or a (kind, value) tuple giving the Dirichlet / Neumann target.

Returns

tuple[HorizontalBC, VerticalBC]

Source code in src/plumax/les_fvm/boundary.py

def build_default_concentration_bc(
    bc_x: (str | tuple[str, str] | tuple[tuple[str, float], tuple[str, float]]) = (
        "dirichlet",
        "outflow",
    ),
    bc_y: (
        str | tuple[str, str] | tuple[tuple[str, float], tuple[str, float]]
    ) = "periodic",
    bc_z: (str | tuple[str, str] | tuple[tuple[str, float], tuple[str, float]]) = (
        "neumann",
        "neumann",
    ),
) -> tuple[HorizontalBC, VerticalBC]:
    """Build ``(HorizontalBC, VerticalBC)`` from user-facing BC specs.

    Each of ``bc_x``, ``bc_y``, ``bc_z`` can be:

    - ``"periodic"``       — periodic on both faces of that axis.
    - ``(west, east)`` (for ``bc_x``) / ``(south, north)`` (``bc_y``) /
      ``(bottom, top)`` (``bc_z``) where each entry is either a BC-kind
      string (``"dirichlet"``, ``"neumann"``, ``"outflow"``) or a
      ``(kind, value)`` tuple giving the Dirichlet / Neumann target.

    Returns
    -------
    tuple[HorizontalBC, VerticalBC]
    """
    w_bc, e_bc = _split_horizontal_spec(bc_x, faces=("west", "east"))
    s_bc, n_bc = _split_horizontal_spec(bc_y, faces=("south", "north"))

    bc_set = BoundaryConditionSet(south=s_bc, north=n_bc, west=w_bc, east=e_bc)
    bot_kind, bot_val, top_kind, top_val = _split_vertical_spec(bc_z)
    return (
        HorizontalBC(bc_set=bc_set),
        VerticalBC(
            bottom_kind=bot_kind,
            bottom_value=bot_val,
            top_kind=top_kind,
            top_value=top_val,
        ),
    )

diffusion_tendency

diffusion_tendency(concentration: Float[Array, 'Nz Ny Nx'], eddy_diffusivity: EddyDiffusivity, plume_grid: PlumeGrid3D, periodic: tuple[bool, bool] = (False, False)) -> Float[Array, 'Nz Ny Nx']

Full 3-D eddy diffusion tendency ∇·(K ∇C) at interior T-points.

Parameters

concentration : Float[Array, "Nz Ny Nx"] Tracer at T-points. eddy_diffusivity : EddyDiffusivity (K_h, K_z) diagonal diffusivity. plume_grid : PlumeGrid3D periodic : tuple[bool, bool], default (False, False) (x_periodic, y_periodic). For a field K_h on a periodic axis the coefficient ghost is wrapped (rather than left edge-padded by :meth:EddyDiffusivity.as_arrays) so the paired periodic seam faces share one diffusivity and their diffusive fluxes cancel. A scalar K_h is uniform and already conserves, so it is untouched.

Returns

Float[Array, "Nz Ny Nx"] Diffusive tendency at T-points, zero on every ghost face.

Notes

The horizontal operator runs in finitevolX wall="open" mode, so the lateral wall diffusive fluxes K_h·(C_edge − C_ghost)/dx are assembled from the ghost ring. The caller must fill the ghost cells with the desired lateral BC (Dirichlet / periodic) before calling — see :func:~plumax.les_fvm.boundary.apply_boundary_conditions. This is what makes horizontal diffusion honor bc_x / bc_y.

Source code in src/plumax/les_fvm/diffusion.py

@eqx.filter_jit
def diffusion_tendency(
    concentration: Float[Array, "Nz Ny Nx"],
    eddy_diffusivity: EddyDiffusivity,
    plume_grid: PlumeGrid3D,
    periodic: tuple[bool, bool] = (False, False),
) -> Float[Array, "Nz Ny Nx"]:
    """Full 3-D eddy diffusion tendency ``∇·(K ∇C)`` at interior T-points.

    Parameters
    ----------
    concentration : Float[Array, "Nz Ny Nx"]
        Tracer at T-points.
    eddy_diffusivity : EddyDiffusivity
        ``(K_h, K_z)`` diagonal diffusivity.
    plume_grid : PlumeGrid3D
    periodic : tuple[bool, bool], default ``(False, False)``
        ``(x_periodic, y_periodic)``.  For a **field** ``K_h`` on a periodic
        axis the coefficient ghost is *wrapped* (rather than left edge-padded
        by :meth:`EddyDiffusivity.as_arrays`) so the paired periodic seam
        faces share one diffusivity and their diffusive fluxes cancel.  A
        scalar ``K_h`` is uniform and already conserves, so it is untouched.

    Returns
    -------
    Float[Array, "Nz Ny Nx"]
        Diffusive tendency at T-points, zero on every ghost face.

    Notes
    -----
    The horizontal operator runs in finitevolX ``wall="open"`` mode, so the
    lateral wall diffusive fluxes ``K_h·(C_edge − C_ghost)/dx`` are assembled
    from the ghost ring.  The caller must fill the ghost cells with the
    desired lateral BC (Dirichlet / periodic) *before* calling — see
    :func:`~plumax.les_fvm.boundary.apply_boundary_conditions`.  This is what
    makes horizontal diffusion honor ``bc_x`` / ``bc_y``.
    """
    k_h, k_z = eddy_diffusivity.as_arrays()
    x_periodic, y_periodic = periodic
    if jnp.ndim(k_h) > 0:
        # Field K_h: wrap the coefficient ghost on periodic axes so a periodic
        # seam uses one shared face diffusivity (scalar K_h needs nothing).
        if x_periodic:
            k_h = k_h.at[:, :, 0].set(k_h[:, :, -2]).at[:, :, -1].set(k_h[:, :, 1])
        if y_periodic:
            k_h = k_h.at[:, 0, :].set(k_h[:, -2, :]).at[:, -1, :].set(k_h[:, 1, :])
    horizontal_op = Diffusion3D(grid=plume_grid.grid)
    horizontal = horizontal_op(concentration, k_h, wall="open")
    vertical = vertical_diffusion_tendency(concentration, k_z, plume_grid.dz)
    return horizontal + vertical

mo_eddy_diffusivity

mo_eddy_diffusivity(grid: PlumeGrid3D, surface_layer: SurfaceLayer, pbl_height: float | Float[Array, 'time y x'], *, it: int, k_h: float | Float[Array, 'nz ny nx']) -> EddyDiffusivity

Monin–Obukhov K-profile for K_z, with a caller-supplied K_h.

K_z(z) = κ u* z (1 − min(z/z_i, 1))² / φ_h(z/L).

The cubic-in-depth profile vanishes at the ground and at the boundary-layer top, and φ_h (:func:plumax.met.surface_layer.phi_h) carries the stability: below 1 in unstable air, above 1 in stable air. The min matters: without it (1 − z/z_i)² grows again above z_i, and the free troposphere would mix.

In neutral air the mid-PBL value agrees with :func:pg_eddy_diffusivity class D to within 30 % only for z_i between 110 m and 300 m, taking u from a 10 m wind over z₀ = 0.1 m and PG's reference distance where σ_z = z_i / 2; the ratio does not depend on the wind speed. Outside that range they disagree by construction: K_z(z_i / 2) = κ u z_i / 8 grows linearly with z_i while the PG constant saturates, so the ratio exceeds 4 at z_i = 1000 m. This is a comparison with an empirical curve, not a validity limit of the profile.

Algorithm

1. Take the height above ground of each T-level, z = (k + ½) Δz, the convention :class:~plumax.met.surface_layer.MoninObukhov uses.
2. Take u*, 1/L and z_i at time index it.
3. Evaluate the profile on every T-point of the interior.

Parameters:

Name	Type	Description	Default
grid	PlumeGrid3D	The analysis grid.	required
surface_layer	SurfaceLayer	u* and 1/L on (time, y, x), from :class:~plumax.met.surface_layer.MoninObukhov.	required
pbl_height	float | Float[Array, 'time y x']	Boundary-layer height z_i [m], a scalar or MetField.pbl_height.	required
it	int	Time index into surface_layer (and a field pbl_height).	required
k_h	float | Float[Array, 'nz ny nx']	Horizontal diffusivity [m²/s], passed through: a constant, or :func:smagorinsky_k_h.	required

Returns:

Type	Description
EddyDiffusivity	EddyDiffusivity with vertical on interior T-points.

Raises:

Type	Description
EquinoxRuntimeError	Under jit too, if pbl_height is negative.

Examples:

Neutral air at mid-PBL gives κ u* z_i / 8:

>>> import jax.numpy as jnp
>>> from plumax.les_fvm.grid import make_grid
>>> from plumax.met.surface_layer import SurfaceLayer
>>> grid = make_grid((0.0, 40.0, 4), (0.0, 40.0, 4), (0.0, 160.0, 4))
>>> layer = SurfaceLayer(
...     u_star=jnp.full((1, 4, 4), 0.5),
...     inverse_obukhov_length=jnp.zeros((1, 4, 4)),
...     w_star=jnp.zeros((1, 4, 4)),
... )
>>> eddy = mo_eddy_diffusivity(grid, layer, 120.0, it=0, k_h=1.0)
>>> eddy.vertical.shape
(4, 4, 4)
>>> bool(jnp.allclose(eddy.vertical[1], 0.4 * 0.5 * 120.0 / 8.0))
True

Above the boundary-layer top nothing mixes:

>>> bool((eddy.vertical[3] == 0.0).all())
True

Source code in src/plumax/les_fvm/diffusion.py

def mo_eddy_diffusivity(
    grid: PlumeGrid3D,
    surface_layer: SurfaceLayer,
    pbl_height: float | Float[Array, "time y x"],
    *,
    it: int,
    k_h: float | Float[Array, "nz ny nx"],
) -> EddyDiffusivity:
    """Monin–Obukhov K-profile for ``K_z``, with a caller-supplied ``K_h``.

    K_z(z) = κ u* z (1 − min(z/z_i, 1))² / φ_h(z/L).

    The cubic-in-depth profile vanishes at the ground and at the boundary-layer
    top, and φ_h (:func:`plumax.met.surface_layer.phi_h`) carries the stability:
    below 1 in unstable air, above 1 in stable air. The ``min`` matters: without
    it (1 − z/z_i)² grows again above z_i, and the free troposphere would mix.

    In neutral air the mid-PBL value agrees with :func:`pg_eddy_diffusivity` class D
    to within 30 % only for z_i between 110 m and 300 m, taking u* from a 10 m wind
    over z₀ = 0.1 m and PG's reference distance where σ_z = z_i / 2; the ratio does
    not depend on the wind speed. Outside that range they disagree by construction:
    K_z(z_i / 2) = κ u* z_i / 8 grows linearly with z_i while the PG constant
    saturates, so the ratio exceeds 4 at z_i = 1000 m. This is a comparison with an
    empirical curve, not a validity limit of the profile.

    Algorithm:
        1. Take the height above ground of each T-level, z = (k + ½) Δz, the
           convention :class:`~plumax.met.surface_layer.MoninObukhov` uses.
        2. Take u*, 1/L and z_i at time index ``it``.
        3. Evaluate the profile on every T-point of the interior.

    Args:
        grid: The analysis grid.
        surface_layer: u* and 1/L on ``(time, y, x)``, from
            :class:`~plumax.met.surface_layer.MoninObukhov`.
        pbl_height: Boundary-layer height z_i [m], a scalar or
            ``MetField.pbl_height``.
        it: Time index into ``surface_layer`` (and a field ``pbl_height``).
        k_h: Horizontal diffusivity [m²/s], passed through: a constant, or
            :func:`smagorinsky_k_h`.

    Returns:
        ``EddyDiffusivity`` with ``vertical`` on interior T-points.

    Raises:
        equinox.EquinoxRuntimeError: Under ``jit`` too, if ``pbl_height`` is
            negative.

    Examples:
        Neutral air at mid-PBL gives κ u* z_i / 8:

        >>> import jax.numpy as jnp
        >>> from plumax.les_fvm.grid import make_grid
        >>> from plumax.met.surface_layer import SurfaceLayer
        >>> grid = make_grid((0.0, 40.0, 4), (0.0, 40.0, 4), (0.0, 160.0, 4))
        >>> layer = SurfaceLayer(
        ...     u_star=jnp.full((1, 4, 4), 0.5),
        ...     inverse_obukhov_length=jnp.zeros((1, 4, 4)),
        ...     w_star=jnp.zeros((1, 4, 4)),
        ... )
        >>> eddy = mo_eddy_diffusivity(grid, layer, 120.0, it=0, k_h=1.0)
        >>> eddy.vertical.shape
        (4, 4, 4)
        >>> bool(jnp.allclose(eddy.vertical[1], 0.4 * 0.5 * 120.0 / 8.0))
        True

        Above the boundary-layer top nothing mixes:

        >>> bool((eddy.vertical[3] == 0.0).all())
        True
    """
    # met.field imports les_fvm.grid, so importing met at module level is a cycle
    from plumax.met.surface_layer import KARMAN, phi_h

    # z = z_k − z_ground = (k + ½) Δz
    height = grid.z - grid.z[0] + 0.5 * grid.dz
    u_star = surface_layer.u_star[it]
    inverse_length = surface_layer.inverse_obukhov_length[it]
    depth = jnp.asarray(pbl_height, dtype=height.dtype)
    depth = eqx.error_if(depth, jnp.any(depth < 0.0), "pbl_height must be ≥ 0")
    if depth.ndim == 3:
        depth = depth[it]
    depth = jnp.broadcast_to(depth, u_star.shape)

    # ζ = z / L
    zeta = einx.multiply("level, y x -> level y x", height, inverse_length)
    # z / z_i, infinite for a collapsed layer so that K_z = 0 there
    relative_height = einx.divide("level, y x -> level y x", height, depth)
    # (1 − min(z/z_i, 1))²
    depth_shape = (1.0 - jnp.minimum(relative_height, 1.0)) ** 2
    # κ u* z
    surface_scale = KARMAN * einx.multiply("level, y x -> level y x", height, u_star)
    # K_z = κ u* z (1 − min(z/z_i, 1))² / φ_h(ζ)
    vertical = surface_scale * depth_shape / phi_h(zeta)
    return EddyDiffusivity(horizontal=k_h, vertical=vertical)

pg_eddy_diffusivity

pg_eddy_diffusivity(stability_class: str, wind_speed: float, reference_distance: float = 200.0) -> EddyDiffusivity

Rough K-theory calibration from Pasquill-Gifford σ curves.

For a neutral puff travelling at u under stability class C, the cross-wind and vertical spreads at a reference distance x_ref are σ_y(x_ref) and σ_z(x_ref). Taylor's identity K = σ² / (2 t) with t = x_ref / u gives a simple constant effective diffusivity. The result is the average K over [0, x_ref]; users wanting a sharper match across a wide range of travel distances should pass a spatially-varying field directly.

Parameters

stability_class : str One of "A".."F". wind_speed : float Representative wind speed u [m/s]. Must be > 0. reference_distance : float, default 200.0 Reference downwind distance x_ref [m] at which the σ curves are sampled.

Returns

EddyDiffusivity Scalar (K_h, K_z) calibrated at x_ref.

Raises

ValueError If wind_speed or reference_distance is non-positive.

Source code in src/plumax/les_fvm/diffusion.py

def pg_eddy_diffusivity(
    stability_class: str,
    wind_speed: float,
    reference_distance: float = 200.0,
) -> EddyDiffusivity:
    """Rough K-theory calibration from Pasquill-Gifford σ curves.

    For a neutral puff travelling at ``u`` under stability class ``C``,
    the cross-wind and vertical spreads at a reference distance
    ``x_ref`` are ``σ_y(x_ref)`` and ``σ_z(x_ref)``.  Taylor's identity
    ``K = σ² / (2 t)`` with ``t = x_ref / u`` gives a simple constant
    effective diffusivity.  The result is the *average* K over
    ``[0, x_ref]``; users wanting a sharper match across a wide range
    of travel distances should pass a spatially-varying field directly.

    Parameters
    ----------
    stability_class : str
        One of ``"A"``..``"F"``.
    wind_speed : float
        Representative wind speed ``u`` [m/s].  Must be > 0.
    reference_distance : float, default 200.0
        Reference downwind distance ``x_ref`` [m] at which the σ curves
        are sampled.

    Returns
    -------
    EddyDiffusivity
        Scalar ``(K_h, K_z)`` calibrated at ``x_ref``.

    Raises
    ------
    ValueError
        If ``wind_speed`` or ``reference_distance`` is non-positive.
    """
    if wind_speed <= 0.0:
        raise ValueError(
            f"pg_eddy_diffusivity: wind_speed must be > 0 (got {wind_speed!r})"
        )
    if reference_distance <= 0.0:
        raise ValueError(
            "pg_eddy_diffusivity: reference_distance must be > 0 "
            f"(got {reference_distance!r})"
        )
    params = get_pg_params(stability_class)
    distance = jnp.asarray(reference_distance, dtype=jnp.float32)
    _, sigma_y, sigma_z = calculate_pg_dispersion(distance, params)
    t_ref = reference_distance / wind_speed
    k_h = float(sigma_y**2 / (2.0 * t_ref))
    k_z = float(sigma_z**2 / (2.0 * t_ref))
    return EddyDiffusivity(horizontal=k_h, vertical=k_z)

smagorinsky_k_h

smagorinsky_k_h(wind: PrescribedWindField, grid: PlumeGrid3D, *, c_s: float = 0.2, t: float | Float[Array, '']) -> Float[Array, 'nz ny nx']

Smagorinsky horizontal diffusivity from the resolved deformation.

K_h = (c_s Δ)² |D|, Δ = √(Δx Δy), |D|² = D_T² + D_S²,

with the tension D_T = ∂u/∂x − ∂v/∂y and the shear D_S = ∂v/∂x + ∂u/∂y. These are the two independent components of the trace-free horizontal strain, so uniform translation, solid-body rotation and isotropic horizontal divergence all give K_h = 0. Vertical shear is left out on purpose: K_h only mixes along levels, and ∂u/∂z in the boundary layer would swamp it on a grid with Δz ≪ Δx.

Algorithm

1. Evaluate the wind at t on its C-grid stagger: u on U-points, v on V-points, level by level.
2. D_T lands on T-points; D_S lands on corners, and D_S² is averaged from the four corners to the T-point, so a checkerboard in D_S cannot cancel.
3. K_h = (c_s Δ)² |D| on interior T-points.

The wind's ghost ring is an edge pad, so its normal gradient across the domain edge reads zero and the outermost interior cells are biased low; cells one further in are exact for a linear wind.

Parameters:

Name	Type	Description	Default
wind	PrescribedWindField	The wind field.	required
grid	PlumeGrid3D	The grid it lives on.	required
c_s	float	Smagorinsky constant c_s (> 0).	0.2
t	float | Float[Array, '']	Time [s].	required

Returns:

Type	Description
Float[Array, 'nz ny nx']	K_h [m²/s] on interior T-points, the layout EddyDiffusivity pads.

Raises:

Type	Description
ValueError	If c_s is not positive.

Examples:

A uniform wind has no deformation, and a shear flow has K_h = (c_s Δ)² S:

>>> import jax.numpy as jnp
>>> from plumax.les_fvm.grid import make_grid
>>> from plumax.les_fvm.wind import uniform_wind_field, wind_field_from_callable
>>> grid = make_grid((0.0, 80.0, 8), (0.0, 80.0, 8), (0.0, 40.0, 4))
>>> uniform = uniform_wind_field(grid, u=5.0, v=1.0)
>>> bool((smagorinsky_k_h(uniform, grid, t=0.0) == 0.0).all())
True
>>> shear = wind_field_from_callable(
...     grid, lambda t, x, y, z: (0.1 * y, jnp.zeros_like(y), jnp.zeros_like(y))
... )
>>> k_h = smagorinsky_k_h(shear, grid, t=0.0)
>>> bool(jnp.allclose(k_h[:, 1:-1, 1:-1], (0.2 * 10.0) ** 2 * 0.1, rtol=1e-4))
True

References

Smagorinsky (1963), Mon. Wea. Rev. 91, 99–164.

Source code in src/plumax/les_fvm/diffusion.py

def smagorinsky_k_h(
    wind: PrescribedWindField,
    grid: PlumeGrid3D,
    *,
    c_s: float = 0.2,
    t: float | Float[Array, ""],
) -> Float[Array, "nz ny nx"]:
    """Smagorinsky horizontal diffusivity from the resolved deformation.

    K_h = (c_s Δ)² |D|,  Δ = √(Δx Δy),  |D|² = D_T² + D_S²,

    with the tension D_T = ∂u/∂x − ∂v/∂y and the shear D_S = ∂v/∂x + ∂u/∂y.
    These are the two independent components of the trace-free horizontal strain,
    so uniform translation, solid-body rotation and isotropic horizontal
    divergence all give K_h = 0. Vertical shear is left out on purpose: K_h only
    mixes along levels, and ∂u/∂z in the boundary layer would swamp it on a grid
    with Δz ≪ Δx.

    Algorithm:
        1. Evaluate the wind at ``t`` on its C-grid stagger: u on U-points, v on
           V-points, level by level.
        2. D_T lands on T-points; D_S lands on corners, and D_S² is averaged from
           the four corners to the T-point, so a checkerboard in D_S cannot
           cancel.
        3. K_h = (c_s Δ)² |D| on interior T-points.

    The wind's ghost ring is an edge pad, so its normal gradient across the
    domain edge reads zero and the outermost interior cells are biased low;
    cells one further in are exact for a linear wind.

    Args:
        wind: The wind field.
        grid: The grid it lives on.
        c_s: Smagorinsky constant c_s (> 0).
        t: Time [s].

    Returns:
        K_h [m²/s] on interior T-points, the layout ``EddyDiffusivity`` pads.

    Raises:
        ValueError: If ``c_s`` is not positive.

    Examples:
        A uniform wind has no deformation, and a shear flow has K_h = (c_s Δ)² S:

        >>> import jax.numpy as jnp
        >>> from plumax.les_fvm.grid import make_grid
        >>> from plumax.les_fvm.wind import uniform_wind_field, wind_field_from_callable
        >>> grid = make_grid((0.0, 80.0, 8), (0.0, 80.0, 8), (0.0, 40.0, 4))
        >>> uniform = uniform_wind_field(grid, u=5.0, v=1.0)
        >>> bool((smagorinsky_k_h(uniform, grid, t=0.0) == 0.0).all())
        True
        >>> shear = wind_field_from_callable(
        ...     grid, lambda t, x, y, z: (0.1 * y, jnp.zeros_like(y), jnp.zeros_like(y))
        ... )
        >>> k_h = smagorinsky_k_h(shear, grid, t=0.0)
        >>> bool(jnp.allclose(k_h[:, 1:-1, 1:-1], (0.2 * 10.0) ** 2 * 0.1, rtol=1e-4))
        True

    References:
        Smagorinsky (1963), Mon. Wea. Rev. 91, 99–164.
    """
    if c_s <= 0.0:
        raise ValueError(f"smagorinsky_k_h: c_s must be > 0 (got {c_s!r})")
    dx, dy = grid.dx, grid.dy
    u, v, _ = wind(jnp.asarray(t, dtype=grid.x.dtype))

    def squared_deformation(
        u_level: Float[Array, "Ny Nx"], v_level: Float[Array, "Ny Nx"]
    ) -> Float[Array, "ny nx"]:
        # D_T = ∂u/∂x − ∂v/∂y at T-points
        tension = tensor_strain(u_level, v_level, dx, dy)[1:-1, 1:-1]
        # D_S = ∂v/∂x + ∂u/∂y at corners
        shear = shear_strain(u_level, v_level, dx, dy)
        # ⟨D_S²⟩ at T-points
        shear_squared = avg_xy_bwd(shear**2)
        # |D|² = D_T² + ⟨D_S²⟩
        return tension**2 + shear_squared

    deformation_squared = jax.vmap(squared_deformation)(u[1:-1], v[1:-1])
    # |D| with a zero gradient at rest: √ has an infinite slope at 0, so take it of
    # 1 there (the double-where trick), keeping uniform-wind adjoints finite
    is_at_rest = deformation_squared == 0.0
    safe_squared = jnp.where(is_at_rest, 1.0, deformation_squared)
    deformation = jnp.where(is_at_rest, 0.0, jnp.sqrt(safe_squared))
    # K_h = (c_s Δ)² |D|, with Δ² = Δx Δy
    return c_s**2 * dx * dy * deformation

build_forward

build_forward(*, domain_x: tuple[float, float, int], domain_y: tuple[float, float, int], domain_z: tuple[float, float, int], save_times: Array, source_location: tuple[float, float, float] | None = None, uniform_wind: tuple[float, float, float], control: SpatialControl | None = None, eddy_diffusivity: float | tuple[float, float] = 1.0, source_radius: float | None = None, advection_scheme: str = 'upwind1', receptor_index: Array | None = None, solver_kwargs: dict | None = None, bc_x: str | tuple = ('dirichlet', 'outflow'), bc_y: str | tuple = 'periodic') -> EulerianForward4DVar

Assemble a differentiable :class:EulerianForward4DVar.

The transport configuration mirrors :func:~plumax.les_fvm.simulate_eulerian_dispersion but is frozen except for the emission time series, which becomes the 4D-Var control vector. A uniform wind and scalar / anisotropic constant diffusivity keep the v1 inversion fast and the twin experiment exactly reproducible.

Parameters:

Name	Type	Description	Default
domain_x	tuple[float, float, int]	(min, max, n_interior) along x.	required
domain_y	tuple[float, float, int]	(min, max, n_interior) along y.	required
domain_z	tuple[float, float, int]	(min, max, n_interior) along z.	required
save_times	Array	Observation/integration times [s], shape (n_t,).	required
source_location	tuple[float, float, float] | None	(x, y, z) release point [m] of the emission time series. Exactly one of this and control is given.	None
uniform_wind	tuple[float, float, float]	Constant (u, v, w) [m/s].	required
control	SpatialControl | None	Per-cell surface flux control, injected through :class:~plumax.les_fvm.source.SurfaceFluxSource.	None
eddy_diffusivity	float | tuple[float, float]	Constant scalar K or (K_h, K_z) [m²/s].	1.0
source_radius	float | None	Gaussian source radius [m] (defaults to 2·max(dx,…)).	None
advection_scheme	str	Reconstruction scheme ("upwind1" is the robust, cheap default for inversion; WENO variants are available).	'upwind1'
receptor_index	Array | None	Optional (n_obs, 2) (iy, ix) receptor cells.	None
solver_kwargs	dict | None	Overrides for the diffrax solve.	None
bc_x	str | tuple	West and east boundary conditions, as :func:~plumax.les_fvm.boundary.build_default_concentration_bc takes them; (("dirichlet", c), "outflow") prescribes an inflow concentration c on the west face.	('dirichlet', 'outflow')
bc_y	str | tuple	South and north boundary conditions, likewise.	'periodic'

Returns:

Type	Description
EulerianForward4DVar	A configured :class:EulerianForward4DVar.

Source code in src/plumax/les_fvm/fourdvar.py

def build_forward(
    *,
    domain_x: tuple[float, float, int],
    domain_y: tuple[float, float, int],
    domain_z: tuple[float, float, int],
    save_times: jax.Array,
    source_location: tuple[float, float, float] | None = None,
    uniform_wind: tuple[float, float, float],
    control: SpatialControl | None = None,
    eddy_diffusivity: float | tuple[float, float] = 1.0,
    source_radius: float | None = None,
    advection_scheme: str = "upwind1",
    receptor_index: jax.Array | None = None,
    solver_kwargs: dict | None = None,
    bc_x: str | tuple = ("dirichlet", "outflow"),
    bc_y: str | tuple = "periodic",
) -> EulerianForward4DVar:
    """Assemble a differentiable :class:`EulerianForward4DVar`.

    The transport configuration mirrors
    :func:`~plumax.les_fvm.simulate_eulerian_dispersion` but is frozen except
    for the emission time series, which becomes the 4D-Var control vector.
    A uniform wind and scalar / anisotropic constant diffusivity keep the v1
    inversion fast and the twin experiment exactly reproducible.

    Args:
        domain_x: ``(min, max, n_interior)`` along x.
        domain_y: ``(min, max, n_interior)`` along y.
        domain_z: ``(min, max, n_interior)`` along z.
        save_times: Observation/integration times [s], shape ``(n_t,)``.
        source_location: ``(x, y, z)`` release point [m] of the emission time
            series. Exactly one of this and ``control`` is given.
        uniform_wind: Constant ``(u, v, w)`` [m/s].
        control: Per-cell surface flux control, injected through
            :class:`~plumax.les_fvm.source.SurfaceFluxSource`.
        eddy_diffusivity: Constant scalar ``K`` or ``(K_h, K_z)`` [m²/s].
        source_radius: Gaussian source radius [m] (defaults to ``2·max(dx,…)``).
        advection_scheme: Reconstruction scheme (``"upwind1"`` is the robust,
            cheap default for inversion; WENO variants are available).
        receptor_index: Optional ``(n_obs, 2)`` ``(iy, ix)`` receptor cells.
        solver_kwargs: Overrides for the diffrax solve.
        bc_x: West and east boundary conditions, as
            :func:`~plumax.les_fvm.boundary.build_default_concentration_bc`
            takes them; ``(("dirichlet", c), "outflow")`` prescribes an inflow
            concentration c on the west face.
        bc_y: South and north boundary conditions, likewise.

    Returns:
        A configured :class:`EulerianForward4DVar`.
    """
    # Validate the time grid up front: t0/t1, SaveAt(ts=...) and jnp.interp all
    # assume ≥ 2 strictly-increasing save times. Catching it here gives a clear
    # error instead of an opaque diffrax / interpolation failure later.
    if (source_location is None) == (control is None):
        raise ValueError(
            "build_forward: pass exactly one of `source_location` and `control`."
        )
    st = np.asarray(save_times, dtype=float)
    if st.ndim != 1 or st.size < 2:
        raise ValueError(
            "build_forward: `save_times` must be a 1-D array with ≥ 2 entries "
            f"(got shape {st.shape})."
        )
    if np.any(np.diff(st) <= 0.0):
        raise ValueError("build_forward: `save_times` must be strictly increasing.")

    # Build the grid (and hence the state, save times and control vector) in the
    # canonical float dtype so they all agree: float64 when JAX x64 is enabled,
    # float32 otherwise. A mismatch makes the diffrax checkpoint buffer reject
    # the (promoted) tendency.
    dtype = jnp.result_type(float)
    plume_grid = make_grid(domain_x, domain_y, domain_z, dtype=dtype)
    wind_field = uniform_wind_field(
        plume_grid=plume_grid,
        u=uniform_wind[0],
        v=uniform_wind[1],
        w=uniform_wind[2],
    )
    eddy = make_eddy_diffusivity(eddy_diffusivity)

    import diffrax

    from plumax.les_fvm.simulate import _max_diffusivity, stable_step_bound

    # Reject an invalid diffusivity regardless of the solver / controller: a
    # negative (anti-diffusion) or non-finite K is unstable or guard-defeating
    # even on the adaptive path, so validate it before any branch.
    _k_h, _k_z = _max_diffusivity(eddy)

    # Stability guard for the fixed-step 4D-Var solve. The default path is
    # Heun + ConstantStepSize + dt0=1 (see `concentration_history`); an
    # over-large step silently blows up *inside* the optimisation loop, where
    # it is hardest to diagnose. The explicit CFL bound only applies to an
    # explicit solver taking a fixed step, so the guard runs only when: (a) the
    # controller is fixed — ConstantStepSize (largest step = dt0) or StepTo
    # (largest step = widest prescribed interval) — and (b) the solver is
    # explicit. Adaptive controllers are error-controlled and implicit solvers
    # are stable at these steps, so both are skipped.
    _kw = dict(solver_kwargs or {})
    _solver = _kw.get("solver", diffrax.Heun())
    _controller = _kw.get("stepsize_controller", diffrax.ConstantStepSize())
    _max_step: float | None = None
    if isinstance(_controller, diffrax.ConstantStepSize):
        _max_step = float(_kw.get("dt0", 1.0))
    elif isinstance(_controller, diffrax.StepTo):
        _ts = np.asarray(_controller.ts, dtype=float)
        _max_step = float(np.max(np.diff(_ts))) if _ts.size >= 2 else 0.0
    if _max_step is not None and not isinstance(
        _solver, diffrax.AbstractImplicitSolver
    ):
        _bound = stable_step_bound(
            dx=plume_grid.dx,
            dy=plume_grid.dy,
            dz=plume_grid.dz,
            max_u=abs(float(uniform_wind[0])),
            max_v=abs(float(uniform_wind[1])),
            max_w=abs(float(uniform_wind[2])),
            k_h=_k_h,
            k_z=_k_z,
        )
        if _max_step > _bound:
            raise ValueError(
                f"build_forward: fixed step {_max_step:g} s exceeds the stable "
                f"step ~{_bound:g} s for the fixed-step 4D-Var solve on this "
                f"grid+wind+K. Reduce the step (solver_kwargs 'dt0' or StepTo "
                f"'ts'), soften K, or pass an adaptive stepsize_controller."
            )

    horizontal_bc, vertical_bc = build_default_concentration_bc(
        bc_x=bc_x, bc_y=bc_y, bc_z=("neumann", "neumann")
    )
    save_times = jnp.asarray(save_times, dtype=plume_grid.x.dtype)

    nz, ny, nx = plume_grid.shape
    if control is not None:
        n_layers, *surface_shape = plume_grid.interior_shape
        _check_control_on_grid(control, tuple(surface_shape), save_times.shape[0])

        def source_factory(q_series: jax.Array) -> SurfaceFluxSource:
            flux_fn = control.flux_function(q_series, save_times)
            return SurfaceFluxSource(flux_fn, plume_grid.dz, n_layers)
    else:
        # Build the source *geometry* once, with a placeholder unit rate. This
        # runs the (eager) source-location bounds check and fixes the static
        # ``density`` profile outside the differentiated path; only the emission
        # rate ``q`` flows through the trace, so the per-call factory just swaps
        # ``emission_fn`` and never re-runs the host-side ``float(...)`` validation.
        base_source = make_gaussian_source(
            plume_grid=plume_grid,
            emission_rate=1.0,
            source_location=source_location,
            source_radius=source_radius,
        )

        def source_factory(q_series: jax.Array) -> GaussianSource:
            return GaussianSource(
                emission_fn=_interpolated_emission(q_series, save_times),
                density=base_source.density,
            )

    def rhs_factory(
        q_series: jax.Array, bc_scale: jax.Array | None = None
    ) -> EulerianDispersionRHS:
        source = source_factory(q_series)
        scaled_bc = (
            horizontal_bc
            if bc_scale is None
            else scale_dirichlet_faces(horizontal_bc, bc_scale)
        )
        return EulerianDispersionRHS(
            plume_grid=plume_grid,
            wind_field=wind_field,
            eddy_diffusivity=eddy,
            source=source,
            horizontal_bc=scaled_bc,
            vertical_bc=vertical_bc,
            advection_scheme=advection_scheme,
        )

    observation_op = ColumnObservationOperator(
        dz=plume_grid.dz, receptor_index=receptor_index
    )
    return EulerianForward4DVar(
        rhs_factory=rhs_factory,
        observation_op=observation_op,
        save_times=save_times,
        initial_concentration=jnp.zeros((nz, ny, nx), dtype=plume_grid.x.dtype),
        t0=float(save_times[0]),
        t1=float(save_times[-1]),
        solver_kwargs=dict(solver_kwargs or {}),
        plume_grid=plume_grid,
        control=control,
    )

build_problem

build_problem(*, forward: EulerianForward4DVar, observations: Array, prior_mean: Array | None = None, prior_covariance: Array | None = None, obs_variance: Array | float, control: SpatialControl | TemporalControl | None = None, background: BackgroundControl | None = None) -> FourDVarProblem

Assemble a :class:FourDVarProblem, factorising the prior covariance.

Parameters:

Name	Type	Description	Default
forward	EulerianForward4DVar	The differentiable transport forward.	required
observations	Array	Column observations y, shape (n_t, n_obs).	required
prior_mean	Array | None	Source prior mean S_b, shape (n_t,), for the emission time series. Give either this and prior_covariance, or control.	None
prior_covariance	Array | None	Source prior covariance B, shape (n_t, n_t) (e.g. a Matérn-3/2 temporal prior from :func:plumax.lagrangian.inversion.matern32_covariance). Cholesky -factorised once here.	None
obs_variance	Array | float	Observation-error variance R (scalar or broadcastable to observations); R = R_retr + R_repr.	required
control	SpatialControl | TemporalControl | None	The control and its prior: the same :class:SpatialControl that forward was built with, or a :class:TemporalControl.	None
background	BackgroundControl | None	Adds the initial concentration and the per-face boundary scaling to the control, with a block-diagonal prior. None pins both, and leaves the problem as it was without this argument.	None

Returns:

Type	Description
FourDVarProblem	The configured :class:FourDVarProblem.

Source code in src/plumax/les_fvm/fourdvar.py

def build_problem(
    *,
    forward: EulerianForward4DVar,
    observations: jax.Array,
    prior_mean: jax.Array | None = None,
    prior_covariance: jax.Array | None = None,
    obs_variance: jax.Array | float,
    control: SpatialControl | TemporalControl | None = None,
    background: BackgroundControl | None = None,
) -> FourDVarProblem:
    """Assemble a :class:`FourDVarProblem`, factorising the prior covariance.

    Args:
        forward: The differentiable transport forward.
        observations: Column observations ``y``, shape ``(n_t, n_obs)``.
        prior_mean: Source prior mean ``S_b``, shape ``(n_t,)``, for the emission
            time series. Give either this and ``prior_covariance``, or ``control``.
        prior_covariance: Source prior covariance ``B``, shape ``(n_t, n_t)``
            (e.g. a Matérn-3/2 temporal prior from
            :func:`plumax.lagrangian.inversion.matern32_covariance`). Cholesky
            -factorised once here.
        obs_variance: Observation-error variance ``R`` (scalar or broadcastable
            to ``observations``); ``R = R_retr + R_repr``.
        control: The control and its prior: the same :class:`SpatialControl`
            that ``forward`` was built with, or a :class:`TemporalControl`.
        background: Adds the initial concentration and the per-face boundary
            scaling to the control, with a block-diagonal prior. ``None`` pins
            both, and leaves the problem as it was without this argument.

    Returns:
        The configured :class:`FourDVarProblem`.
    """
    if isinstance(control, SpatialControl):
        problem = _build_spatial_problem(
            forward, observations, prior_mean, prior_covariance, obs_variance, control
        )
    else:
        problem = _build_temporal_problem(
            forward, observations, prior_mean, prior_covariance, obs_variance, control
        )
    return problem if background is None else _with_background(problem, background)

laplace_sample

laplace_sample(problem: FourDVarProblem, whitened_map: Array, key: Array, n_samples: int) -> jax.Array

Draw source-space samples from the Laplace posterior N(S*, P_S).

The draws come from vardax.Posterior.sample on the Gauss–Newton Laplace posterior at the MAP source, S = S* + P_S^½ ξ with ξ ~ N(0, I); P_S is the covariance of :func:posterior_covariance.

Parameters:

Name	Type	Description	Default
problem	FourDVarProblem	The assembled 4D-Var problem.	required
whitened_map	Array	The MAP whitened control χ*, shape (n_state,).	required
key	Array	A jax.random PRNG key.	required
n_samples	int	Number of posterior draws to return.	required

Returns:

Type	Description
Array	Source-space samples, shape (n_samples, n_t), or the flux fields
Array	(n_samples, K, ny, nx) of a :class:SpatialControl problem.

Source code in src/plumax/les_fvm/fourdvar.py

def laplace_sample(
    problem: FourDVarProblem,
    whitened_map: jax.Array,
    key: jax.Array,
    n_samples: int,
) -> jax.Array:
    """Draw source-space samples from the Laplace posterior ``N(S*, P_S)``.

    The draws come from ``vardax.Posterior.sample`` on the Gauss–Newton
    Laplace posterior at the MAP source, ``S = S* + P_S^½ ξ`` with
    ``ξ ~ N(0, I)``; ``P_S`` is the covariance of :func:`posterior_covariance`.

    Args:
        problem: The assembled 4D-Var problem.
        whitened_map: The MAP whitened control ``χ*``, shape ``(n_state,)``.
        key: A ``jax.random`` PRNG key.
        n_samples: Number of posterior draws to return.

    Returns:
        Source-space samples, shape ``(n_samples, n_t)``, or the flux fields
        ``(n_samples, K, ny, nx)`` of a :class:`SpatialControl` problem.
    """
    if not _is_traced(n_samples) and int(n_samples) <= 0:
        raise ValueError("laplace_sample: `n_samples` must be a positive integer.")
    chi_star = jnp.asarray(whitened_map, dtype=float)
    source = problem.source_from_whitened(chi_star)
    posterior = _vardax_laplace(problem, source)
    # S⁽ᵉ⁾ = S* + P_S^½ ξ⁽ᵉ⁾, ξ⁽ᵉ⁾ ~ N(0, I), on the (1, n_state) analysis state
    draws = posterior.sample(key, int(n_samples))
    source_draws = draws[:, 0, : problem.n_source]
    if problem.control is not None:
        return problem.control.to_field(source_draws)
    return source_draws

posterior_covariance

posterior_covariance(problem: FourDVarProblem, whitened_map: Array) -> PosteriorCovariance

Gauss-Newton Laplace posterior covariance around a 4D-Var MAP.

vardax's LaplaceCovariance of the whitened control at whitened_map: the Gauss–Newton precision I + J̃ᵀ R⁻¹ J̃ in χ, mapped to the source as P_S = L P_χ Lᵀ for the prior Cholesky factor L.

Parameters:

Name	Type	Description	Default
problem	FourDVarProblem	The assembled 4D-Var problem.	required
whitened_map	Array	The MAP whitened control χ* (e.g. :attr:FourDVarResult.whitened), shape (n_state,).	required

Returns:

Name	Type	Description
The	PosteriorCovariance	class:PosteriorCovariance holding P_χ and P_S.

Raises:

Type	Description
ImportError	Without the assim extra.

Source code in src/plumax/les_fvm/fourdvar.py

def posterior_covariance(
    problem: FourDVarProblem, whitened_map: jax.Array
) -> PosteriorCovariance:
    """Gauss-Newton Laplace posterior covariance around a 4D-Var MAP.

    vardax's ``LaplaceCovariance`` of the whitened control at
    ``whitened_map``: the Gauss–Newton precision ``I + J̃ᵀ R⁻¹ J̃`` in ``χ``,
    mapped to the source as ``P_S = L P_χ Lᵀ`` for the prior Cholesky factor
    ``L``.

    Args:
        problem: The assembled 4D-Var problem.
        whitened_map: The MAP whitened control ``χ*`` (e.g.
            :attr:`FourDVarResult.whitened`), shape ``(n_state,)``.

    Returns:
        The :class:`PosteriorCovariance` holding ``P_χ`` and ``P_S``.

    Raises:
        ImportError: Without the ``assim`` extra.
    """
    chi = jnp.asarray(whitened_map, dtype=float)
    return _laplace_posterior(problem, problem.source_from_whitened(chi))

solve_4dvar

solve_4dvar(problem: FourDVarProblem, *, method: Literal['lbfgs', 'gauss_newton'] = 'lbfgs', initial_source: Array | None = None, rtol: float = 1e-06, atol: float = 1e-08, max_steps: int = 100, solver: AbstractMinimiser | None = None, compute_posterior: bool = False, outer_iters: int = 5, inner_iters: int = 50) -> FourDVarResult

Minimise the 4D-Var cost in whitened space, via vardax.

method="lbfgs" runs :meth:FourDVarProblem.analysis: vardax.ThreeDVar over the whitened control χ, so the prior is the identity Gaussian and the gradient is the exact discrete adjoint of the FV solver.

method="gauss_newton" runs :meth:FourDVarProblem.incremental_analysis: vardax.IncrementalFourDVar, Gauss–Newton outer iterations with CG inner iterations over the same whitened cost. The forward is affine in the control, so the linearised cost is the cost itself and one outer iteration with a converged CG solve reaches the MAP. Each outer iteration forms the whole Jacobian, one VJP per observation, so it costs more than L-BFGS when the observations are many and the problem is well conditioned.

The observation-space dual (PSAS) is not offered: it belongs to vardax (jejjohnson/vardax#72).

Parameters:

Name	Type	Description	Default
problem	FourDVarProblem	The assembled 4D-Var problem.	required
method	Literal['lbfgs', 'gauss_newton']	The minimiser.	'lbfgs'
initial_source	Array | None	Optional first-guess emission signal S₀ (defaults to the prior mean, i.e. χ₀ = 0), for L-BFGS only: IncrementalFourDVar always starts from the prior mean.	None
rtol	float	Relative tolerance of the optimiser, or of the CG solves.	1e-06
atol	float	Absolute tolerance of the optimiser, or of the CG solves.	1e-08
max_steps	int	Maximum L-BFGS iterations.	100
solver	AbstractMinimiser | None	Optional optimistix minimiser (defaults to optx.LBFGS), for L-BFGS only.	None
compute_posterior	bool	When True, also evaluate the Gauss-Newton Laplace :func:posterior_covariance at the MAP and attach it to the result.	False
outer_iters	int	Gauss–Newton outer iterations.	5
inner_iters	int	CG iterations per Gauss–Newton outer iteration.	50

Returns:

Name	Type	Description
The	FourDVarResult	class:FourDVarResult with the MAP source and diagnostics.
	FourDVarResult	n_steps counts L-BFGS iterations or Gauss–Newton outer iterations.
	FourDVarResult	converged is optimistix's verdict for L-BFGS, and for Gauss–Newton
	FourDVarResult	whether ‖∇𝒥(χ*)‖ ≤ atol + rtol ‖∇𝒥(0)‖, because vardax reports none.

Raises:

Type	Description
ValueError	For an unknown method, "psas" included, or an option the method cannot honour.
ImportError	Without the assim extra (pip install 'plumax[assim]').

Source code in src/plumax/les_fvm/fourdvar.py

def solve_4dvar(
    problem: FourDVarProblem,
    *,
    method: Literal["lbfgs", "gauss_newton"] = "lbfgs",
    initial_source: jax.Array | None = None,
    rtol: float = 1e-6,
    atol: float = 1e-8,
    max_steps: int = 100,
    solver: optx.AbstractMinimiser | None = None,
    compute_posterior: bool = False,
    outer_iters: int = 5,
    inner_iters: int = 50,
) -> FourDVarResult:
    """Minimise the 4D-Var cost in whitened space, via vardax.

    ``method="lbfgs"`` runs :meth:`FourDVarProblem.analysis`:
    ``vardax.ThreeDVar`` over the whitened control ``χ``, so the prior is the
    identity Gaussian and the gradient is the exact discrete adjoint of the FV
    solver.

    ``method="gauss_newton"`` runs :meth:`FourDVarProblem.incremental_analysis`:
    ``vardax.IncrementalFourDVar``, Gauss–Newton outer iterations with CG inner
    iterations over the same whitened cost. The forward is affine in the
    control, so the linearised cost is the cost itself and one outer iteration
    with a converged CG solve reaches the MAP. Each outer iteration forms the
    whole Jacobian, one VJP per observation, so it costs more than L-BFGS when
    the observations are many and the problem is well conditioned.

    The observation-space dual (PSAS) is not offered: it belongs to vardax
    (jejjohnson/vardax#72).

    Args:
        problem: The assembled 4D-Var problem.
        method: The minimiser.
        initial_source: Optional first-guess emission signal ``S₀`` (defaults to
            the prior mean, i.e. ``χ₀ = 0``), for L-BFGS only:
            ``IncrementalFourDVar`` always starts from the prior mean.
        rtol: Relative tolerance of the optimiser, or of the CG solves.
        atol: Absolute tolerance of the optimiser, or of the CG solves.
        max_steps: Maximum L-BFGS iterations.
        solver: Optional optimistix minimiser (defaults to ``optx.LBFGS``), for
            L-BFGS only.
        compute_posterior: When ``True``, also evaluate the Gauss-Newton Laplace
            :func:`posterior_covariance` at the MAP and attach it to the result.
        outer_iters: Gauss–Newton outer iterations.
        inner_iters: CG iterations per Gauss–Newton outer iteration.

    Returns:
        The :class:`FourDVarResult` with the MAP source and diagnostics.
        ``n_steps`` counts L-BFGS iterations or Gauss–Newton outer iterations.
        ``converged`` is optimistix's verdict for L-BFGS, and for Gauss–Newton
        whether ‖∇𝒥(χ*)‖ ≤ atol + rtol ‖∇𝒥(0)‖, because vardax reports none.

    Raises:
        ValueError: For an unknown ``method``, ``"psas"`` included, or an option
            the method cannot honour.
        ImportError: Without the ``assim`` extra (``pip install
            'plumax[assim]'``).
    """
    if method == "gauss_newton":
        if initial_source is not None or solver is not None:
            raise ValueError(
                "solve_4dvar: `initial_source` and `solver` apply to L-BFGS only; "
                "vardax.IncrementalFourDVar starts from the prior mean and solves "
                "with CG."
            )
        flat_map, chi_star, cost, n_steps, converged = _solve_gauss_newton(
            problem, outer_iters, inner_iters, rtol, atol
        )
    elif method == "lbfgs":
        flat_map, chi_star, cost, n_steps, converged = _solve_lbfgs(
            problem, initial_source, rtol, atol, max_steps, solver
        )
    elif method == "psas":
        raise ValueError(
            "solve_4dvar: dual PSAS belongs to vardax and is not available yet "
            "(https://github.com/jejjohnson/vardax/issues/72)."
        )
    else:
        raise ValueError(f"solve_4dvar: unknown method {method!r}.")
    finite = bool(jnp.all(jnp.isfinite(flat_map)))
    # `throw=False` hides the optimiser outcome, so surface optimistix's own
    # verdict (`result == successful`). L-BFGS on these stiff transport
    # problems often exhausts `max_steps` while still returning a good, finite
    # MAP, so a bare `converged is False` is common and not alarming on its
    # own; a non-finite MAP is the hazard to flag.
    if not finite:
        warnings.warn(
            "solve_4dvar: the optimiser returned a non-finite source; the "
            "solve diverged. The posterior is skipped (it would be NaN). Check "
            "the prior conditioning, `obs_variance`, and `initial_source`.",
            RuntimeWarning,
            stacklevel=2,
        )
    # Skip the posterior on a non-finite MAP: the Gauss-Newton precision at a
    # NaN control would attach a NaN covariance.
    posterior = (
        _laplace_posterior(problem, flat_map) if compute_posterior and finite else None
    )
    n_source = problem.n_source
    source_chol = problem.prior_chol[:n_source, :n_source]
    return FourDVarResult(
        source=flat_map[:n_source],
        whitened=chi_star,
        cost=cost,
        n_steps=n_steps,
        converged=converged,
        posterior=posterior,
        prior_mean=problem.prior_mean[:n_source],
        prior_covariance=source_chol @ source_chol.T,
        control=problem.control,
        background=_background_analysis(problem, flat_map),
    )

make_grid

make_grid(domain_x: tuple[float, float, int], domain_y: tuple[float, float, int], domain_z: tuple[float, float, int], dtype: dtype = jnp.float32) -> PlumeGrid3D

Build a :class:PlumeGrid3D from (x_min, x_max, nx_interior) triples.

Parameters

domain_x : tuple (x_min, x_max, nx_interior) Physical x-range + number of interior cells in x. domain_y : tuple (y_min, y_max, ny_interior) Physical y-range + number of interior cells in y. domain_z : tuple (z_min, z_max, nz_interior) Physical z-range + number of interior cells in z. z_min is treated as the ground level (Neumann-wall BC target below). dtype : jnp.dtype, default float32 Coordinate-array dtype.

Returns

PlumeGrid3D

Raises

ValueError If any interior cell count is < 4 (WENO5 stencils need at least four interior cells in every horizontal direction) or if any physical extent is non-positive.

Source code in src/plumax/les_fvm/grid.py

def make_grid(
    domain_x: tuple[float, float, int],
    domain_y: tuple[float, float, int],
    domain_z: tuple[float, float, int],
    dtype: jnp.dtype = jnp.float32,
) -> PlumeGrid3D:
    """Build a :class:`PlumeGrid3D` from ``(x_min, x_max, nx_interior)`` triples.

    Parameters
    ----------
    domain_x : tuple (x_min, x_max, nx_interior)
        Physical x-range + number of interior cells in x.
    domain_y : tuple (y_min, y_max, ny_interior)
        Physical y-range + number of interior cells in y.
    domain_z : tuple (z_min, z_max, nz_interior)
        Physical z-range + number of interior cells in z.  ``z_min`` is
        treated as the ground level (Neumann-wall BC target below).
    dtype : jnp.dtype, default float32
        Coordinate-array dtype.

    Returns
    -------
    PlumeGrid3D

    Raises
    ------
    ValueError
        If any interior cell count is < 4 (WENO5 stencils need at least
        four interior cells in every horizontal direction) or if any
        physical extent is non-positive.
    """
    x_min, x_max, nx_int = domain_x
    y_min, y_max, ny_int = domain_y
    z_min, z_max, nz_int = domain_z

    for label, count in (("x", nx_int), ("y", ny_int), ("z", nz_int)):
        if count < 4:
            raise ValueError(
                f"make_grid: need at least 4 interior cells in {label} "
                f"for WENO5 stencils (got {count})"
            )

    Lx = float(x_max - x_min)
    Ly = float(y_max - y_min)
    Lz = float(z_max - z_min)
    for label, L, expr in (
        ("x", Lx, "x_max - x_min"),
        ("y", Ly, "y_max - y_min"),
        ("z", Lz, "z_max - z_min"),
    ):
        if L <= 0.0:
            raise ValueError(
                f"make_grid: {label}-extent must be positive (got {expr} = {L})"
            )

    grid = CartesianGrid3D.from_interior(
        nx_interior=int(nx_int),
        ny_interior=int(ny_int),
        nz_interior=int(nz_int),
        Lx=Lx,
        Ly=Ly,
        Lz=Lz,
    )
    # Interior cell centres — offset by half a cell from the lower edge,
    # matching the finitevolX convention of co-locating T-points at
    # cell centres and drop ghost cells at i = 0 and i = Nx - 1.
    x = jnp.asarray(x_min + (jnp.arange(nx_int, dtype=dtype) + 0.5) * (Lx / nx_int))
    y = jnp.asarray(y_min + (jnp.arange(ny_int, dtype=dtype) + 0.5) * (Ly / ny_int))
    z = jnp.asarray(z_min + (jnp.arange(nz_int, dtype=dtype) + 0.5) * (Lz / nz_int))
    return PlumeGrid3D(grid=grid, x=x, y=y, z=z)

simulate_eulerian_dispersion

simulate_eulerian_dispersion(*, domain_x: tuple[float, float, int], domain_y: tuple[float, float, int], domain_z: tuple[float, float, int], t_start: float, t_end: float, save_interval: float, emission_rate: float | Callable[[Float[Array, '']], Float[Array, '']], source_location: tuple[float, float, float], source_radius: float | None = None, wind_schedule: WindSchedule | None = None, uniform_wind: tuple[float, float, float] | None = None, wind_field: PrescribedWindField | None = None, eddy_diffusivity: float | tuple[float, float] | str | EddyDiffusivity = 'pg', stability_class: str | None = None, pg_reference_distance: float = 200.0, advection_scheme: AdvectionScheme = 'weno5', bc_x: str | tuple[str, str] | tuple[tuple[str, float], tuple[str, float]] = ('dirichlet', 'outflow'), bc_y: str | tuple[str, str] | tuple[tuple[str, float], tuple[str, float]] = 'periodic', bc_z: str | tuple[str, str] | tuple[tuple[str, float], tuple[str, float]] = ('neumann', 'neumann'), solver: SolverName = 'tsit5', dt0: float = 1.0, rtol: float = 0.001, atol: float = 1e-08, max_steps: int = 100000, max_wind_speed: float | None = None, initial_concentration: Float[Array, 'Nz Ny Nx'] | None = None, seed: int = 0) -> xr.Dataset

Simulate 3-D Eulerian tracer transport under a prescribed wind.

Parameters

domain_x, domain_y, domain_z : tuple (min, max, n_interior) Physical extent + number of interior cells per axis. t_start, t_end : float Simulation window [s]. save_interval : float Save cadence [s]. Saved times are np.arange(t_start, t_end + save_interval, save_interval). emission_rate : float or callable Methane source rate [kg/s] (constant or t -> q(t)). source_location : tuple (x, y, z) Release point [m]. source_radius : float, optional Gaussian source radius [m]. Defaults to 2 · max(dx, dy, dz). wind_schedule : WindSchedule, optional uniform_wind : (u, v, w), optional wind_field : PrescribedWindField, optional One of these three must be provided — they are mutually exclusive. eddy_diffusivity : float | (K_h, K_z) | "pg" | EddyDiffusivity K-theory eddy diffusivity spec. The string "pg" pulls a rough calibration from the PG σ curves — requires stability_class and a non-zero mean wind speed. stability_class : str, optional Required when eddy_diffusivity="pg". pg_reference_distance : float, default 200.0 Reference downwind distance for PG → K-theory conversion. advection_scheme : {"upwind1", "weno3", "weno5", ...}, default "weno5" Horizontal reconstruction scheme. bc_x, bc_y, bc_z : BC specs See :func:boundary.build_default_concentration_bc. solver : {"tsit5", "dopri5", "ssprk3", "heun"}, default "tsit5" diffrax time integrator. "ssprk3" and "heun" are explicit-fixed-step and ignore rtol/atol. dt0 : float Initial / fixed step size [s]. rtol, atol : float Adaptive-step controller tolerances (ignored for fixed-step solvers). max_steps : int Upper bound on diffrax internal steps. max_wind_speed : float, optional Conservative upper bound on each wind component |u|, |v|, |w| [m/s], used only by the fixed-step stability guard. Supply it for a callable wind whose peak the guard cannot otherwise sample reliably; when None the peak is estimated from the field (dense time grid + any WindSchedule knots in the window). initial_concentration : ndarray, optional Interior-shaped (nz, ny, nx) initial tracer field. Defaults to zero. seed : int Currently unused (reserved for future stochastic variants).

Returns

xr.Dataset concentration(time, z, y, x) [kg/m³] and the diagnostic column_concentration(time, y, x) [kg/m²]. The dataset carries model metadata in ds.attrs.

Raises

ValueError For inconsistent wind/source/emission specs.

Source code in src/plumax/les_fvm/simulate.py

def simulate_eulerian_dispersion(
    *,
    # Domain and resolution
    domain_x: tuple[float, float, int],
    domain_y: tuple[float, float, int],
    domain_z: tuple[float, float, int],
    # Time
    t_start: float,
    t_end: float,
    save_interval: float,
    # Source
    emission_rate: float | Callable[[Float[Array, ""]], Float[Array, ""]],
    source_location: tuple[float, float, float],
    source_radius: float | None = None,
    # Flow
    wind_schedule: WindSchedule | None = None,
    uniform_wind: tuple[float, float, float] | None = None,
    wind_field: PrescribedWindField | None = None,
    # SGS / eddy diffusivity
    eddy_diffusivity: (float | tuple[float, float] | str | EddyDiffusivity) = "pg",
    stability_class: str | None = None,
    pg_reference_distance: float = 200.0,
    # Numerics
    advection_scheme: AdvectionScheme = "weno5",
    bc_x: (str | tuple[str, str] | tuple[tuple[str, float], tuple[str, float]]) = (
        "dirichlet",
        "outflow",
    ),
    bc_y: (
        str | tuple[str, str] | tuple[tuple[str, float], tuple[str, float]]
    ) = "periodic",
    bc_z: (str | tuple[str, str] | tuple[tuple[str, float], tuple[str, float]]) = (
        "neumann",
        "neumann",
    ),
    solver: SolverName = "tsit5",
    dt0: float = 1.0,
    rtol: float = 1e-3,
    atol: float = 1e-8,
    max_steps: int = 100_000,
    max_wind_speed: float | None = None,
    # Initial condition
    initial_concentration: Float[Array, "Nz Ny Nx"] | None = None,
    seed: int = 0,
) -> xr.Dataset:
    """Simulate 3-D Eulerian tracer transport under a prescribed wind.

    Parameters
    ----------
    domain_x, domain_y, domain_z : tuple (min, max, n_interior)
        Physical extent + number of interior cells per axis.
    t_start, t_end : float
        Simulation window [s].
    save_interval : float
        Save cadence [s].  Saved times are
        ``np.arange(t_start, t_end + save_interval, save_interval)``.
    emission_rate : float or callable
        Methane source rate [kg/s] (constant or ``t -> q(t)``).
    source_location : tuple (x, y, z)
        Release point [m].
    source_radius : float, optional
        Gaussian source radius [m].  Defaults to ``2 · max(dx, dy, dz)``.
    wind_schedule : WindSchedule, optional
    uniform_wind : (u, v, w), optional
    wind_field : PrescribedWindField, optional
        One of these three must be provided — they are mutually exclusive.
    eddy_diffusivity : float | (K_h, K_z) | "pg" | EddyDiffusivity
        K-theory eddy diffusivity spec.  The string ``"pg"`` pulls a rough
        calibration from the PG σ curves — requires ``stability_class`` and
        a non-zero mean wind speed.
    stability_class : str, optional
        Required when ``eddy_diffusivity="pg"``.
    pg_reference_distance : float, default 200.0
        Reference downwind distance for PG → K-theory conversion.
    advection_scheme : {"upwind1", "weno3", "weno5", ...}, default "weno5"
        Horizontal reconstruction scheme.
    bc_x, bc_y, bc_z : BC specs
        See :func:`boundary.build_default_concentration_bc`.
    solver : {"tsit5", "dopri5", "ssprk3", "heun"}, default "tsit5"
        diffrax time integrator.  ``"ssprk3"`` and ``"heun"`` are
        explicit-fixed-step and ignore ``rtol``/``atol``.
    dt0 : float
        Initial / fixed step size [s].
    rtol, atol : float
        Adaptive-step controller tolerances (ignored for fixed-step solvers).
    max_steps : int
        Upper bound on diffrax internal steps.
    max_wind_speed : float, optional
        Conservative upper bound on each wind component ``|u|, |v|, |w|``
        [m/s], used only by the fixed-step stability guard. Supply it for a
        callable wind whose peak the guard cannot otherwise sample reliably;
        when ``None`` the peak is estimated from the field (dense time grid +
        any ``WindSchedule`` knots in the window).
    initial_concentration : ndarray, optional
        Interior-shaped ``(nz, ny, nx)`` initial tracer field.  Defaults
        to zero.
    seed : int
        Currently unused (reserved for future stochastic variants).

    Returns
    -------
    xr.Dataset
        ``concentration(time, z, y, x)`` [kg/m³] and the diagnostic
        ``column_concentration(time, y, x)`` [kg/m²].  The dataset
        carries model metadata in ``ds.attrs``.

    Raises
    ------
    ValueError
        For inconsistent wind/source/emission specs.
    """
    del seed  # Reserved for future stochastic branch.

    plume_grid = make_grid(domain_x, domain_y, domain_z)

    wf = _resolve_wind_field(
        plume_grid=plume_grid,
        wind_schedule=wind_schedule,
        uniform_wind=uniform_wind,
        wind_field=wind_field,
    )

    source = make_gaussian_source(
        plume_grid=plume_grid,
        emission_rate=emission_rate,
        source_location=source_location,
        source_radius=source_radius,
    )

    eddy = _resolve_eddy_diffusivity(
        eddy_diffusivity=eddy_diffusivity,
        stability_class=stability_class,
        wind_schedule=wind_schedule,
        uniform_wind=uniform_wind,
        pg_reference_distance=pg_reference_distance,
        t_start=t_start,
        t_end=t_end,
    )
    # Validate inputs independently of the solver: a negative / non-finite K is
    # unstable (or guard-defeating) on the adaptive path too, and a malformed
    # `max_wind_speed` would silently disable the fixed-step guard.
    _max_diffusivity(eddy)
    if max_wind_speed is not None and (
        not np.isfinite(max_wind_speed) or max_wind_speed < 0.0
    ):
        raise ValueError(
            f"max_wind_speed must be finite and non-negative (got {max_wind_speed!r})."
        )

    horizontal_bc, vertical_bc = build_default_concentration_bc(
        bc_x=bc_x, bc_y=bc_y, bc_z=bc_z
    )

    rhs = EulerianDispersionRHS(
        plume_grid=plume_grid,
        wind_field=wf,
        eddy_diffusivity=eddy,
        source=source,
        horizontal_bc=horizontal_bc,
        vertical_bc=vertical_bc,
        advection_scheme=advection_scheme,
    )

    concentration0 = _build_initial_concentration(
        plume_grid=plume_grid, initial_concentration=initial_concentration
    )
    save_times = _build_save_times(
        t_start=t_start, t_end=t_end, save_interval=save_interval
    )
    # In-window WindSchedule knots — a piecewise-linear schedule attains its
    # speed extrema at knots, so feed them to the fixed-step stability guard so
    # a fast knot between the dense samples is not missed.
    knot_times: tuple[float, ...] = ()
    if wind_schedule is not None:
        knots = np.asarray(wind_schedule.times, dtype=float)
        knot_times = tuple(float(t) for t in knots if t_start <= t <= t_end)
    solution = _solve(
        rhs=rhs,
        t_start=t_start,
        t_end=t_end,
        dt0=dt0,
        y0=concentration0,
        save_times=save_times,
        solver_name=solver,
        rtol=rtol,
        atol=atol,
        max_steps=max_steps,
        extra_sample_times=knot_times,
        max_wind_speed=max_wind_speed,
    )
    return _to_dataset(
        plume_grid=plume_grid,
        times=np.asarray(solution.ts),
        concentration_history=np.asarray(solution.ys),
        attrs={
            "title": "Eulerian 3-D methane dispersion (finitevolX)",
            "model": "L2: prescribed wind + K-theory diffusivity",
            "advection_scheme": advection_scheme,
            "solver": solver,
        },
    )

stable_step_bound

stable_step_bound(*, dx: float, dy: float, dz: float, max_u: float, max_v: float, max_w: float, k_h: float, k_z: float, cfl_safety: float = 1.0) -> float

Maximum stable step for an explicit fixed-step transport solve.

Advection and diffusion are advanced in the same explicit step, so their tendencies add and the conservative forward-Euler / Heun limit is set by the summed rate::

adv_rate  = |u|/dx + |v|/dy + |w|/dz
diff_rate = 2·(K_h/dx² + K_h/dy² + K_z/dz²)
bound     = cfl_safety / (adv_rate + diff_rate)

Using min(1/adv_rate, 1/diff_rate) instead would permit up to a 2× step when the two rates are comparable. The bound is inf only when the field is both motionless and diffusion-free.

Parameters

dx, dy, dz : float Grid spacing [m]. max_u, max_v, max_w : float Component-wise maximum wind speed |u|, |v|, |w| over the run [m/s]. k_h, k_z : float Maximum horizontal / vertical eddy diffusivity [m²/s]. cfl_safety : float, default 1.0 Multiplier on the raw stability bound. 1.0 is the forward-Euler sum-CFL / FTCS-diffusion limit; the RK schemes used here are at least as stable, so this is a conservative guard.

Returns

float The maximum stable dt0 [s].

Notes

This is a first-order (upwind / forward-Euler) sum-CFL estimate — a conservative safety net that catches gross step-size errors (the dt0-orders-of-magnitude-too-large failure mode). It is not a rigorous stability certificate for the higher-order or low-dissipation reconstructions (e.g. the default weno5, or naive): those semidiscrete operators have different stability regions, and an explicit RK step can amplify their weakly-damped modes even below this bound. Keep dt0 comfortably under the reported value when using such schemes.

Source code in src/plumax/les_fvm/simulate.py

def stable_step_bound(
    *,
    dx: float,
    dy: float,
    dz: float,
    max_u: float,
    max_v: float,
    max_w: float,
    k_h: float,
    k_z: float,
    cfl_safety: float = 1.0,
) -> float:
    """Maximum stable step for an explicit fixed-step transport solve.

    Advection and diffusion are advanced in the *same* explicit step, so
    their tendencies add and the conservative forward-Euler / Heun limit is
    set by the summed rate::

        adv_rate  = |u|/dx + |v|/dy + |w|/dz
        diff_rate = 2·(K_h/dx² + K_h/dy² + K_z/dz²)
        bound     = cfl_safety / (adv_rate + diff_rate)

    Using ``min(1/adv_rate, 1/diff_rate)`` instead would permit up to a 2×
    step when the two rates are comparable. The bound is ``inf`` only when
    the field is both motionless and diffusion-free.

    Parameters
    ----------
    dx, dy, dz : float
        Grid spacing [m].
    max_u, max_v, max_w : float
        Component-wise maximum wind speed |u|, |v|, |w| over the run [m/s].
    k_h, k_z : float
        Maximum horizontal / vertical eddy diffusivity [m²/s].
    cfl_safety : float, default 1.0
        Multiplier on the raw stability bound.  ``1.0`` is the forward-Euler
        sum-CFL / FTCS-diffusion limit; the RK schemes used here are at
        least as stable, so this is a conservative guard.

    Returns
    -------
    float
        The maximum stable ``dt0`` [s].

    Notes
    -----
    This is a first-order (upwind / forward-Euler) sum-CFL estimate — a
    conservative safety net that catches gross step-size errors (the
    ``dt0``-orders-of-magnitude-too-large failure mode). It is **not** a
    rigorous stability certificate for the higher-order or low-dissipation
    reconstructions (e.g. the default ``weno5``, or ``naive``): those
    semidiscrete operators have different stability regions, and an explicit
    RK step can amplify their weakly-damped modes even below this bound. Keep
    ``dt0`` comfortably under the reported value when using such schemes.
    """
    adv_rate = max_u / dx + max_v / dy + max_w / dz
    diff_rate = 2.0 * (k_h / dx**2 + k_h / dy**2 + k_z / dz**2)
    # Advection + diffusion act in the same explicit step, so the stable-step
    # limit is set by the summed rate — 1/(adv_rate + diff_rate) — not by the
    # looser min(1/adv_rate, 1/diff_rate).
    total_rate = adv_rate + diff_rate
    return float("inf") if total_rate <= 0.0 else cfl_safety / total_rate

make_gaussian_source

make_gaussian_source(plume_grid: PlumeGrid3D, emission_rate: float | Callable[[Float[Array, '']], Float[Array, '']], source_location: tuple[float, float, float], source_radius: float | None = None) -> GaussianSource

Build a :class:GaussianSource on plume_grid.

Parameters

plume_grid : PlumeGrid3D emission_rate : float or callable Either a constant rate [kg/s], or a pure function t -> q(t) that returns an instantaneous rate as a scalar JAX array. source_location : tuple (x, y, z) Release point [m]. Must lie inside the interior domain. source_radius : float, optional Isotropic Gaussian radius [m]. Defaults to 2 · max(dx, dy, dz) — wide enough to avoid single-cell aliasing without smearing the plume.

Returns

GaussianSource

Raises

ValueError If emission_rate is a constant that is negative, or if the source location lies outside the interior domain.

Source code in src/plumax/les_fvm/source.py

def make_gaussian_source(
    plume_grid: PlumeGrid3D,
    emission_rate: float | Callable[[Float[Array, ""]], Float[Array, ""]],
    source_location: tuple[float, float, float],
    source_radius: float | None = None,
) -> GaussianSource:
    """Build a :class:`GaussianSource` on ``plume_grid``.

    Parameters
    ----------
    plume_grid : PlumeGrid3D
    emission_rate : float or callable
        Either a constant rate [kg/s], or a pure function ``t -> q(t)``
        that returns an instantaneous rate as a scalar JAX array.
    source_location : tuple (x, y, z)
        Release point [m].  Must lie inside the interior domain.
    source_radius : float, optional
        Isotropic Gaussian radius [m].  Defaults to ``2 · max(dx, dy, dz)``
        — wide enough to avoid single-cell aliasing without smearing the
        plume.

    Returns
    -------
    GaussianSource

    Raises
    ------
    ValueError
        If ``emission_rate`` is a constant that is negative, or if the
        source location lies outside the interior domain.
    """
    x_src, y_src, z_src = (float(c) for c in source_location)
    x_min, x_max = float(plume_grid.x[0]), float(plume_grid.x[-1])
    y_min, y_max = float(plume_grid.y[0]), float(plume_grid.y[-1])
    z_min, z_max = float(plume_grid.z[0]), float(plume_grid.z[-1])
    if not (
        x_min <= x_src <= x_max and y_min <= y_src <= y_max and z_min <= z_src <= z_max
    ):
        raise ValueError(
            "make_gaussian_source: source_location "
            f"{source_location!r} lies outside the interior domain "
            f"x∈[{x_min}, {x_max}], y∈[{y_min}, {y_max}], "
            f"z∈[{z_min}, {z_max}]"
        )

    if callable(emission_rate):
        emission_fn = emission_rate
    else:
        q_const = float(emission_rate)
        if q_const < 0.0:
            raise ValueError(
                "make_gaussian_source: constant `emission_rate` must be "
                f"≥ 0 (got {emission_rate!r})"
            )
        q_const_arr = jnp.asarray(q_const, dtype=plume_grid.x.dtype)

        def emission_fn(t):
            del t
            return q_const_arr

    if source_radius is None:
        source_radius = 2.0 * max(plume_grid.dx, plume_grid.dy, plume_grid.dz)
    else:
        source_radius = float(source_radius)
        if source_radius <= 0.0:
            raise ValueError(
                "make_gaussian_source: `source_radius` must be > 0 "
                f"(got {source_radius!r})"
            )

    X, Y, Z = coord_arrays_from_grid(plume_grid)
    r2 = (X - x_src) ** 2 + (Y - y_src) ** 2 + (Z - z_src) ** 2
    profile = jnp.exp(-r2 / (2.0 * source_radius**2))
    # Normalise so ∫ density dV = 1 over the interior (discrete sum × Δx Δy Δz).
    cell_volume = plume_grid.dx * plume_grid.dy * plume_grid.dz
    total = jnp.sum(profile) * cell_volume
    density = profile / total
    return GaussianSource(emission_fn=emission_fn, density=density)

uniform_wind_field

uniform_wind_field(plume_grid: PlumeGrid3D, u: float = 0.0, v: float = 0.0, w: float = 0.0) -> PrescribedWindField

Space- and time-constant wind field.

Parameters

plume_grid : PlumeGrid3D u, v, w : float, default 0.0 Cartesian wind velocity components [m/s].

Returns

PrescribedWindField

Raises

ValueError If any of u, v, w is non-finite — a NaN/inf wind would make the CFL bound NaN (disabling the fixed-step guard) and corrupt the solve.

Source code in src/plumax/les_fvm/wind.py

def uniform_wind_field(
    plume_grid: PlumeGrid3D,
    u: float = 0.0,
    v: float = 0.0,
    w: float = 0.0,
) -> PrescribedWindField:
    """Space- and time-constant wind field.

    Parameters
    ----------
    plume_grid : PlumeGrid3D
    u, v, w : float, default 0.0
        Cartesian wind velocity components [m/s].

    Returns
    -------
    PrescribedWindField

    Raises
    ------
    ValueError
        If any of ``u``, ``v``, ``w`` is non-finite — a NaN/inf wind would make
        the CFL bound NaN (disabling the fixed-step guard) and corrupt the solve.
    """
    for name, comp in (("u", u), ("v", v), ("w", w)):
        if not math.isfinite(float(comp)):
            raise ValueError(
                f"uniform_wind_field: {name} must be finite (got {comp!r})."
            )
    nz, ny, nx = plume_grid.interior_shape
    u_arr = jnp.full((nz, ny, nx), u, dtype=plume_grid.x.dtype)
    v_arr = jnp.full((nz, ny, nx), v, dtype=plume_grid.x.dtype)
    w_arr = jnp.full((nz, ny, nx), w, dtype=plume_grid.x.dtype)

    def interior_fn(t):
        del t
        return u_arr, v_arr, w_arr

    return PrescribedWindField(plume_grid=plume_grid, interior_fn=interior_fn)

wind_field_from_callable

wind_field_from_callable(plume_grid: PlumeGrid3D, fn: Callable[[Float[Array, ''], Float[Array, 'nz ny nx'], Float[Array, 'nz ny nx'], Float[Array, 'nz ny nx']], tuple[Float[Array, 'nz ny nx'], Float[Array, 'nz ny nx'], Float[Array, 'nz ny nx']]]) -> PrescribedWindField

General-purpose wind field from a user-supplied (t, X, Y, Z) function.

Each velocity component is sampled at its C-grid stagger point so the field the advection operator consumes is correctly located:

• u at U-points — T-point coordinates shifted +dx/2 in x,
• v at V-points — T-point coordinates shifted +dy/2 in y,
• w at T-points (collocated, per the les_fvm convention).

fn is therefore evaluated three times — once per stagger — and only the matching component is kept from each call. A divergence-free analytic wind sampled this way stays discretely divergence-free (sampling all three at T-points would displace u/v by half a cell and introduce spurious tracer sources).

Parameters

plume_grid : PlumeGrid3D fn : callable Pure function (t, X, Y, Z) -> (u, v, w) where each of X, Y, Z is an interior-shaped coordinate array. Each output must have the same interior shape.

Returns

PrescribedWindField

Source code in src/plumax/les_fvm/wind.py

def wind_field_from_callable(
    plume_grid: PlumeGrid3D,
    fn: Callable[
        [
            Float[Array, ""],
            Float[Array, "nz ny nx"],
            Float[Array, "nz ny nx"],
            Float[Array, "nz ny nx"],
        ],
        tuple[
            Float[Array, "nz ny nx"],
            Float[Array, "nz ny nx"],
            Float[Array, "nz ny nx"],
        ],
    ],
) -> PrescribedWindField:
    """General-purpose wind field from a user-supplied ``(t, X, Y, Z)`` function.

    Each velocity component is sampled at its C-grid stagger point so the
    field the advection operator consumes is correctly located:

    - ``u`` at U-points — T-point coordinates shifted ``+dx/2`` in x,
    - ``v`` at V-points — T-point coordinates shifted ``+dy/2`` in y,
    - ``w`` at T-points (collocated, per the ``les_fvm`` convention).

    ``fn`` is therefore evaluated three times — once per stagger — and
    only the matching component is kept from each call.  A divergence-free
    analytic wind sampled this way stays discretely divergence-free
    (sampling all three at T-points would displace ``u``/``v`` by half a
    cell and introduce spurious tracer sources).

    Parameters
    ----------
    plume_grid : PlumeGrid3D
    fn : callable
        Pure function ``(t, X, Y, Z) -> (u, v, w)`` where each of
        ``X, Y, Z`` is an interior-shaped coordinate array.  Each output
        must have the same interior shape.

    Returns
    -------
    PrescribedWindField
    """
    X, Y, Z = coord_arrays_from_grid(plume_grid)
    half_dx = 0.5 * plume_grid.dx
    half_dy = 0.5 * plume_grid.dy
    # Stagger the sampling coordinates: U-points sit half a cell east of
    # the T-point in x, V-points half a cell north in y.
    X_u = X + half_dx
    Y_v = Y + half_dy

    def interior_fn(t):
        u = fn(t, X_u, Y, Z)[0]
        v = fn(t, X, Y_v, Z)[1]
        w = fn(t, X, Y, Z)[2]
        return u, v, w

    return PrescribedWindField(plume_grid=plume_grid, interior_fn=interior_fn)

wind_field_from_schedule

wind_field_from_schedule(plume_grid: PlumeGrid3D, schedule: WindSchedule, w: float = 0.0) -> PrescribedWindField

Time-varying but spatially uniform wind field.

u(t) and v(t) are interpolated from the WindSchedule (shared with gauss_puff); w remains constant in time and space (0 by default).

Parameters

plume_grid : PlumeGrid3D schedule : WindSchedule w : float, default 0.0 Constant vertical velocity [m/s].

Returns

PrescribedWindField

Source code in src/plumax/les_fvm/wind.py

def wind_field_from_schedule(
    plume_grid: PlumeGrid3D,
    schedule: WindSchedule,
    w: float = 0.0,
) -> PrescribedWindField:
    """Time-varying but spatially uniform wind field.

    ``u(t)`` and ``v(t)`` are interpolated from the ``WindSchedule``
    (shared with ``gauss_puff``); ``w`` remains constant in time and
    space (``0`` by default).

    Parameters
    ----------
    plume_grid : PlumeGrid3D
    schedule : WindSchedule
    w : float, default 0.0
        Constant vertical velocity [m/s].

    Returns
    -------
    PrescribedWindField
    """
    nz, ny, nx = plume_grid.interior_shape
    ones = jnp.ones((nz, ny, nx), dtype=plume_grid.x.dtype)
    w_arr = w * ones

    def interior_fn(t):
        u_t, v_t = schedule.wind_at(t)
        return u_t * ones, v_t * ones, w_arr

    return PrescribedWindField(plume_grid=plume_grid, interior_fn=interior_fn)