Gaussian puff

Time-resolved Gaussian puff (Pasquill–Gifford), diffrax wind integration, and Ornstein–Uhlenbeck sub-grid turbulence. Tier I.

gauss_puff

Gaussian puff model (JAX + diffrax + NumPyro).

Submodules

• dispersion — Pasquill-Gifford σ(s) coefficients + shared Briggs dispatch
• wind — piecewise-linear WindSchedule + diffrax cumulative integrals
• turbulence — Ornstein-Uhlenbeck sub-grid position disturbances on puff centres
• puff — single-puff kernel, evolve_puffs, simulate_puff (xarray)
• inference — NumPyro Bayesian inference for Q (constant) and Q_i (random walk)

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

puff_concentration_vmap module-attribute

puff_concentration_vmap = vmap(puff_concentration, in_axes=(None, None, None, 0, 0, 0, 0, 0, 0, 0), out_axes=0)

Vectorised over puffs (per-puff params batched along axis 0, receptor shared).

PuffState

Bases: Module

Container for the Lagrangian state of a puff ensemble at a single time.

Attributes

release_times : Float[Array, "N"] Puff release times [s]. x, y, z : Float[Array, "N"] Puff-centre positions at the evaluation time [m]. travel_distance : Float[Array, "N"] Cumulative travel distance since release [m]. mass : Float[Array, "N"] Puff mass [kg]. Inactive puffs (not yet released) have mass 0.

OUTurbulence dataclass

OUTurbulence(sigma_fluctuations: float = 0.5, correlation_time: float = 60.0)

Ornstein-Uhlenbeck turbulence for puff-release position disturbances.

Attributes

sigma_fluctuations : float OU fluctuation amplitude [m/√s]. Together with correlation_time it sets the steady-state RMS offset σ_∞ = σ_f · √(τ_c / 2) [m]. correlation_time : float OU correlation time [s]. Offsets for puffs released within τ_c of each other are strongly correlated; puffs released more than a few τ_c apart are effectively independent.

Notes

Gorroño et al. (2023) calibrated comparable parameters against large-eddy simulations of methane plumes; defaults in the Eucalyptus release are (σ_f, τ_c) = (0.5 m/s, 60 s) giving a steady-state RMS offset of about 2.7 m — appropriate for 20 m-GSD Sentinel-2 pixels.

stationary_std property

stationary_std: float

Steady-state RMS per-component offset σ_f · √(τ_c / 2) [m].

WindSchedule

Bases: Module

Piecewise-linear time-varying wind field.

Attributes

times : Float[Array, "T"] Monotone increasing time knots [s]. u_wind : Float[Array, "T"] x-component of wind velocity at the knots [m/s]. v_wind : Float[Array, "T"] y-component of wind velocity at the knots [m/s].

Notes

The schedule behaves as a piecewise-linear interpolant between knots. Queries outside [times[0], times[-1]] linearly extrapolate — consumers that need strict clamping should range-check times before calling.

from_speed_direction classmethod

from_speed_direction(times: ndarray | ndarray, wind_speed: ndarray | ndarray, wind_direction: ndarray | ndarray) -> WindSchedule

Build a schedule from (speed, direction) in meteorological "from" convention.

wind_direction = 270° is a wind from the west (u > 0); wind_direction = 0° is a wind from the north (v < 0).

u = −|V| · sin(θ),   v = −|V| · cos(θ),   θ = direction (radians).

Parameters

times : array-like, shape (T,) Monotone time knots [s]. wind_speed : array-like, shape (T,) Wind speed at each knot [m/s]. Must be ≥ 0. wind_direction : array-like, shape (T,) Wind direction at each knot [degrees from North].

Returns

WindSchedule

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

@classmethod
def from_speed_direction(
    cls,
    times: np.ndarray | jnp.ndarray,
    wind_speed: np.ndarray | jnp.ndarray,
    wind_direction: np.ndarray | jnp.ndarray,
) -> WindSchedule:
    """Build a schedule from (speed, direction) in meteorological "from" convention.

    ``wind_direction = 270°`` is a wind *from* the west (u > 0);
    ``wind_direction = 0°`` is a wind *from* the north (v < 0).

        u = −|V| · sin(θ),   v = −|V| · cos(θ),   θ = direction (radians).

    Parameters
    ----------
    times : array-like, shape (T,)
        Monotone time knots [s].
    wind_speed : array-like, shape (T,)
        Wind speed at each knot [m/s]. Must be ≥ 0.
    wind_direction : array-like, shape (T,)
        Wind direction at each knot [degrees from North].

    Returns
    -------
    WindSchedule
    """
    times_j = jnp.asarray(times, dtype=jnp.float32)
    speed_j = jnp.asarray(wind_speed, dtype=jnp.float32)
    theta = jnp.deg2rad(jnp.asarray(wind_direction, dtype=jnp.float32))
    u = -speed_j * jnp.sin(theta)
    v = -speed_j * jnp.cos(theta)
    return cls(times=times_j, u_wind=u, v_wind=v)

wind_at

wind_at(t: Float[Array, '']) -> tuple[Float[Array, ''], Float[Array, '']]

Interpolate the wind velocity components at time t.

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

def wind_at(self, t: Float[Array, ""]) -> tuple[Float[Array, ""], Float[Array, ""]]:
    """Interpolate the wind velocity components at time ``t``."""
    return self._interp_u().evaluate(t), self._interp_v().evaluate(t)

speed_at

speed_at(t: Float[Array, '']) -> Float[Array, '']

Interpolate the wind speed |V(t)| at time t.

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

def speed_at(self, t: Float[Array, ""]) -> Float[Array, ""]:
    """Interpolate the wind speed |V(t)| at time ``t``."""
    u, v = self.wind_at(t)
    return jnp.sqrt(u**2 + v**2)

calculate_briggs_dispersion_xyz

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

Evaluate σ_x(s), σ_y(s), σ_z(s) from Briggs parameters.

Thin wrapper around the plume model's Briggs evaluator that additionally exposes σ_x = σ_y so this routine is signature-compatible with :func:calculate_pg_dispersion. Lets callers swap dispersion schemes without changing downstream code.

Parameters

distance : jnp.ndarray Puff travel distance [m]. params : jnp.ndarray, shape (6,) Briggs parameters [a_y, b_y, c_y, a_z, b_z, c_z].

Returns

