Gaussian plume

Steady-state Gaussian plume with Briggs dispersion, plus NumPyro emission-rate inference. Tier I.

gauss_plume

Steady-state Gaussian plume model (JAX + NumPyro).

Submodules

• dispersion — Briggs σ_y(x), σ_z(x) coefficients (A-F stability classes)
• plume — coordinate rotation + JIT forward model + xarray wrapper
• plume_rise — Briggs plume rise Δh(x) and the effective release height
• inference — NumPyro Bayesian inference for Q (and optional stability)
• observation — column + averaging-kernel observation shared by every tier
• likelihoods — Gaussian, Student-t and heteroscedastic observation noise
• multi_source — masked-K source slots and their posterior

The inference submodule is imported lazily on first attribute access so that from plumax import gauss_plume does not pull in NumPyro for consumers of only the forward model.

plume_concentration_vmap module-attribute

plume_concentration_vmap = vmap(plume_concentration, in_axes=(0, 0, 0, None, None, None, None, None, None, None), out_axes=0)

Vectorised over receptor locations — x, y, z batched along axis 0.

GaussianNoise

Bases: Noise

Homoscedastic Gaussian noise, ε ~ 𝒩(0, σ²).

log p(y | ŷ) = −½ ((y − ŷ)/σ)² − log σ − ½ log 2π.

The Tier I models default to this with σ = obs_noise_std, so their numbers are unchanged.

Attributes:

Name	Type	Description
std	float | Float[Array, '*#obs']	Noise standard deviation σ > 0, in the units of y; a scalar or an array that broadcasts against the observations.

Raises:

Type	Description
ParameterError	If std is not positive.

Examples:

The log-density at the mean of a unit Gaussian is −½ log 2π:

>>> import jax.numpy as jnp
>>> noise = GaussianNoise(std=1.0)
>>> value = noise.log_prob(jnp.zeros(3), jnp.zeros(3))
>>> value.shape
(3,)
>>> bool(jnp.allclose(value, -0.5 * jnp.log(2.0 * jnp.pi)))
True

A non-positive scale is rejected:

>>> GaussianNoise(std=0.0)
Traceback (most recent call last):
...
ParameterError: std must be > 0, got 0.0

HeteroscedasticNoise

Bases: Noise

Per-observation Gaussian noise, ε_i ~ 𝒩(0, σ_i²), from a retrieval-error map.

log p(y_i | ŷ_i) = −½ ((y_i − ŷ_i)/σ_i)² − log σ_i − ½ log 2π.

Unlike :class:GaussianNoise, σ must give one value per observation, and a mismatch with the observations is an error rather than a broadcast.

Attributes:

Name	Type	Description
std	Float[Array, ' obs']	Per-observation standard deviation σ_i > 0, in the units of y.

Raises:

Type	Description
ParameterError	If any std is not positive, or (from log_prob and the NumPyro site) if its shape differs from the observations'.

Examples:

Doubling σ_i halves the weight of a residual:

>>> import jax.numpy as jnp
>>> noise = HeteroscedasticNoise(std=jnp.array([1.0, 2.0]))
>>> value = noise.log_prob(jnp.array([2.0, 4.0]), jnp.zeros(2))
>>> bool(jnp.allclose(value[0] - value[1], jnp.log(2.0)))
True

Build it from variances, e.g. Instrument.observation_variance():

>>> HeteroscedasticNoise.from_variance(jnp.array([4.0, 9.0])).std.tolist()
[2.0, 3.0]

from_variance classmethod

from_variance(variance: Float[Array, ' obs']) -> HeteroscedasticNoise

The noise with σ_i = √(variance_i).

Source code in src/plumax/gauss_plume/_src/likelihoods.py

@classmethod
def from_variance(cls, variance: Float[Array, " obs"]) -> HeteroscedasticNoise:
    """The noise with σ_i = √(variance_i)."""
    # σ_i = √(R_ii)
    std = jnp.sqrt(jnp.asarray(variance))
    return cls(std=std)

Noise

Bases: Module

Observation-noise model p(y | ŷ) for a Tier I likelihood.

Subclasses define :meth:distribution (the NumPyro distribution, for models) and :meth:log_prob (plain JAX, for losses); they must agree.

distribution abstractmethod

distribution(mean: Float[Array, '*obs']) -> dist.Distribution

The NumPyro distribution of y given the prediction ŷ = mean.

Source code in src/plumax/gauss_plume/_src/likelihoods.py

@abc.abstractmethod
def distribution(self, mean: Float[Array, "*obs"]) -> dist.Distribution:
    """The NumPyro distribution of y given the prediction ŷ = ``mean``."""

log_prob abstractmethod

log_prob(obs: Float[Array, '*obs'], mean: Float[Array, '*obs']) -> Float[Array, '*obs']

Pointwise log-density log p(y | ŷ), without NumPyro.

Source code in src/plumax/gauss_plume/_src/likelihoods.py

@abc.abstractmethod
def log_prob(
    self, obs: Float[Array, "*obs"], mean: Float[Array, "*obs"]
) -> Float[Array, "*obs"]:
    """Pointwise log-density log p(y | ŷ), without NumPyro."""

StudentTNoise

Bases: Noise

Heavy-tailed Student-t noise, ε ~ t_ν(0, s), for outliers near the detection floor.

log p(y | ŷ) = log Γ((ν+1)/2) − log Γ(ν/2) − ½ log(ν π) − log s − (ν+1)/2 · log(1 + ((y − ŷ)/s)² / ν).

ν ≈ 5 is the Tier I design's default for low-SNR retrievals; as ν → ∞ it tends to :class:GaussianNoise with σ = s.

Attributes:

Name	Type	Description
scale	float | Float[Array, '*#obs']	Scale s > 0, in the units of y; a scalar or an array that broadcasts against the observations. The standard deviation is s √(ν/(ν − 2)) for ν > 2.
df	float	Degrees of freedom ν > 0.

Raises:

Type	Description
ParameterError	If scale or df is not positive.

Examples:

A 10-scale outlier costs far less than under a Gaussian:

>>> import jax.numpy as jnp
>>> obs, mean = jnp.array([10.0]), jnp.array([0.0])
>>> student = StudentTNoise(scale=1.0).log_prob(obs, mean)
>>> gaussian = GaussianNoise(std=1.0).log_prob(obs, mean)
>>> bool(student[0] > gaussian[0] + 30.0)
True

Large ν recovers the Gaussian:

>>> near = StudentTNoise(scale=1.0, df=1.0e8).log_prob(obs, mean)
>>> bool(jnp.allclose(near, gaussian, rtol=1.0e-6))
True

GolderDispersion

GolderDispersion(inverse_obukhov_length: Float[Array, ''] | float, z0: float, interpolate: bool = False)

Bases: Module

PG σ_y, σ_z with the class chosen from Monin–Obukhov scales (Golder 1972).

σ_i(x) = σ_i^PG(x; class(1/L, z₀)), i ∈ {y, z}, with the Briggs–PG curves of :func:~plumax.gauss_plume.calculate_briggs_dispersion. With interpolate=True the class index s = :func:golder_stability_index is continuous and σ_i(x) = (1 − w) σ_i^PG(x; ⌊s⌋) + w σ_i^PG(x; ⌊s⌋ + 1), w = s − ⌊s⌋, which is differentiable in L (a plumax modelling choice, not part of Golder). Pass it as dispersion_params to :func:~plumax.gauss_plume.plume_concentration or as stability_class to :func:~plumax.gauss_plume.simulate_plume.

Attributes:

Name	Type	Description
inverse_obukhov_length	Float[Array, '']	1/L [m⁻¹] of the column; zero is neutral.
z0	float	Roughness length z₀ > 0 [m]; clamped to :data:GOLDER_Z0_RANGE.
interpolate	bool	Blend neighbouring PG classes linearly in 1/L (default False: the hard class lookup, as in Golder and MMIF).

Raises:

Type	Description
ParameterError	If z0 is not positive.

Examples:

Neutral air reproduces PG class D exactly:

>>> import jax.numpy as jnp
>>> from plumax.gauss_plume import calculate_briggs_dispersion
>>> from plumax.gauss_plume import get_dispersion_params
>>> distance = jnp.array([100.0, 1000.0])
>>> neutral = GolderDispersion(inverse_obukhov_length=0.0, z0=0.1)
>>> sigma_y, sigma_z = neutral(distance)
>>> pg_y, pg_z = calculate_briggs_dispersion(distance, get_dispersion_params("D"))
>>> bool(jnp.allclose(sigma_y, pg_y) & jnp.allclose(sigma_z, pg_z))
True

The interpolated σ moves smoothly with L inside a class:

>>> import jax
>>> def sigma_z_at_1km(inverse_length):
...     dispersion = GolderDispersion(inverse_length, z0=0.1, interpolate=True)
...     return dispersion(jnp.array(1000.0))[1]
>>> bool(jax.grad(sigma_z_at_1km)(-1.0 / 50.0) < 0.0)
True

Build it from a met column's surface-layer diagnosis:

>>> from plumax.les_fvm.grid import make_grid
>>> from plumax.met import MetField, MoninObukhov
>>> grid = make_grid(
...     (0.0, 400.0, 4), (0.0, 400.0, 4), (0.0, 200.0, 4), dtype=jnp.float64
... )
>>> shape = (1, 4, 4, 4)
>>> field = MetField(
...     plume_grid=grid, times=jnp.zeros(1), u=jnp.full(shape, 5.0),
...     v=jnp.zeros(shape), w=jnp.zeros(shape),
...     temperature=jnp.full(shape, 300.0), pressure=jnp.full(shape, 1.0e5),
...     pbl_height=jnp.full((1, 4, 4), 800.0),
... )
>>> layer = MoninObukhov()(field, z0=0.1, sensible_heat_flux=0.0)
>>> GolderDispersion.from_surface_layer(layer, z0=0.1, it=0, iy=1, ix=2).z0
0.1

References

Golder, D. (1972). Relations among stability parameters in the surface layer. Boundary-Layer Meteorology 3, 47–58. US EPA (2024). MMIF v4.1.1, pasquill_gifford.f90 (pg_golder).

Source code in src/plumax/gauss_plume/_src/mo_dispersion.py

def __init__(
    self,
    inverse_obukhov_length: Float[Array, ""] | float,
    z0: float,
    interpolate: bool = False,
) -> None:
    self.inverse_obukhov_length = jnp.asarray(inverse_obukhov_length)
    self.z0 = float(z0)
    self.interpolate = bool(interpolate)

from_surface_layer classmethod

from_surface_layer(layer: SurfaceLayer, *, z0: float, it: int, iy: int, ix: int, interpolate: bool = False) -> GolderDispersion

Golder dispersion at one column of a :class:plumax.met.SurfaceLayer.

Parameters:

Name	Type	Description	Default
layer	SurfaceLayer	Surface-layer scales from :class:plumax.met.MoninObukhov; only its 1/L is used.	required
z0	float	The roughness length z₀ [m] the layer was diagnosed with.	required
it	int	Met time index.	required
iy	int	Analysis y index.	required
ix	int	Analysis x index.	required
interpolate	bool	See the class attribute.	False

Returns:

Type	Description
GolderDispersion	The dispersion model for that column.

Source code in src/plumax/gauss_plume/_src/mo_dispersion.py

@classmethod
def from_surface_layer(
    cls,
    layer: SurfaceLayer,
    *,
    z0: float,
    it: int,
    iy: int,
    ix: int,
    interpolate: bool = False,
) -> GolderDispersion:
    """Golder dispersion at one column of a :class:`plumax.met.SurfaceLayer`.

    Args:
        layer: Surface-layer scales from :class:`plumax.met.MoninObukhov`;
            only its 1/L is used.
        z0: The roughness length z₀ [m] the layer was diagnosed with.
        it: Met time index.
        iy: Analysis y index.
        ix: Analysis x index.
        interpolate: See the class attribute.

    Returns:
        The dispersion model for that column.
    """
    inverse_length = layer.inverse_obukhov_length[it, iy, ix]
    return cls(inverse_length, z0=z0, interpolate=interpolate)

MaskedSources

Bases: Module

K_max source slots with emission rates, locations and activity indicators.

c(𝐬) = Σ_k a_k Q_k c₁(𝐬; x_k, y_k).

Attributes:

Name	Type	Description
emission_rate	Float[Array, ' source']	Q_k [kg s⁻¹] per slot.
x	Float[Array, ' source']	Source x_k [m] per slot.
y	Float[Array, ' source']	Source y_k [m] per slot.
active	Float[Array, ' source']	Activity a_k per slot, 1 on and 0 off (a float, so a relaxed indicator in [0, 1] works too).

Examples:

Two slots, one on, one off: the field is the first source's alone.

