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 emission term.
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 Methane source term. horizontal_bc : HorizontalBC vertical_bc : VerticalBC advection_scheme : str, default "weno5" Horizontal reconstruction scheme for :func:advection.advection_tendency.

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.

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[[Array], EulerianDispersionRHS], observation_op: ColumnObservationOperator, save_times: Array, initial_concentration: Array, t0: float, t1: float, solver_kwargs: dict = dict())

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 — 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[[Array], EulerianDispersionRHS]`	`q_series -> EulerianDispersionRHS` closing over the static transport configuration; the emission rate is the only free input (linearly interpolated between save times).
`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).

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

    rhs = self.rhs_factory(jnp.asarray(q_series))
    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=self.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)

A strong-constraint 4D-Var problem over a scalar emission time series.

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.

source_from_whitened ¶

source_from_whitened(chi: Array) -> jax.Array

Map whitened control χ back to the emission signal S.

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 emission signal ``S``."""
    return self.prior_mean + self.prior_chol @ chi

cost ¶

cost(chi: Array) -> jax.Array

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

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

FourDVarResult `dataclass` ¶

FourDVarResult(source: Array, whitened: Array, cost: float, n_steps: int, posterior: PosteriorCovariance | None = None, converged: bool = True)

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 `χ*`, shape `(n_t,)`.
`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`	optx's strict success flag (`sol.result == successful`). 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.

PosteriorCovariance `dataclass` ¶

PosteriorCovariance(whitened_covariance: Array, source_covariance: Array)

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

Attributes:

Name	Type	Description
`whitened_covariance`	`Array`	`P_χ`, shape `(n_t, n_t)`.
`source_covariance`	`Array`	`P_S`, shape `(n_t, n_t)`.

source_std `property` ¶

source_std: Array

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

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

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

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)

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], uniform_wind: tuple[float, float, float], 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) -> 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]`	`(x, y, z)` release point [m].	required
`uniform_wind`	`tuple[float, float, float]`	Constant `(u, v, w)` [m/s].	required
`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`

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],
    uniform_wind: tuple[float, float, float],
    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,
) -> 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].
        uniform_wind: Constant ``(u, v, w)`` [m/s].
        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.

    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.
    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=("dirichlet", "outflow"), bc_y="periodic", bc_z=("neumann", "neumann")
    )
    save_times = jnp.asarray(save_times, dtype=plume_grid.x.dtype)

    # 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 rhs_factory(q_series: jax.Array) -> EulerianDispersionRHS:
        source = GaussianSource(
            emission_fn=_interpolated_emission(q_series, save_times),
            density=base_source.density,
        )
        return EulerianDispersionRHS(
            plume_grid=plume_grid,
            wind_field=wind_field,
            eddy_diffusivity=eddy,
            source=source,
            horizontal_bc=horizontal_bc,
            vertical_bc=vertical_bc,
            advection_scheme=advection_scheme,
        )

    nz, ny, nx = plume_grid.shape
    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 {}),
    )

build_problem ¶

build_problem(*, forward: EulerianForward4DVar, observations: Array, prior_mean: Array, prior_covariance: Array, obs_variance: Array | float) -> 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`	Source prior mean `S_b`, shape `(n_t,)`.	required
`prior_covariance`	`Array`	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.	required
`obs_variance`	`Array \| float`	Observation-error variance `R` (scalar or broadcastable to `observations`); `R = R_retr + R_repr`.	required

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,
    prior_covariance: jax.Array,
    obs_variance: jax.Array | float,
) -> 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,)``.
        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``.

    Returns:
        The configured :class:`FourDVarProblem`.
    """
    sb = jnp.asarray(prior_mean, dtype=float)
    b = jnp.asarray(prior_covariance, dtype=float)
    y = jnp.asarray(observations, dtype=float)
    if sb.ndim != 1:
        raise ValueError("build_problem: `prior_mean` must be 1-D (n_t,).")
    n_t = sb.shape[0]
    if b.shape != (n_t, n_t):
        raise ValueError(
            f"build_problem: `prior_covariance` must be ({n_t}, {n_t}), got {b.shape}."
        )
    if y.ndim != 2 or y.shape[0] != n_t:
        raise ValueError(
            f"build_problem: `observations` must be (n_t={n_t}, n_obs), got {y.shape}."
        )
    # Reject non-finite inputs: a NaN would flow through the cost / posterior as
    # a finite-`source` solve with a NaN objective and covariance (the source-
    # only divergence check in `solve_4dvar` cannot catch that).
    for _name, _arr in (
        ("prior_mean", sb),
        ("prior_covariance", b),
        ("observations", y),
    ):
        if not bool(jnp.all(jnp.isfinite(_arr))):
            raise ValueError(
                f"build_problem: `{_name}` must be finite (contains NaN/inf)."
            )
    # Adaptive-jitter Cholesky: gx.safe_cholesky escalates the jitter on failure
    # (JIT-compatible), succeeding for badly-scaled Matérn priors where a fixed
    # 1e-9 kick would not.
    chol = gx.safe_cholesky(
        lx.MatrixLinearOperator(
            b, tags=frozenset({lx.symmetric_tag, lx.positive_semidefinite_tag})
        )
    )

    # Resolve obs_variance to R, the (n_t, n_obs) diagonal. A length-n_t 1-D
    # input is treated as one variance *per overpass* (R_t), so reshape it to a
    # column before broadcasting — otherwise numpy would align it with the
    # trailing n_obs axis. A length-n_obs 1-D input stays per-receptor, and a
    # scalar broadcasts to everything. When n_t == n_obs a bare 1-D vector is
    # genuinely ambiguous between the two, so reject it and require an explicit
    # 2-D shape rather than silently picking one and changing the objective.
    r_in = np.asarray(obs_variance, dtype=float)
    n_obs = y.shape[1]
    if r_in.ndim == 1 and n_t == n_obs and r_in.shape[0] == n_t:
        raise ValueError(
            f"build_problem: `obs_variance` of length {n_t} is ambiguous when "
            f"n_t == n_obs == {n_t} (per-overpass vs per-receptor). Pass an "
            f"explicit ({n_t}, 1) column for per-overpass R_t, a (1, {n_obs}) "
            f"row for per-receptor, or a full ({n_t}, {n_obs}) array."
        )
    if r_in.ndim == 1 and r_in.shape[0] == n_t:
        r_in = r_in[:, None]
    if not np.all(np.isfinite(r_in)) or np.any(r_in <= 0.0):
        # NaN passes ``<= 0`` (every comparison is False), so check finiteness
        # explicitly — a NaN R would give a NaN cost and posterior.
        raise ValueError(
            "build_problem: `obs_variance` entries must be finite and > 0."
        )
    try:
        r = jnp.broadcast_to(jnp.asarray(r_in), y.shape)
    except (ValueError, TypeError) as exc:
        raise ValueError(
            f"build_problem: `obs_variance` shape {np.shape(obs_variance)} is not "
            f"broadcastable to observations {y.shape}."
        ) from exc
    return FourDVarProblem(
        forward=forward,
        prior_mean=sb,
        prior_chol=chol,
        observations=y,
        obs_variance=r,
    )

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