sigma_x, sigma_y, sigma_z : jnp.ndarray Dispersion coefficients [m]. σ_x = σ_y.

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

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

    Thin wrapper around the plume model's Briggs evaluator that additionally
    exposes ``σ_x = σ_y`` so this routine is signature-compatible with
    :func:`calculate_pg_dispersion`. Lets callers swap dispersion schemes
    without changing downstream code.

    Parameters
    ----------
    distance : jnp.ndarray
        Puff travel distance [m].
    params : jnp.ndarray, shape (6,)
        Briggs parameters ``[a_y, b_y, c_y, a_z, b_z, c_z]``.

    Returns
    -------
    sigma_x, sigma_y, sigma_z : jnp.ndarray
        Dispersion coefficients [m]. ``σ_x = σ_y``.
    """
    sigma_y, sigma_z = calculate_briggs_dispersion(distance, params)
    return sigma_y, sigma_y, sigma_z

calculate_pg_dispersion

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

Evaluate σ_x(s), σ_y(s), σ_z(s) from Pasquill-Gifford parameters.

σ_i(s) = exp(a_i + b_i · ln(s) + c_i · (ln s)²) for i ∈ {y, z}; σ_x := σ_y.

Parameters

distance : jnp.ndarray Puff travel distance [m]. Non-positive values are clamped to 1 m so that ln(s) stays finite. params : jnp.ndarray, shape (6,) [a_y, b_y, c_y, a_z, b_z, c_z].

Returns

sigma_x, sigma_y : jnp.ndarray Horizontal dispersion coefficients [m]. σ_x = σ_y. sigma_z : jnp.ndarray Vertical dispersion coefficient [m].

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

@jit
def calculate_pg_dispersion(
    distance: jnp.ndarray,
    params: jnp.ndarray,
) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
    """Evaluate σ_x(s), σ_y(s), σ_z(s) from Pasquill-Gifford parameters.

    σ_i(s) = exp(a_i + b_i · ln(s) + c_i · (ln s)²) for i ∈ {y, z};
    σ_x := σ_y.

    Parameters
    ----------
    distance : jnp.ndarray
        Puff travel distance [m]. Non-positive values are clamped to 1 m
        so that ``ln(s)`` stays finite.
    params : jnp.ndarray, shape (6,)
        ``[a_y, b_y, c_y, a_z, b_z, c_z]``.

    Returns
    -------
    sigma_x, sigma_y : jnp.ndarray
        Horizontal dispersion coefficients [m]. ``σ_x = σ_y``.
    sigma_z : jnp.ndarray
        Vertical dispersion coefficient [m].
    """
    ay, by, cy, az, bz, cz = params
    s = jnp.maximum(distance, 1.0)
    ln_s = jnp.log(s)
    sigma_y = jnp.exp(ay + by * ln_s + cy * ln_s**2)
    sigma_z = jnp.exp(az + bz * ln_s + cz * ln_s**2)
    return sigma_y, sigma_y, sigma_z

get_dispersion_scheme

get_dispersion_scheme(scheme: str) -> tuple[dict[str, jnp.ndarray], DispersionFn]

Return (params_dict, calculator_fn) for a named dispersion scheme.

Parameters

scheme : str One of 'pg' (Pasquill-Gifford, default) or 'briggs'.

Raises

ValueError If scheme is not a registered key.

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

def get_dispersion_scheme(
    scheme: str,
) -> tuple[dict[str, jnp.ndarray], DispersionFn]:
    """Return ``(params_dict, calculator_fn)`` for a named dispersion scheme.

    Parameters
    ----------
    scheme : str
        One of ``'pg'`` (Pasquill-Gifford, default) or ``'briggs'``.

    Raises
    ------
    ValueError
        If ``scheme`` is not a registered key.
    """
    if scheme not in DISPERSION_SCHEMES:
        raise ValueError(
            f"get_dispersion_scheme: `scheme` must be one of "
            f"{tuple(DISPERSION_SCHEMES)}, got {scheme!r}"
        )
    return DISPERSION_SCHEMES[scheme]

get_pg_params

get_pg_params(stability_class: str) -> jnp.ndarray

Return the Pasquill-Gifford 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_puff/dispersion.py

def get_pg_params(stability_class: str) -> jnp.ndarray:
    """Return the Pasquill-Gifford 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 PG_DISPERSION_PARAMS:
        raise ValueError(
            f"get_pg_params: stability_class must be one of "
            f"{STABILITY_CLASSES}, got {stability_class!r}"
        )
    return PG_DISPERSION_PARAMS[stability_class]

evolve_puffs

evolve_puffs(schedule: WindSchedule, release_times: Float[Array, N], current_time: Float[Array, ''], source_location: tuple[float, float, float], puff_mass: Float[Array, N] | float, position_disturbance: tuple[Float[Array, N], Float[Array, N]] | None = None) -> PuffState

Advance a puff ensemble to current_time under a time-varying wind.

Uses :func:cumulative_wind_integrals to compute (I_u, I_v, S) at the release times and at current_time in a single diffrax solve, then derives per-puff positions and travel distances as differences.

Puffs with release_times > current_time are "not yet released" — their mass is set to 0 and position is placed at the source, so they contribute nothing to the downstream field.

Parameters

schedule : WindSchedule Time-varying wind field. release_times : Float[Array, "N"] Monotone increasing puff release times [s]. current_time : Float[Array, ""] Evaluation time [s]. source_location : tuple of float, (3,) (x, y, z) source coordinates [m]. z ≥ 0. puff_mass : Float[Array, "N"] or float Per-puff mass [kg], or a scalar applied to all puffs. position_disturbance : tuple of Float[Array, "N"], optional Per-puff (Δx, Δy) offsets [m] applied on top of the wind advection. Intended for OU-process turbulence samples (see :func:plumax.gauss_puff.turbulence.sample_ou_offsets), but any caller-provided offset array is accepted. Offsets for not-yet-released puffs are ignored.

Returns

puff_state : PuffState Puff positions, travel distances, and masses at current_time.

Source code in src/plumax/gauss_puff/puff.py

@eqx.filter_jit
def evolve_puffs(
    schedule: WindSchedule,
    release_times: Float[Array, N],
    current_time: Float[Array, ""],
    source_location: tuple[float, float, float],
    puff_mass: Float[Array, N] | float,
    position_disturbance: tuple[Float[Array, N], Float[Array, N]] | None = None,
) -> PuffState:
    """Advance a puff ensemble to ``current_time`` under a time-varying wind.

    Uses :func:`cumulative_wind_integrals` to compute ``(I_u, I_v, S)`` at the
    release times and at ``current_time`` in a single diffrax solve, then
    derives per-puff positions and travel distances as differences.

    Puffs with ``release_times > current_time`` are "not yet released" — their
    mass is set to 0 and position is placed at the source, so they contribute
    nothing to the downstream field.

    Parameters
    ----------
    schedule : WindSchedule
        Time-varying wind field.
    release_times : Float[Array, "N"]
        Monotone increasing puff release times [s].
    current_time : Float[Array, ""]
        Evaluation time [s].
    source_location : tuple of float, (3,)
        ``(x, y, z)`` source coordinates [m]. ``z ≥ 0``.
    puff_mass : Float[Array, "N"] or float
        Per-puff mass [kg], or a scalar applied to all puffs.
    position_disturbance : tuple of Float[Array, "N"], optional
        Per-puff ``(Δx, Δy)`` offsets [m] applied on top of the wind advection.
        Intended for OU-process turbulence samples (see
        :func:`plumax.gauss_puff.turbulence.sample_ou_offsets`),
        but any caller-provided offset array is accepted. Offsets for
        not-yet-released puffs are ignored.

    Returns
    -------
    puff_state : PuffState
        Puff positions, travel distances, and masses at ``current_time``.
    """
    src_x, src_y, src_z = source_location

    save_at = jnp.concatenate(
        [release_times, jnp.atleast_1d(jnp.asarray(current_time))]
    )
    sort_idx = jnp.argsort(save_at)
    save_at_sorted = save_at[sort_idx]

    I_u_s, I_v_s, S_s = cumulative_wind_integrals(schedule, save_at_sorted)

    unsort_idx = jnp.argsort(sort_idx)
    I_u = I_u_s[unsort_idx]
    I_v = I_v_s[unsort_idx]
    S = S_s[unsort_idx]

    I_u_rel, I_v_rel, S_rel = I_u[:-1], I_v[:-1], S[:-1]
    I_u_now, I_v_now, S_now = I_u[-1], I_v[-1], S[-1]

    active = release_times <= current_time
    x = src_x + jnp.where(active, I_u_now - I_u_rel, 0.0)
    y = src_y + jnp.where(active, I_v_now - I_v_rel, 0.0)
    if position_disturbance is not None:
        dx_puff, dy_puff = position_disturbance
        x = x + jnp.where(active, dx_puff, 0.0)
        y = y + jnp.where(active, dy_puff, 0.0)
    z = jnp.full_like(release_times, src_z)
    s = jnp.where(active, S_now - S_rel, 0.0)

    mass_arr = jnp.broadcast_to(jnp.asarray(puff_mass), release_times.shape)
    mass = jnp.where(active, mass_arr, 0.0)

    return PuffState(
        release_times=release_times,
        x=x,
        y=y,
        z=z,
        travel_distance=s,
        mass=mass,
    )

frequency_to_release_interval

frequency_to_release_interval(release_frequency: float) -> float

Convert a release frequency f [Hz] into a release interval Δt [s].

Δt = 1 / f.

Source code in src/plumax/gauss_puff/puff.py

def frequency_to_release_interval(release_frequency: float) -> float:
    """Convert a release frequency f [Hz] into a release interval Δt [s].

    ``Δt = 1 / f``.
    """
    if not (release_frequency > 0.0):
        raise ValueError(
            f"frequency_to_release_interval: `release_frequency` must be > 0 "
            f"(got {release_frequency!r})"
        )
    return 1.0 / release_frequency

make_release_times

make_release_times(t_start: float, t_end: float, release_frequency: float) -> jnp.ndarray

Evenly spaced puff release times in [t_start, t_end).

Parameters

t_start, t_end : float Simulation time window [s]. t_start < t_end required. release_frequency : float Puff release rate [Hz]. Must be > 0.

