GeoStack overview

1The decision¶

GeoStack is one operator algebra across three stacks: GeoML (numerics and inference), GeoModels (forward models) and GeoStack proper (data, pipelines, evaluation). None of them merges into another. The JAX libraries stay pipeline-free, and they only reach pipekit through thin adapters/pipekit.py modules.

The work splits into six workstreams. Each one makes one decision, and its letter tags every proposed (X) call on the examples page and every gate in Table 6.

• A · One framework. Today every library invents its own pipeline glue. pipekit replaces it with one Operator protocol, Sequential and Graph composition, config round-tripping and sandboxed YAML.

• B · Indexed data at both ends of the lifecycle. Observations can be catalogued whereas model states are scattered files. geocatalog indexes both: L0 to L2 observations as rows queried with GeoSlices, and L3 to L5 model states through its state profile (run time, member, producer and parent columns, in geocatalog.states). The readers split by level, with georeader below L3 and xrreader from L3 up.

• C · Two carriers, one algebra. geotoolz (GeoTensor) and xrtoolz (xarray) share the operator algebra and the patcher, and only the substrate differs.

• D · Data assimilation as a cycle. The ForwardModel, ObservationOperator and AnalysisStep protocols of pipekit-cycle are the only contract. filterax, vardax and plumax satisfy them through adapters, so an emulator can replace a physics model without touching the cycle.

• E · Training with lineage. pipekit-train wraps Equinox + Optax today, with NumPyro and BlackJAX for Bayesian fits, and Lightning and Keras 3 are planned. pipekit-experiment adds a tracker protocol and a content-addressed ModelRegistry.

• F · Evaluation as infrastructure. pipekit-evaluate scores every cell of the benchmark cube against an Oracle, with the metrics from xrtoolz.metrics. The Unit × Lens × Stage tags and a content-addressed evaluation report are proposed (Report 14).

End state. GeoML owns linear algebra, kernels and inference, and it never imports pipekit outside an adapter. GeoModels owns the physics and only exposes it through adapters. pipekit owns composition, cycles, training and evaluation. geocatalog owns the one index, for observations and (through its state profile) for model states, whereas geotoolz and xrtoolz own the carrier-specific operators. Orchestration and serving stay external.

2The elevator pitch¶

A vertically integrated, JAX-native research software stack for geophysical modeling, inference and data assimilation: from structured linear algebra primitives, through composable operator graphs, to indexed multi-tier data products, with ML pluggable at every level.

Great JAX numerics and great xarray pipelines exist today. However, what is missing is the operator-graph abstraction which connects them, with content-addressed data and models throughout.

3The shape of the problem¶

There are five gaps which motivate everything below. The first four are old, whereas the fifth has only become visible recently.

The stack addresses all five.

4Three organising axes¶

Three axes structure the framework (Table 1). They are orthogonal, so every operation sits somewhere on all three.

4.1The modeling cycle (axis 1)¶

4.2The L0–L5 data lifecycle (axis 2)¶

Table 2 names the six levels.

The science axis runs left to right. Until retrieval it is one stream, as the same L1 radiances feed every geoscience. Retrieval is where it fans out into ocean, atmosphere, land and ice variables, and each domain then carries its own analysis and reanalysis or forecast. L5 holds the products derived from those states, such as extremes (marine heatwaves, droughts, storms) or impact indicators. Which products belong in L5 is still an open question.

The engineering axis runs top to bottom, from computing a level down to storing it, and it maps every level to the GeoStack tool that handles it.

• Compute: what performs the step into each level. radtrax retrieves (L1 → L2). GeoML grids and gap-fills (L2 → L3: pyrox, kernellib, the optimal interpolation of vardax). GeoModels with data assimilation model and assimilate (L3 → L4: somax and plumax with filterax and vardax), which is the inner loop of Figure 1. Lastly, xtremax derives extremes (L4 → L5).

• Compose: pipekit wires every level into operator graphs, cycles and training loops, whatever the carrier.

The rows below split where the carrier changes:

• Observations (L0 to L2): georeader reads them into a GeoTensor, geocatalog indexes them as GeoSlices and geotoolz transforms them.

• Model states and derived products (L3 to L5): xrreader reads them into xarray, geocatalog.states indexes them in the same catalogs as state rows (its state profile), and xrtoolz transforms and verifies them.

