tracerax: design details

The implementation reference behind the tracerax design: the core types, what callers pass in, the algorithms with their pitfalls, JAX behaviour, tests, dependencies, and a survey of the code that moves in.

Core types¶

The state is a pytree of arrays; the closure is an abstract module; results are small modules that carry their units. Every name is proposed.

# Shapes: N particles; nx × ny × nz bins; nx × ny surface cells

import abc

import equinox as eqx
from jaxtyping import Array, Bool, Float


# proposed (tracerax)
class ParticleState(eqx.Module):
    """An ensemble of N particles; every field is a pytree leaf."""

    # xₚ = (x, y, z) in metres, in a local frame around the source
    position: Float[Array, "particle 3"]
    # u'ₚ, the turbulent velocity fluctuation (zeros for pure advection)
    velocity: Float[Array, "particle 3"]
    # mₚ in kg; a release of rate Q over T gives mₚ = Q T / N
    mass: Float[Array, "particle"]
    # seconds since release; negative before a staggered release
    age: Float[Array, "particle"]
    # False before release and after leaving the domain
    active: Bool[Array, "particle"]


# proposed (tracerax)
class TurbulenceModel(eqx.Module):
    """σ(z) and τ_L(z), the closure a Langevin step reads."""

    @abc.abstractmethod
    def at(
        self, height: Float[Array, "particle"]
    ) -> tuple[Float[Array, "particle 3"], Float[Array, "particle 3"]]:
        """Return (σ, τ_L) per particle, one column per velocity component."""


# proposed (tracerax)
class Reflect(eqx.Module):
    """Reflecting walls in z: the ground and, optionally, the PBL top."""

    ground: float = 0.0
    top: float | None = None


# proposed (tracerax)
class Residence(eqx.Module):
    """Particle-seconds per bin: Rⱼ = ∑ₚ ∑ₖ Δtₖ 1[xₚ(tₖ) ∈ j]."""

    seconds: Float[Array, "x y z"]
    n_particles: int = eqx.field(static=True)


# proposed (tracerax)
class Footprint(eqx.Module):
    """Fⱼ in s m² kg⁻¹: sensitivity of one receptor to a surface flux in j."""

    sensitivity: Float[Array, "x y"]
    surface_layer_depth: float

The bins keep plumax’s (x, y, z) axis order so results compare directly with plumax’s; xrops relabels them when it wraps the output.

Types to adhere to¶

Callers pass three things in: a wind, a turbulence closure, and bin edges.

• The wind is a function of positions and time, wind(position, time) -> velocity, with shapes (N, 3), () → (N, 3). A gridded wind is a functax interpolation plan closed over the field. plumax’s wind today is a function of time only, t → (u, v, w), uniform in space; trx.advect.uniform(wind_of_time) adapts it.

• The closure subclasses TurbulenceModel. plumax’s HomogeneousTurbulence and HannaTurbulence (the profiles of Hanna (1982)) already satisfy the shape of at as frozen dataclasses; under tracerax they become eqx.Module subclasses, or are wrapped, so that jax.grad of σ_w² traces through them.

• The bins are edge arrays per axis. They come from a gridax analysis geometry, from a finitevolX model grid for a model run, or are written by hand.

# Shapes: N particles

from plumax.lagrangian import HannaTurbulence


# user
# u(x, t) at every particle, from a functax plan over the gridded wind
def wind_at(position, time): ...  # (N, 3), () → (N, 3)


# today (plumax)
# σ(z), τ_L(z) from Hanna (1982) boundary-layer scaling
turbulence = HannaTurbulence(u_star=0.4, pbl_height=1000.0, obukhov_length=-50.0)

Algorithms¶

Deterministic advection¶

Each particle follows the resolved wind:

1. Wrap wind_at as a diffrax ODETerm over the (N, 3) position array, so the ensemble is one state.

2. Step it with an integrix stepper: a fixed-step SSP Runge–Kutta for gridded winds, an adaptive pair for analytic ones.

3. Save at the requested times, or only the final state.

4. Deactivate particles that leave the domain: they keep their last position and active = False.

Cost: K steps × s stages × one wind evaluation of N points. Pitfalls: a linearly interpolated wind is only continuous, not smooth, so an adaptive error estimate stalls at cell faces; use fixed steps or a cubic plan there. A particle outside the wind’s grid must be clamped before interpolation, or it reads garbage.

The Langevin step, well mixed, with reflection¶

For Gaussian turbulence with σ(z) and τ_L(z) per component, the velocity fluctuation is an Ornstein–Uhlenbeck process plus, for the vertical component, Thomson’s (1987) well-mixed drift. The linear part is integrated with its exact transition over a step Δt, with ξ a standard normal draw:

The vertical drift is added with an Euler step, and the position follows the mean wind plus the fluctuation:

1. Read σ, τ_L from the closure at each particle’s height; floor σ at a small positive value.

2. Draw ξ from this step’s key and apply the exact OU transition to all three components.

3. Compute ∂σ_w²/∂z by jax.vmap(jax.grad(...)) of the closure, and add the drift to u’_w.

4. Move the particles.

5. Reflect: below the ground, z → −z and u’_w → −u’_w; above the top h, z → 2h − z and u’_w → −u’_w.

6. Shorten the last step so the run ends exactly at the horizon (plumax’s step_durations).

Cost: O(N) per step, plus one gradient of σ_w² per particle. Pitfalls:

• The exact OU transition makes homogeneous runs well mixed for any Δt, but the drift term is first order, so Δt must stay small against τ_L near the ground, where τ_L is shortest.

• Convective profiles carry a (z/h)^(2/3) factor whose derivative is singular at the ground; plumax floors z at 0.1 m before evaluating the closure.