Returns

release_times : jnp.ndarray, shape (N,) Times [t_start, t_start + Δt, t_start + 2Δt, ...] with Δt = 1 / release_frequency, all strictly less than t_end.

Notes

The count is ceil((t_end − t_start) / Δt), so when the window is not an exact multiple of Δt the final, partial interval still gets a release (e.g. [0, 0.5) at 1 Hz yields one puff at t = 0, and [0, 10.5) yields eleven). That trailing puff represents a shorter slice of the continuous source, so its emitted mass must be weighted by its actual interval duration — see :func:release_durations and how simulate_puff assigns puff_mass. All release times remain strictly < t_end.

Source code in src/plumax/gauss_puff/puff.py

def make_release_times(
    t_start: float,
    t_end: float,
    release_frequency: float,
) -> jnp.ndarray:
    """Evenly spaced puff release times in ``[t_start, t_end)``.

    Parameters
    ----------
    t_start, t_end : float
        Simulation time window [s]. ``t_start < t_end`` required.
    release_frequency : float
        Puff release rate [Hz]. Must be > 0.

    Returns
    -------
    release_times : jnp.ndarray, shape (N,)
        Times ``[t_start, t_start + Δt, t_start + 2Δt, ...]`` with
        ``Δt = 1 / release_frequency``, all strictly less than ``t_end``.

    Notes
    -----
    The count is ``ceil((t_end − t_start) / Δt)``, so when the window is not an
    exact multiple of ``Δt`` the **final, partial interval still gets a
    release** (e.g. ``[0, 0.5)`` at 1 Hz yields one puff at ``t = 0``, and
    ``[0, 10.5)`` yields eleven). That trailing puff represents a shorter slice
    of the continuous source, so its emitted mass must be weighted by its actual
    interval duration — see :func:`release_durations` and how ``simulate_puff``
    assigns ``puff_mass``. All release times remain strictly ``< t_end``.
    """
    if not (t_end > t_start):
        raise ValueError(
            f"make_release_times: `t_end` must be > `t_start` "
            f"(got t_start={t_start!r}, t_end={t_end!r})"
        )
    dt = frequency_to_release_interval(release_frequency)
    n = int(np.ceil((t_end - t_start) / dt))
    return jnp.asarray(t_start + np.arange(n) * dt, dtype=jnp.float32)

puff_concentration

puff_concentration(x: ndarray, y: ndarray, z: ndarray, puff_x: float, puff_y: float, puff_z: float, sigma_x: float, sigma_y: float, sigma_z: float, puff_mass: float) -> jnp.ndarray

Evaluate the concentration from a single Gaussian puff (JIT).

Includes ground-reflection via an image source at z = −puff_z.

Parameters

x, y, z : jnp.ndarray Receptor coordinates [m]. Broadcast-compatible. puff_x, puff_y, puff_z : float Puff-centre position [m]. puff_z ≥ 0. sigma_x, sigma_y, sigma_z : float Dispersion coefficients [m]. All > 0. puff_mass : float Mass of the puff [kg]. Must be ≥ 0.

Returns

concentration : jnp.ndarray Mass concentration [kg/m³], same shape as the broadcast of the input coordinates.

Source code in src/plumax/gauss_puff/puff.py

@jit
def puff_concentration(
    x: jnp.ndarray,
    y: jnp.ndarray,
    z: jnp.ndarray,
    puff_x: float,
    puff_y: float,
    puff_z: float,
    sigma_x: float,
    sigma_y: float,
    sigma_z: float,
    puff_mass: float,
) -> jnp.ndarray:
    """Evaluate the concentration from a single Gaussian puff (JIT).

    Includes ground-reflection via an image source at ``z = −puff_z``.

    Parameters
    ----------
    x, y, z : jnp.ndarray
        Receptor coordinates [m]. Broadcast-compatible.
    puff_x, puff_y, puff_z : float
        Puff-centre position [m]. ``puff_z ≥ 0``.
    sigma_x, sigma_y, sigma_z : float
        Dispersion coefficients [m]. All > 0.
    puff_mass : float
        Mass of the puff [kg]. Must be ≥ 0.

    Returns
    -------
    concentration : jnp.ndarray
        Mass concentration [kg/m³], same shape as the broadcast of the input
        coordinates.
    """
    dx = x - puff_x
    dy = y - puff_y
    dz_direct = z - puff_z
    dz_reflected = z + puff_z

    normalization = jnp.power(2.0 * jnp.pi, 1.5) * sigma_x * sigma_y * sigma_z

    exp_x = jnp.exp(-0.5 * jnp.square(dx / sigma_x))
    exp_y = jnp.exp(-0.5 * jnp.square(dy / sigma_y))
    exp_z_direct = jnp.exp(-0.5 * jnp.square(dz_direct / sigma_z))
    exp_z_reflected = jnp.exp(-0.5 * jnp.square(dz_reflected / sigma_z))

    return (
        (puff_mass / normalization) * exp_x * exp_y * (exp_z_direct + exp_z_reflected)
    )

release_durations

release_durations(release_times: ndarray, t_end: float) -> jnp.ndarray

Per-puff emission-interval durations [s] for a continuous source.

Each puff carries the mass emitted over its interval, i.e. up to the next release (or t_end for the last puff). The durations therefore sum to t_end − release_times[0] exactly: interior intervals equal the release spacing Δt while the final, possibly partial interval is the remainder. Multiply by the emission rate Q to get the per-puff mass.

Parameters

release_times : jnp.ndarray, shape (N,) Monotone increasing release times [s], as from :func:make_release_times. t_end : float End of the emission window [s]; must exceed the last release time.

Returns

durations : jnp.ndarray, shape (N,) Per-puff interval durations [s], all > 0.

Source code in src/plumax/gauss_puff/puff.py

def release_durations(
    release_times: jnp.ndarray,
    t_end: float,
) -> jnp.ndarray:
    """Per-puff emission-interval durations [s] for a continuous source.

    Each puff carries the mass emitted over its interval, i.e. up to the next
    release (or ``t_end`` for the last puff). The durations therefore sum to
    ``t_end − release_times[0]`` exactly: interior intervals equal the release
    spacing ``Δt`` while the final, possibly partial interval is the remainder.
    Multiply by the emission rate ``Q`` to get the per-puff mass.

    Parameters
    ----------
    release_times : jnp.ndarray, shape (N,)
        Monotone increasing release times [s], as from
        :func:`make_release_times`.
    t_end : float
        End of the emission window [s]; must exceed the last release time.

    Returns
    -------
    durations : jnp.ndarray, shape (N,)
        Per-puff interval durations [s], all > 0.
    """
    rel = np.asarray(release_times, dtype=np.float64)
    if rel.size == 0:
        return jnp.zeros((0,), dtype=jnp.float32)
    edges = np.append(rel, float(t_end))
    return jnp.asarray(np.diff(edges), dtype=jnp.float32)

release_interval_to_frequency

release_interval_to_frequency(release_interval: float) -> float

Convert a release interval Δt [s] into a release frequency f [Hz].

f = 1 / Δt. Both Δt and f are common ways to parametrise puff cadence; this pair of helpers makes the conversion explicit.

Source code in src/plumax/gauss_puff/puff.py

def release_interval_to_frequency(release_interval: float) -> float:
    """Convert a release interval Δt [s] into a release frequency f [Hz].

    ``f = 1 / Δt``. Both Δt and f are common ways to parametrise puff cadence;
    this pair of helpers makes the conversion explicit.
    """
    if not (release_interval > 0.0):
        raise ValueError(
            f"release_interval_to_frequency: `release_interval` must be > 0 "
            f"(got {release_interval!r})"
        )
    return 1.0 / release_interval

simulate_puff

simulate_puff(emission_rate: float | ndarray, source_location: tuple[float, float, float], wind_speed: ndarray, wind_direction: ndarray, stability_class: str, domain_x: tuple[float, float, int], domain_y: tuple[float, float, int], domain_z: tuple[float, float, int], time_array: ndarray, release_frequency: float = 1.0, scheme: str = 'pg', background_conc: float = 0.0, turbulence: OUTurbulence | None = None, turbulence_seed: int | Generator | None = None) -> xr.Dataset

Simulate a time-resolved Gaussian-puff dispersion field on a 3-D grid.