• Both carriers: geopatcher patches fields whatever the carrier, and geocloud stores both on object storage (COG for observations, Zarr or NetCDF for states).

ML enters at every transition.

4.3Two streams converge¶

The lifecycle is not one chain. A top-down satellite stream and a bottom-up in-situ stream converge between L2 and L3 (Figure 4).

The matchup pattern is canonical and it deserves its own operator family (stage 2 runs it, and where it lives is an open question).

4.4Three layers of transformation¶

A third lens joins the cycle and the lifecycle, which is what we do to the data on its way into a model (Figure 5).

4.5The six-step cycle (axis 3)¶

The six-step cycle is the journey from research to production (Figure 6). Every domain library follows this cycle implicitly, and the framework makes it explicit and composable.

In symbols, let $\gvect{\theta} \in \mathbb{R}^{D_\theta}$ be the quantity we want (an emission rate, a model state), $f$ the physical forward model, $h$ the observation operator, $\vect{y} \in \mathbb{R}^{D_y}$ the observations and $\gvect{\varepsilon} \in \mathbb{R}^{D_y}$ their error, with covariance $\gmat{\Sigma}_\varepsilon \in \mathbb{R}^{D_y \times D_y}$. In the six-step cycle (and in its figures), $\gvect{\theta}$ stands for whatever we infer, so it covers both the state $\vect{x}$ and the parameters of Figure 1. Step 1 is the generative story:

where the error is Gaussian,

Step 2 inverts it with classical, model-based inference. Bayes’ rule with a prior $p(\gvect{\theta})$ defines the target:

Computing (3) is only exact in tractable cases, such as a linear-Gaussian model solved with the Kalman filter or a GP. The usual step-2 methods (a finite-ensemble EnKF, 3D- or 4D-Var) approximate it, and slowly. So step 2 is the oracle because it is the most trusted reference we have, and not because it is exact.

Step 3 trains an emulator $\hat{f}(\cdot; \gvect{\theta}_\text{emu})$ with weights $\gvect{\theta}_\text{emu} \in \mathbb{R}^{D_\text{emu}}$ on step-1 simulations by (4), and step 4 reruns step 2 with $\hat{f}$ in place of $f$:

Step 5 skips the inversion. It trains a conditional density $q(\gvect{\theta} \mid \vect{y}; \gvect{\nu})$ with weights $\gvect{\nu} \in \mathbb{R}^{D_\nu}$, which maps observations straight to an approximate posterior, on pairs $(\gvect{\theta}, \vect{y})$ drawn from the generative model (1), by (5):

so that the trained density approximates the posterior,

Table 3 says where each step lives in the stack. There are two structural commitments which fall out of the cycle:

4.6How the axes relate¶

The axes form a cube of operations. At any point you are somewhere on the modeling cycle (research stage), somewhere on the lifecycle (data maturity) and somewhere on the six-step cycle (production maturity). Figure 7 crosses the lifecycle transitions with the steps.

Every cell is a benchmarkable artifact, scored against the earlier step that computes the same thing (an emulator against the simulator it replaces, and emulator-based and amortized inference against the posterior of the step-2 oracle). The benchmark ladder of pipekit-evaluate ships this rule (Report 17). The whole lifecycle is a nested set of six-step cycles, each gated by benchmarks. This is what makes the framework improvable in practice, since improving a cell re-runs its downstream column, the lineage is content-addressed, and nothing is opaque.

5The big picture stack¶

We show the stack at two zoom levels: the layers first, then the packages inside GeoStack.

5.1Four layers¶

5.2GeoStack, zoomed in¶

Most of the current work happens in the GeoStack layer (Figure 9).

5.3The dependency graph¶

Figure 10 shows which repositories import which.

5.4Ownership rules¶

A package only imports from the layers below it, so there are no circular dependencies. geotoolz and xrtoolz re-export the common framework symbols, so user code usually imports a single package. Each of the three stacks of the decision has its own rules. The GeoStack rules (pipekit and its packages, geocatalog, geopatcher, geotoolz and xrtoolz) are in Report 16, and the GeoModels rules (somax, plumax, radtrax, on top of the GeoNumerics layer) are on the package boundaries page. The GeoML rules are only stated here: the JAX libraries touch pipekit only through an optional adapter module, and never in their core.

1. gaussx owns structured linear algebra and Gaussian primitives. It imports no GeoStack package.

