integrix: integration for JAX

integrix is the stack’s integration library: one integral, I = ∫ f dμ, computed three ways (a rule, a scan or a march), with shared adaptivity and adjoints. It sits at Core tier 1 on functax’s rules and gaussx’s solves, extends diffrax with the stack’s steppers, and takes over finitevolX’s steppers and pipekit-jax’s adjoints. Status: proposed, at design-sketch depth.

At a glance¶

integrix sits one tier above the rules it sums with, so every model, numerics library and assimilation library that steps or integrates can import it. This page is the design; core types, algorithms and the code survey are in integrix: design details.

The problem¶

Every library in the stack integrates something, and each writes its own integrator. Turning one satellite methane plume into an emission rate takes six integrals; four are coded in different places today, and two are not written yet:

The rate is then Q = IME · U / L, with U the wind speed and L the plume length. Each row is a sum of weights against samples, or a step that repeats one; each library has rebuilt the weights, the stepping or the adjoint on its own.

Three problems follow. The same hand-written trapezoid sums sit in plumax, radtrax and xrtoolz. finitevolX holds eleven generic stepper classes that have nothing to do with finite volumes, so a spectral or particle model that wants an SSP Runge–Kutta step has to import a finite-volume library. And the adjoints that make 4D-Var and plume inversions affordable (checkpointed and truncated backward sweeps) live in pipekit-jax, the pipeline library, rather than next to the steppers they differentiate.

One integral, three operations¶

Every row of the table is an integral of a quantity f over a domain Ω with a measure μ:

Here Ω is a time interval, a column, a band, a set of directions, an image or a probability space, and μ is length, area, a spectral response or a distribution. The problems differ, but they reduce to three operations:

A march is still a rule in time. A Butcher tableau’s weights bᵢ are quadrature weights over one step, and Gauss–Legendre collocation is an implicit Runge–Kutta method. That is why quadrature procedures and time steppers belong in one library.

Two pieces of machinery are shared by all three:

• Adaptivity. Adaptive quadrature compares a Gauss–Kronrod pair; adaptive time stepping compares an embedded Runge–Kutta pair. Both estimate the error from two rules and refine where it is large.

• Adjoints. A rule’s adjoint is its weights, a scan’s is a reverse scan, and a march’s is a backward march with checkpointing.

The weights w of a rule come from functax: they are a property of a basis on a domain. integrix owns what is done with them: summing, scanning, marching, adapting and differentiating.

Modules and boundaries¶

Six modules: one per operation, plus adaptivity, adjoints and probability. integrix owns the machinery, not the domains: spectral responses stay in radtrax and plumax, mass enhancement in plumax and exceedance events in xtremax; they pass in weights, measures or integrands.

What integrix owns, and what it hands off¶

The test for every name: does it sum, scan, march or adapt over a domain, or differentiate one of those? Then integrix owns it.

What moves in¶

In use¶

The methane example end to end: a rule for band radiance and integrated mass, a scan for optical depth, and a march for transport with a checkpointed adjoint. Every name is proposed; the rules come from functax and the Runge–Kutta steppers subclass diffrax’s.

# Shapes: W wavelengths; Z layers; Ny × Nx plume pixels; K steps; C state size

import diffrax as dfx
import functax as fx
import integrix as ix

# --- rule: band radiance ---
# proposed (functax)
# wₗ for ∫ g(λ) dλ on the sampled wavelengths
band_rule = fx.quad.trapezoid(wavelength)  # (W,) → rule
# proposed (integrix)
# L_b = ∫ L(λ) R(λ) dλ / ∫ R(λ) dλ
band_radiance = ix.rule.weighted_mean(  # (…, W) → (…,)
    band_rule, radiance, weight=response
)

# --- scan: optical depth ---
# proposed (functax)
layer_rule = fx.quad.trapezoid(altitude)  # (Z,) → rule
# proposed (integrix)
# τ(z) = ∫₀ᶻ k(z′) n(z′) dz′
optical_depth = ix.scan.cumulative(layer_rule, absorption * density)  # (Z,) → (Z,)

# --- rule: integrated mass enhancement ---
# proposed (functax)
pixel_rule = fx.quad.cell_area(plume_grid)  # (Ny, Nx) → rule
# proposed (integrix)
# IME = ∬ ΔΩ(x, y) dA
mass = ix.rule.integrate(pixel_rule, enhancement)  # (Ny, Nx) → ()

# --- march: plume transport, differentiable ---
# not pinned (diffrax)
terms = dfx.ODETerm(advection_diffusion)
# proposed (integrix)
# yₙ₊₁ = Φ(yₙ, tₙ, h), three-stage SSP Runge–Kutta
stepper = ix.march.RK3SSP()
# proposed (integrix)
# backward sweep, storing every 50th state
adjoint = ix.adjoint.Checkpointed(every=50)
# not pinned (diffrax)
solution = dfx.diffeqsolve(  # (C,) → (K, C)
    terms,
    stepper,
    t0=0.0,
    t1=duration,
    dt0=step,
    y0=initial,
    adjoint=adjoint,
    saveat=dfx.SaveAt(ts=times),
)

