tracerax: Lagrangian particles and trajectories for JAX

tracerax carries positions forward in time and analyses where they went: particle seeding, deterministic and stochastic advection, binning to grids, footprints, finite-time Lyapunov exponents, residence times and connectivity. It sits at GeoNumerics tier 2, on gridax’s grids, functax’s interpolation and integrix’s stepping, and takes over the particle machinery now inside plumax. Status: proposed, at design-sketch depth.

At a glance¶

The line with finitevolX, settled on the layer map: a scheme that discards its departure points after each step is an operator on a field, and stays with the Eulerian libraries; anything whose state is the positions themselves is tracerax’s. This page is the design; core types, algorithms and the code survey are in tracerax: design details.

The problem¶

Two questions about a methane plume are Lagrangian. Forward: where does the gas from this source go, and what concentration does it leave in each pixel? Backward: which surface cells could have produced the enhancement this satellite pixel saw? Both release particles, move them with the wind plus turbulence, and count where they end up:

The particle machinery is generic, but it sits inside a plume model, so the ocean side (sesmax’s drifters, larval connectivity, retention in eddies) and xrops’ planned Lagrangian views would have to import plumax to use it, or write it again.

Positions carried forward¶

The state of a Lagrangian calculation is a set of N particles: positions xₚ, masses mₚ, ages, an active flag, and for a stochastic model a velocity fluctuation u’ₚ. Everything in tracerax either moves that state or reads it.

Deterministic advection. Each particle follows the resolved wind, dxₚ/dt = u(xₚ, t), with u at the particle’s position interpolated from a grid by functax and the step taken by an integrix stepper.

Stochastic advection. Unresolved turbulence enters as a Markov-1 Langevin model for the velocity fluctuations: an Ornstein–Uhlenbeck update on all three components, as plumax does it, with the well-mixed drift of Thomson (1987) acting in the vertical, with σ_w(z) the velocity standard deviation and τ_L(z) the Lagrangian time scale supplied by a turbulence closure:

Particles reflect at the ground and, when it is given, at the top of the boundary layer.

Binning. The concentration in cell j of volume Vⱼ is the particle mass inside it:

Footprints. Released backward in time from a receptor, particles that spend time near the ground say how sensitive the receptor is to a surface flux in each cell. With h_s the depth of the surface layer and ρ the air density:

Trajectory analysis. With Φ the flow map that carries a position from t₀ to t₁, the finite-time Lyapunov exponent measures how fast neighbours separate; its ridges are transport barriers:

In JAX the gradient of the flow map is jax.jacfwd of the integration itself, so no finite-difference grid of seeds is needed. Residence time counts how long particles stay in a region; a connectivity matrix Cᵢⱼ counts the fraction released in region i that reaches region j.

Modules and boundaries¶

Five modules. The particle state is an Equinox module of arrays, so it is a pytree that jax.lax.scan, jit and vmap carry as plumax’s already is.

What tracerax owns, and what it hands off¶

The test for every name: is its state a set of positions carried forward in time, or does it analyse such positions? Then tracerax owns it.

What moves in¶

In use¶

The plume both ways: release particles at the source and bin them into a concentration field; release them backward from a satellite pixel and read its surface footprint. Every tracerax name is proposed.

# Shapes: N particles; K steps; Ny × Nx surface cells

import jax
import tracerax as trx
from plumax.lagrangian import HannaTurbulence

release_key, forward_key, backward_key = jax.random.split(jax.random.key(0), 3)


# user
# u(x, t) at particle positions, wrapping a functax interpolation plan
def wind_at(positions, time): ...  # (N, 3) → (N, 3)


# today* (plumax)
# σ_w(z), τ_L(z) from boundary-layer scaling
turbulence = HannaTurbulence(...)

