plumax

Where is the plume, how much is coming out, where does it come from, how often does it happen. Differentiable transport models (Gaussian plume and puff, Lagrangian particles, Eulerian finite volumes) and the Bayesian and variational machinery to invert them, for any passive tracer: methane, CO₂, NO₂, SO₂, smoke or heat. Everything is JAX: jit, vmap and grad through the physics.

Four questions¶

✅ ships today · 🚧 designed, not shipped. The roadmap has the status of every tier.

Where plumax sits¶

Turning radiance into a column of gas is radiative transfer, and that lives in radtrax. plumax takes the column, or any concentration field, and reasons about the plume: its shape, its transport and its source. plumax imports radtrax (the coupled tier maps a simulated plume to radiance through it); radtrax never imports plumax.

What it stands on¶

plumax sits in GeoStack’s GeoModels layer beside radtrax. It reaches down to finitevolX for the Eulerian grid, to gaussx for every posterior covariance, and to pipekit for the Operator interface. Inference sits to one side: every tier exposes the same forward(params, met) → observations, so the samplers and assimilation schemes are written once against it. NumPyro does that for sampling; vardax (variational, the assim extra) runs the Tier II optimal interpolation and the Tier III 3D-Var; filterax (ensemble Kalman) is planned.

The modeling cycle¶

plumax is organised around a five-tier data-driven modeling cycle. Each tier shares a fixed forward interface, so emulation and amortized inference are substitutions rather than rewrites:

1. Simple model: a generative story you can simulate from.

2. Model-based inference: slow but exact ground truth.

3. Model emulator: a fast surrogate (skip it if the model is cheap).

4. Emulator-based inference: step 2 in seconds.

5. Amortized inference: a predictor that learns the posterior map.

6. Improve: upgrade any component; the previous step is its ground truth.

The Roadmap & Architecture page has the tier breakdown (Gaussian → Lagrangian → Eulerian → Coupled → Population).

Installation¶

uv add plumax

gaussx, pipekit and radtrax are not yet on PyPI; from a checkout they resolve from git via [tool.uv.sources]. NumPyro (Bayesian inference), vardax (the Tier II and Tier III source inversions) and the notebook dependencies are optional extras:

uv add "plumax[inference,assim]"

Quickstart: simulate a plume, then recover its emission rate¶

A single stack 20 m tall releases a passive tracer at $Q = 0.5$ kg s⁻¹ into a 4 m s⁻¹ westerly, under Pasquill stability class C. The steady Gaussian plume gives the concentration $C$ [kg m⁻³] at a point $(x, y, z)$, with $x$ downwind, $y$ crosswind, $z$ up, source height $h$, wind speed $u$, and the Briggs widths $\sigma_y(x)$, $\sigma_z(x)$ of the stability class; the second exponential reflects the plume off the ground:

$K$ sensors at 5 m record $o_k = C(x_k, y_k, z_k) + b + \varepsilon_k$, with a background $b$ and noise $\varepsilon_k \sim \mathcal{N}(0, s^2)$. NUTS samples the posterior over $(Q, b)$; $S$ is the number of draws.

import jax
import jax.numpy as jnp
from plumax.gauss_plume import infer_emission_rate, simulate_plume

# Shapes: X = 100, Y = 80, Z = 13 grid points; K = 200 sensors; S = 500 draws

# today (plumax)
# C(x, y, z) on the grid, Q = 0.5 kg/s
truth = simulate_plume(                                  # → Dataset (X, Y, Z)
    emission_rate=0.5,
    source_location=(0.0, 0.0, 20.0),
    wind_speed=4.0,
    wind_direction=270.0,
    stability_class="C",
    domain_x=(0.0, 1000.0, 100),
    domain_y=(-200.0, 200.0, 80),
    domain_z=(0.0, 60.0, 13),
)

# every 40th grid column at z = 5 m
sensors = truth.concentration.isel(z=1).stack(k=("x", "y"))[::40]  # → (K,)

# oₖ = C(xₖ, yₖ, zₖ) + εₖ,  εₖ ~ 𝒩(0, s²)
noise = 2e-7 * jax.random.normal(jax.random.key(0), sensors.shape)  # (K,)
observed = jnp.asarray(sensors.values) + noise           # (K,) → (K,)
sensor_xyz = (
    jnp.asarray(sensors.x.values),
    jnp.asarray(sensors.y.values),
    jnp.full(sensors.shape, 5.0),
)

# today (plumax)
# p(Q, b | o) ∝ 𝒩(o; C_Q + b, s²I) · p(Q) · p(b)
posterior = infer_emission_rate(                         # → {Q: (S,), b: (S,)}
    observed,
    sensor_xyz,
    source_location=(0.0, 0.0, 20.0),
    wind_speed=4.0,
    wind_direction=270.0,
    stability_class="C",
    prior_mean=0.3,
    prior_std=0.3,
    obs_noise_std=2e-7,
    num_warmup=300,
    num_samples=500,
)

# 𝔼[Q | o], √Var[Q | o]
q = posterior["emission_rate"]                           # (S,)
print(f"Q = {q.mean():.4f} ± {q.std():.4f} kg/s")      # → Q = 0.5000 ± 0.0004 kg/s

simulate_plume and infer_emission_rate in the API reference list every parameter.

Not just methane¶

The transport and the inference treat a passive tracer: an emission rate in kg s⁻¹ goes in, a concentration in kg m⁻³ comes out. Three places still assume methane: coupled.rtm converts column mass to ΔVMR with CH4_MOLAR_MASS_KG_PER_MOL; les_fvm labels its outputs “methane”; and the population intensity kernels are named for methane sectors, although their maths is sector-agnostic.

Links¶

• Roadmap & Architecture

• API Reference

• Changelog

• GitHub