Wraps :func:evolve_puffs + :func:simulate_puff_field with grid construction, a diffrax-driven time-varying wind solve, and optional column integration. Returns an :class:xarray.Dataset with variables on a (time, x, y, z) grid.

Parameters

emission_rate : float or ndarray Continuous emission rate Q [kg/s]. Scalar constant Q, or an array of length n_puffs specifying Q at each puff release time. The puff mass is Q · Δt_release with Δt_release = 1/release_frequency. Entries must be ≥ 0 (scalar or array); Q = 0 yields a zero field. source_location : tuple of float, (3,) (x, y, z) source coordinates [m]. z ≥ 0. wind_speed : ndarray, shape (T_w,) Wind speed [m/s] at each entry in time_array. Must be same length. wind_direction : ndarray, shape (T_w,) Wind direction [deg from North]. Same length as time_array. 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. time_array : ndarray, shape (T,) Monotone increasing evaluation times [s]. release_frequency : float Puff release rate [Hz]. Δt_release = 1/release_frequency. Default 1 Hz. scheme : str Dispersion scheme: 'pg' (Pasquill-Gifford, default) or 'briggs' (sharing coefficients with the steady plume model). background_conc : float Additive background concentration [kg/m³]. Default 0. turbulence : OUTurbulence, optional Ornstein-Uhlenbeck sub-grid turbulence applied as per-puff (Δx, Δy) offsets at release time. Adds realistic meander on top of the wind-driven advection. Defaults to None (deterministic puffs). turbulence_seed : int or numpy.random.Generator, optional Seed for the OU sampler; required for bit-exact reproducibility when turbulence is set.

Returns

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

Raises

ValueError On invalid stability class, non-positive emission rate, shape mismatches, unknown scheme, or malformed domain tuples.

Source code in src/plumax/gauss_puff/puff.py

def simulate_puff(
    emission_rate: float | np.ndarray,
    source_location: tuple[float, float, float],
    wind_speed: np.ndarray,
    wind_direction: np.ndarray,
    stability_class: str,
    domain_x: tuple[float, float, int],
    domain_y: tuple[float, float, int],
    domain_z: tuple[float, float, int],
    time_array: np.ndarray,
    release_frequency: float = 1.0,
    scheme: str = "pg",
    background_conc: float = 0.0,
    turbulence: OUTurbulence | None = None,
    turbulence_seed: int | np.random.Generator | None = None,
) -> xr.Dataset:
    """Simulate a time-resolved Gaussian-puff dispersion field on a 3-D grid.

    Wraps :func:`evolve_puffs` + :func:`simulate_puff_field` with grid
    construction, a diffrax-driven time-varying wind solve, and optional
    column integration. Returns an :class:`xarray.Dataset` with variables
    on a ``(time, x, y, z)`` grid.

    Parameters
    ----------
    emission_rate : float or ndarray
        Continuous emission rate Q [kg/s]. Scalar constant Q, or an array of
        length ``n_puffs`` specifying Q at each puff release time. The puff
        mass is ``Q · Δt_release`` with ``Δt_release = 1/release_frequency``.
        Entries must be ≥ 0 (scalar or array); ``Q = 0`` yields a zero field.
    source_location : tuple of float, (3,)
        ``(x, y, z)`` source coordinates [m]. ``z ≥ 0``.
    wind_speed : ndarray, shape (T_w,)
        Wind speed [m/s] at each entry in ``time_array``. Must be same length.
    wind_direction : ndarray, shape (T_w,)
        Wind direction [deg from North]. Same length as ``time_array``.
    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.
    time_array : ndarray, shape (T,)
        Monotone increasing evaluation times [s].
    release_frequency : float
        Puff release rate [Hz]. Δt_release = 1/release_frequency. Default 1 Hz.
    scheme : str
        Dispersion scheme: ``'pg'`` (Pasquill-Gifford, default) or
        ``'briggs'`` (sharing coefficients with the steady plume model).
    background_conc : float
        Additive background concentration [kg/m³]. Default 0.
    turbulence : OUTurbulence, optional
        Ornstein-Uhlenbeck sub-grid turbulence applied as per-puff ``(Δx, Δy)``
        offsets at release time. Adds realistic meander on top of the
        wind-driven advection. Defaults to ``None`` (deterministic puffs).
    turbulence_seed : int or numpy.random.Generator, optional
        Seed for the OU sampler; required for bit-exact reproducibility
        when ``turbulence`` is set.

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

    Raises
    ------
    ValueError
        On invalid stability class, non-positive emission rate, shape
        mismatches, unknown scheme, or malformed domain tuples.
    """
    import xarray as xr  # lazy — only this entry point requires xarray

    scheme_params_dict, dispersion_fn = get_dispersion_scheme(scheme)

    if stability_class not in scheme_params_dict:
        raise ValueError(
            f"simulate_puff: `stability_class` must be one of "
            f"{STABILITY_CLASSES} (got {stability_class!r})"
        )
    if len(source_location) != 3:
        raise ValueError(
            "simulate_puff: `source_location` must be (x, y, z); "
            f"got length {len(source_location)}"
        )
    if source_location[2] < 0.0:
        raise ValueError(
            f"simulate_puff: source height must be ≥ 0 (got z={source_location[2]!r})"
        )
    time_array = np.asarray(time_array, dtype=np.float32)
    wind_speed = np.asarray(wind_speed, dtype=np.float32)
    wind_direction = np.asarray(wind_direction, dtype=np.float32)
    if time_array.ndim != 1 or time_array.size < 2:
        raise ValueError("simulate_puff: `time_array` must be 1-D with ≥ 2 entries")
    if wind_speed.shape != time_array.shape:
        raise ValueError(
            f"simulate_puff: `wind_speed` shape {wind_speed.shape} must match "
            f"`time_array` shape {time_array.shape}"
        )
    if wind_direction.shape != time_array.shape:
        raise ValueError(
            f"simulate_puff: `wind_direction` shape {wind_direction.shape} "
            f"must match `time_array` shape {time_array.shape}"
        )
    if np.any(wind_speed < 0.0):
        raise ValueError("simulate_puff: `wind_speed` entries must be ≥ 0")

    for name, domain in (
        ("domain_x", domain_x),
        ("domain_y", domain_y),
        ("domain_z", domain_z),
    ):
        if len(domain) != 3:
            raise ValueError(
                f"simulate_puff: `{name}` must be (start, stop, n_points); "
                f"got length {len(domain)}"
            )
        start, stop, n = domain
        if not (start < stop):
            raise ValueError(
                f"simulate_puff: `{name}` requires start < stop "
                f"(got start={start!r}, stop={stop!r})"
            )
        if int(n) < 2:
            raise ValueError(
                f"simulate_puff: `{name}` requires n_points ≥ 2 (got n={n!r})"
            )
    if not (release_frequency > 0.0):
        raise ValueError(
            f"simulate_puff: `release_frequency` must be > 0 "
            f"(got {release_frequency!r})"
        )

    dispersion_params = scheme_params_dict[stability_class]

    x_grid = np.linspace(domain_x[0], domain_x[1], int(domain_x[2]))
    y_grid = np.linspace(domain_y[0], domain_y[1], int(domain_y[2]))
    z_grid = np.linspace(domain_z[0], domain_z[1], int(domain_z[2]))
    n_x, n_y, n_z = len(x_grid), len(y_grid), len(z_grid)
    n_t = len(time_array)

    X, Y, Z = np.meshgrid(x_grid, y_grid, z_grid, indexing="ij")
    x_flat = jnp.asarray(X.ravel(), dtype=jnp.float32)
    y_flat = jnp.asarray(Y.ravel(), dtype=jnp.float32)
    z_flat = jnp.asarray(Z.ravel(), dtype=jnp.float32)

    dt_release = frequency_to_release_interval(release_frequency)
    release_times = make_release_times(
        float(time_array[0]), float(time_array[-1]), release_frequency
    )
    n_puffs = int(release_times.shape[0])
    # Per-puff emission-interval durations: full Δt for interior puffs, the
    # remainder for the final (possibly partial) interval. Weighting mass by
    # these durations conserves the total emitted mass Q·(t_end − t_start) even
    # when the window is not an exact multiple of the release interval.
    durations = release_durations(release_times, float(time_array[-1]))

    if np.ndim(emission_rate) == 0:
        q_scalar = float(emission_rate)
        if q_scalar < 0.0:
            raise ValueError(
                "simulate_puff: scalar `emission_rate` must be ≥ 0 "
                f"(got {emission_rate!r})"
            )
        puff_mass = jnp.asarray(q_scalar * durations, dtype=jnp.float32)
    else:
        Q = np.asarray(emission_rate, dtype=np.float32)
        if Q.shape != (n_puffs,):
            raise ValueError(
                f"simulate_puff: array `emission_rate` must have shape "
                f"({n_puffs},) matching the number of releases in "
                f"[t_start, t_end) at release_frequency={release_frequency}, "
                f"got {Q.shape}"
            )
        if np.any(Q < 0.0):
            raise ValueError("simulate_puff: `emission_rate` entries must be ≥ 0")
        puff_mass = jnp.asarray(Q * np.asarray(durations), dtype=jnp.float32)

    schedule = WindSchedule.from_speed_direction(time_array, wind_speed, wind_direction)

    if turbulence is not None:
        from plumax.gauss_puff.turbulence import sample_ou_offsets

        dx_np, dy_np = sample_ou_offsets(
            turbulence,
            np.asarray(release_times),
            seed=turbulence_seed,
        )
        position_disturbance = (
            jnp.asarray(dx_np, dtype=jnp.float32),
            jnp.asarray(dy_np, dtype=jnp.float32),
        )
    else:
        position_disturbance = None

    concentration = np.zeros((n_t, n_x, n_y, n_z), dtype=np.float32)
    for i_t, t_current in enumerate(time_array):
        puff_state = evolve_puffs(
            schedule,
            release_times,
            jnp.asarray(float(t_current), dtype=jnp.float32),
            source_location,
            puff_mass,
            position_disturbance=position_disturbance,
        )
        field_flat = simulate_puff_field(
            (x_flat, y_flat, z_flat),
            puff_state,
            dispersion_params,
            dispersion_fn,
        )
        concentration[i_t] = np.asarray(field_flat).reshape((n_x, n_y, n_z))

    concentration = concentration + background_conc

    # Trapezoidal vertical integration over the (endpoint-inclusive) z grid;
    # a plain ``sum * dz`` over-counts a vertically uniform component (such as
    # the background) by one grid interval. See gauss_plume.simulate_plume.
    column = np.trapezoid(concentration, x=z_grid, axis=3)

    ds = xr.Dataset(
        data_vars={
            "concentration": (["time", "x", "y", "z"], concentration),
            "column_concentration": (["time", "x", "y"], column),
            "wind_speed": (["time"], np.asarray(wind_speed)),
            "wind_direction": (["time"], np.asarray(wind_direction)),
        },
        coords={
            "time": np.asarray(time_array),
            "x": x_grid,
            "y": y_grid,
            "z": z_grid,
        },
        attrs={
            "title": "Gaussian puff dispersion (JAX + diffrax)",
            "source_x": float(source_location[0]),
            "source_y": float(source_location[1]),
            "source_z": float(source_location[2]),
            "stability_class": stability_class,
            "dispersion_scheme": scheme,
            "release_frequency": release_frequency,
            "release_interval": dt_release,
            "n_puffs": n_puffs,
            "background_concentration": background_conc,
            "ou_sigma_fluctuations": (
                float(turbulence.sigma_fluctuations) if turbulence else 0.0
            ),
            "ou_correlation_time": (
                float(turbulence.correlation_time) if turbulence else 0.0
            ),
        },
    )
    ds["concentration"].attrs = {"long_name": "Mass concentration", "units": "kg/m^3"}
    ds["column_concentration"].attrs = {
        "long_name": "Column-integrated concentration",
        "units": "kg/m^2",
    }
    ds["wind_speed"].attrs = {"long_name": "Wind speed", "units": "m/s"}
    ds["wind_direction"].attrs = {
        "long_name": "Wind direction",
        "units": "degrees from North (meteorological)",
    }
    return ds