The rules are weights from functax, so the same ix.rule call serves a band, a column and an image. The march is ordinary diffrax: integrix supplies the stepper and the adjoint, not a new solve loop.

Integration across the stack¶

Seven libraries would import integrix, pipekit-jax would hand code to it, and spectraldiffx’s examples could use it. Most changes are an import: the steppers keep their diffrax interface, and a rule replaces a hand-written sum.

The largest change is finitevolX’s, and it is an import swap:

# Today:
# today (finitevolX)
from finitevolx import RK3SSP

# After:
# proposed (integrix)
from integrix.march import RK3SSP

The class keeps its diffrax interface, so call sites do not change; finitevolX keeps a deprecated alias for one release.

Decisions and open questions¶

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

Decisions 10 to 13, on integrals over space, are in integrix: integrals over space; all four were decided on 10 Nov 2026.

1. Name: integrix¶

Decided (5 Nov 2026). integrix, free on PyPI when checked. integrax is taken; quadrax, marchax, cumulax and evolvax were the other candidates.

2. Scope: all integration¶

Decided (5 Nov 2026). Time, space, spectrum, angle, exceedance and probability, not only quadrature and time stepping. All six reduce to the rule, scan and march above, and share adaptivity and adjoints.

3. Rules in functax, procedures here¶

Decided (5 Nov 2026). The nodes and weights of a rule are a property of a basis on a domain, so functax owns them. integrix owns what sums, scans, adapts or marches with them. Gaussian rules (Gauss–Hermite, unscented, sigma points) stay in gaussx.

4. Layer: Core tier 1¶

Decided (5 Nov 2026, on the layer map). integrix imports functax and gaussx, both Core tier 0, and extends diffrax. Every library that steps or integrates sits above it.

5. Adaptive quadrature: written here, on quadax’s design¶

Decided (5 Nov 2026). Adaptive Gauss–Kronrod, tanh–sinh, Romberg and cumulative integrals are written in integrix, drawing on quadax’s design, with no dependency on quadax. quadax’s results serve as references in tests.

6. Semi-Lagrangian advection stays in finitevolX¶

Decided (5 Nov 2026). It discards its departure points after each step, so it is a field scheme, not a stepper. 11 of finitevolX’s 12 stepper classes move (0.1.0 added AB3Solver); SemiLagrangianSolver stays and takes its interpolation from functax.

7. vardax’s adjoints: argminax’s, not integrix’s¶

Decided (5 Nov 2026), after reading vardax. vardax’s KStepAdjoint and OneStepAdjoint subclass optimistix.AbstractAdjoint: they differentiate through an inner minimiser (warm-up iterations under stop_gradient, then k live steps), not backward through time. They go to argminax, as the layer map says. integrix keeps the time adjoints of the forward run that 4D-Var’s cost contains: checkpointed backward marches and truncated scans.

8. Steppers: extend diffrax or own a loop?¶

Leaning: extend diffrax, with a small loop for multistep schemes. Eight of the eleven steppers moving from finitevolX (AB3Solver among them) already subclass diffrax’s solvers, so diffeqsolve, its step-size controllers and its adjoints keep working. The other three, AB2Solver, LeapfrogRAFSolver and SplitExplicitRKSolver, are plain eqx.Modules today, because multistep and split-explicit schemes carry state between steps that diffrax’s solver protocol handles awkwardly. integrix keeps them as its own stepping loop, with the same adjoint options.

9. Monte Carlo here or in pyrox?¶

Open. Plain Monte Carlo and quasi-Monte Carlo expectations are integration and fit integrix.expect. Sampling from a posterior (MCMC, SVI) stays in NumPyro and pyrox; integrix only averages what they draw.

Roadmap¶

Six waves, each a GitHub issue epic in the integrix repository once it exists. integrix needs functax’s rules first, so wave 0 waits on functax’s wave 0 (geometry) and its rules wave.

Waves 1 and 3 can run side by side. Wave 5 is one issue per consumer, filed in that consumer’s repository and linked to the integrix epic.

Sources¶

• Diffrax documentation: the solver and adjoint protocols integrix extends.

• quadax: adaptive quadrature in JAX, the design integrix’s adaptive module draws on.

• SciPy integrate reference: the test references for rules and adaptive quadrature.

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

• functax: the rules integrix sums with.

• Code read on 5 Nov 2026 at geostack’s submodule pins: finitevolX _src/timestepping/diffrax_solvers.py and semi_lagrangian.py; pipekit-jax adjoints.py; somax _src/solvers.py; plumax, radtrax and xrtoolz SRF and quadrature call sites; spectraldiffx scripts/qg_model.py.