>>> import jax.numpy as jnp
>>> from plumax.gauss_plume.dispersion import BRIGGS_DISPERSION_PARAMS
>>> sources = MaskedSources(
...     emission_rate=jnp.array([0.1, 5.0]), x=jnp.array([0.0, 50.0]),
...     y=jnp.array([0.0, 50.0]), active=jnp.array([1.0, 0.0]),
... )
>>> sources.active_count().tolist()
1.0
>>> receptors = (jnp.array([300.0]), jnp.array([0.0]), jnp.array([1.0]))
>>> params = BRIGGS_DISPERSION_PARAMS["D"]
>>> field = sources.concentration(
...     receptors, source_z=2.0, wind_u=5.0, wind_v=0.0, dispersion_params=params
... )
>>> single = plume_concentration(*receptors, 0.0, 0.0, 2.0, 5.0, 0.0, 0.1, params)
>>> bool(jnp.allclose(field, single))
True

k_max property

k_max: int

Number of slots K_max.

active_count

active_count() -> Float[Array, '']

K = Σ_k a_k.

Source code in src/plumax/gauss_plume/_src/multi_source.py

def active_count(self) -> Float[Array, ""]:
    """K = Σ_k a_k."""
    # K = Σ_k a_k
    return einx.sum("source ->", self.active)

unit_responses

unit_responses(receptor_coords: tuple[Float[Array, ' obs'], Float[Array, ' obs'], Float[Array, ' obs']], *, source_z: float | Float[Array, ''], wind_u: float | Float[Array, ''], wind_v: float | Float[Array, ''], dispersion_params: Float[Array, ' param']) -> Float[Array, 'source obs']

Unit-rate plume c₁(𝐬_i; x_k, y_k) of every slot at every receptor.

Source code in src/plumax/gauss_plume/_src/multi_source.py

def unit_responses(
    self,
    receptor_coords: tuple[
        Float[Array, " obs"], Float[Array, " obs"], Float[Array, " obs"]
    ],
    *,
    source_z: float | Float[Array, ""],
    wind_u: float | Float[Array, ""],
    wind_v: float | Float[Array, ""],
    dispersion_params: Float[Array, " param"],
) -> Float[Array, "source obs"]:
    """Unit-rate plume c₁(𝐬_i; x_k, y_k) of every slot at every receptor."""
    x_obs, y_obs, z_obs = receptor_coords

    def one_source(source_x, source_y):
        return plume_concentration(
            x_obs,
            y_obs,
            z_obs,
            source_x,
            source_y,
            source_z,
            wind_u,
            wind_v,
            1.0,
            dispersion_params,
        )

    return jax.vmap(one_source)(self.x, self.y)

concentration

concentration(receptor_coords: tuple[Float[Array, ' obs'], Float[Array, ' obs'], Float[Array, ' obs']], *, source_z: float | Float[Array, ''], wind_u: float | Float[Array, ''], wind_v: float | Float[Array, ''], dispersion_params: Float[Array, ' param']) -> Float[Array, ' obs']

c(𝐬_i) = Σ_k a_k Q_k c₁(𝐬_i; x_k, y_k) [kg m⁻³].

Source code in src/plumax/gauss_plume/_src/multi_source.py

def concentration(
    self,
    receptor_coords: tuple[
        Float[Array, " obs"], Float[Array, " obs"], Float[Array, " obs"]
    ],
    *,
    source_z: float | Float[Array, ""],
    wind_u: float | Float[Array, ""],
    wind_v: float | Float[Array, ""],
    dispersion_params: Float[Array, " param"],
) -> Float[Array, " obs"]:
    """c(𝐬_i) = Σ_k a_k Q_k c₁(𝐬_i; x_k, y_k) [kg m⁻³]."""
    responses = self.unit_responses(
        receptor_coords,
        source_z=source_z,
        wind_u=wind_u,
        wind_v=wind_v,
        dispersion_params=dispersion_params,
    )
    # a_k Q_k
    weights = self.active * self.emission_rate
    # Σ_k a_k Q_k c₁(𝐬_i; x_k, y_k)
    return einx.dot("source, source obs -> obs", weights, responses)

MultiSourcePosterior dataclass

MultiSourcePosterior(emission_rate: ndarray, x: ndarray, y: ndarray, active: ndarray, background: ndarray)

Posterior draws of a masked-K source model.

Attributes:

Name	Type	Description
emission_rate	ndarray	Q_k per draw and slot [kg s⁻¹], shape (sample, source).
x	ndarray	x_k per draw and slot [m].
y	ndarray	y_k per draw and slot [m].
active	ndarray	a_k per draw and slot (bool).
background	ndarray	Background per draw [kg m⁻³], shape (sample,).

Examples:

Three draws of two slots:

>>> import numpy as np
>>> posterior = MultiSourcePosterior(
...     emission_rate=np.array([[0.1, 0.5], [0.2, 0.6], [0.3, 0.7]]),
...     x=np.zeros((3, 2)), y=np.zeros((3, 2)),
...     active=np.array([[True, False], [True, True], [True, False]]),
...     background=np.zeros(3),
... )
>>> posterior.count_probabilities().round(3).tolist()
[0.0, 0.667, 0.333]
>>> summary = posterior.per_source(1)
>>> summary.probability_active, summary.emission_rate
(0.3333333333333333, 0.6)

k_max property

k_max: int

Number of slots K_max.

active_count

active_count() -> np.ndarray

K = Σ_k a_k per draw, shape (sample,).

Source code in src/plumax/gauss_plume/_src/multi_source.py

def active_count(self) -> np.ndarray:
    """K = Σ_k a_k per draw, shape ``(sample,)``."""
    # K = Σ_k a_k
    return einx.sum("sample source -> sample", np.asarray(self.active, dtype=int))

count_probabilities

count_probabilities() -> np.ndarray

P(K = k | y) for k = 0, …, K_max, shape (K_max + 1,).

Source code in src/plumax/gauss_plume/_src/multi_source.py

def count_probabilities(self) -> np.ndarray:
    """P(K = k | y) for k = 0, …, K_max, shape ``(K_max + 1,)``."""
    counts = np.bincount(self.active_count(), minlength=self.k_max + 1)
    return counts / counts.sum()

saturated_fraction

saturated_fraction() -> float

P(K = K_max | y): the mass at the cap, which drives the RJMCMC promotion.

Source code in src/plumax/gauss_plume/_src/multi_source.py

def saturated_fraction(self) -> float:
    """P(K = K_max | y): the mass at the cap, which drives the RJMCMC promotion."""
    return float(self.count_probabilities()[-1])

per_source

per_source(k: int) -> SourceSummary

Summary of slot k over the draws where it is active.

Raises:

Type	Description
ParameterError	If k is not a slot index.

Source code in src/plumax/gauss_plume/_src/multi_source.py

def per_source(self, k: int) -> SourceSummary:
    """Summary of slot ``k`` over the draws where it is active.

    Raises:
        ParameterError: If ``k`` is not a slot index.
    """
    if not 0 <= k < self.k_max:
        raise ParameterError(f"slot k must be in [0, {self.k_max}), got {k}")
    on = np.asarray(self.active[:, k], dtype=bool)
    probability = float(on.mean())
    if not on.any():
        nan = float("nan")
        return SourceSummary(nan, nan, nan, nan, probability)
    rates = np.asarray(self.emission_rate[on, k])
    return SourceSummary(
        emission_rate=float(rates.mean()),
        emission_std=float(rates.std()),
        x=float(np.asarray(self.x[on, k]).mean()),
        y=float(np.asarray(self.y[on, k]).mean()),
        probability_active=probability,
    )

SourceSummary dataclass

SourceSummary(emission_rate: float, emission_std: float, x: float, y: float, probability_active: float)

Posterior summary of one slot, conditioned on the slot being active.

It satisfies :class:plumax.population.PerEventPosterior, so event_from_posterior(summary, x=summary.x, y=summary.y, …) builds the slot's catalogue event.

Attributes:

Name	Type	Description
emission_rate	float	Posterior mean of Q_k over draws with a_k = 1 [kg s⁻¹]; NaN when the slot is never active.
emission_std	float	Posterior standard deviation of Q_k over those draws.
x	float	Posterior mean of x_k over those draws [m].
y	float	Posterior mean of y_k over those draws [m].
probability_active	float	Posterior probability P(a_k = 1 | y).

ColumnObservation

Bases: Module

Averaging-kernel-weighted column at a set of receptors.

y_i = A_i Σ_k h_k c_{k,i}.

Attributes:

Name	Type	Description
layer_weights	Float[Array, ' level']	Level weights h_k, the layer thicknesses [m] for a column integral.
averaging_kernel	Float[Array, ' obs'] | None	Per-receptor scalar averaging kernel A_i. None is the identity kernel.
receptor_index	Int[Array, 'obs 2'] | None	(iy, ix) interior-cell index of each receptor on the field :meth:from_field takes. None observes every column, flattened row-major over (y, x).

Examples:

Two levels 10 m thick, every column of a uniform field:

>>> import jax.numpy as jnp
>>> observation = ColumnObservation(layer_weights=jnp.full(2, 10.0))
>>> observation.from_field(jnp.ones((2, 3, 4))).shape
(12,)

Two receptors with an averaging kernel:

>>> observation = ColumnObservation(
...     layer_weights=jnp.full(2, 10.0),
...     averaging_kernel=jnp.array([1.0, 0.5]),
...     receptor_index=jnp.array([[0, 0], [2, 3]]),
... )
>>> observation.from_field(jnp.ones((2, 3, 4))).tolist()
[20.0, 10.0]

Trapezoid weights for profiles on irregular levels:

>>> levels = jnp.array([0.0, 10.0, 30.0])
>>> ColumnObservation.trapezoid(levels).layer_weights.tolist()
[5.0, 15.0, 10.0]

trapezoid classmethod

trapezoid(levels: Float[Array, ' level'], averaging_kernel: Float[Array, ' obs'] | None = None) -> ColumnObservation

Column over profiles on strictly increasing levels, trapezoid rule.

h_0 = Δz_0 / 2, h_k = (Δz_{k-1} + Δz_k) / 2, h_{K-1} = Δz_{K-2} / 2,

so Σ_k h_k c_k is jnp.trapezoid(c, x=levels).

Parameters:

Name	Type	Description	Default
levels	Float[Array, ' level']	Level heights z_k [m].	required
averaging_kernel	Float[Array, ' obs'] | None	Per-receptor scalar averaging kernel A_i.	None

Returns:

Type	Description
ColumnObservation	The observation, with no receptor index (it takes profiles).

Source code in src/plumax/gauss_plume/_src/observation.py

@classmethod
def trapezoid(
    cls,
    levels: Float[Array, " level"],
    averaging_kernel: Float[Array, " obs"] | None = None,
) -> ColumnObservation:
    """Column over profiles on strictly increasing levels, trapezoid rule.

    h_0 = Δz_0 / 2,  h_k = (Δz_{k-1} + Δz_k) / 2,  h_{K-1} = Δz_{K-2} / 2,

    so Σ_k h_k c_k is ``jnp.trapezoid(c, x=levels)``.

    Args:
        levels: Level heights z_k [m].
        averaging_kernel: Per-receptor scalar averaging kernel A_i.

    Returns:
        The observation, with no receptor index (it takes profiles).
    """
    # Δz_k = z_{k+1} − z_k
    spacing = jnp.diff(levels)
    # h_k = (Δz_{k-1} + Δz_k) / 2, with Δz_{-1} = Δz_{K-1} = 0
    layer_weights = 0.5 * (jnp.pad(spacing, (0, 1)) + jnp.pad(spacing, (1, 0)))
    return cls(layer_weights=layer_weights, averaging_kernel=averaging_kernel)

for_grid classmethod

for_grid(plume_grid: PlumeGrid3D, *, instrument: Instrument | None = None, receptor_index: Int[Array, 'obs 2'] | None = None) -> ColumnObservation

Column over a les_fvm analysis grid, midpoint rule (h_k = Δz).

An instrument places each receptor in the T-cell that contains it and brings its averaging kernel. Without one, receptor_index (or every column) is observed with the identity kernel.

Parameters:

Name	Type	Description	Default
plume_grid	PlumeGrid3D	The analysis grid the field lives on.	required
instrument	Instrument | None	Receptors (x, y) [m] and their averaging kernel.	None
receptor_index	Int[Array, 'obs 2'] | None	(iy, ix) interior-cell indices, used when no instrument is given.	None

Returns:

Type	Description
ColumnObservation	The observation for interior fields of shape (nz, ny, nx).

Raises:

Type	Description
DomainError	If a receptor lies outside the grid.
ValueError	If both instrument and receptor_index are given.

Source code in src/plumax/gauss_plume/_src/observation.py