simulate_puff_field

simulate_puff_field(receptor_coords: tuple[Float[Array, P], Float[Array, P], Float[Array, P]], puff_state: PuffState, dispersion_params: Float[Array, 6], dispersion_fn: Callable) -> Float[Array, P]

Sum the contributions of all puffs at each receptor.

Parameters

receptor_coords : tuple of Float[Array, "P"] (x, y, z) receptor coordinates [m]. puff_state : PuffState Per-puff positions, travel distances, and masses. dispersion_params : Float[Array, "6"] Dispersion coefficient vector for the chosen scheme and stability class. dispersion_fn : callable A function with signature (distance, params) → (σ_x, σ_y, σ_z). See :func:plumax.gauss_puff.dispersion.calculate_pg_dispersion and :func:calculate_briggs_dispersion_xyz.

Returns

concentration : Float[Array, "P"] Total concentration [kg/m³] at each receptor.

Source code in src/plumax/gauss_puff/puff.py

def simulate_puff_field(
    receptor_coords: tuple[Float[Array, P], Float[Array, P], Float[Array, P]],
    puff_state: PuffState,
    dispersion_params: Float[Array, 6],
    dispersion_fn: Callable,
) -> Float[Array, P]:
    """Sum the contributions of all puffs at each receptor.

    Parameters
    ----------
    receptor_coords : tuple of Float[Array, "P"]
        ``(x, y, z)`` receptor coordinates [m].
    puff_state : PuffState
        Per-puff positions, travel distances, and masses.
    dispersion_params : Float[Array, "6"]
        Dispersion coefficient vector for the chosen scheme and stability class.
    dispersion_fn : callable
        A function with signature ``(distance, params) → (σ_x, σ_y, σ_z)``.
        See :func:`plumax.gauss_puff.dispersion.calculate_pg_dispersion`
        and :func:`calculate_briggs_dispersion_xyz`.

    Returns
    -------
    concentration : Float[Array, "P"]
        Total concentration [kg/m³] at each receptor.
    """
    x_r, y_r, z_r = receptor_coords

    sigma_x, sigma_y, sigma_z = dispersion_fn(
        puff_state.travel_distance, dispersion_params
    )

    per_puff = puff_concentration_vmap(
        x_r,
        y_r,
        z_r,
        puff_state.x,
        puff_state.y,
        puff_state.z,
        sigma_x,
        sigma_y,
        sigma_z,
        puff_state.mass,
    )  # (N_puffs, P)
    return jnp.sum(per_puff, axis=0)

sample_ou_offsets

sample_ou_offsets(turbulence: OUTurbulence, release_times: Float[ndarray, N] | ndarray, seed: int | Generator | None = None) -> tuple[np.ndarray, np.ndarray]

Draw per-puff (Δx, Δy) offsets from a 2-D OU process at release times.

Integrates the OU SDE via its exact transition density — the two-point Gaussian update that is exact for constant coefficients — so the result is statistically correct even for coarse release_times spacing and does not accumulate Euler-Maruyama bias.

Parameters

turbulence : OUTurbulence Turbulence parameters. release_times : ndarray, shape (N,) Puff release times [s]. Must be monotone non-decreasing. The first puff starts at the OU stationary distribution so the ensemble is already mixed (no warm-up transient visible in the field). seed : int or numpy.random.Generator, optional Seed or generator for reproducibility. If None a fresh default generator is used and results will vary across calls.

Returns

dx, dy : ndarray, shape (N,) Per-puff position offsets [m], to be added to the advected puff centres in :func:~plumax.gauss_puff.puff.evolve_puffs.

Source code in src/plumax/gauss_puff/turbulence.py