2. geonnax owns basis functions, neural layers and coordinate transforms. somax and the other models compose them and never reimplement them.

3. kernellib owns covariance kernels and tapers. It imports no inference library.

4. optax_bayes owns the Bayesian learning rule as optax transforms. It imports no model library.

5. spectraldiffx and finitevolx own discretisations (spectral and finite-volume operators, no physics).

6. pyrox owns GP and Bayesian NN building blocks for NumPyro. Its pipekit surface (KrigingOperator) only lives in pyrox.adapters.pipekit.

7. filterax owns ensemble Kalman methods and vardax owns 3D/4D-Var. Each one only exposes an AnalysisStep through adapters/pipekit.py.

8. gauss_flows owns Gaussianization and normalizing flows, and xtremax owns extreme-value and point-process models. Neither imports pipekit outside an adapter.

6Walkthrough: methane attribution¶

This is the canonical example, a complete pipeline from satellite scenes to source emission estimates, with ML at several stages.

For the data, we use EMIT imaging-spectrometer scenes (L1 radiances) over an area of interest, TCCON column measurements as in-situ truth, and CAMS reanalysis and aircraft campaigns as evaluation references. The dimensions we use below are $B$ spectral bands, $N_\text{pix}$ pixels per scene (H × W in the code sketches), $M$ observations, $N$ grid cells, $N_\mathcal{E} = 40$ ensemble members $e \in \mathcal{E}$ (N_e in the code) and $T = 24$ hourly steps.

The eight steps of Figure 11 group into five stages. Each stage states the problem it solves before the code does, and stages 3 and 4 have closed-form special cases which double as test oracles (Table 6).

6.1Stage 1 · Retrieve methane (steps 1–2: geocatalog → georeader → geotoolz)¶

TL;DR. Find every EMIT scene over the area, mask clouds, and turn each pixel’s spectrum into a methane enhancement.

Problem. For pixel $n$, $\vect{y}_n \in \mathbb{R}^B$ is its radiance spectrum, $\gvect{\mu}_y \in \mathbb{R}^B$ is the scene-background mean, $\gmat{\Sigma}_y \in \mathbb{R}^{B \times B}$ is the background covariance and $\vect{J}_h \in \mathbb{R}^B$ is the methane target signature (unit absorption times $\gvect{\mu}_y$). The matched-filter enhancement estimate $\hat{\alpha}_n \in \mathbb{R}$ is the least-squares amplitude of $\vect{J}_h$ in the whitened residual:

The matched filter (7) is the fast classical enhancement estimate, and its derivation is in radtrax · 1. The full-physics retrieval (optimal estimation on a radiative-transfer model) is step 2 of the six-step cycle for the L1→L2 cell (Figure 7). A learned retrieval takes the same slot in the pipeline, and optimal estimation stays as its oracle.

6.2Stage 2 · Match with in-situ (step 3: geocatalog matchup)¶

TL;DR. Pair each TCCON measurement with the one satellite retrieval nearest to it in space and time; measurements with nothing nearby are dropped.

Problem. In-situ observation $j$ sits at location $\vect{s}_j \in \mathbb{R}^2$ and time $t_j \in \mathbb{R}$, satellite retrieval $r$ sits at $\vect{s}_r$ and $t_r$, and $d(\cdot, \cdot) \in \mathbb{R}_{\ge 0}$ is the great-circle distance. The candidates for $j$ are the retrievals inside both tolerances, (8):

If $\mathcal{C}_j$ is empty, then $j$ has no matchup. Otherwise $j$ keeps the single candidate with the smallest scaled distance (9), so no TCCON record is counted twice:

Ties go to the earlier retrieval and then to the smaller catalog index, so the selection is deterministic.

The pairs serve two jobs, so we split TCCON once, by station, before anything is fitted. The pairs from the training sites are the (label, feature) data which a learned retrieval trains on, whereas the pairs from the held-out sites validate stage 1 and are the only TCCON references we use in stage 5. No site appears in both, so the held-out scores never include a label the model saw. How much of the disagreement between the pairs is representativeness error is an open question.

6.3Stage 3 · Grid and gap-fill (step 4: pyrox through pipekit)¶

TL;DR. Turn scattered retrievals into a regular map with an uncertainty for every grid cell.

