Multi-source

Masked-K multi-source plumes: up to K_max source slots, each with an emission rate, a location and an on/off indicator. MaskedSources evaluates the summed field and MultiSourcePosterior turns posterior draws into a posterior over the active count and per-source summaries for the population catalogue. The NumPyro model and runner are multi_source_plume_model and infer_multi_source in Gaussian plume. Tier I, reused by the Tier IV coupled version.

multi_source

Masked-K source slots and their posterior (implementation in plumax.gauss_plume._src.multi_source).

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