def sample_ou_offsets(
    turbulence: OUTurbulence,
    release_times: Float[np.ndarray, N] | np.ndarray,
    seed: int | np.random.Generator | None = None,
) -> tuple[np.ndarray, np.ndarray]:
    """Draw per-puff (Δx, Δy) offsets from a 2-D OU process at release times.

    Integrates the OU SDE via its *exact* transition density — the two-point
    Gaussian update that is exact for constant coefficients — so the result
    is statistically correct even for coarse ``release_times`` spacing and
    does not accumulate Euler-Maruyama bias.

    Parameters
    ----------
    turbulence : OUTurbulence
        Turbulence parameters.
    release_times : ndarray, shape (N,)
        Puff release times [s]. Must be monotone non-decreasing. The first
        puff starts at the OU stationary distribution so the ensemble is
        already mixed (no warm-up transient visible in the field).
    seed : int or numpy.random.Generator, optional
        Seed or generator for reproducibility. If ``None`` a fresh default
        generator is used and results will vary across calls.

    Returns
    -------
    dx, dy : ndarray, shape (N,)
        Per-puff position offsets [m], to be added to the advected puff
        centres in :func:`~plumax.gauss_puff.puff.evolve_puffs`.
    """
    rel = np.asarray(release_times, dtype=np.float64).reshape(-1)
    if rel.size == 0:
        return np.empty((0,), dtype=np.float64), np.empty((0,), dtype=np.float64)
    if np.any(np.diff(rel) < 0.0):
        raise ValueError(
            "sample_ou_offsets: `release_times` must be monotone non-decreasing."
        )

    rng = seed if isinstance(seed, np.random.Generator) else np.random.default_rng(seed)

    sigma_f = float(turbulence.sigma_fluctuations)
    tau_c = float(turbulence.correlation_time)
    sigma_inf = turbulence.stationary_std  # per-component stationary std

    n = rel.size
    dx = np.empty(n, dtype=np.float64)
    dy = np.empty(n, dtype=np.float64)

    # Initial puff drawn from the stationary distribution so the trajectory
    # is already mixed — no need to run burn-in before the first release.
    z0 = rng.standard_normal(2) * sigma_inf
    dx[0] = z0[0]
    dy[0] = z0[1]

    # σ_f = 0 → deterministic OU, stays at the initial draw (which will
    # itself be exactly zero since σ_inf = 0). Shortcut to avoid spurious
    # zero-variance noise draws.
    if sigma_f == 0.0:
        dx[:] = 0.0
        dy[:] = 0.0
        return dx, dy

    for i in range(1, n):
        dt = rel[i] - rel[i - 1]
        decay = np.exp(-dt / tau_c)
        # Conditional variance of p(t+Δt) | p(t) under the exact OU transition.
        step_var = sigma_inf**2 * (1.0 - decay**2)
        step_std = np.sqrt(max(step_var, 0.0))
        eps = rng.standard_normal(2)
        dx[i] = decay * dx[i - 1] + step_std * eps[0]
        dy[i] = decay * dy[i - 1] + step_std * eps[1]

    return dx, dy

cumulative_wind_integrals

cumulative_wind_integrals(schedule: WindSchedule, save_at: Float[Array, M], rtol: float = 1e-06, atol: float = 1e-08) -> tuple[Float[Array, M], Float[Array, M], Float[Array, M]]

Compute cumulative wind integrals via a single diffrax ODE solve.

Integrates d(I_u, I_v, S)/dt = (u(t), v(t), √(u² + v²)) from schedule.times[0] to max(save_at[-1], schedule.times[-1]) with initial condition (0, 0, 0) using :class:diffrax.Tsit5 with an adaptive :class:diffrax.PIDController, saving at the requested times.

Parameters

schedule : WindSchedule Time-varying wind field. save_at : Float[Array, "M"] Monotone non-decreasing times [s] at which to return the integrals. Save times outside [schedule.times[0], schedule.times[-1]] are evaluated using the linear extrapolation supported by :attr:WindSchedule — use the schedule knots to bound the physically meaningful range when that matters. rtol, atol : float Relative and absolute tolerances for the adaptive step controller.

Returns

I_u, I_v, S : Float[Array, "M"] Cumulative integrals of u, v, and wind speed, each evaluated at save_at.

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

@eqx.filter_jit
def cumulative_wind_integrals(
    schedule: WindSchedule,
    save_at: Float[Array, M],
    rtol: float = 1e-6,
    atol: float = 1e-8,
) -> tuple[Float[Array, M], Float[Array, M], Float[Array, M]]:
    """Compute cumulative wind integrals via a single diffrax ODE solve.

    Integrates ``d(I_u, I_v, S)/dt = (u(t), v(t), √(u² + v²))`` from
    ``schedule.times[0]`` to ``max(save_at[-1], schedule.times[-1])`` with
    initial condition ``(0, 0, 0)`` using :class:`diffrax.Tsit5` with an
    adaptive :class:`diffrax.PIDController`, saving at the requested times.

    Parameters
    ----------
    schedule : WindSchedule
        Time-varying wind field.
    save_at : Float[Array, "M"]
        **Monotone non-decreasing** times [s] at which to return the integrals.
        Save times outside ``[schedule.times[0], schedule.times[-1]]`` are
        evaluated using the linear extrapolation supported by
        :attr:`WindSchedule` — use the schedule knots to bound the physically
        meaningful range when that matters.
    rtol, atol : float
        Relative and absolute tolerances for the adaptive step controller.

    Returns
    -------
    I_u, I_v, S : Float[Array, "M"]
        Cumulative integrals of u, v, and wind speed, each evaluated at
        ``save_at``.
    """
    interp_u = schedule._interp_u()
    interp_v = schedule._interp_v()

    def rhs(t: Float[Array, ""], y: Float[Array, 3], args) -> Float[Array, 3]:
        u = interp_u.evaluate(t)
        v = interp_v.evaluate(t)
        speed = jnp.sqrt(u**2 + v**2)
        return jnp.stack([u, v, speed])

    t0 = schedule.times[0]
    t1 = jnp.maximum(save_at[-1], schedule.times[-1])

    sol = diffrax.diffeqsolve(
        diffrax.ODETerm(rhs),
        diffrax.Tsit5(),
        t0=t0,
        t1=t1,
        dt0=None,
        y0=jnp.zeros(3, dtype=schedule.u_wind.dtype),
        saveat=diffrax.SaveAt(ts=save_at),
        stepsize_controller=diffrax.PIDController(rtol=rtol, atol=atol),
        max_steps=None,
    )
    ys = sol.ys
    return ys[:, 0], ys[:, 1], ys[:, 2]

gaussian_puff_model

gaussian_puff_model(observations: ndarray | None = None, receptor_coords: tuple[ndarray, ndarray, ndarray] | None = None, observation_times: ndarray | None = None, source_location: tuple[float, float, float] | None = None, schedule: WindSchedule | None = None, release_times: ndarray | None = None, release_interval: float | None = None, stability_class: str = 'C', scheme: str = 'pg', prior_emission_rate_mean: float = 0.1, prior_emission_rate_std: float = 0.05, background_prior_std: float = 5e-07, obs_noise_std: float = 5e-07, puff_durations: ndarray | None = None) -> jnp.ndarray

NumPyro model for the Gaussian puff forward with a constant Q prior.

Priors

emission_rate ~ LogNormal(μ_log, σ_log) [kg/s] background ~ HalfNormal(σ_bg) [kg/m³]

Likelihood

obs ~ Normal(f_puff(emission_rate, ...) + background, σ_obs)

where (μ_log, σ_log) are the log-space parameters derived from the requested real-scale (prior_emission_rate_mean, prior_emission_rate_std).

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,) [m]. observation_times : jnp.ndarray, shape (N,), optional Time at which each observation was taken [s]. Matches observations in length. source_location : tuple of float, (3,), optional schedule : WindSchedule, optional release_times : jnp.ndarray, shape (N_puffs,), optional release_interval : float, optional Δt_release [s]. Required (as a presence flag) when a forward evaluation is needed. Used for the per-puff mass m = Q · Δt only as a fallback when puff_durations is not supplied. puff_durations : jnp.ndarray, shape (N_puffs,), optional Per-puff emission-interval durations [s] (from :func:plumax.gauss_puff.puff.release_durations). When given, the per-puff mass is m_i = Q · duration_i — so the trailing, possibly partial interval carries only its actual duration, matching simulate_puff and conserving emitted mass. Falls back to Q · release_interval for every puff when omitted. stability_class : str scheme : str Dispersion scheme: 'pg' or 'briggs'. prior_emission_rate_mean, prior_emission_rate_std : float background_prior_std, obs_noise_std : float

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