Sampling happens in whitened space — χ ~ N(χ*, P_χ) via a Cholesky factor of P_χ — and each draw is mapped through :meth:FourDVarProblem.source_from_whitened, which is exactly the affine map that turns N(χ*, P_χ) into N(S*, P_S).

Parameters:

Name	Type	Description	Default
`problem`	`FourDVarProblem`	The assembled 4D-Var problem.	required
`whitened_map`	`Array`	The MAP whitened control `χ*`, shape `(n_t,)`.	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)`.

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

    Sampling happens in whitened space — ``χ ~ N(χ*, P_χ)`` via a Cholesky
    factor of ``P_χ`` — and each draw is mapped through
    :meth:`FourDVarProblem.source_from_whitened`, which is exactly the affine
    map that turns ``N(χ*, P_χ)`` into ``N(S*, P_S)``.

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

    Returns:
        Source-space samples, shape ``(n_samples, n_t)``.
    """
    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)
    post = posterior_covariance(problem, chi_star)
    n_t = chi_star.shape[0]
    # Adaptive-jitter Cholesky against round-off in the whitened covariance.
    chol_chi = gx.safe_cholesky(
        lx.MatrixLinearOperator(
            post.whitened_covariance,
            tags=frozenset({lx.symmetric_tag, lx.positive_semidefinite_tag}),
        )
    )
    noise = jax.random.normal(key, (int(n_samples), n_t), dtype=chi_star.dtype)
    chi_samples = chi_star[None, :] + noise @ chol_chi.T
    return jax.vmap(problem.source_from_whitened)(chi_samples)

posterior_covariance ¶

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

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

Builds the Gauss-Newton Hessian H_GN = I + J̃ᵀ R⁻¹ J̃ of the whitened cost at whitened_map and inverts it with :mod:gaussx (the Hessian is wrapped as a symmetric positive-semidefinite :class:lineax.MatrixLinearOperator and inverted via :func:gaussx.inv). The whitened posterior covariance is P_χ = H_GN⁻¹ and the source-space covariance is 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_t,)`.	required

