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
inference — NumPyro Bayesian inference for Q (and optional stability)

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.

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]

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) -> 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,) [a_y, b_y, c_y, a_z, b_z, c_z] (Briggs parameters).

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,
) -> 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,)
        ``[a_y, b_y, c_y, a_z, b_z, c_z]`` (Briggs parameters).

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

    z_diff = z - source_z
    z_sum = z + source_z
    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, domain_x: tuple[float, float, int], domain_y: tuple[float, float, int], domain_z: tuple[float, float, int], background_conc: float = 0.0) -> 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 One of 'A', ..., 'F'. 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.

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,
    domain_x: tuple[float, float, int],
    domain_y: tuple[float, float, int],
    domain_z: tuple[float, float, int],
    background_conc: float = 0.0,
) -> 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
        One of ``'A'``, ..., ``'F'``.
    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``.

    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})"
        )
    if 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})"
            )

    dispersion_params = get_dispersion_params(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,
    )
    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_class,
            "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

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) -> 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 ~ Normal(f_plume(emission_rate, ...) + 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³].

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,
) -> 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 ~ Normal(f_plume(emission_rate, ...) + 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³].

    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

    return numpyro.sample(
        "obs",
        dist.Normal(predicted, obs_noise_std),
        obs=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) -> dict[str, NDArray]

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.

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. background_prior_std : float HalfNormal scale for the background-concentration prior [kg/m³].

Returns¶

samples : dict[str, ndarray] Posterior draws keyed by site name: 'emission_rate', 'background', and (if infer_stability) 'stability_idx'.

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,
) -> dict[str, NDArray]:
    """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.

    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.
    background_prior_std : float
        HalfNormal scale for the background-concentration prior [kg/m³].

    Returns
    -------
    samples : dict[str, ndarray]
        Posterior draws keyed by site name: ``'emission_rate'``,
        ``'background'``, and (if ``infer_stability``) ``'stability_idx'``.

    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,
    )
    if print_summary:
        mcmc.print_summary()

    return {k: np.asarray(v) for k, v in mcmc.get_samples().items()}

Gaussian plume¶

gauss_plume ¶

plume_concentration_vmap module-attribute ¶

calculate_briggs_dispersion ¶

Parameters¶

Returns¶

get_dispersion_params ¶

Parameters¶

Returns¶

Raises¶

plume_concentration ¶

Parameters¶

Returns¶

rotate_to_wind_frame ¶

Parameters¶

Returns¶

simulate_plume ¶

Parameters¶

Returns¶

Raises¶

gaussian_plume_model ¶

Parameters¶

Raises¶

infer_emission_rate ¶

Parameters¶

Returns¶

Raises¶

plume_concentration_vmap `module-attribute` ¶