def gaussian_puff_model(
    observations: jnp.ndarray | None = None,
    receptor_coords: tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray] | None = None,
    observation_times: jnp.ndarray | None = None,
    source_location: tuple[float, float, float] | None = None,
    schedule: WindSchedule | None = None,
    release_times: jnp.ndarray | None = None,
    release_interval: float | None = None,
    stability_class: str = "C",
    scheme: str = "pg",
    prior_emission_rate_mean: float = 0.1,
    prior_emission_rate_std: float = 0.05,
    background_prior_std: float = 5e-7,
    obs_noise_std: float = 5e-7,
    puff_durations: jnp.ndarray | None = None,
) -> jnp.ndarray:
    """NumPyro model for the Gaussian puff forward with a constant Q prior.

    Priors:
        emission_rate ~ LogNormal(μ_log, σ_log)                     [kg/s]
        background    ~ HalfNormal(σ_bg)                            [kg/m³]

    Likelihood:
        obs ~ Normal(f_puff(emission_rate, ...) + background, σ_obs)

    where ``(μ_log, σ_log)`` are the log-space parameters derived from the
    requested real-scale ``(prior_emission_rate_mean, prior_emission_rate_std)``.

    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,) [m].
    observation_times : jnp.ndarray, shape (N,), optional
        Time at which each observation was taken [s]. Matches ``observations``
        in length.
    source_location : tuple of float, (3,), optional
    schedule : WindSchedule, optional
    release_times : jnp.ndarray, shape (N_puffs,), optional
    release_interval : float, optional
        ``Δt_release`` [s]. Required (as a presence flag) when a forward
        evaluation is needed. Used for the per-puff mass ``m = Q · Δt`` only as
        a fallback when ``puff_durations`` is not supplied.
    puff_durations : jnp.ndarray, shape (N_puffs,), optional
        Per-puff emission-interval durations [s] (from
        :func:`plumax.gauss_puff.puff.release_durations`). When given, the
        per-puff mass is ``m_i = Q · duration_i`` — so the trailing, possibly
        partial interval carries only its actual duration, matching
        ``simulate_puff`` and conserving emitted mass. Falls back to
        ``Q · release_interval`` for every puff when omitted.
    stability_class : str
    scheme : str
        Dispersion scheme: ``'pg'`` or ``'briggs'``.
    prior_emission_rate_mean, prior_emission_rate_std : float
    background_prior_std, obs_noise_std : float
    """
    _require_forward_inputs(
        observations,
        receptor_coords,
        observation_times,
        source_location,
        schedule,
        release_times,
        release_interval,
    )
    scheme_params_dict, dispersion_fn = get_dispersion_scheme(scheme)
    if stability_class not in scheme_params_dict:
        raise ValueError(
            f"gaussian_puff_model: `stability_class` must be one of "
            f"{STABILITY_CLASSES}, got {stability_class!r}"
        )
    dispersion_params = scheme_params_dict[stability_class]

    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))
    background = numpyro.sample("background", dist.HalfNormal(background_prior_std))

    # The forward model is evaluated whenever its inputs are present —
    # independently of whether `observations` were provided. This keeps
    # NumPyro's `Predictive` path (observations=None + full inputs) producing
    # model-driven draws, while prior-predictive calls without forward inputs
    # fall through to a background-only draw.
    forward_ready = _forward_inputs_present(
        receptor_coords,
        observation_times,
        source_location,
        schedule,
        release_times,
        release_interval,
    )

    if forward_ready:
        if puff_durations is not None:
            puff_mass = emission_rate * jnp.asarray(puff_durations)
        else:
            puff_mass = emission_rate * release_interval * jnp.ones_like(release_times)
        predicted = _predict_observations(
            puff_mass,
            release_times,
            receptor_coords,
            observation_times,
            source_location,
            schedule,
            dispersion_params,
            dispersion_fn,
        )
        predicted = predicted + background
    else:
        predicted = background

    return numpyro.sample(
        "obs",
        dist.Normal(predicted, obs_noise_std),
        obs=observations,
    )

gaussian_puff_rw_model

gaussian_puff_rw_model(observations: ndarray, receptor_coords: tuple[ndarray, ndarray, ndarray], observation_times: ndarray, source_location: tuple[float, float, float], schedule: WindSchedule, release_times: ndarray, release_interval: float, stability_class: str = 'C', scheme: str = 'pg', prior_emission_rate_mean: float = 0.1, prior_emission_rate_std: float = 0.05, rw_step_std: float = 0.02, background_prior_std: float = 5e-07, obs_noise_std: float = 5e-07, puff_durations: ndarray | None = None) -> jnp.ndarray

Random-walk state-space model for a time-varying emission rate Q_i.

Each puff i has its own emission rate Q_i. A Gaussian random walk on ln Q_i gives a smoothness prior, with the first step anchored by the LogNormal(prior_mean, prior_std) prior:

ln Q_0            ~ Normal(μ_log, σ_log)
ln Q_i | ln Q_{i−1} ~ Normal(ln Q_{i−1}, rw_step_std)        for i ≥ 1
Q_i              = exp(ln Q_i)                              [kg/s]

Likelihood is the same as :func:gaussian_puff_model.

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

def gaussian_puff_rw_model(
    observations: jnp.ndarray,
    receptor_coords: tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray],
    observation_times: jnp.ndarray,
    source_location: tuple[float, float, float],
    schedule: WindSchedule,
    release_times: jnp.ndarray,
    release_interval: float,
    stability_class: str = "C",
    scheme: str = "pg",
    prior_emission_rate_mean: float = 0.1,
    prior_emission_rate_std: float = 0.05,
    rw_step_std: float = 0.02,
    background_prior_std: float = 5e-7,
    obs_noise_std: float = 5e-7,
    puff_durations: jnp.ndarray | None = None,
) -> jnp.ndarray:
    """Random-walk state-space model for a time-varying emission rate Q_i.

    Each puff ``i`` has its own emission rate ``Q_i``. A Gaussian random walk
    on ``ln Q_i`` gives a smoothness prior, with the first step anchored by
    the LogNormal(``prior_mean``, ``prior_std``) prior:

        ln Q_0            ~ Normal(μ_log, σ_log)
        ln Q_i | ln Q_{i−1} ~ Normal(ln Q_{i−1}, rw_step_std)        for i ≥ 1
        Q_i              = exp(ln Q_i)                              [kg/s]

    Likelihood is the same as :func:`gaussian_puff_model`.
    """
    scheme_params_dict, dispersion_fn = get_dispersion_scheme(scheme)
    if stability_class not in scheme_params_dict:
        raise ValueError(
            f"gaussian_puff_rw_model: `stability_class` must be one of "
            f"{STABILITY_CLASSES}, got {stability_class!r}"
        )
    dispersion_params = scheme_params_dict[stability_class]

    n_puffs = release_times.shape[0]
    if n_puffs == 0:
        raise ValueError(
            "gaussian_puff_rw_model: `release_times` must contain ≥ 1 puff; "
            "got n_puffs=0. Check that t_end − t_start ≥ 1/release_frequency."
        )

    mu_log, sigma_log = _lognormal_from_moments(
        prior_emission_rate_mean, prior_emission_rate_std
    )

    innovations = numpyro.sample(
        "innovations",
        dist.Normal(0.0, 1.0).expand([n_puffs]).to_event(1),
    )
    # Build ln Q_i by cumulative sum of scaled innovations, first anchored
    # to the LogNormal prior location.
    first = mu_log + sigma_log * innovations[0]
    increments = rw_step_std * innovations[1:]
    log_q = jnp.concatenate([first[None], first + jnp.cumsum(increments)])
    emission_series = numpyro.deterministic("emission_rate", jnp.exp(log_q))

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

    if puff_durations is not None:
        puff_mass = emission_series * jnp.asarray(puff_durations)
    else:
        puff_mass = emission_series * release_interval
    predicted = _predict_observations(
        puff_mass,
        release_times,
        receptor_coords,
        observation_times,
        source_location,
        schedule,
        dispersion_params,
        dispersion_fn,
    )
    predicted = 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], observation_times: NDArray, source_location: tuple[float, float, float], wind_times: NDArray, wind_speed: NDArray, wind_direction: NDArray, release_frequency: float, t_start: float, t_end: float, stability_class: str = 'C', scheme: str = 'pg', prior_mean: float = 0.1, prior_std: float = 0.05, obs_noise_std: float = 5e-07, num_warmup: int = 500, num_samples: int = 1000, num_chains: int = 1, seed: int = 0, progress_bar: bool = False, print_summary: bool = False) -> dict[str, NDArray]