Returns:

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

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.

    Builds the Gauss-Newton Hessian ``H_GN = I + J̃ᵀ R⁻¹ J̃`` of the whitened
    cost at ``whitened_map`` and inverts it with :mod:`gaussx` (the Hessian is
    wrapped as a symmetric positive-semidefinite
    :class:`lineax.MatrixLinearOperator` and inverted via :func:`gaussx.inv`).
    The whitened posterior covariance is ``P_χ = H_GN⁻¹`` and the source-space
    covariance is ``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_t,)``.

    Returns:
        The :class:`PosteriorCovariance` holding ``P_χ`` and ``P_S``.
    """
    chi = jnp.asarray(whitened_map, dtype=float)
    h_gn, _ = _gauss_newton_hessian(problem, chi)
    hessian_op = lx.MatrixLinearOperator(
        h_gn, tags=frozenset({lx.symmetric_tag, lx.positive_semidefinite_tag})
    )
    p_chi = gx.inv(hessian_op).as_matrix()
    # Symmetrise to clean up asymmetric round-off before the congruence map.
    p_chi = gx.symmetrize(p_chi)
    chol = problem.prior_chol  # L with B = L Lᵀ
    p_source = chol @ p_chi @ chol.T
    return PosteriorCovariance(
        whitened_covariance=p_chi,
        source_covariance=gx.symmetrize(p_source),
    )

solve_4dvar ¶

solve_4dvar(problem: FourDVarProblem, *, 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) -> FourDVarResult

Minimise the 4D-Var cost with L-BFGS in whitened space.

Optimises χ (whitened control), so the prior is the identity Gaussian and the gradient — the exact discrete adjoint of the FV solver — flows through jax.value_and_grad of :meth:FourDVarProblem.cost.

Parameters:

Name	Type	Description	Default
`problem`	`FourDVarProblem`	The assembled 4D-Var problem.	required
`initial_source`	`Array \| None`	Optional first-guess emission signal `S₀` (defaults to the prior mean, i.e. `χ₀ = 0`).	`None`
`rtol`	`float`	Relative optimiser tolerance.	`1e-06`
`atol`	`float`	Absolute optimiser tolerance.	`1e-08`
`max_steps`	`int`	Maximum optimiser iterations.	`100`
`solver`	`AbstractMinimiser \| None`	Optional optimistix minimiser (defaults to `optx.LBFGS`).	`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`

Returns:

Name	Type	Description
`The`	`FourDVarResult`	class:`FourDVarResult` with the MAP source and diagnostics.

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

def solve_4dvar(
    problem: FourDVarProblem,
    *,
    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,
) -> FourDVarResult:
    """Minimise the 4D-Var cost with L-BFGS in whitened space.

    Optimises ``χ`` (whitened control), so the prior is the identity Gaussian
    and the gradient — the exact discrete adjoint of the FV solver — flows
    through ``jax.value_and_grad`` of :meth:`FourDVarProblem.cost`.

    Args:
        problem: The assembled 4D-Var problem.
        initial_source: Optional first-guess emission signal ``S₀`` (defaults to
            the prior mean, i.e. ``χ₀ = 0``).
        rtol: Relative optimiser tolerance.
        atol: Absolute optimiser tolerance.
        max_steps: Maximum optimiser iterations.
        solver: Optional optimistix minimiser (defaults to ``optx.LBFGS``).
        compute_posterior: When ``True``, also evaluate the Gauss-Newton Laplace
            :func:`posterior_covariance` at the MAP and attach it to the result.

    Returns:
        The :class:`FourDVarResult` with the MAP source and diagnostics.
    """
    import optimistix as optx

    n_t = problem.prior_mean.shape[0]
    if initial_source is None:
        chi0 = jnp.zeros(n_t)
    else:
        s0 = jnp.asarray(initial_source, dtype=float)
        # χ₀ = L⁻¹ (S₀ − S_b) — solve the triangular system rather than invert.
        chi0 = jax.scipy.linalg.solve_triangular(
            problem.prior_chol, s0 - problem.prior_mean, lower=True
        )

    if solver is None:
        solver = optx.LBFGS(rtol=rtol, atol=atol)
    sol = optx.minimise(
        lambda chi, _args: problem.cost(chi),
        solver,
        chi0,
        max_steps=max_steps,
        throw=False,
    )
    # `throw=False` suppresses the failure exception, so surface the outcome
    # instead of trusting `sol.value` blindly. `converged` is optx's strict
    # success flag; L-BFGS on these stiff transport problems frequently
    # exhausts `max_steps` while still returning a good, finite MAP, so a bare
    # `converged is False` is common and not alarming on its own. The genuine
    # hazard the caller must be warned about loudly is a *diverged* solve whose
    # MAP is non-finite and would silently poison `posterior_covariance`.
    converged = bool(sol.result == optx.RESULTS.successful)
    chi_star = sol.value
    source = problem.source_from_whitened(chi_star)
    finite = bool(jnp.all(jnp.isfinite(source)))
    if not finite:
        warnings.warn(
            f"solve_4dvar: the optimiser returned a non-finite source "
            f"(result={sol.result}); the solve diverged. The posterior is "
            f"skipped (it would be NaN). Check the prior conditioning, "
            f"`obs_variance`, and `initial_source`.",
            RuntimeWarning,
            stacklevel=2,
        )
    # Skip the posterior on a non-finite MAP: evaluating the Gauss-Newton
    # Jacobian + inverse at a NaN control would raise or attach a NaN covariance
    # — the exact downstream poisoning the finiteness check exists to prevent.
    posterior = (
        posterior_covariance(problem, chi_star)
        if compute_posterior and finite
        else None
    )
    return FourDVarResult(
        source=source,
        whitened=chi_star,
        cost=float(problem.cost(chi_star)),
        n_steps=int(sol.stats.get("num_steps", -1)),
        converged=converged,
        posterior=posterior,
    )

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)

Eulerian finite volume¶

les_fvm ¶

Submodules¶

HorizontalBC ¶

advection_variant ¶

VerticalBC ¶

Parameters¶

advection_variant ¶

EddyDiffusivity ¶

Attributes¶

as_arrays ¶

EulerianDispersionRHS ¶

Parameters¶

ColumnObservationOperator dataclass ¶

EulerianForward4DVar dataclass ¶

concentration_history ¶

predict ¶

FourDVarProblem dataclass ¶

source_from_whitened ¶

cost ¶

value_and_grad ¶

FourDVarResult dataclass ¶

PosteriorCovariance dataclass ¶

source_std property ¶

PlumeGrid3D ¶

Attributes¶

Notes¶

shape property ¶

interior_shape property ¶

GaussianSource ¶

Attributes¶

PrescribedWindField ¶

Parameters¶

advection_tendency ¶

Parameters¶

Returns¶

Notes¶

apply_boundary_conditions ¶

build_default_concentration_bc ¶

Returns¶

diffusion_tendency ¶

Parameters¶

Returns¶

Notes¶

pg_eddy_diffusivity ¶

Parameters¶

Returns¶

Raises¶

build_forward ¶

build_problem ¶

laplace_sample ¶

posterior_covariance ¶

solve_4dvar ¶

make_grid ¶

Parameters¶

Returns¶

Raises¶

simulate_eulerian_dispersion ¶

Parameters¶

Returns¶

Raises¶

stable_step_bound ¶

Parameters¶

Returns¶

Notes¶

make_gaussian_source ¶

Parameters¶

Returns¶

Raises¶

uniform_wind_field ¶

Parameters¶

Returns¶

Raises¶

wind_field_from_callable ¶

Parameters¶

Returns¶

wind_field_from_schedule ¶

Parameters¶

Returns¶

ColumnObservationOperator `dataclass` ¶

EulerianForward4DVar `dataclass` ¶

FourDVarProblem `dataclass` ¶

FourDVarResult `dataclass` ¶

PosteriorCovariance `dataclass` ¶

source_std `property` ¶

shape `property` ¶

interior_shape `property` ¶