• Perfect reflection keeps the well-mixed state only when σ_w is uniform near the wall Wilson & Flesch, 1993; with strong gradients at the wall it biases the near-ground density, which is exactly where a footprint is counted.

Binning by scatter-add¶

With edges e₀ < … < eₙ per axis, a particle’s cell index is i = searchsorted(e, x) − 1, kept when 0 ≤ i < n. The three indices flatten to one, and a scatter-add accumulates the weight:

For a continuous source of rate Q released as N particles at t = 0 in a steady flow, the steady concentration is cⱼ = Q Rⱼ / (N Vⱼ), with Vⱼ the cell volume. Cost: O(N) per step with a single .at[].add. Pitfalls: a histogram is piecewise constant in the positions, so its gradient with respect to them is zero almost everywhere. Gradients with respect to the wind or the closure’s parameters need a smooth deposit instead, such as cloud-in-cell (linear) weights, which spread each particle over the 2ᵈ nearest cells.

Backward footprints¶

Particles released at a receptor and integrated backward in time give the receptor’s sensitivity to a surface flux. With h_s = f h the surface-layer depth and ρ the air density:

1. Reverse the sign of the mean wind.

2. Reverse the order of the forward step durations, so the partial step comes first, and sample the wind at each interval’s start on the receptor’s clock.

3. Run the Langevin step; accumulate residence of the particles below h_s per surface cell.

4. Divide by N ρ h_s. There is no division by cell area: F pairs with a flux in kg m⁻² s⁻¹.

Pitfall: reusing the forward integrator with the wind reversed is exact for the OU part when the turbulence is stationary, because that process is time-reversible. In inhomogeneous turbulence the backward-time drift is not the forward one with the wind reversed Flesch et al., 1995; this needs a test before the footprints are trusted near the ground.

Finite-time Lyapunov exponents¶

With Φ the flow map of the deterministic advection from t₀ to t₁ and ∇Φ its Jacobian at a seed x₀:

1. Write the flow map for one seed as a function x₀ ↦ x(t₁).

2. Take jax.jacfwd of it and jax.vmap over the seeds.

3. Form C with einx (einx.dot("... a b, ... a c -> ... b c", J, J)) and take the largest eigenvalue with jnp.linalg.eigvalsh.

Cost: about d times one integration for d dimensions, against 2d extra trajectories per seed for central differences. Pitfalls: ∇Φ grows exponentially with |t₁ − t₀|, so long windows overflow in float32; split the window and multiply the Jacobians, accumulating the logarithm. FTLE is defined for the deterministic flow only: no turbulence draw.

Residence time and connectivity¶

With regions labelled on the bins, residence in region A is τ_A = ∑ₖ Δtₖ 1[xₚ(tₖ) ∈ A] per particle, and the connectivity from release region i to region j is the fraction of i’s particles found in j at the end (or ever, for a visiting matrix):

Both reuse the binning: a label lookup per particle, then a one-hot contraction over particles with einx, O(N × regions) per step.

JAX compatibility¶

Keys: one key per run, split up front into one key per step (jr.split(key, n_steps)), so a scan step never reuses a key and the draws do not depend on how the run is chunked. The step function takes its key as an argument; nothing reads numpy.random.

Scan carry: (ParticleState, accumulator), where the accumulator is the residence array of the bins. Saving trajectories turns the per-step position into a scan output instead.

Memory: O(N + bins) for an accumulating run; O(K N) when the trajectory is saved, which is the usual way a long run runs out of memory. Reverse-mode gradients through a long scan store every step’s carry; use jax.checkpoint on the step, or forward mode when the parameters are few.

Precision: positions in metres far from the origin lose resolution in float32 (a 1 cm step at 10⁶ m is below float32’s resolution). Run in a local frame with the origin near the source (xrops’ LocalFrame), or in float64.

Testing¶

• plumax oracle. For homogeneous and Hanna turbulence with a uniform wind and the same keys, the moved Langevin step, binning and footprint match plumax 28f257e to round-off. This test is deleted when plumax re-exports tracerax.

• Well-mixed stays well-mixed. Particles started uniform in [0, h] with Hanna profiles stay uniform after many τ_L: a Kolmogorov–Smirnov test on the heights, with the threshold from the KS statistic’s sampling distribution at the chosen N, not a tuned tolerance.

• Taylor (1921). In homogeneous turbulence, the variance of the displacement matches the closed form, within a few standard errors of the sample variance:

• Mass conservation. With reflecting walls and a domain larger than the plume, total residence equals N times the horizon; with an open boundary it equals the sum over active particles.

• Forward and backward duality. For a uniform wind and homogeneous turbulence, the forward concentration at a receptor from a unit surface source equals the backward footprint summed over that source’s cells, within the Monte Carlo error.

• Deterministic advection. Solid-body rotation returns particles to their start after one period; the double gyre’s FTLE ridges match Shadden et al. (2005), and the jacfwd Jacobian matches central finite differences.

• Transforms. Every public function gives the same result eagerly, under jit and under vmap over a batch of one; gradients with respect to closure parameters match finite differences at a fixed key.

External dependencies¶

Where it stands¶

Nothing named tracerax exists yet. The code that moves in is plumax’s lagrangian subpackage at plumax 28f257e:

What the move adds rather than carries over: a wind that varies in space (plumax’s wind is a function of time only, t → (u, v, w)), deterministic advection through integrix, FTLE, residence time and connectivity, staggered releases, and the active mask. None of these exist anywhere in the stack today.

Neighbours checked and left alone: finitevolX 82090a3 has semi_lagrangian_step (field in, field out by map_coordinates at departure points), which is Eulerian transport and stays there; xrtoolz b3855d7 has xrops.metrics._src.lagrangian, a six-line stub reserved for a later view, and nothing to move.