Problem. $\vect{y} \in \mathbb{R}^M$ are the L2 retrievals, $\sigma_\varepsilon^2 \in \mathbb{R}_{>0}$ is their noise variance and $k$ is a Matérn-5/2 kernel with a 50 km length scale. $\mat{K}_{SS} \in \mathbb{R}^{M \times M}$ is $k$ between the $M$ observation locations $\mat{S}$, $\mat{K}_{*S} \in \mathbb{R}^{N \times M}$ is $k$ between the $N$ grid cells and the observation locations (with $\mat{K}_{S*} = \mat{K}_{*S}^\top$), and $\mat{K}_{**} \in \mathbb{R}^{N \times N}$ is $k$ between the grid cells. The GP posterior mean $\gvect{\mu}_* \in \mathbb{R}^N$ and variance $\vect{v}_* \in \mathbb{R}^N$ on the grid are:

Kriging is the oracle for the L2→L3 cell, and DINEOF and neural gap-fillers are judged against it. Stage 4 assimilates this map, with $\gvect{\mu}_*$ as its observation and $\op{diag}(\vect{v}_*)$ as its observation-error covariance.

6.4Stage 4 · Assimilate hourly (step 5: pipekit-cycle → plumax → filterax)¶

TL;DR. Step a transport model forward each hour and pull it towards the observations.

Problem. $\vect{x}_t \in \mathbb{R}^N$ is the methane state at hour $t$, $f$ is one plumax transport step of 3,600 s, $\mat{H} \in \mathbb{R}^{N \times N}$ is the column observation operator, $\vect{y}_t \in \mathbb{R}^N$ is the stage-3 map for hour $t$, and $\gmat{\Sigma}_\varepsilon \in \mathbb{R}^{N \times N}$ is the observation-error covariance (here $\gmat{\Sigma}_\varepsilon = \op{diag}(\vect{v}_*)$ from the stage-3 kriging). The background covariance $\gmat{\Sigma}_{x,t \mid t-1} \in \mathbb{R}^{N \times N}$ is estimated from the $N_\mathcal{E}$ members. Each cycle forecasts every member $e = 1, \dots, N_\mathcal{E}$ by (12):

and then updates each member against its own perturbed copy of the observations (the stochastic EnKF), with $\gvect{\varepsilon}^{(e)} \sim \mathcal{N}(\vect{0}, \gmat{\Sigma}_\varepsilon)$ drawn independently per member:

where the gain $\mat{G}_t \in \mathbb{R}^{N \times N}$ is

The perturbations matter, because if we update every member with the same unperturbed $\vect{y}_t$, the anomalies shrink by $(\mat{I} - \mat{G}_t \mat{H})$ and the analysis ensemble is too narrow, which would bias the CRPS in stage 5. $\gmat{\Sigma}_{x,t \mid t-1}$ is inflated by 1.05 and localised with a Gaspari–Cohn taper. The localisation trades exactness for removing spurious long-range correlations, and the perturbations make the estimate random. However, they are drawn from one key seeded from the config of the run, so a rerun reproduces them (gate G1). A square-root (ETKF) variant replaces the perturbations with a deterministic anomaly transform. With a linear $f$ and $\mat{H}$, no inflation and no localisation, its analysis mean equals the Kalman update applied to the sample mean and covariance of the ensemble itself, which is gate D2.

6.5Stage 5 · Index, evaluate, register (steps 6–8: geocatalog.states → pipekit-evaluate → pipekit-experiment)¶

TL;DR. Store every analysis, score it against three independent references, and attach the score to the model that made it.

Problem. For a forecast distribution with CDF $F$, a reference value $y \in \mathbb{R}$, and independent draws $y', y'' \in \mathbb{R}$ from $F$, the continuous ranked probability score is given below. Lower is better, and with one member it reduces to the absolute error. Mass closure checks that the total methane is conserved to within 1 %.

The same pipeline, written as YAML, becomes the regulatory artifact (Report 9, use case 9). The same operators run in a notebook (case 1), an ETL job (case 2), a FastAPI endpoint (case 5) or a Fused UDF (the Fused integration design).

7Per-layer briefs¶

Here we give a compact summary per package, by layer (Figure 8). Where a report or a GeoModels page covers a package, the brief is a link to it.

8Key code patterns¶

Each pattern has runnable-looking code on the overview examples page.

9Open questions¶

None of these block the work, but all of them deserve some thought before we commit to a decision.

10A plan¶