NUTS inference of a constant emission rate from puff observations.

Parameters

observations : ndarray, shape (N,) Observed concentrations [kg/m³]. observation_coords : tuple of ndarray, each shape (N,) (x, y, z) receptor coordinates [m]. observation_times : ndarray, shape (N,) Evaluation time for each observation [s]. source_location : tuple of float, (3,) wind_times, wind_speed, wind_direction : ndarray, shape (T_w,) Wind schedule. Direction follows the meteorological "from" convention. release_frequency : float Puff release rate [Hz]. t_start, t_end : float Simulation window [s]. Release times are laid down on [t_start, t_end). stability_class : str scheme : str Dispersion scheme: 'pg' or 'briggs'. prior_mean, prior_std : float Real-scale prior moments for Q. num_warmup, num_samples, num_chains, seed : int progress_bar, print_summary : bool

Returns

samples : dict[str, ndarray] Posterior draws keyed by site name: 'emission_rate', 'background'.

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

def infer_emission_rate(
    observations: NDArray,
    observation_coords: tuple[NDArray, NDArray, NDArray],
    observation_times: NDArray,
    source_location: tuple[float, float, float],
    wind_times: NDArray,
    wind_speed: NDArray,
    wind_direction: NDArray,
    release_frequency: float,
    t_start: float,
    t_end: float,
    stability_class: str = "C",
    scheme: str = "pg",
    prior_mean: float = 0.1,
    prior_std: float = 0.05,
    obs_noise_std: float = 5e-7,
    num_warmup: int = 500,
    num_samples: int = 1000,
    num_chains: int = 1,
    seed: int = 0,
    progress_bar: bool = False,
    print_summary: bool = False,
) -> dict[str, NDArray]:
    """NUTS inference of a **constant** emission rate from puff observations.

    Parameters
    ----------
    observations : ndarray, shape (N,)
        Observed concentrations [kg/m³].
    observation_coords : tuple of ndarray, each shape (N,)
        ``(x, y, z)`` receptor coordinates [m].
    observation_times : ndarray, shape (N,)
        Evaluation time for each observation [s].
    source_location : tuple of float, (3,)
    wind_times, wind_speed, wind_direction : ndarray, shape (T_w,)
        Wind schedule. Direction follows the meteorological "from" convention.
    release_frequency : float
        Puff release rate [Hz].
    t_start, t_end : float
        Simulation window [s]. Release times are laid down on ``[t_start, t_end)``.
    stability_class : str
    scheme : str
        Dispersion scheme: ``'pg'`` or ``'briggs'``.
    prior_mean, prior_std : float
        Real-scale prior moments for Q.
    num_warmup, num_samples, num_chains, seed : int
    progress_bar, print_summary : bool

    Returns
    -------
    samples : dict[str, ndarray]
        Posterior draws keyed by site name: ``'emission_rate'``, ``'background'``.
    """
    obs, coords, times = _validate_inference_inputs(
        "infer_emission_rate",
        observations,
        observation_coords,
        observation_times,
        wind_times,
        wind_speed,
        wind_direction,
        source_location,
        release_frequency,
        t_start,
        t_end,
        prior_mean,
        prior_std,
    )

    schedule = WindSchedule.from_speed_direction(wind_times, wind_speed, wind_direction)
    release_times = make_release_times(t_start, t_end, release_frequency)
    release_interval = frequency_to_release_interval(release_frequency)
    # Per-puff durations conserve emitted mass: the trailing partial interval
    # carries only its actual duration (matches simulate_puff).
    puff_durations = release_durations(release_times, t_end)

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

    mcmc = MCMC(
        NUTS(gaussian_puff_model),
        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,
        observation_times=times_jax,
        source_location=source_location,
        schedule=schedule,
        release_times=release_times,
        release_interval=release_interval,
        puff_durations=puff_durations,
        stability_class=stability_class,
        scheme=scheme,
        prior_emission_rate_mean=prior_mean,
        prior_emission_rate_std=prior_std,
        obs_noise_std=obs_noise_std,
    )
    if print_summary:
        mcmc.print_summary()
    return {k: np.asarray(v) for k, v in mcmc.get_samples().items()}

infer_emission_timeseries

infer_emission_timeseries(observations: NDArray, observation_coords: tuple[NDArray, NDArray, NDArray], observation_times: NDArray, source_location: tuple[float, float, float], wind_times: NDArray, wind_speed: NDArray, wind_direction: NDArray, release_frequency: float, t_start: float, t_end: float, stability_class: str = 'C', scheme: str = 'pg', prior_mean: float = 0.1, prior_std: float = 0.05, rw_step_std: float = 0.02, obs_noise_std: float = 5e-07, num_warmup: int = 500, num_samples: int = 1000, num_chains: int = 1, seed: int = 0, progress_bar: bool = False, print_summary: bool = False) -> dict[str, NDArray]

NUTS inference of a time-varying emission rate series Q_i.

Uses a Gaussian random-walk prior on ln Q along puff indices, anchored by the LogNormal(prior_mean, prior_std) prior on ln Q_0. See :func:gaussian_puff_rw_model.

Returns posterior draws with key 'emission_rate' of shape (num_samples, n_puffs).

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

def infer_emission_timeseries(
    observations: NDArray,
    observation_coords: tuple[NDArray, NDArray, NDArray],
    observation_times: NDArray,
    source_location: tuple[float, float, float],
    wind_times: NDArray,
    wind_speed: NDArray,
    wind_direction: NDArray,
    release_frequency: float,
    t_start: float,
    t_end: float,
    stability_class: str = "C",
    scheme: str = "pg",
    prior_mean: float = 0.1,
    prior_std: float = 0.05,
    rw_step_std: float = 0.02,
    obs_noise_std: float = 5e-7,
    num_warmup: int = 500,
    num_samples: int = 1000,
    num_chains: int = 1,
    seed: int = 0,
    progress_bar: bool = False,
    print_summary: bool = False,
) -> dict[str, NDArray]:
    """NUTS inference of a **time-varying** emission rate series Q_i.

    Uses a Gaussian random-walk prior on ``ln Q`` along puff indices, anchored
    by the LogNormal(``prior_mean``, ``prior_std``) prior on ``ln Q_0``. See
    :func:`gaussian_puff_rw_model`.

    Returns posterior draws with key ``'emission_rate'`` of shape
    ``(num_samples, n_puffs)``.
    """
    obs, coords, times = _validate_inference_inputs(
        "infer_emission_timeseries",
        observations,
        observation_coords,
        observation_times,
        wind_times,
        wind_speed,
        wind_direction,
        source_location,
        release_frequency,
        t_start,
        t_end,
        prior_mean,
        prior_std,
    )
    if not (rw_step_std > 0.0):
        raise ValueError(
            "infer_emission_timeseries: `rw_step_std` must be > 0 "
            f"(got {rw_step_std!r})"
        )
    if not (obs_noise_std > 0.0):
        raise ValueError(
            "infer_emission_timeseries: `obs_noise_std` must be > 0 "
            f"(got {obs_noise_std!r})"
        )

    schedule = WindSchedule.from_speed_direction(wind_times, wind_speed, wind_direction)
    release_times = make_release_times(t_start, t_end, release_frequency)
    release_interval = frequency_to_release_interval(release_frequency)
    puff_durations = release_durations(release_times, t_end)

    mcmc = MCMC(
        NUTS(gaussian_puff_rw_model),
        num_warmup=num_warmup,
        num_samples=num_samples,
        num_chains=num_chains,
        progress_bar=progress_bar,
    )
    mcmc.run(
        jax.random.PRNGKey(seed),
        observations=jnp.asarray(obs),
        receptor_coords=tuple(jnp.asarray(c) for c in coords),
        observation_times=jnp.asarray(times),
        source_location=source_location,
        schedule=schedule,
        release_times=release_times,
        release_interval=release_interval,
        puff_durations=puff_durations,
        stability_class=stability_class,
        scheme=scheme,
        prior_emission_rate_mean=prior_mean,
        prior_emission_rate_std=prior_std,
        rw_step_std=rw_step_std,
        obs_noise_std=obs_noise_std,
    )
    if print_summary:
        mcmc.print_summary()
    return {k: np.asarray(v) for k, v in mcmc.get_samples().items()}