@classmethod
def for_grid(
    cls,
    plume_grid: PlumeGrid3D,
    *,
    instrument: Instrument | None = None,
    receptor_index: Int[Array, "obs 2"] | None = None,
) -> ColumnObservation:
    """Column over a ``les_fvm`` analysis grid, midpoint rule (h_k = Δz).

    An instrument places each receptor in the T-cell that contains it and
    brings its averaging kernel. Without one, ``receptor_index`` (or every
    column) is observed with the identity kernel.

    Args:
        plume_grid: The analysis grid the field lives on.
        instrument: Receptors ``(x, y)`` [m] and their averaging kernel.
        receptor_index: ``(iy, ix)`` interior-cell indices, used when no
            instrument is given.

    Returns:
        The observation for interior fields of shape ``(nz, ny, nx)``.

    Raises:
        DomainError: If a receptor lies outside the grid.
        ValueError: If both ``instrument`` and ``receptor_index`` are given.
    """
    nz = plume_grid.interior_shape[0]
    layer_weights = jnp.full(nz, plume_grid.dz, dtype=plume_grid.z.dtype)
    if instrument is None:
        return cls(layer_weights=layer_weights, receptor_index=receptor_index)
    if receptor_index is not None:
        raise ValueError(
            "ColumnObservation.for_grid: pass `instrument` or "
            "`receptor_index`, not both."
        )
    receptors = np.asarray(instrument.receptors, dtype=float)
    ix = _containing_cell(receptors[:, 0], np.asarray(plume_grid.x), "x")
    iy = _containing_cell(receptors[:, 1], np.asarray(plume_grid.y), "y")
    index = jnp.asarray(np.column_stack([iy, ix]))
    return cls(
        layer_weights=layer_weights,
        averaging_kernel=instrument.ak,
        receptor_index=index,
    )

column

column(profiles: Float[Array, 'obs level']) -> Float[Array, ' obs']

Apply the observation to per-receptor profiles c_{k,i}.

Parameters:

Name	Type	Description	Default
profiles	Float[Array, 'obs level']	Concentration at each receptor and level.	required

Returns:

Type	Description
Float[Array, ' obs']	The observation y_i.

Source code in src/plumax/gauss_plume/_src/observation.py

def column(self, profiles: Float[Array, "obs level"]) -> Float[Array, " obs"]:
    """Apply the observation to per-receptor profiles c_{k,i}.

    Args:
        profiles: Concentration at each receptor and level.

    Returns:
        The observation y_i.
    """
    # Σ_k h_k c_{k,i}
    column = einx.dot("obs level, level -> obs", profiles, self.layer_weights)
    if self.averaging_kernel is None:
        return column
    # y_i = A_i Σ_k h_k c_{k,i}
    return self.averaging_kernel * column

from_field

from_field(concentration: Float[Array, 'level y x']) -> Float[Array, ' obs']

Apply the observation to an interior field on the analysis grid.

Parameters:

Name	Type	Description	Default
concentration	Float[Array, 'level y x']	Interior concentration c, ghost rings stripped.	required

Returns:

Type	Description
Float[Array, ' obs']	The observation y_i at each receptor.

Source code in src/plumax/gauss_plume/_src/observation.py

def from_field(
    self, concentration: Float[Array, "level y x"]
) -> Float[Array, " obs"]:
    """Apply the observation to an interior field on the analysis grid.

    Args:
        concentration: Interior concentration c, ghost rings stripped.

    Returns:
        The observation y_i at each receptor.
    """
    if self.receptor_index is None:
        profiles = einx.rearrange("level y x -> (y x) level", concentration)
    else:
        iy = self.receptor_index[:, 0]
        ix = self.receptor_index[:, 1]
        # c_{k,i} = c[k, iy_i, ix_i]
        gathered = concentration[:, iy, ix]
        profiles = einx.rearrange("level obs -> obs level", gathered)
    return self.column(profiles)

BriggsPlumeRise

Bases: Module

Briggs plume rise Δh(x) of a stack plume, from its buoyancy and momentum.

Δh(x) = max(Δh_b(x), Δh_m(x)), with the buoyant and momentum rises of the module docstring.

Algorithm

1. Fluxes F_b = g v_s d² (T_s − T_a)/(4 T_s), floored at 0, and F_m = v_s² d² T_a/(4 T_s).
2. Buoyant rise: the 2/3 law 1.6 F_b^{1/3} x^{2/3}/u, capped at the class's final rise.
3. Momentum rise: the jet law 1.44 d (v_s/u)^{2/3} (x/d)^{1/3}, capped at 3 d v_s/u, and in stable air also at 1.5 (F_m/u)^{1/3} s^{−1/6}.
4. Take the larger of the two.

Attributes:

Name	Type	Description
exit_velocity	float | Float[Array, '']	Stack exit velocity v_s ≥ 0 [m s⁻¹].
exit_temperature	float | Float[Array, '']	Stack gas temperature T_s > 0 [K].
stack_diameter	float | Float[Array, '']	Inside stack-top diameter d > 0 [m].
ambient_temperature	float | Float[Array, '']	Ambient air temperature T_a > 0 [K].
stability_class	str	Pasquill–Gifford class 'A' … 'F'; E and F use the stable formulas with ∂θ/∂z from :data:STABLE_THETA_GRADIENT.

Raises:

Type	Description
ParameterError	If the class is not A–F or a parameter is out of range.

Examples:

A hot stack in neutral air: past x_f the rise is the final rise.

>>> import jax.numpy as jnp
>>> rise = BriggsPlumeRise(
...     exit_velocity=10.0, exit_temperature=400.0, stack_diameter=2.0,
...     ambient_temperature=290.0, stability_class="D",
... )
>>> heights = rise(jnp.array([100.0, 5000.0, 20000.0]), wind_speed=5.0)
>>> heights.shape
(3,)
>>> bool(heights[1] == heights[2])
True

A cold, still stack does not rise:

>>> still = BriggsPlumeRise(
...     exit_velocity=0.0, exit_temperature=290.0, stack_diameter=2.0,
...     ambient_temperature=290.0,
... )
>>> still(jnp.array([500.0]), wind_speed=5.0).tolist()
[0.0]

Classes outside A–F are rejected:

>>> BriggsPlumeRise(
...     exit_velocity=1.0, exit_temperature=300.0, stack_diameter=1.0,
...     ambient_temperature=290.0, stability_class="G",
... )
Traceback (most recent call last):
...
ParameterError: stability_class must be one of A-F, got 'G'

References

Briggs, G. A. (1969). Plume Rise. USAEC TID-25075 (Eqs. 4.19b, 4.28, 4.33, 5.2). U.S. EPA, User's Guide for the Industrial Source Complex (ISC3) Dispersion Models, Vol. II, §1.1.4 (after Briggs 1971, 1975).

is_stable property

is_stable: bool

Whether the class is E or F (the stable formulas).

buoyancy_flux

buoyancy_flux() -> Float[Array, '']

F_b = g v_s d² (T_s − T_a)/(4 T_s) [m⁴ s⁻³], floored at 0 (no sinking).

Source code in src/plumax/gauss_plume/_src/plume_rise.py

def buoyancy_flux(self) -> Float[Array, ""]:
    """F_b = g v_s d² (T_s − T_a)/(4 T_s) [m⁴ s⁻³], floored at 0 (no sinking)."""
    # T_s − T_a
    excess = self.exit_temperature - self.ambient_temperature
    # g v_s d² (T_s − T_a) / (4 T_s)
    flux = (
        GRAVITY
        * self.exit_velocity
        * self.stack_diameter**2
        * excess
        / (4.0 * self.exit_temperature)
    )
    return jnp.maximum(flux, 0.0)

momentum_flux

momentum_flux() -> Float[Array, '']

F_m = v_s² d² T_a/(4 T_s) [m⁴ s⁻²].

Source code in src/plumax/gauss_plume/_src/plume_rise.py

def momentum_flux(self) -> Float[Array, ""]:
    """F_m = v_s² d² T_a/(4 T_s) [m⁴ s⁻²]."""
    # v_s² d² T_a / (4 T_s)
    return jnp.asarray(
        self.exit_velocity**2
        * self.stack_diameter**2
        * self.ambient_temperature
        / (4.0 * self.exit_temperature)
    )

stability_parameter

stability_parameter() -> Float[Array, '']

s = (g / T_a) ∂θ/∂z [s⁻²] of the stable classes; 0 for A–D.

Source code in src/plumax/gauss_plume/_src/plume_rise.py

def stability_parameter(self) -> Float[Array, ""]:
    """s = (g / T_a) ∂θ/∂z [s⁻²] of the stable classes; 0 for A–D."""
    gradient = STABLE_THETA_GRADIENT.get(self.stability_class, 0.0)
    # s = g ∂θ/∂z / T_a
    return jnp.asarray(GRAVITY * gradient / self.ambient_temperature)

EmissionRateSamples

EmissionRateSamples(draws: dict[str, NDArray], *, prior_mean: float, prior_std: float)

Bases: dict

NUTS draws by site name, with the per-event payload Tier V reads.

A dict of posterior arrays (samples["emission_rate"] is unchanged) that also satisfies population.PerEventPosterior: the posterior mean and standard deviation of Q, resampled draws of Q, and the log-density of the emission-rate prior the draws were made under, Q ~ LogNormal(μ, σ²) moment-matched to (prior_mean, prior_std).

Parameters:

Name	Type	Description	Default
draws	dict[str, NDArray]	Posterior arrays keyed by site name.	required
prior_mean	float	Real-scale prior mean of Q [kg/s].	required
prior_std	float	Real-scale prior standard deviation of Q [kg/s].	required

Source code in src/plumax/gauss_plume/inference.py

def __init__(
    self, draws: dict[str, NDArray], *, prior_mean: float, prior_std: float
) -> None:
    super().__init__(draws)
    self.prior_mean = prior_mean
    self.prior_std = prior_std

emission_rate property

emission_rate: NDArray

Posterior mean of Q [kg/s].

emission_std property

emission_std: NDArray

Posterior standard deviation of Q [kg/s].

samples

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

n draws of Q, resampled with replacement from the NUTS draws.

Source code in src/plumax/gauss_plume/inference.py

def samples(self, key: jax.Array, n: int) -> jax.Array:
    """n draws of Q, resampled with replacement from the NUTS draws."""
    return EmpiricalEmission(self["emission_rate"]).sample(key, n)

prior_logpdf

prior_logpdf(q: Array) -> jax.Array

log π(q) of the emission-rate prior of :func:gaussian_plume_model.

Source code in src/plumax/gauss_plume/inference.py

def prior_logpdf(self, q: jax.Array) -> jax.Array:
    """log π(q) of the emission-rate prior of :func:`gaussian_plume_model`."""
    prior = LognormalEmission.from_moments(self.prior_mean, self.prior_std)
    return prior.log_prob(q)

calculate_briggs_dispersion

calculate_briggs_dispersion(distance: ndarray, params: ndarray) -> tuple[jnp.ndarray, jnp.ndarray]

Evaluate σ_y(x), σ_z(x) from Briggs parameters.

σ_i(x) = a_i · x · (1 + b_i · x)^{c_i} for i ∈ {y, z}.

Parameters

distance : jnp.ndarray Downwind distance from source [m]. Negative / zero values are clamped to 1 m so that the coefficient evaluation stays finite; upwind points are masked elsewhere in the plume model. params : jnp.ndarray, shape (6,) [a_y, b_y, c_y, a_z, b_z, c_z].

Returns

sigma_y : jnp.ndarray Lateral dispersion coefficient [m], same shape as distance. sigma_z : jnp.ndarray Vertical dispersion coefficient [m], same shape as distance.

Source code in src/plumax/gauss_plume/dispersion.py

@jit
def calculate_briggs_dispersion(
    distance: jnp.ndarray,
    params: jnp.ndarray,
) -> tuple[jnp.ndarray, jnp.ndarray]:
    """Evaluate σ_y(x), σ_z(x) from Briggs parameters.

    σ_i(x) = a_i · x · (1 + b_i · x)^{c_i} for i ∈ {y, z}.

    Parameters
    ----------
    distance : jnp.ndarray
        Downwind distance from source [m]. Negative / zero values are clamped
        to 1 m so that the coefficient evaluation stays finite; upwind points
        are masked elsewhere in the plume model.
    params : jnp.ndarray, shape (6,)
        ``[a_y, b_y, c_y, a_z, b_z, c_z]``.

    Returns
    -------
    sigma_y : jnp.ndarray
        Lateral dispersion coefficient [m], same shape as ``distance``.
    sigma_z : jnp.ndarray
        Vertical dispersion coefficient [m], same shape as ``distance``.
    """
    ay, by, cy, az, bz, cz = params
    x = jnp.maximum(distance, 1.0)
    sigma_y = ay * x * jnp.power(1.0 + by * x, cy)
    sigma_z = az * x * jnp.power(1.0 + bz * x, cz)
    return sigma_y, sigma_z

get_dispersion_params

get_dispersion_params(stability_class: str) -> jnp.ndarray

Return the Briggs parameter vector for a stability class.

Parameters

stability_class : str One of 'A', 'B', 'C', 'D', 'E', 'F'.

Returns

params : jnp.ndarray, shape (6,) [a_y, b_y, c_y, a_z, b_z, c_z].

Raises

ValueError If stability_class is not one of A-F.

Source code in src/plumax/gauss_plume/dispersion.py

def get_dispersion_params(stability_class: str) -> jnp.ndarray:
    """Return the Briggs parameter vector for a stability class.

    Parameters
    ----------
    stability_class : str
        One of ``'A'``, ``'B'``, ``'C'``, ``'D'``, ``'E'``, ``'F'``.

    Returns
    -------
    params : jnp.ndarray, shape (6,)
        ``[a_y, b_y, c_y, a_z, b_z, c_z]``.

    Raises
    ------
    ValueError
        If ``stability_class`` is not one of A-F.
    """
    if stability_class not in BRIGGS_DISPERSION_PARAMS:
        raise ValueError(
            f"get_dispersion_params: stability_class must be one of "
            f"{STABILITY_CLASSES}, got {stability_class!r}"
        )
    return BRIGGS_DISPERSION_PARAMS[stability_class]