# --- forward: where the gas goes ---
# proposed (tracerax)
particles = trx.particles.release_at(source, n=N, mass=rate * duration, key=release_key)
# proposed (tracerax)
# du' = −u'/τ_L dt + ½ ∂σ_w²/∂z (1 + u'²/σ_w²) dt + (2σ_w²/τ_L)^½ dW
final = trx.advect.langevin(  # (N, 3) → (N, 3)
    particles,
    wind_at,
    turbulence,
    times,
    reflect=trx.advect.Reflect(ground=0.0, top=pbl_height),
    key=forward_key,
)
# proposed (tracerax)
# c_j = ∑ₚ mₚ 1[xₚ ∈ j] / V_j
concentration = trx.bin.concentration(final, grid)  # → (Ny, Nx)

# --- backward: what the pixel could have seen ---
# proposed (tracerax)
# F_j = ∑ₚ ∫ 1[xₚ(t) ∈ j, zₚ(t) < h_s] dt / (N ρ h_s)
footprint = trx.footprint.surface(  # → (Ny, Nx)
    receptor, wind_at, turbulence, grid, duration=600.0, key=backward_key
)

The footprint times a prior flux map is the pixel’s predicted enhancement, which is what an inversion in plumax fits. Because the run is a JAX function, the same footprint is differentiable in the wind and the closure’s parameters.

Integration across the stack¶

For plumax the change is an import:

# Today:
# today (plumax)
from plumax.lagrangian.particles import ParticleState, langevin_step

# After:
# proposed (tracerax)
from tracerax.particles import ParticleState

# proposed (tracerax)
from tracerax.advect import langevin_step

Decisions and open questions¶

Three decided, two leaning, two open. Each entry gives the position first, then the trade-off.

1. Layer: GeoNumerics tier 2¶

Decided (5 Nov 2026, on the layer map). tracerax imports gridax (tier 1), geophysix (tier 0), and functax and integrix (Core). plumax, sesmax-ocean and xrops sit above it.

2. The line with the Eulerian libraries¶

Decided (5 Nov 2026). A scheme that traces back to departure points and then discards them, such as finitevolX’s semi-Lagrangian advection, is an operator on a field and stays with finitevolX and spectraldiffx. tracerax owns everything whose state is the positions.

3. Turbulence closures stay in plumax¶

Decided (5 Nov 2026, on the layer map). Particles take σ_w(z) and τ_L(z) through a protocol, and plumax’s closures satisfy it. Open within this: whether tracerax defines the protocol (so that sesmax-ocean can write its own closures without plumax) or plumax keeps it. Leaning: tracerax defines it.

4. Stepping through integrix or a loop of its own¶

Open. Deterministic advection is an ordinary ODE and fits integrix’s steppers. The Langevin step reflects particles at boundaries, a discrete event that diffrax’s SDE solvers handle awkwardly, and plumax today runs it in a lax.scan of its own. Leaning: deterministic advection through integrix, the stochastic step as tracerax’s own scan with the same adjoint options.

5. Trajectories stored or accumulated¶

Leaning: accumulate by default. A footprint or a concentration needs only running sums, which plumax already accumulates inside its scan; storing K × N × 3 positions is optional, at chosen times.

6. Randomness¶

Leaning. One key per run, split once into a key per stage (release, forward, backward) and per step inside the scan, as the house style asks; no global seeds.

7. Coastlines and other boundaries¶

Open. Reflection at the ground and the boundary-layer top is simple. Reflection at a coastline needs the model’s mask (finitevolX’s Mask2D) and the exact crossing point, which is costly for many particles. Leaning: nearest-wet-cell projection first, exact reflection when an ocean consumer needs it.

Roadmap¶

Five waves, each a GitHub issue epic in the tracerax repository once it exists.

Sources¶

• The GeoStack layer map: tracerax’s layer, tier, imports and moves.

• Thomson (1987): the well-mixed condition.

• Lin et al. (2003): backward footprints.

• Haller (2015): FTLE and transport barriers.

• Code read on 5 Nov 2026 at geostack’s submodule pins and xrtoolz b3855d7: plumax lagrangian/ (particles.py, concentration.py, footprint.py, turbulence.py, inversion.py); xrops metrics/_src/lagrangian.py; finitevolX’s semi-Lagrangian advection.