Each phase produces packages which are useful on their own (Table 4).

10.1Phases 4 and 5, month by month¶

10.2Deferred until a project needs it¶

Table 5 lists what waits for a project, and what pulls it in.

Phase 7 covers the documentation site, a research journal, the first external benchmark adoption (OceanBench, WeatherBench-Ocean?) and capstone tutorials per domain.

10.3Gates¶

A month is done when its gates pass, and not when its packages exist (Figure 12). Every pass condition is a number or a closed form, and several of them reuse the closed-form stages of the walkthrough as oracles (Table 6).

10.4Risks¶

• The schedule slips. Six months for Phases 4 and 5 is optimistic for one person, and eight to twelve is realistic. Fallback: every month ships a usable package, so nothing depends on the whole plan finishing.

• Phase 6 is needed early. A project may need pipekit-jax or in-situ catalogs before the framework is ready. Fallback: we wrap jitted functions in pipekit-jax’s JaxModelOp until the full backend lands, and we jump phases, since the sequence is a default and not a contract.

• The Slice refactor ripples late. The longer GeoSlice stays concrete, the more code binds to it. Fallback: we introduce a Slice protocol in month 2 with GeoSlice as its only implementation, and add profile and trajectory slices later.

• pipekit-train becomes three half-backends. Equal status for Lightning, Equinox and Keras 3 triples the maintenance. Fallback: one default backend per use case, and the others are community-supported.

• Nobody else adopts it. Adoption takes years. Fallback: GeoBench is the first consumer, so the substrate earns its keep without outside users.

11Where to read next¶

11.1How the pages are written¶

Each page checks a design, a scoping report or a source note against the code pinned in references/, and the code takes precedence. A Differs box marks each place where the code (or a check the page runs) disagrees with the design. The gaps that matter become workstreams with proposals, and the rest go to open questions or upcoming work. Every number on a page comes from a notebook which asserts it. Pseudocode follows one convention: 1) a shapes header names the dimensions, 2) the equation sits in a comment above the line that computes it, 3) shapes go on the right, and 4) every call into the stack carries a status marker (markers).

11.2The pages¶

We suggest you start with the capability index and the overview examples, and then read by question.

Across the stack

• Report 9: twenty use cases mapped onto the libraries as they ship.

• Report 16: the geo data stack, the ownership rules and the realignment plan across 18 distributions.

pipekit

• Report 1: the toolz lineage, and why pipekit’s operators are classes; now pipekit’s Background page.

• Report 2: what pipekit 0.0.2 ships, in fourteen groups (A to N); the tour is now pipekit’s tour of the core, and the planning stays here.

• Report 5: jit, vmap and grad through pipelines, and what pipekit-jax ships.

• Report 10: pipekit-cycle, time-stepping and data assimilation.

• Report 11: pipekit-train, training for emulators and amortized inference.

• Report 12: pipekit-experiment, the registry, trackers and reproducibility artifacts.

• Report 14: pipekit-evaluate, multidimensional evaluation.

geotoolz and xrtoolz

• Report 3: the sister libraries, pipekit-array, geotoolz and xrtoolz.

• Report 6: geocatalog, the spatiotemporal index.

• Report 7: geopatcher, the four-axis patcher.

• Report 13: L3 to L5 indexing, as geocatalog’s state profile.

Benchmarking

• Report 15: benchmarks as content-addressed contracts.

• Report 17: the six-step cycle as the benchmark ladder, step 2 as oracle.

• Benchmarks gallery: eighteen draft benchmark entries in five domains.

GeoModels

• Package boundaries: what moves from geotoolz and plumax into plumax and a new radtrax.

• radtrax · 1: the Beer–Lambert column retrieval, term by term.

• radtrax · 2: the SBMP, MBSP and MBMP family of multi-band retrievals.

• radtrax · 3: radtrax’s design and roadmap (formerly rtmx).

• plumax · 1: emission rate from first principles, from the plume to the IME.

12In summary¶

GeoStack is one operator algebra from structured linear algebra to indexed multi-tier data products, which is organised by three orthogonal axes and split into the three stacks of the decision. Classical methods stay as oracles, benchmarks are content-addressed contracts, and differentiability only lives in the JAX packages. The plan is a default sequence which projects will reshape. So pick a page from Where to read next and come back with questions, since no part of the stack requires the rest.

13Glossary¶