golder_class_boundaries

golder_class_boundaries(z0: float) -> Float[Array, ' boundary']

Inverse Obukhov lengths at Golder's five PG class boundaries.

(1/L)_k = (ln z₀ − d_k) / c_k for the A/B, B/C, C/D, D/E and E/F boundaries, with z₀ clamped to :data:GOLDER_Z0_RANGE.

Parameters:

Name	Type	Description	Default
z0	float	Roughness length z₀ > 0 [m].	required

Returns:

Type	Description
Float[Array, ' boundary']	The five boundaries [m⁻¹], increasing from A/B (most negative) to E/F.

Examples:

At z₀ = 0.1 m the unstable boundaries sit at L ≈ −10.5, −30.4 and −104 m, the stable ones at L ≈ 112 and 26.9 m:

>>> import jax.numpy as jnp
>>> boundaries = golder_class_boundaries(0.1)
>>> [round(float(1.0 / b), 1) for b in boundaries]
[-10.5, -30.4, -104.1, 111.6, 26.9]

z₀ beyond 0.5 m is clamped:

>>> bool(jnp.allclose(golder_class_boundaries(2.0), golder_class_boundaries(0.5)))
True

Source code in src/plumax/gauss_plume/_src/mo_dispersion.py

def golder_class_boundaries(z0: float) -> Float[Array, " boundary"]:
    """Inverse Obukhov lengths at Golder's five PG class boundaries.

    (1/L)_k = (ln z₀ − d_k) / c_k for the A/B, B/C, C/D, D/E and E/F
    boundaries, with z₀ clamped to :data:`GOLDER_Z0_RANGE`.

    Args:
        z0: Roughness length z₀ > 0 [m].

    Returns:
        The five boundaries [m⁻¹], increasing from A/B (most negative) to E/F.

    Examples:
        At z₀ = 0.1 m the unstable boundaries sit at L ≈ −10.5, −30.4 and
        −104 m, the stable ones at L ≈ 112 and 26.9 m:

        >>> import jax.numpy as jnp
        >>> boundaries = golder_class_boundaries(0.1)
        >>> [round(float(1.0 / b), 1) for b in boundaries]
        [-10.5, -30.4, -104.1, 111.6, 26.9]

        z₀ beyond 0.5 m is clamped:

        >>> bool(jnp.allclose(golder_class_boundaries(2.0), golder_class_boundaries(0.5)))
        True
    """
    roughness = jnp.asarray(z0)
    roughness = eqx.error_if(roughness, roughness <= 0.0, "z0 must be positive")
    clamped = jnp.clip(roughness, *GOLDER_Z0_RANGE)
    log_z0 = jnp.log(clamped)
    scale = jnp.asarray(GOLDER_BOUNDARY_SCALE, dtype=log_z0.dtype)
    offset = jnp.asarray(GOLDER_BOUNDARY_OFFSET, dtype=log_z0.dtype)
    # (1/L)_k = (ln z₀ − d_k) / c_k
    return (log_z0 - offset) / scale

golder_stability_class

golder_stability_class(inverse_obukhov_length: Float[Array, '*shape'] | float, z0: float) -> Int[Array, '*shape']

Golder's (1972) PG class for 1/L and z₀, as an index 0 (A) … 5 (F).

The class is the number of boundaries (1/L)_k strictly below 1/L, so a value on a boundary goes to the less stable class and 1/L = 0 is D.

Parameters:

Name	Type	Description	Default
inverse_obukhov_length	Float[Array, '*shape'] | float	1/L [m⁻¹]; zero is neutral, negative unstable.	required
z0	float	Roughness length z₀ > 0 [m].	required

Returns:

Type	Description
Int[Array, '*shape']	The class index into :data:~plumax.gauss_plume.STABILITY_CLASSES.

Examples:

Over z₀ = 0.1 m, L = −20 m is B, neutral is D and L = 50 m is E:

>>> import jax.numpy as jnp
>>> inverse = jnp.array([-1.0 / 20.0, 0.0, 1.0 / 50.0])
>>> golder_stability_class(inverse, 0.1).tolist()
[1, 3, 4]

Source code in src/plumax/gauss_plume/_src/mo_dispersion.py

def golder_stability_class(
    inverse_obukhov_length: Float[Array, "*shape"] | float, z0: float
) -> Int[Array, "*shape"]:
    """Golder's (1972) PG class for 1/L and z₀, as an index 0 (A) … 5 (F).

    The class is the number of boundaries (1/L)_k strictly below 1/L, so a
    value on a boundary goes to the less stable class and 1/L = 0 is D.

    Args:
        inverse_obukhov_length: 1/L [m⁻¹]; zero is neutral, negative unstable.
        z0: Roughness length z₀ > 0 [m].

    Returns:
        The class index into :data:`~plumax.gauss_plume.STABILITY_CLASSES`.

    Examples:
        Over z₀ = 0.1 m, L = −20 m is B, neutral is D and L = 50 m is E:

        >>> import jax.numpy as jnp
        >>> inverse = jnp.array([-1.0 / 20.0, 0.0, 1.0 / 50.0])
        >>> golder_stability_class(inverse, 0.1).tolist()
        [1, 3, 4]
    """
    inverse_obukhov_length = jnp.asarray(inverse_obukhov_length)
    boundaries = golder_class_boundaries(z0)
    above = inverse_obukhov_length[..., None] > boundaries
    return jnp.sum(above, axis=-1)

golder_stability_index

golder_stability_index(inverse_obukhov_length: Float[Array, '*shape'] | float, z0: float) -> Float[Array, '*shape']

Continuous PG class index from 1/L and z₀, piecewise linear in 1/L.

s(1/L) interpolates linearly through the knots (1/L, s) = ((1/L)_AB − Δ_A, 0), ((1/L)_AB, 0.5), ((1/L)_BC, 1.5), ((1/L)_CD, 2.5), (0, 3), ((1/L)_DE, 3.5), ((1/L)_EF, 4.5), ((1/L)_EF + Δ_F, 5), and is clamped to [0, 5] outside them, with Δ_A = ((1/L)_BC − (1/L)_AB)/2 and Δ_F = ((1/L)_EF − (1/L)_DE)/2 (the half-width of the neighbouring class).

This is a plumax modelling choice, not part of Golder (1972): Golder gives only the class boundaries. Each boundary maps to the half-integer between its two classes, so round of the index (ties to the less stable class) is the hard class of :func:golder_stability_class, and neutral maps to exactly 3 (class D). Its derivative in 1/L is non-zero everywhere except deep inside A (1/L < (1/L)_AB − Δ_A) and F (1/L > (1/L)_EF + Δ_F).

Parameters:

Name	Type	Description	Default
inverse_obukhov_length	Float[Array, '*shape'] | float	1/L [m⁻¹]; zero is neutral, negative unstable.	required
z0	float	Roughness length z₀ > 0 [m] (a scalar).	required

Returns:

Type	Description
Float[Array, '*shape']	The class index s ∈ [0, 5], 0 for A and 5 for F.

Examples:

Boundaries land on half-integers and neutral on D:

>>> import jax.numpy as jnp
>>> boundaries = golder_class_boundaries(0.1)
>>> golder_stability_index(boundaries, 0.1).tolist()
[0.5, 1.5, 2.5, 3.5, 4.5]
>>> float(golder_stability_index(0.0, 0.1))
3.0

Source code in src/plumax/gauss_plume/_src/mo_dispersion.py

def golder_stability_index(
    inverse_obukhov_length: Float[Array, "*shape"] | float, z0: float
) -> Float[Array, "*shape"]:
    """Continuous PG class index from 1/L and z₀, piecewise linear in 1/L.

    s(1/L) interpolates linearly through the knots (1/L, s) =
    ((1/L)_AB − Δ_A, 0), ((1/L)_AB, 0.5), ((1/L)_BC, 1.5), ((1/L)_CD, 2.5),
    (0, 3), ((1/L)_DE, 3.5), ((1/L)_EF, 4.5), ((1/L)_EF + Δ_F, 5), and is
    clamped to [0, 5] outside them, with Δ_A = ((1/L)_BC − (1/L)_AB)/2 and
    Δ_F = ((1/L)_EF − (1/L)_DE)/2 (the half-width of the neighbouring class).

    This is a plumax modelling choice, not part of Golder (1972): Golder gives
    only the class boundaries. Each boundary maps to the half-integer between
    its two classes, so ``round`` of the index (ties to the less stable class)
    is the hard class of :func:`golder_stability_class`, and neutral maps to
    exactly 3 (class D). Its derivative in 1/L is non-zero everywhere except
    deep inside A (1/L < (1/L)_AB − Δ_A) and F (1/L > (1/L)_EF + Δ_F).

    Args:
        inverse_obukhov_length: 1/L [m⁻¹]; zero is neutral, negative unstable.
        z0: Roughness length z₀ > 0 [m] (a scalar).

    Returns:
        The class index s ∈ [0, 5], 0 for A and 5 for F.

    Examples:
        Boundaries land on half-integers and neutral on D:

        >>> import jax.numpy as jnp
        >>> boundaries = golder_class_boundaries(0.1)
        >>> golder_stability_index(boundaries, 0.1).tolist()
        [0.5, 1.5, 2.5, 3.5, 4.5]
        >>> float(golder_stability_index(0.0, 0.1))
        3.0
    """
    inverse_obukhov_length = jnp.asarray(inverse_obukhov_length)
    boundaries = golder_class_boundaries(z0)
    # Δ_A = ((1/L)_BC − (1/L)_AB) / 2
    half_width_a = 0.5 * (boundaries[1] - boundaries[0])
    # Δ_F = ((1/L)_EF − (1/L)_DE) / 2
    half_width_f = 0.5 * (boundaries[4] - boundaries[3])
    neutral = jnp.zeros((1,), dtype=boundaries.dtype)
    knots = jnp.concatenate(
        [
            (boundaries[0] - half_width_a)[None],
            boundaries[:3],
            neutral,
            boundaries[3:],
            (boundaries[4] + half_width_f)[None],
        ]
    )
    index_knots = jnp.asarray(_INDEX_KNOTS, dtype=boundaries.dtype)
    return jnp.interp(inverse_obukhov_length, knots, index_knots)

plume_concentration

plume_concentration(x: ndarray, y: ndarray, z: ndarray, source_x: float, source_y: float, source_z: float, wind_u: float, wind_v: float, emission_rate: float, dispersion_params: ndarray | GolderDispersion, plume_rise: BriggsPlumeRise | None = None) -> jnp.ndarray

Evaluate the steady-state Gaussian plume concentration.

Computes the mass concentration [kg/m³] at receptor points (x, y, z) in the fixed reference frame. Coordinates are rotated into the wind-aligned frame internally, upwind points are masked to zero, and ground reflection is applied.

Parameters

x, y, z : jnp.ndarray Receptor coordinates in the fixed frame [m]. Broadcast-compatible. source_x, source_y, source_z : float Source location [m]. source_z is the emission height above ground. wind_u, wind_v : float Wind velocity components [m/s]. The total wind speed is clamped below to :data:MIN_WIND_SPEED for numerical stability. emission_rate : float Continuous emission rate Q [kg/s]. dispersion_params : jnp.ndarray, shape (6,), or GolderDispersion [a_y, b_y, c_y, a_z, b_z, c_z] (Briggs parameters of one PG class), or a :class:~plumax.gauss_plume.GolderDispersion that picks the PG class from the Obukhov length and roughness length. plume_rise : BriggsPlumeRise, optional Plume rise Δh(x'). When given, the release height at each receptor is the effective height source_z + Δh(x') with Δh evaluated at the receptor's downwind distance and the (clamped) wind speed; it is differentiable in the stack parameters. Default None keeps the fixed height source_z.

Returns

concentration : jnp.ndarray Mass concentration [kg/m³], same shape as the broadcast of the input coordinates. Zero at upwind locations (x' ≤ 0).

Source code in src/plumax/gauss_plume/plume.py

@jit
def plume_concentration(
    x: jnp.ndarray,
    y: jnp.ndarray,
    z: jnp.ndarray,
    source_x: float,
    source_y: float,
    source_z: float,
    wind_u: float,
    wind_v: float,
    emission_rate: float,
    dispersion_params: jnp.ndarray | GolderDispersion,
    plume_rise: BriggsPlumeRise | None = None,
) -> jnp.ndarray:
    """Evaluate the steady-state Gaussian plume concentration.

    Computes the mass concentration [kg/m³] at receptor points ``(x, y, z)``
    in the fixed reference frame. Coordinates are rotated into the
    wind-aligned frame internally, upwind points are masked to zero, and
    ground reflection is applied.

    Parameters
    ----------
    x, y, z : jnp.ndarray
        Receptor coordinates in the fixed frame [m]. Broadcast-compatible.
    source_x, source_y, source_z : float
        Source location [m]. ``source_z`` is the emission height above ground.
    wind_u, wind_v : float
        Wind velocity components [m/s]. The total wind speed is clamped
        below to :data:`MIN_WIND_SPEED` for numerical stability.
    emission_rate : float
        Continuous emission rate Q [kg/s].
    dispersion_params : jnp.ndarray, shape (6,), or GolderDispersion
        ``[a_y, b_y, c_y, a_z, b_z, c_z]`` (Briggs parameters of one PG
        class), or a :class:`~plumax.gauss_plume.GolderDispersion` that picks
        the PG class from the Obukhov length and roughness length.
    plume_rise : BriggsPlumeRise, optional
        Plume rise Δh(x'). When given, the release height at each receptor is
        the effective height ``source_z + Δh(x')`` with Δh evaluated at the
        receptor's downwind distance and the (clamped) wind speed; it is
        differentiable in the stack parameters. Default ``None`` keeps the
        fixed height ``source_z``.

    Returns
    -------
    concentration : jnp.ndarray
        Mass concentration [kg/m³], same shape as the broadcast of the
        input coordinates. Zero at upwind locations (x' ≤ 0).
    """
    wind_speed = jnp.sqrt(wind_u**2 + wind_v**2)
    wind_speed = jnp.maximum(wind_speed, MIN_WIND_SPEED)

    x_downwind, y_crosswind = rotate_to_wind_frame(
        x, y, source_x, source_y, wind_u, wind_v
    )

    x_positive = jnp.maximum(x_downwind, 1.0)
    if isinstance(dispersion_params, GolderDispersion):
        sigma_y, sigma_z = dispersion_params(x_positive)
    else:
        sigma_y, sigma_z = calculate_briggs_dispersion(x_positive, dispersion_params)

    normalization = 2.0 * jnp.pi * wind_speed * sigma_y * sigma_z

    exp_y = jnp.exp(-0.5 * jnp.square(y_crosswind / sigma_y))

    if plume_rise is None:
        release_height = source_z
    else:
        # H_eff = H + Δh(x')
        release_height = source_z + plume_rise(x_positive, wind_speed=wind_speed)

    z_diff = z - release_height
    z_sum = z + release_height
    exp_z_direct = jnp.exp(-0.5 * jnp.square(z_diff / sigma_z))
    exp_z_reflected = jnp.exp(-0.5 * jnp.square(z_sum / sigma_z))

    concentration = (
        (emission_rate / normalization) * exp_y * (exp_z_direct + exp_z_reflected)
    )
    return jnp.where(x_downwind > 0.0, concentration, 0.0)

rotate_to_wind_frame

rotate_to_wind_frame(x: ndarray, y: ndarray, source_x: float, source_y: float, wind_u: float, wind_v: float) -> tuple[jnp.ndarray, jnp.ndarray]

Rotate fixed-frame coordinates into the wind-aligned frame.

The wind-aligned frame has its origin at the source, its x'-axis pointing downwind, and its y'-axis pointing crosswind (90° counter-clockwise from x', following the right-hand rule). The rotation angle is θ = atan2(wind_v, wind_u) [rad].

Parameters

x, y : jnp.ndarray Receptor coordinates in the fixed frame [m]. Must have matching shapes (or be broadcast-compatible). source_x, source_y : float Source location in the fixed frame [m]. wind_u, wind_v : float Wind velocity components [m/s].

Returns

x_wind : jnp.ndarray Downwind coordinate x' [m] (> 0 downwind, < 0 upwind of source). y_wind : jnp.ndarray Crosswind coordinate y' [m].

Source code in src/plumax/gauss_plume/plume.py

@jit
def rotate_to_wind_frame(
    x: jnp.ndarray,
    y: jnp.ndarray,
    source_x: float,
    source_y: float,
    wind_u: float,
    wind_v: float,
) -> tuple[jnp.ndarray, jnp.ndarray]:
    """Rotate fixed-frame coordinates into the wind-aligned frame.

    The wind-aligned frame has its origin at the source, its x'-axis pointing
    downwind, and its y'-axis pointing crosswind (90° counter-clockwise from
    x', following the right-hand rule). The rotation angle is
    ``θ = atan2(wind_v, wind_u)`` [rad].

    Parameters
    ----------
    x, y : jnp.ndarray
        Receptor coordinates in the fixed frame [m]. Must have matching
        shapes (or be broadcast-compatible).
    source_x, source_y : float
        Source location in the fixed frame [m].
    wind_u, wind_v : float
        Wind velocity components [m/s].

    Returns
    -------
    x_wind : jnp.ndarray
        Downwind coordinate x' [m] (> 0 downwind, < 0 upwind of source).
    y_wind : jnp.ndarray
        Crosswind coordinate y' [m].
    """
    dx = x - source_x
    dy = y - source_y

    theta = jnp.arctan2(wind_v, wind_u)
    cos_theta = jnp.cos(theta)
    sin_theta = jnp.sin(theta)

    x_wind = dx * cos_theta + dy * sin_theta
    y_wind = -dx * sin_theta + dy * cos_theta
    return x_wind, y_wind

simulate_plume

simulate_plume(emission_rate: float, source_location: tuple[float, float, float], wind_speed: float, wind_direction: float, stability_class: str | GolderDispersion, domain_x: tuple[float, float, int], domain_y: tuple[float, float, int], domain_z: tuple[float, float, int], background_conc: float = 0.0, plume_rise: BriggsPlumeRise | None = None) -> xr.Dataset

Simulate a steady-state Gaussian plume on a 3-D grid.

Wraps :func:plume_concentration with grid construction, wind-from-angle to (u, v) conversion, and column integration, returning an :class:xarray.Dataset suitable for plotting / I/O.

Wind-direction convention follows the meteorological "from" convention: wind_direction = 270° is a wind from the west (flowing east); wind_direction = 0° is a wind from the north (flowing south).

u = -wind_speed · sin(θ),     v = -wind_speed · cos(θ),
    with θ = wind_direction in radians.

Parameters

emission_rate : float Continuous emission rate Q [kg/s]. Must be strictly positive. source_location : tuple of float, (3,) (x, y, z) source coordinates [m]. z ≥ 0. wind_speed : float Wind speed magnitude [m/s]. Must be strictly positive. wind_direction : float Wind direction [degrees from North, meteorological "from" convention]. stability_class : str or GolderDispersion One of 'A', ..., 'F', or a :class:~plumax.gauss_plume.GolderDispersion that picks the class from the Obukhov length and roughness length. domain_x, domain_y, domain_z : tuple of (float, float, int) (start, stop, n_points) for each axis. start < stop and n_points ≥ 2. background_conc : float Additive background concentration [kg/m³]. Default 0.0. plume_rise : BriggsPlumeRise, optional Plume rise added to the source height z (see :func:plume_concentration). Default None: no rise.

Returns

ds : xarray.Dataset Variables: concentration (x, y, z) [kg/m³], column_concentration (x, y) [kg/m²].

Raises

ValueError On invalid stability class, non-positive emission rate, non-positive wind speed, malformed domain tuples, or negative source height.

Source code in src/plumax/gauss_plume/plume.py

def simulate_plume(
    emission_rate: float,
    source_location: tuple[float, float, float],
    wind_speed: float,
    wind_direction: float,
    stability_class: str | GolderDispersion,
    domain_x: tuple[float, float, int],
    domain_y: tuple[float, float, int],
    domain_z: tuple[float, float, int],
    background_conc: float = 0.0,
    plume_rise: BriggsPlumeRise | None = None,
) -> xr.Dataset:
    """Simulate a steady-state Gaussian plume on a 3-D grid.

    Wraps :func:`plume_concentration` with grid construction, wind-from-angle
    to (u, v) conversion, and column integration, returning an
    :class:`xarray.Dataset` suitable for plotting / I/O.

    Wind-direction convention follows the meteorological "from" convention:
    ``wind_direction = 270°`` is a wind *from* the west (flowing east);
    ``wind_direction = 0°`` is a wind *from* the north (flowing south).

        u = -wind_speed · sin(θ),     v = -wind_speed · cos(θ),
            with θ = wind_direction in radians.

    Parameters
    ----------
    emission_rate : float
        Continuous emission rate Q [kg/s]. Must be strictly positive.
    source_location : tuple of float, (3,)
        (x, y, z) source coordinates [m]. ``z ≥ 0``.
    wind_speed : float
        Wind speed magnitude [m/s]. Must be strictly positive.
    wind_direction : float
        Wind direction [degrees from North, meteorological "from" convention].
    stability_class : str or GolderDispersion
        One of ``'A'``, ..., ``'F'``, or a
        :class:`~plumax.gauss_plume.GolderDispersion` that picks the class from
        the Obukhov length and roughness length.
    domain_x, domain_y, domain_z : tuple of (float, float, int)
        ``(start, stop, n_points)`` for each axis. ``start < stop`` and
        ``n_points ≥ 2``.
    background_conc : float
        Additive background concentration [kg/m³]. Default ``0.0``.
    plume_rise : BriggsPlumeRise, optional
        Plume rise added to the source height ``z`` (see
        :func:`plume_concentration`). Default ``None``: no rise.

    Returns
    -------
    ds : xarray.Dataset
        Variables: ``concentration`` (x, y, z) [kg/m³],
        ``column_concentration`` (x, y) [kg/m²].

    Raises
    ------
    ValueError
        On invalid stability class, non-positive emission rate, non-positive
        wind speed, malformed domain tuples, or negative source height.
    """
    import xarray as xr  # lazy — only this entry point requires xarray

    if not (emission_rate > 0.0):
        raise ValueError(
            f"simulate_plume: `emission_rate` must be > 0 (got {emission_rate!r})"
        )
    if not (wind_speed > 0.0):
        raise ValueError(
            f"simulate_plume: `wind_speed` must be > 0 (got {wind_speed!r})"
        )
    is_golder = isinstance(stability_class, GolderDispersion)
    if not is_golder and stability_class not in BRIGGS_DISPERSION_PARAMS:
        raise ValueError(
            f"simulate_plume: `stability_class` must be one of "
            f"{STABILITY_CLASSES} (got {stability_class!r})"
        )
    if len(source_location) != 3:
        raise ValueError(
            "simulate_plume: `source_location` must be (x, y, z); "
            f"got length {len(source_location)}"
        )
    if source_location[2] < 0.0:
        raise ValueError(
            f"simulate_plume: source height must be ≥ 0 (got z={source_location[2]!r})"
        )
    for name, domain in (
        ("domain_x", domain_x),
        ("domain_y", domain_y),
        ("domain_z", domain_z),
    ):
        if len(domain) != 3:
            raise ValueError(
                f"simulate_plume: `{name}` must be (start, stop, n_points); "
                f"got length {len(domain)}"
            )
        start, stop, n = domain
        if not (start < stop):
            raise ValueError(
                f"simulate_plume: `{name}` requires start < stop "
                f"(got start={start!r}, stop={stop!r})"
            )
        if int(n) < 2:
            raise ValueError(
                f"simulate_plume: `{name}` requires n_points ≥ 2 (got n={n!r})"
            )

    if isinstance(stability_class, GolderDispersion):
        dispersion_params = stability_class
        stability_attr = "golder"
    else:
        dispersion_params = get_dispersion_params(stability_class)
        stability_attr = stability_class

    x = np.linspace(domain_x[0], domain_x[1], int(domain_x[2]))
    y = np.linspace(domain_y[0], domain_y[1], int(domain_y[2]))
    z = np.linspace(domain_z[0], domain_z[1], int(domain_z[2]))
    n_x, n_y, n_z = len(x), len(y), len(z)

    theta_rad = np.deg2rad(wind_direction)
    wind_u = -wind_speed * np.sin(theta_rad)
    wind_v = -wind_speed * np.cos(theta_rad)

    src_x, src_y, src_z = source_location

    X, Y, Z = np.meshgrid(x, y, z, indexing="ij")
    conc_flat = plume_concentration(
        jnp.array(X.ravel()),
        jnp.array(Y.ravel()),
        jnp.array(Z.ravel()),
        src_x,
        src_y,
        src_z,
        float(wind_u),
        float(wind_v),
        emission_rate,
        dispersion_params,
        plume_rise=plume_rise,
    )
    concentration = np.asarray(conc_flat).reshape((n_x, n_y, n_z))
    concentration = concentration + background_conc

    # Trapezoidal vertical integration: the z grid from ``np.linspace`` includes
    # both endpoints, so a plain ``sum * dz`` over-counts a vertically uniform
    # component (e.g. a background) by one grid interval (n_z / (n_z - 1)).
    column = np.trapezoid(concentration, x=z, axis=2)

    ds = xr.Dataset(
        data_vars={
            "concentration": (["x", "y", "z"], concentration),
            "column_concentration": (["x", "y"], column),
        },
        coords={"x": x, "y": y, "z": z},
        attrs={
            "title": "Steady-state Gaussian plume (JAX)",
            "emission_rate": emission_rate,
            "emission_rate_units": "kg/s",
            "source_x": src_x,
            "source_y": src_y,
            "source_z": src_z,
            "wind_speed": wind_speed,
            "wind_speed_units": "m/s",
            "wind_direction": wind_direction,
            "wind_direction_units": "degrees from North (meteorological)",
            "wind_u": float(wind_u),
            "wind_v": float(wind_v),
            "stability_class": stability_attr,
            "background_concentration": background_conc,
        },
    )
    ds["concentration"].attrs = {"long_name": "Mass concentration", "units": "kg/m^3"}
    ds["column_concentration"].attrs = {
        "long_name": "Column-integrated concentration",
        "units": "kg/m^2",
    }
    return ds

briggs_plume_rise

briggs_plume_rise(x: Float[Array, '*obs'], *, exit_velocity: float | Float[Array, ''], exit_temperature: float | Float[Array, ''], stack_diameter: float | Float[Array, ''], ambient_temperature: float | Float[Array, ''], wind_speed: float | Float[Array, ''], stability_class: str = 'D') -> Float[Array, '*obs']

Briggs plume rise Δh(x) [m]; see :class:BriggsPlumeRise.

Parameters:

Name	Type	Description	Default
x	Float[Array, '*obs']	Downwind distance [m].	required
exit_velocity	float | Float[Array, '']	Stack exit velocity v_s [m s⁻¹].	required
exit_temperature	float | Float[Array, '']	Stack gas temperature T_s [K].	required
stack_diameter	float | Float[Array, '']	Stack diameter d [m].	required
ambient_temperature	float | Float[Array, '']	Ambient temperature T_a [K].	required
wind_speed	float | Float[Array, '']	Wind speed u at the stack top [m s⁻¹].	required
stability_class	str	Pasquill–Gifford class 'A' … 'F'.	'D'

Returns:

Type	Description
Float[Array, '*obs']	Δh(x), with the final rise past x_f.

Examples:

The final rise of a buoyant plume in neutral air (F_b < 55):

>>> import jax.numpy as jnp
>>> rise = briggs_plume_rise(
...     jnp.array([1.0e5]), exit_velocity=10.0, exit_temperature=400.0,
...     stack_diameter=2.0, ambient_temperature=290.0, wind_speed=5.0,
... )
>>> flux = 9.80665 * 10.0 * 4.0 * 110.0 / 1600.0
>>> final = 1.6 * flux ** (1 / 3) * (49.0 * flux ** 0.625) ** (2 / 3) / 5.0
>>> bool(jnp.allclose(rise, final))
True

Source code in src/plumax/gauss_plume/_src/plume_rise.py

def briggs_plume_rise(
    x: Float[Array, "*obs"],
    *,
    exit_velocity: float | Float[Array, ""],
    exit_temperature: float | Float[Array, ""],
    stack_diameter: float | Float[Array, ""],
    ambient_temperature: float | Float[Array, ""],
    wind_speed: float | Float[Array, ""],
    stability_class: str = "D",
) -> Float[Array, "*obs"]:
    """Briggs plume rise Δh(x) [m]; see :class:`BriggsPlumeRise`.

    Args:
        x: Downwind distance [m].
        exit_velocity: Stack exit velocity v_s [m s⁻¹].
        exit_temperature: Stack gas temperature T_s [K].
        stack_diameter: Stack diameter d [m].
        ambient_temperature: Ambient temperature T_a [K].
        wind_speed: Wind speed u at the stack top [m s⁻¹].
        stability_class: Pasquill–Gifford class ``'A'`` … ``'F'``.

    Returns:
        Δh(x), with the final rise past x_f.

    Examples:
        The final rise of a buoyant plume in neutral air (F_b < 55):

        >>> import jax.numpy as jnp
        >>> rise = briggs_plume_rise(
        ...     jnp.array([1.0e5]), exit_velocity=10.0, exit_temperature=400.0,
        ...     stack_diameter=2.0, ambient_temperature=290.0, wind_speed=5.0,
        ... )
        >>> flux = 9.80665 * 10.0 * 4.0 * 110.0 / 1600.0
        >>> final = 1.6 * flux ** (1 / 3) * (49.0 * flux ** 0.625) ** (2 / 3) / 5.0
        >>> bool(jnp.allclose(rise, final))
        True
    """
    plume_rise = BriggsPlumeRise(
        exit_velocity=exit_velocity,
        exit_temperature=exit_temperature,
        stack_diameter=stack_diameter,
        ambient_temperature=ambient_temperature,
        stability_class=stability_class,
    )
    return plume_rise(x, wind_speed=wind_speed)

effective_height

effective_height(x: Float[Array, '*obs'], *, stack_height: float | Float[Array, ''], exit_velocity: float | Float[Array, ''], exit_temperature: float | Float[Array, ''], stack_diameter: float | Float[Array, ''], ambient_temperature: float | Float[Array, ''], wind_speed: float | Float[Array, ''], stability_class: str = 'D') -> Float[Array, '*obs']

Effective release height H_eff(x) = H_stack + Δh(x) [m].

Parameters:

Name	Type	Description	Default
x	Float[Array, '*obs']	Downwind distance [m].	required
stack_height	float | Float[Array, '']	Physical stack height H_stack [m].	required
exit_velocity	float | Float[Array, '']	Stack exit velocity v_s [m s⁻¹].	required
exit_temperature	float | Float[Array, '']	Stack gas temperature T_s [K].	required
stack_diameter	float | Float[Array, '']	Stack diameter d [m].	required
ambient_temperature	float | Float[Array, '']	Ambient temperature T_a [K].	required
wind_speed	float | Float[Array, '']	Wind speed u at the stack top [m s⁻¹].	required
stability_class	str	Pasquill–Gifford class 'A' … 'F'.	'D'

Returns:

Type	Description
Float[Array, '*obs']	H_eff at each x.

Examples:

At the stack itself the plume has not risen yet:

>>> import jax.numpy as jnp
>>> effective_height(
...     jnp.array([0.0]), stack_height=30.0, exit_velocity=10.0,
...     exit_temperature=400.0, stack_diameter=2.0,
...     ambient_temperature=290.0, wind_speed=5.0,
... ).tolist()
[30.0]

Source code in src/plumax/gauss_plume/_src/plume_rise.py

def effective_height(
    x: Float[Array, "*obs"],
    *,
    stack_height: float | Float[Array, ""],
    exit_velocity: float | Float[Array, ""],
    exit_temperature: float | Float[Array, ""],
    stack_diameter: float | Float[Array, ""],
    ambient_temperature: float | Float[Array, ""],
    wind_speed: float | Float[Array, ""],
    stability_class: str = "D",
) -> Float[Array, "*obs"]:
    """Effective release height H_eff(x) = H_stack + Δh(x) [m].

    Args:
        x: Downwind distance [m].
        stack_height: Physical stack height H_stack [m].
        exit_velocity: Stack exit velocity v_s [m s⁻¹].
        exit_temperature: Stack gas temperature T_s [K].
        stack_diameter: Stack diameter d [m].
        ambient_temperature: Ambient temperature T_a [K].
        wind_speed: Wind speed u at the stack top [m s⁻¹].
        stability_class: Pasquill–Gifford class ``'A'`` … ``'F'``.

    Returns:
        H_eff at each x.

    Examples:
        At the stack itself the plume has not risen yet:

        >>> import jax.numpy as jnp
        >>> effective_height(
        ...     jnp.array([0.0]), stack_height=30.0, exit_velocity=10.0,
        ...     exit_temperature=400.0, stack_diameter=2.0,
        ...     ambient_temperature=290.0, wind_speed=5.0,
        ... ).tolist()
        [30.0]
    """
    rise = briggs_plume_rise(
        x,
        exit_velocity=exit_velocity,
        exit_temperature=exit_temperature,
        stack_diameter=stack_diameter,
        ambient_temperature=ambient_temperature,
        wind_speed=wind_speed,
        stability_class=stability_class,
    )
    # H_eff = H_stack + Δh
    return stack_height + rise

gaussian_plume_model

gaussian_plume_model(observations: ndarray | None = None, receptor_coords: tuple[ndarray, ndarray, ndarray] | None = None, source_location: tuple[float, float, float] | None = None, wind_u: float | None = None, wind_v: float | None = None, stability_class: str = 'C', prior_emission_rate_mean: float = 0.1, prior_emission_rate_std: float = 0.05, infer_stability: bool = False, background_prior_std: float = 5e-07, obs_noise_std: float = 5e-07, noise: Noise | None = None) -> jnp.ndarray

NumPyro model for the steady-state Gaussian plume.

Priors

emission_rate ~ LogNormal(μ_log, σ_log) [kg/s] stability_idx ~ Categorical(1/6, ..., 1/6) (optional) [A..F] background ~ HalfNormal(σ_bg) [kg/m³]

where (μ_log, σ_log) are the log-space parameters derived from the requested real-scale (prior_emission_rate_mean, prior_emission_rate_std) via the method-of-moments inversion in :func:_lognormal_from_moments.

Likelihood

obs ~ noise(f_plume(emission_rate, ...) + background)

with noise defaulting to GaussianNoise(obs_noise_std), i.e. obs ~ Normal(f_plume + background, σ_obs).

Parameters

observations : jnp.ndarray, shape (N,), optional Observed concentrations [kg/m³]. If None, runs in prior-predictive mode. receptor_coords : tuple of jnp.ndarray, optional (x, y, z), each shape (N,), in the fixed frame [m]. Required whenever observations is not None. source_location : tuple of float, optional (x, y, z) source coordinates [m]. Required whenever observations is not None. wind_u, wind_v : float, optional Wind velocity components [m/s]. Required whenever observations is not None. stability_class : str Stability class used if infer_stability=False. Default 'C'. prior_emission_rate_mean, prior_emission_rate_std : float Real-scale prior mean / std for the emission rate [kg/s]. The implementation converts these to log-space parameters internally, so the realised prior has the requested mean and variance. infer_stability : bool If True, sample a categorical stability_idx latent over A-F. background_prior_std : float HalfNormal scale for the background concentration [kg/m³]. obs_noise_std : float Observation noise σ [kg/m³] of the default Gaussian noise. Ignored when noise is given. noise : Noise, optional Observation-noise model, e.g. StudentTNoise(scale, df=5.0) or HeteroscedasticNoise(std) from :mod:plumax.gauss_plume.likelihoods. Default None is GaussianNoise(obs_noise_std).

Raises

ValueError If observations is provided but any of receptor_coords, source_location, wind_u, or wind_v is None — that combination disconnects the likelihood from the forward model and would produce a meaningless posterior. Prior-predictive runs (observations=None) are allowed with any subset of forward-model inputs.

Source code in src/plumax/gauss_plume/inference.py

def gaussian_plume_model(
    observations: jnp.ndarray | None = None,
    receptor_coords: tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray] | None = None,
    source_location: tuple[float, float, float] | None = None,
    wind_u: float | None = None,
    wind_v: float | None = None,
    stability_class: str = "C",
    prior_emission_rate_mean: float = 0.1,
    prior_emission_rate_std: float = 0.05,
    infer_stability: bool = False,
    background_prior_std: float = 5e-7,
    obs_noise_std: float = 5e-7,
    noise: Noise | None = None,
) -> jnp.ndarray:
    """NumPyro model for the steady-state Gaussian plume.

    Priors:
        emission_rate ~ LogNormal(μ_log, σ_log)                    [kg/s]
        stability_idx ~ Categorical(1/6, ..., 1/6)  (optional)     [A..F]
        background    ~ HalfNormal(σ_bg)                            [kg/m³]

    where ``(μ_log, σ_log)`` are the log-space parameters derived from the
    requested real-scale ``(prior_emission_rate_mean, prior_emission_rate_std)``
    via the method-of-moments inversion in :func:`_lognormal_from_moments`.

    Likelihood:
        obs ~ noise(f_plume(emission_rate, ...) + background)

    with ``noise`` defaulting to ``GaussianNoise(obs_noise_std)``, i.e.
    ``obs ~ Normal(f_plume + background, σ_obs)``.

    Parameters
    ----------
    observations : jnp.ndarray, shape (N,), optional
        Observed concentrations [kg/m³]. If None, runs in prior-predictive mode.
    receptor_coords : tuple of jnp.ndarray, optional
        ``(x, y, z)``, each shape (N,), in the fixed frame [m]. Required
        whenever ``observations`` is not None.
    source_location : tuple of float, optional
        ``(x, y, z)`` source coordinates [m]. Required whenever ``observations``
        is not None.
    wind_u, wind_v : float, optional
        Wind velocity components [m/s]. Required whenever ``observations``
        is not None.
    stability_class : str
        Stability class used if ``infer_stability=False``. Default ``'C'``.
    prior_emission_rate_mean, prior_emission_rate_std : float
        **Real-scale** prior mean / std for the emission rate [kg/s]. The
        implementation converts these to log-space parameters internally, so
        the realised prior has the requested mean and variance.
    infer_stability : bool
        If True, sample a categorical ``stability_idx`` latent over A-F.
    background_prior_std : float
        HalfNormal scale for the background concentration [kg/m³].
    obs_noise_std : float
        Observation noise σ [kg/m³] of the default Gaussian noise. Ignored
        when ``noise`` is given.
    noise : Noise, optional
        Observation-noise model, e.g. ``StudentTNoise(scale, df=5.0)`` or
        ``HeteroscedasticNoise(std)`` from :mod:`plumax.gauss_plume.likelihoods`.
        Default ``None`` is ``GaussianNoise(obs_noise_std)``.

    Raises
    ------
    ValueError
        If ``observations`` is provided but any of ``receptor_coords``,
        ``source_location``, ``wind_u``, or ``wind_v`` is ``None`` — that
        combination disconnects the likelihood from the forward model and
        would produce a meaningless posterior. Prior-predictive runs
        (``observations=None``) are allowed with any subset of forward-model
        inputs.
    """
    forward_inputs_present = (
        receptor_coords is not None
        and source_location is not None
        and wind_u is not None
        and wind_v is not None
    )
    if observations is not None and not forward_inputs_present:
        missing = [
            name
            for name, value in (
                ("receptor_coords", receptor_coords),
                ("source_location", source_location),
                ("wind_u", wind_u),
                ("wind_v", wind_v),
            )
            if value is None
        ]
        raise ValueError(
            "gaussian_plume_model: `observations` were provided but the "
            "forward model cannot be evaluated — missing inputs: "
            f"{', '.join(missing)}. Either pass all of "
            "(receptor_coords, source_location, wind_u, wind_v) or set "
            "observations=None for prior-predictive mode."
        )

    mu_log, sigma_log = _lognormal_from_moments(
        prior_emission_rate_mean, prior_emission_rate_std
    )
    emission_rate = numpyro.sample("emission_rate", dist.LogNormal(mu_log, sigma_log))

    if infer_stability:
        stability_probs = jnp.ones(len(STABILITY_CLASSES)) / len(STABILITY_CLASSES)
        stability_idx = numpyro.sample(
            "stability_idx", dist.Categorical(probs=stability_probs)
        )
        all_params = jnp.stack([BRIGGS_DISPERSION_PARAMS[c] for c in STABILITY_CLASSES])
        dispersion_params = all_params[stability_idx]
    else:
        dispersion_params = get_dispersion_params(stability_class)

    background = numpyro.sample("background", dist.HalfNormal(background_prior_std))

    if forward_inputs_present:
        x_obs, y_obs, z_obs = receptor_coords
        src_x, src_y, src_z = source_location
        predicted = plume_concentration(
            x_obs,
            y_obs,
            z_obs,
            src_x,
            src_y,
            src_z,
            wind_u,
            wind_v,
            emission_rate,
            dispersion_params,
        )
        predicted = predicted + background
    else:
        # Prior-predictive mode (observations is None and no forward inputs).
        predicted = background

    noise_model = GaussianNoise(std=obs_noise_std) if noise is None else noise
    return noise_model("obs", predicted, observations)

infer_emission_rate

infer_emission_rate(observations: NDArray, observation_coords: tuple[NDArray, NDArray, NDArray], source_location: tuple[float, float, float], wind_speed: float, wind_direction: float, stability_class: str = 'C', infer_stability: bool = False, prior_mean: float = 0.1, prior_std: float = 0.05, num_warmup: int = 500, num_samples: int = 1000, num_chains: int = 1, seed: int = 0, progress_bar: bool = False, print_summary: bool = False, obs_noise_std: float = 5e-07, background_prior_std: float = 5e-07, noise: Noise | None = None) -> EmissionRateSamples

Infer the plume emission rate via NUTS.

Wraps :func:gaussian_plume_model with input validation, wind-from-angle to (u, v) conversion, and posterior unpacking into NumPy arrays.

Identifiability target: three overpasses with varied wind direction should give a posterior CV(Q) ≤ 50 %, and a single overpass CV(Q) ≤ 100 % under a tight met wind prior (Tier I design, "Identifiability — quantitative target"). A wider posterior means the data do not constrain Q yet.

Parameters

observations : ndarray, shape (N,) Observed concentrations [kg/m³]. observation_coords : tuple of ndarray (x, y, z), each shape (N,), in the fixed frame [m]. source_location : tuple of float, (3,) (x, y, z) source coordinates [m]. wind_speed : float Wind speed magnitude [m/s]. Must be > 0. wind_direction : float Wind direction [degrees from North, meteorological "from" convention]. stability_class : str Stability class used if infer_stability=False. Default 'C'. infer_stability : bool If True, jointly sample a categorical stability-class latent. prior_mean, prior_std : float Prior parameters for the emission rate [kg/s]. num_warmup, num_samples, num_chains : int NUTS sampler configuration. seed : int PRNG seed. progress_bar : bool If True, show the NumPyro progress bar. Default False. print_summary : bool If True, call mcmc.print_summary() after the run. obs_noise_std : float Observation-noise σ for the Gaussian likelihood [kg/m³]. Set this to the sensor noise so the posterior width reflects the instrument rather than the hard-coded default. Ignored when noise is given. background_prior_std : float HalfNormal scale for the background-concentration prior [kg/m³]. noise : Noise, optional Observation-noise model (see :func:gaussian_plume_model). Default None is GaussianNoise(obs_noise_std), today's likelihood.

Returns

samples : EmissionRateSamples Posterior draws keyed by site name: 'emission_rate', 'background', and (if infer_stability) 'stability_idx'. A dict that also carries the Tier V per-event payload (emission_rate / emission_std attributes, samples and prior_logpdf).

Raises

ValueError On invalid stability class, non-positive wind speed, shape mismatches between observations and observation_coords, or empty data.

Source code in src/plumax/gauss_plume/inference.py

def infer_emission_rate(
    observations: NDArray,
    observation_coords: tuple[NDArray, NDArray, NDArray],
    source_location: tuple[float, float, float],
    wind_speed: float,
    wind_direction: float,
    stability_class: str = "C",
    infer_stability: bool = False,
    prior_mean: float = 0.1,
    prior_std: float = 0.05,
    num_warmup: int = 500,
    num_samples: int = 1000,
    num_chains: int = 1,
    seed: int = 0,
    progress_bar: bool = False,
    print_summary: bool = False,
    obs_noise_std: float = 5e-7,
    background_prior_std: float = 5e-7,
    noise: Noise | None = None,
) -> EmissionRateSamples:
    """Infer the plume emission rate via NUTS.

    Wraps :func:`gaussian_plume_model` with input validation, wind-from-angle
    to (u, v) conversion, and posterior unpacking into NumPy arrays.

    Identifiability target: three overpasses with varied wind direction
    should give a posterior CV(Q) ≤ 50 %, and a single overpass CV(Q) ≤ 100 %
    under a tight met wind prior (Tier I design, "Identifiability —
    quantitative target"). A wider posterior means the data do not constrain
    Q yet.

    Parameters
    ----------
    observations : ndarray, shape (N,)
        Observed concentrations [kg/m³].
    observation_coords : tuple of ndarray
        ``(x, y, z)``, each shape (N,), in the fixed frame [m].
    source_location : tuple of float, (3,)
        ``(x, y, z)`` source coordinates [m].
    wind_speed : float
        Wind speed magnitude [m/s]. Must be > 0.
    wind_direction : float
        Wind direction [degrees from North, meteorological "from" convention].
    stability_class : str
        Stability class used if ``infer_stability=False``. Default ``'C'``.
    infer_stability : bool
        If True, jointly sample a categorical stability-class latent.
    prior_mean, prior_std : float
        Prior parameters for the emission rate [kg/s].
    num_warmup, num_samples, num_chains : int
        NUTS sampler configuration.
    seed : int
        PRNG seed.
    progress_bar : bool
        If True, show the NumPyro progress bar. Default False.
    print_summary : bool
        If True, call ``mcmc.print_summary()`` after the run.
    obs_noise_std : float
        Observation-noise σ for the Gaussian likelihood [kg/m³]. Set this to
        the sensor noise so the posterior width reflects the instrument rather
        than the hard-coded default. Ignored when ``noise`` is given.
    background_prior_std : float
        HalfNormal scale for the background-concentration prior [kg/m³].
    noise : Noise, optional
        Observation-noise model (see :func:`gaussian_plume_model`). Default
        ``None`` is ``GaussianNoise(obs_noise_std)``, today's likelihood.

    Returns
    -------
    samples : EmissionRateSamples
        Posterior draws keyed by site name: ``'emission_rate'``,
        ``'background'``, and (if ``infer_stability``) ``'stability_idx'``.
        A ``dict`` that also carries the Tier V per-event payload
        (``emission_rate`` / ``emission_std`` attributes, ``samples`` and
        ``prior_logpdf``).

    Raises
    ------
    ValueError
        On invalid stability class, non-positive wind speed, shape mismatches
        between ``observations`` and ``observation_coords``, or empty data.
    """
    if not (wind_speed > 0.0):
        raise ValueError(
            f"infer_emission_rate: `wind_speed` must be > 0 (got {wind_speed!r})"
        )
    if not infer_stability and stability_class not in BRIGGS_DISPERSION_PARAMS:
        raise ValueError(
            f"infer_emission_rate: `stability_class` must be one of "
            f"{STABILITY_CLASSES}, got {stability_class!r}"
        )
    obs = np.asarray(observations)
    if obs.size == 0:
        raise ValueError("infer_emission_rate: `observations` must contain ≥ 1 point")
    if len(observation_coords) != 3:
        raise ValueError(
            "infer_emission_rate: `observation_coords` must be (x, y, z); "
            f"got length {len(observation_coords)}"
        )
    coords = tuple(np.asarray(c) for c in observation_coords)
    for axis_name, arr in zip("xyz", coords, strict=False):
        if arr.shape != obs.shape:
            raise ValueError(
                f"infer_emission_rate: `observation_coords` axis '{axis_name}' "
                f"has shape {arr.shape} ≠ observations shape {obs.shape}"
            )
    if not (prior_mean > 0.0):
        raise ValueError(
            f"infer_emission_rate: `prior_mean` must be > 0 (got {prior_mean!r})"
        )
    if not (prior_std > 0.0):
        raise ValueError(
            f"infer_emission_rate: `prior_std` must be > 0 (got {prior_std!r})"
        )
    if not (obs_noise_std > 0.0):
        raise ValueError(
            f"infer_emission_rate: `obs_noise_std` must be > 0 (got {obs_noise_std!r})"
        )
    if not (background_prior_std > 0.0):
        raise ValueError(
            "infer_emission_rate: `background_prior_std` must be > 0 "
            f"(got {background_prior_std!r})"
        )

    theta_rad = np.deg2rad(wind_direction)
    wind_u = float(-wind_speed * np.sin(theta_rad))
    wind_v = float(-wind_speed * np.cos(theta_rad))

    obs_jax = jnp.asarray(obs)
    coords_jax = tuple(jnp.asarray(c) for c in coords)

    inner = NUTS(gaussian_plume_model)
    # Discrete latents can't be handled by NUTS on their own without
    # `funsor` enumeration. `DiscreteHMCGibbs` wraps NUTS with a Gibbs
    # step over the categorical stability-class site, avoiding the extra
    # dependency while giving a valid joint sampler.
    kernel = DiscreteHMCGibbs(inner) if infer_stability else inner
    mcmc = MCMC(
        kernel,
        num_warmup=num_warmup,
        num_samples=num_samples,
        num_chains=num_chains,
        progress_bar=progress_bar,
    )
    mcmc.run(
        jax.random.PRNGKey(seed),
        observations=obs_jax,
        receptor_coords=coords_jax,
        source_location=source_location,
        wind_u=wind_u,
        wind_v=wind_v,
        stability_class=stability_class,
        prior_emission_rate_mean=prior_mean,
        prior_emission_rate_std=prior_std,
        infer_stability=infer_stability,
        obs_noise_std=obs_noise_std,
        background_prior_std=background_prior_std,
        noise=noise,
    )
    if print_summary:
        mcmc.print_summary()

    draws = {k: np.asarray(v) for k, v in mcmc.get_samples().items()}
    return EmissionRateSamples(draws, prior_mean=prior_mean, prior_std=prior_std)

infer_multi_source

infer_multi_source(observations: NDArray, observation_coords: tuple[NDArray, NDArray, NDArray], *, wind_speed: float, wind_direction: float, k_max: int = 10, candidate_locations: NDArray | None = None, source_bounds: tuple[tuple[float, float], tuple[float, float]] | None = None, source_height: float = 0.0, stability_class: str = 'C', prior_mean: float = 0.1, prior_std: float = 0.05, prior_active: float = 0.3, obs_noise_std: float = 5e-07, background_prior_std: float = 5e-07, noise: Noise | None = None, num_warmup: int = 500, num_samples: int = 1000, num_chains: int = 1, seed: int = 0, progress_bar: bool = False) -> MultiSourcePosterior

Infer up to K_max sources and their on/off indicators.

Wraps :func:multi_source_plume_model: NUTS for the continuous sites, and a Gibbs step (DiscreteHMCGibbs) over the activity indicators, as for the stability class in :func:infer_emission_rate.

Parameters

observations : ndarray, shape (N,) Observed concentrations [kg/m³]. observation_coords : tuple of ndarray (x, y, z), each shape (N,), in the fixed frame [m]. wind_speed : float Wind speed [m/s], > 0. wind_direction : float Wind direction [degrees from North, meteorological "from" convention]. k_max, candidate_locations, source_bounds, source_height, stability_class Slot layout; see :func:multi_source_plume_model. prior_mean, prior_std : float Real-scale prior moments of each Q_k [kg/s]. prior_active, obs_noise_std, background_prior_std, noise See :func:multi_source_plume_model. num_warmup, num_samples, num_chains, seed : int Sampler configuration. progress_bar : bool Show NumPyro's progress bar.

Returns

posterior : MultiSourcePosterior Draws of every slot; count_probabilities() is the posterior over the active count and per_source(k) the slot's summary.

Raises

ValueError On a non-positive wind speed, an unknown stability class, mismatched shapes, non-positive prior or noise scales, or an invalid slot layout.

Source code in src/plumax/gauss_plume/inference.py

def infer_multi_source(
    observations: NDArray,
    observation_coords: tuple[NDArray, NDArray, NDArray],
    *,
    wind_speed: float,
    wind_direction: float,
    k_max: int = 10,
    candidate_locations: NDArray | None = None,
    source_bounds: tuple[tuple[float, float], tuple[float, float]] | None = None,
    source_height: float = 0.0,
    stability_class: str = "C",
    prior_mean: float = 0.1,
    prior_std: float = 0.05,
    prior_active: float = 0.3,
    obs_noise_std: float = 5e-7,
    background_prior_std: float = 5e-7,
    noise: Noise | None = None,
    num_warmup: int = 500,
    num_samples: int = 1000,
    num_chains: int = 1,
    seed: int = 0,
    progress_bar: bool = False,
) -> MultiSourcePosterior:
    """Infer up to K_max sources and their on/off indicators.

    Wraps :func:`multi_source_plume_model`: NUTS for the continuous sites, and
    a Gibbs step (``DiscreteHMCGibbs``) over the activity indicators, as for
    the stability class in :func:`infer_emission_rate`.

    Parameters
    ----------
    observations : ndarray, shape (N,)
        Observed concentrations [kg/m³].
    observation_coords : tuple of ndarray
        ``(x, y, z)``, each shape (N,), in the fixed frame [m].
    wind_speed : float
        Wind speed [m/s], > 0.
    wind_direction : float
        Wind direction [degrees from North, meteorological "from" convention].
    k_max, candidate_locations, source_bounds, source_height, stability_class
        Slot layout; see :func:`multi_source_plume_model`.
    prior_mean, prior_std : float
        Real-scale prior moments of each Q_k [kg/s].
    prior_active, obs_noise_std, background_prior_std, noise
        See :func:`multi_source_plume_model`.
    num_warmup, num_samples, num_chains, seed : int
        Sampler configuration.
    progress_bar : bool
        Show NumPyro's progress bar.

    Returns
    -------
    posterior : MultiSourcePosterior
        Draws of every slot; ``count_probabilities()`` is the posterior over
        the active count and ``per_source(k)`` the slot's summary.

    Raises
    ------
    ValueError
        On a non-positive wind speed, an unknown stability class, mismatched
        shapes, non-positive prior or noise scales, or an invalid slot layout.
    """
    if not (wind_speed > 0.0):
        raise ValueError(
            f"infer_multi_source: `wind_speed` must be > 0 (got {wind_speed!r})"
        )
    if stability_class not in BRIGGS_DISPERSION_PARAMS:
        raise ValueError(
            f"infer_multi_source: `stability_class` must be one of "
            f"{STABILITY_CLASSES}, got {stability_class!r}"
        )
    obs = np.asarray(observations)
    if obs.size == 0:
        raise ValueError("infer_multi_source: `observations` must contain ≥ 1 point")
    if len(observation_coords) != 3:
        raise ValueError(
            "infer_multi_source: `observation_coords` must be (x, y, z); "
            f"got length {len(observation_coords)}"
        )
    coords = tuple(np.asarray(c) for c in observation_coords)
    for axis_name, arr in zip("xyz", coords, strict=True):
        if arr.shape != obs.shape:
            raise ValueError(
                f"infer_multi_source: `observation_coords` axis '{axis_name}' "
                f"has shape {arr.shape} ≠ observations shape {obs.shape}"
            )
    for name, value in (
        ("prior_mean", prior_mean),
        ("prior_std", prior_std),
        ("obs_noise_std", obs_noise_std),
        ("background_prior_std", background_prior_std),
    ):
        if not (value > 0.0):
            raise ValueError(
                f"infer_multi_source: `{name}` must be > 0 (got {value!r})"
            )
    _check_multi_source_layout(k_max, candidate_locations, source_bounds, prior_active)

    theta_rad = np.deg2rad(wind_direction)
    wind_u = float(-wind_speed * np.sin(theta_rad))
    wind_v = float(-wind_speed * np.cos(theta_rad))

    return _run_masked_k(
        multi_source_plume_model,
        jnp.asarray(obs),
        k_max=k_max,
        candidate_locations=candidate_locations,
        prior_active=prior_active,
        num_warmup=num_warmup,
        num_samples=num_samples,
        num_chains=num_chains,
        seed=seed,
        progress_bar=progress_bar,
        receptor_coords=tuple(jnp.asarray(c) for c in coords),
        wind_u=wind_u,
        wind_v=wind_v,
        source_bounds=source_bounds,
        source_height=source_height,
        stability_class=stability_class,
        prior_emission_rate_mean=prior_mean,
        prior_emission_rate_std=prior_std,
        background_prior_std=background_prior_std,
        obs_noise_std=obs_noise_std,
        noise=noise,
    )

multi_source_plume_model

multi_source_plume_model(observations: ndarray | None = None, *, receptor_coords: tuple[ndarray, ndarray, ndarray], wind_u: float, wind_v: float, k_max: int = 10, candidate_locations: ndarray | None = None, source_bounds: tuple[tuple[float, float], tuple[float, float]] | None = None, source_height: float = 0.0, stability_class: str = 'C', prior_emission_rate_mean: float = 0.1, prior_emission_rate_std: float = 0.05, prior_active: float = 0.3, background_prior_std: float = 5e-07, obs_noise_std: float = 5e-07, noise: Noise | None = None) -> jnp.ndarray

Masked-K NumPyro model: up to K_max plumes, each switched on or off.

Priors, for each slot k = 1 … K_max: a_k ~ Bernoulli(prior_active) [on/off] Q_k ~ LogNormal(μ_log, σ_log) [kg/s] (x_k, y_k) = candidate k, or x ordered uniform and y uniform in source_bounds [m] background ~ HalfNormal(σ_bg) [kg/m³]

Likelihood

obs ~ noise(Σ_k a_k Q_k c₁(x_k, y_k) + background)

with noise defaulting to GaussianNoise(obs_noise_std). The active count K = Σ_k a_k is the posterior over the number of sources.

Label switching: with free locations the slots are exchangeable, so their x positions are drawn as the order statistics of K_max uniforms on [x_low, x_high] (x_low + (x_high − x_low) · cumsum(w)[:K_max] with w ~ Dirichlet(1, …, 1) of length K_max + 1), which pins slot k to the k-th source from the west. Fixed candidate locations (known facilities) have no symmetry to break.

Promotion to RJMCMC (Tier I / IV design): masked-K caps K at K_max. Track the posterior mass at the cap, P(K = K_max | y) (:meth:MultiSourcePosterior.saturated_fraction), per overpass; when more than 5 % of overpasses saturate K_max, the cap is binding and the trans-dimensional (RJMCMC) model is due.

Parameters

observations : jnp.ndarray, shape (N,), optional Observed concentrations [kg/m³]; None for predictive draws. receptor_coords : tuple of jnp.ndarray (x, y, z), each shape (N,), in the fixed frame [m]. wind_u, wind_v : float Wind velocity components [m/s]. k_max : int Number of slots K_max (default 10, the design's v1 cap). candidate_locations : jnp.ndarray, shape (k_max, 2), optional Fixed (x, y) per slot, e.g. inventory facilities. source_bounds : ((x_low, x_high), (y_low, y_high)), optional Box the free slot locations are drawn from. Give exactly one of this and candidate_locations. source_height : float Release height shared by every slot [m]. stability_class : str Pasquill–Gifford class 'A' … 'F'. prior_emission_rate_mean, prior_emission_rate_std : float Real-scale prior moments of each Q_k [kg/s]. prior_active : float Prior probability that a slot is on, in (0, 1]. 1.0 keeps every slot on (no discrete sites), which with K_max = 1 is the single-source model. background_prior_std : float HalfNormal scale of the background [kg/m³]. obs_noise_std : float σ of the default Gaussian noise; ignored when noise is given. noise : Noise, optional Observation-noise model from :mod:plumax.gauss_plume.likelihoods.

Raises

ValueError On an invalid slot layout, k_max < 1 or prior_active outside (0, 1].

Source code in src/plumax/gauss_plume/inference.py

def multi_source_plume_model(
    observations: jnp.ndarray | None = None,
    *,
    receptor_coords: tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray],
    wind_u: float,
    wind_v: float,
    k_max: int = 10,
    candidate_locations: jnp.ndarray | None = None,
    source_bounds: tuple[tuple[float, float], tuple[float, float]] | None = None,
    source_height: float = 0.0,
    stability_class: str = "C",
    prior_emission_rate_mean: float = 0.1,
    prior_emission_rate_std: float = 0.05,
    prior_active: float = 0.3,
    background_prior_std: float = 5e-7,
    obs_noise_std: float = 5e-7,
    noise: Noise | None = None,
) -> jnp.ndarray:
    """Masked-K NumPyro model: up to K_max plumes, each switched on or off.

    Priors, for each slot k = 1 … K_max:
        a_k ~ Bernoulli(prior_active)                               [on/off]
        Q_k ~ LogNormal(μ_log, σ_log)                               [kg/s]
        (x_k, y_k) = candidate k, or x ordered uniform and y uniform in
                     ``source_bounds``                               [m]
        background ~ HalfNormal(σ_bg)                               [kg/m³]

    Likelihood:
        obs ~ noise(Σ_k a_k Q_k c₁(x_k, y_k) + background)

    with ``noise`` defaulting to ``GaussianNoise(obs_noise_std)``. The active
    count K = Σ_k a_k is the posterior over the number of sources.

    Label switching: with free locations the slots are exchangeable, so their
    x positions are drawn as the order statistics of K_max uniforms on
    ``[x_low, x_high]`` (``x_low + (x_high − x_low) · cumsum(w)[:K_max]`` with
    ``w ~ Dirichlet(1, …, 1)`` of length K_max + 1), which pins slot k to the
    k-th source from the west. Fixed candidate locations (known facilities)
    have no symmetry to break.

    Promotion to RJMCMC (Tier I / IV design): masked-K caps K at K_max. Track
    the posterior mass at the cap, P(K = K_max | y)
    (:meth:`MultiSourcePosterior.saturated_fraction`), per overpass; when more
    than 5 % of overpasses saturate K_max, the cap is binding and the
    trans-dimensional (RJMCMC) model is due.

    Parameters
    ----------
    observations : jnp.ndarray, shape (N,), optional
        Observed concentrations [kg/m³]; ``None`` for predictive draws.
    receptor_coords : tuple of jnp.ndarray
        ``(x, y, z)``, each shape (N,), in the fixed frame [m].
    wind_u, wind_v : float
        Wind velocity components [m/s].
    k_max : int
        Number of slots K_max (default 10, the design's v1 cap).
    candidate_locations : jnp.ndarray, shape (k_max, 2), optional
        Fixed ``(x, y)`` per slot, e.g. inventory facilities.
    source_bounds : ((x_low, x_high), (y_low, y_high)), optional
        Box the free slot locations are drawn from. Give exactly one of this
        and ``candidate_locations``.
    source_height : float
        Release height shared by every slot [m].
    stability_class : str
        Pasquill–Gifford class ``'A'`` … ``'F'``.
    prior_emission_rate_mean, prior_emission_rate_std : float
        Real-scale prior moments of each Q_k [kg/s].
    prior_active : float
        Prior probability that a slot is on, in (0, 1]. ``1.0`` keeps every
        slot on (no discrete sites), which with K_max = 1 is the single-source
        model.
    background_prior_std : float
        HalfNormal scale of the background [kg/m³].
    obs_noise_std : float
        σ of the default Gaussian noise; ignored when ``noise`` is given.
    noise : Noise, optional
        Observation-noise model from :mod:`plumax.gauss_plume.likelihoods`.

    Raises
    ------
    ValueError
        On an invalid slot layout, ``k_max < 1`` or ``prior_active`` outside
        (0, 1].
    """
    _check_multi_source_layout(k_max, candidate_locations, source_bounds, prior_active)
    dispersion_params = get_dispersion_params(stability_class)

    sources, background = _masked_source_prior(
        k_max=k_max,
        candidate_locations=candidate_locations,
        source_bounds=source_bounds,
        prior_emission_rate_mean=prior_emission_rate_mean,
        prior_emission_rate_std=prior_emission_rate_std,
        prior_active=prior_active,
        background_prior_std=background_prior_std,
    )
    predicted = sources.concentration(
        receptor_coords,
        source_z=source_height,
        wind_u=wind_u,
        wind_v=wind_v,
        dispersion_params=dispersion_params,
    )
    predicted = predicted + background

    noise_model = GaussianNoise(std=obs_noise_std) if noise is None else noise
    return noise_model("obs", predicted, observations)