GeoStack overview

Abstract¶

There is no Keras for geophysics. GeoStack is the connective tissue that is missing: one operator-graph abstraction that runs from structured linear algebra up to end-to-end pipelines, with content-addressed data and models throughout and machine learning pluggable at every step. This page is the map. It ties together the fifteen scoping reports, the GeoModels pages and the benchmarks gallery (Where to read next), and it ends with the open questions and a plan.

Keywords:geophysicsdata assimilationJAXoperator graphsremote sensing¶

How to read this page

Five minutes: read the elevator pitch, look at Figure 2 and Figure 8, and skim the plan.
Building something: start with the walkthrough and the overview examples, then search the capability index before writing new code.
Reviewing the design: go straight to the decision, the ownership rules, the gates and the open questions.

1The decision¶

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

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

A · One framework. Every library invents its own pipeline glue today. 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; model states are scattered files. geocatalog indexes both: L0–L2 observations as rows queried with GeoSlices, and L3–L5 model states through its state profile (run time, member, producer and parent columns, in geocatalog.states). The readers split by level: georeader below L3, xrreader from L3 up.
C · Two carriers, one algebra. geotoolz (GeoTensor) and xrtoolz (xarray) share the operator algebra and the patcher; only the substrate differs.
D · Data assimilation as a cycle. pipekit-cycle’s ForwardModel, ObservationOperator and AnalysisStep protocols are the only contract. filterax, vardax and plumax satisfy them through adapters, so an emulator replaces 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 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 never imports pipekit outside an adapter. GeoModels owns physics and exposes it only through adapters. pipekit owns composition, cycles, training and evaluation. geocatalog owns the one index, for observations and, through its state profile, for model states; 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; what is missing is the operator-graph abstraction that connects them, with content-addressed data and models throughout.

3The shape of the problem¶

Five gaps motivate everything below. The first four are old; the fifth has only become visible recently.

1 · Scattered numerics

The same linear algebra, kernel methods and filtering routines are reimplemented in every project. There are no shared primitives.

2 · ML tooling ignores classical numerics

Keras and PyTorch layers assume supervised learning on tensors. They do not express linear operators, PDE solvers, state estimation or uncertainty propagation.

3 · Monolithic domain software

Ocean models, retrieval algorithms and DA systems are dense, tightly coupled and hard to recombine.

4 · No connective tissue

Even with good numerics and good ML, there is no Keras-like pipeline to compose them for real geophysical work.

5 · No indexed model-state tier

Observations are catalogable (STAC, GeoParquet). Forecasts, analyses and reanalyses are scattered files with bespoke conventions: the L3–L5 half of the lifecycle has nothing like the catalogs that organise L0–L2.

The stack addresses all five.

4Three organising axes¶

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

Table 1:The three axes and what each one orders.

Axis	Orders	Figure
1 · Modeling cycle	the research loop: model, observe, compare, refine	Figure 1
2 · Data lifecycle	data maturity, from raw telemetry (L0) to derived products (L5)	Figure 2, Figure 3
3 · Six-step cycle	the journey from a simple model to amortized inference	Figure 6

4.1The modeling cycle (axis 1)¶

The real-world system feeds two tracks. Modeling runs conceptual model, equations and code, predictions. Observing runs sensor, measurements, observations. Both meet at compare. A random misfit goes to estimate (state by data assimilation, parameters by calibration), whose inner loop updates the predictions; a systematic misfit goes to refine the model, whose outer loop returns to the conceptual model and, dashed, to the sensor. — Figure 1:**The modeling cycle runs on two parallel tracks and two loops.**

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

Figure 2:The data lifecycle, science axis.

One stream of radiances fans out into ocean, atmosphere, land and ice at retrieval, and ends in derived products such as extremes (L5, dashed: examples still to be defined). Figure 3 maps the same levels to the tools that handle them.

Table 2:The six levels of the data lifecycle.

Level	Name	What it holds	Carrier
L0	Unstructured observations	raw instrument output: telemetry, counts, packets	GeoTensor
L1	Structured observations	calibrated, georeferenced radiances	GeoTensor
L2	Retrieved observations	per-pixel geophysical variables, with gaps where retrieval fails	GeoTensor
L3	Analysis	a regular, gap-filled grid: the “world right now”	xarray
L4	Reanalysis and forecast	dynamical model output: data-constrained, fused or projected	xarray
L5	Derived products	extremes, indicators and impacts computed from L4	xarray

Table 2 names the six levels.

The science axis runs left to right. Until retrieval it is one stream: 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 open.

The same six levels across the top. Below, a compute row names what performs the step into each level (radtrax at L2, GeoML at L3, GeoModels with data assimilation at L4, xtremax at L5), then rows for compose (pipekit, all levels), transform (geoops, xrtoolz), patch (geopatcher), carry (GeoTensor, xarray), read (georeader, xrreader), index (geocatalog, geocatalog.states) and store (geocloud, all levels), split by a dashed line where the carrier switches from GeoTensor to xarray. — Figure 3:**The data lifecycle, engineering axis.**

The engineering axis runs top to bottom, from computing a level down to storing it, and 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, vardax’s optimal interpolation); GeoModels with data assimilation model and assimilate (L3 → L4: somax and plumax with filterax and vardax), the inner loop of Figure 1; 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–L2): georeader reads them into a GeoTensor, geocatalog indexes them as GeoSlices and geotoolz transforms them.
Model states and derived products (L3–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 satellite chain L0 to L5 on the left; in-situ networks for ocean, atmosphere, land and ice at the bottom right. L2 and in-situ meet at a matchup node that feeds L3 and yields matchup pairs for validation and bias correction; in-situ also enters L4 directly through assimilation. — Figure 4:**Satellite and in-situ observations converge through matchup.**

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

4.4Three layers of transformation¶

The cycle and the lifecycle are joined by a third lens: what we do to the data on its way into a model (Figure 5).

Figure 5:Scientific data, ML-ready data, embeddings, and who owns each.

The first step is data engineering (GeoStack), the second representation learning (GeoML). The middle layer is the handoff that makes training reproducible: cached, CF-compliant tensors with content-addressed splits. Decoders run the chain backwards, which is what emulators and generative models do.

4.5The six-step cycle (axis 3)¶

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

Five steps in a row, from simple model to amortized inference, each with its defining equation. Step 1 simulates (θ, y) pairs that train steps 3 and 5; step 2 validates steps 3 to 5. A shared improve step under steps 2 to 5 lists each step's ground truth, the step before it. An axis below shows the cost per query falling from hours to milliseconds. — Figure 6:**From a simple model to amortized inference, with one shared improve step.**

In symbols: let θ be the quantity we want (an emission rate, a model state), f the physical forward model, H the observation operator, y the observations and ε their error with covariance R. Step 1 is the generative story:

y = H\left(f(\theta)\right) + \varepsilon, ~~~ \varepsilon \sim \mathcal{N}(0, R)

(1)

Step 2 inverts it with classical, model-based inference. Bayes’ rule with a prior p(θ) defines the target:

p(\theta \mid y) \propto p(y \mid \theta) ~ p(\theta)

(2)

Computing (2) is exact only 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, slowly. Step 2 is the oracle because it is the most trusted reference available, not because it is exact.

Step 3 trains an emulator $\hat f_\phi$ with weights φ on step-1 simulations by (3), and step 4 reruns step 2 with $\hat f_\phi$ in place of f:

\phi^\star = \arg\min_\phi \mathbb{E}_{\theta \sim p(\theta)} \left\| f(\theta) - \hat f_\phi(\theta) \right\|^2

(3)

Step 5 skips the inversion: it trains a conditional density $q_\psi$ with weights ψ that maps observations straight to an approximate posterior, on pairs (θ, y) drawn from the generative model (1), by (4):

\psi^\star = \arg\min_\psi \mathbb{E}_{(\theta, y) \sim p(\theta, y)} \left[ -\log q_\psi(\theta \mid y) \right], ~~~ q_{\psi^\star}(\theta \mid y) \approx p(\theta \mid y)

(4)

Table 3 says where each step lives in the stack. Two structural commitments fall out of the cycle:

Step 2 is the oracle

Classical methods (filterax, vardax, pyrox) are not legacy code waiting to be replaced by neural emulators. They become reference generators that validate steps 3–5. The classical path stays alive, promoted to oracle.

Benchmarks formalise understanding

Writing a benchmark contract (carriers, references, metrics, splits, baselines, expected failure modes) is formalising what you understand about the problem. Every transition between steps is a benchmarkable contract.

Table 3:Where each step lives in the stack.

Step	Lives in
1 · Simple model	GeoModels (`somax`, `plumax`, `radtrax`)
2 · Model-based inference	`pipekit-cycle.DACycle` with classical analysis steps
3 · Emulator	`pipekit-train.SimulationDataset`, registered in `pipekit-experiment.ModelRegistry`
4 · Emulator-based inference	`pipekit-cycle.DACycle` with `NeuralForward`
5 · Amortized inference	`pipekit-train` with `gauss_flows` or `pyrox`
6 · Improve	`pipekit-evaluate.Benchmark` reruns, compared row by row, with lineage in `pipekit-experiment`

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.

A grid with the five lifecycle steps as rows, from read and calibrate (detector counts to radiance) to derive (heatwave return levels), each with a running example, and steps 1 to 5 as columns headed by their equations. Each cell names a method and what it does on the example; the step-2 column is highlighted as the oracle. Arrows above show step 3 scored against step 1, and steps 4 and 5 against step 2. — Figure 7:**Every lifecycle step crossed with every six-step stage is a concrete method.**

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 step-2 oracle’s posterior. pipekit-evaluate’s benchmark ladder ships this rule (Report 17). The whole lifecycle is a nested set of six-step cycles, each gated by benchmarks. That is what makes the framework improvable in practice: improving a cell re-runs its downstream column, lineage is content-addressed, and nothing is opaque.

5The big picture stack¶

Two zoom levels: the layers first, then the packages inside GeoStack.

5.1Four layers¶

Four owned layers between external foundations at the bottom and external orchestration and serving at the top. GeoStack (data) and GeoNumerics (numerics) sit side by side at the bottom, GeoML (learning and inference) on both, GeoModels (physics) on top; imports point down. georeader is marked third-party and read-only. — Figure 8:**The stack in four layers.**

5.2GeoStack, zoomed in¶

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

The GeoStack layer. The pipekit core (pipekit, pipekit-array, pipekit-jax) at the bottom; above it two carrier columns, GeoTensor for observations and xarray for states, in rows by job: read (georeader, geoproducts; xrreader), index (geocatalog; the proposed geocatalog.states), transform (geoops; xrtoolz, xrcore, xrgrad, xreinx, xrsklearn), patch (geopatcher, either carrier) and store (geocloud, either carrier). The carrier-neutral ML loops (pipekit-cycle, -train, -experiment, -evaluate) on top. — Figure 9:**One pipekit core, two carriers, the same jobs on both.**

5.3The dependency graph¶

Figure 10 shows which repositories import which.

The 18 repositories in four layer bands: GeoModels on top, GeoML below, GeoNumerics and GeoStack side by side at the bottom. Edges point down from each repository to those it depends on or imports; dashed edges are optional extras; the geotoolz and pipekit cycle is drawn in red. georeader is marked third-party. — Figure 10:**How the repositories import each other.**

5.4Ownership rules¶

A package imports only from 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. The three stacks of the decision each have their rules: GeoStack’s (pipekit and its packages, geocatalog, geopatcher, geotoolz and xrtoolz) are Report 16’s, and GeoModels’ (somax, plumax, radtrax, on top of the GeoNumerics layer) are on the package boundaries page. GeoML’s rules are stated only here: the JAX libraries touch pipekit only through an optional adapter module, never in their core.

gaussx owns structured linear algebra and Gaussian primitives. It imports no GeoStack package.
geonnax owns basis functions, neural layers and coordinate transforms. somax and the other models compose them and never reimplement them.
kernellib owns covariance kernels and tapers. It imports no inference library.
optax_bayes owns the Bayesian learning rule as optax transforms. It imports no model library.
spectraldiffx and finitevolx own discretisations: spectral and finite-volume operators, no physics.
pyrox owns GP and Bayesian NN building blocks for NumPyro. Its pipekit surface (KrigingOperator) lives only in pyrox.adapters.pipekit.
filterax owns ensemble Kalman methods and vardax owns 3D/4D-Var. Each exposes an AnalysisStep only through adapters/pipekit.py.
gauss_flows owns Gaussianization and normalizing flows; xtremax owns extreme-value and point-process models. Neither imports pipekit outside an adapter.

6Walkthrough: methane attribution¶

The canonical example: a complete pipeline from satellite scenes to source emission estimates, with ML at several stages.

Figure 11:Eight steps from cataloguing scenes to serving estimates.

Stages 1–4 are the science, stage 5 the engineering around it. Each arrow names what flows to the next step; the dashed arrow is the inner loop of Figure 1, the misfit that updates the state and parameters.

The data: 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 used below: B spectral bands, H × W pixels per scene, M observations, N grid cells, $N_e = 40$ ensemble members 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; stages 3 and 4 have closed-form special cases that 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 i, $x_i \in \mathbb{R}^B$ is its radiance spectrum, μ the scene-background mean, $\Sigma \in \mathbb{R}^{B \times B}$ the background covariance and t the methane target signature (unit absorption times μ). The matched filter’s enhancement estimate $\hat\alpha_i$ is the least-squares amplitude of t in the whitened residual:

\hat\alpha_i = \frac{(x_i - \mu)^\top \Sigma^{-1} t}{t^\top \Sigma^{-1} t}

(5)

The matched filter (5) is the fast classical enhancement estimate; 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 $s_j$ and time $t_j$ ; satellite retrieval k at $s_k$ , $t_k$ ; d is great-circle distance. The candidates for j are the retrievals inside both tolerances, (6):

W_j = \left\{ k : |t_j - t_k| \le 3~\mathrm{h} ~~\text{and}~~ d(s_j, s_k) \le 25~\mathrm{km} \right\}

(6)

If $W_j$ is empty, j has no matchup. Otherwise j keeps the single candidate with the smallest scaled distance (7), so no TCCON record is counted twice:

k^\star(j) = \arg\min_{k \in W_j} \left(\frac{d(s_j, s_k)}{25~\mathrm{km}}\right)^2 + \left(\frac{|t_j - t_k|}{3~\mathrm{h}}\right)^2

(7)

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

The pairs serve two jobs, so TCCON is split once, by station, before anything is fitted. Pairs from the training sites are the (label, feature) data a learned retrieval trains on; pairs from the held-out sites validate stage 1 and are the only TCCON references used in stage 5. No site appears in both, so held-out scores never include a label the model saw. How much of the pairs’ disagreement 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. $y \in \mathbb{R}^M$ are the L2 retrievals, σ² their noise variance and k a Matérn-5/2 kernel with a 50 km length scale. $K_{ff}$ is k between the M observation locations, $K_{*f}$ between the N grid cells and the observations, and $K_{**}$ between grid cells. The GP posterior mean m and variance v on the grid are:

m = K_{*f} \left(K_{ff} + \sigma^2 I_M\right)^{-1} y

(8)

v = \operatorname{diag}\left(K_{**} - K_{*f} \left(K_{ff} + \sigma^2 I_M\right)^{-1} K_{f*}\right)

(9)

Kriging is the oracle for the L2→L3 cell; DINEOF and neural gap-fillers are judged against it. Stage 4 assimilates this map: m is its observation and $\mathrm{diag}(v)$ 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. $x_t \in \mathbb{R}^N$ is the methane state at hour t, 𝓜 one plumax transport step of 3,600 s, H the column observation operator, R the observation-error covariance, here $R = \mathrm{diag}(v)$ from the stage-3 kriging, and $P^b$ the background covariance estimated from the $N_e$ members. Each cycle forecasts every member $i = 1, \dots, N_e$ by (10):

x_t^{b,(i)} = \mathcal{M}\left(x_{t-1}^{a,(i)}\right)

(10)

then updates each member against its own perturbed copy of the observations, $\varepsilon^{(i)} \sim \mathcal{N}(0, R)$ drawn independently per member (the stochastic EnKF):

K = P^b H^\top \left(H P^b H^\top + R\right)^{-1}, ~~~ x_t^{a,(i)} = x_t^{b,(i)} + K \left(y_t + \varepsilon^{(i)} - H x_t^{b,(i)}\right)

(11)

The perturbations matter: updating every member with the same unperturbed $y_t$ shrinks the anomalies by $(I - KH)$ and leaves the analysis ensemble too narrow, which would bias the CRPS in stage 5. $P^b$ is inflated by 1.05 and localised with a Gaspari–Cohn taper; localisation trades exactness for removing spurious long-range correlations, and the perturbations make the estimate random. They are drawn from one key seeded from the run’s config, so a rerun reproduces them (gate G1). A square-root (ETKF) variant replaces the perturbations with a deterministic anomaly transform. With a linear 𝓜 and H, no inflation and no localisation, its analysis mean equals the Kalman update applied to the ensemble’s own sample mean and covariance, 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 F, a reference value y, and independent draws X, X′ from F, the continuous ranked probability score is below; lower is better, and with one member it reduces to absolute error. Mass closure checks that total methane is conserved to within 1 %.

\mathrm{CRPS}(F, y) = \mathbb{E}_F |X - y| - \frac{1}{2} \mathbb{E}_F |X - X'|

(12)

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¶

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.

External

GeoStack

GeoNumerics

GeoML

GeoModels

The JAX ecosystem: JAX, Equinox, lineax, optax, diffrax, NumPyro and matfree.

The xarray ecosystem: DataArray, Dataset and DataTree, plus rioxarray, xvec, cf-xarray, regionmask, xskillscore, dask and intake.

georeader (third-party, spaceml-org’s; read-only here) owns GeoTensor, the raster carrier, and the reader infrastructure (RasterioReader, AsyncGeoTIFFReader). Everything raster-shaped builds on it.

Package	What it does	Read
`pipekit`	The operator framework: fourteen groups (A–N)	Report 2, Report 1
`pipekit-array`	Array-API operators	Report 3
`pipekit-jax`	A model wrapper and a diffrax forward model; the operator layer is proposed	Report 5
`pipekit-cycle`	Time stepping and DA cycles	Report 10
`pipekit-train`	Training pipelines	Report 11
`pipekit-experiment`	Model registry and tracking	Report 12
`pipekit-evaluate`	The benchmark ladder	Report 14, Report 15, Report 17
`geotoolz`	GeoTensor operators for remote sensing	Report 3
`geocatalog` (ships as `geotoolz-catalog`)	Find, index, join and load files	Report 6
`geopatcher` (ships as `geotoolz-patcher`)	Split, process and stitch patches	Report 7
`geocatalog.states` (proposed state profile)	A catalog for model states	Report 13
`xrtoolz`	xarray operators for ocean, atmosphere and reanalysis	Report 3

geotoolz and xrtoolz share the same pipekit operator algebra; only the carrier differs.

Package	What it does
`gaussx`	Structured linear algebra, Gaussian primitives, exponential family; wraps lineax and matfree
`spectraldiffx`	Pseudospectral differentiation, filters, PDE solvers
`finitevolx`	Arakawa C-grid finite-volume operators

Pure math with no domain knowledge: all eqx.Module-based, jittable and JAX-traceable. See the gaussx example.

Package	What it does
`geonnax`	Basis functions, neural layers and coordinate transforms (`lonlat_to_cartesian3d`)
`kernellib`	Covariance kernels and scalable kernel methods
`pyrox`	GP and Bayesian deep-learning building blocks for NumPyro models
`vardax`	Learnable variational DA (3D- and 4D-Var)
`filterax`	Differentiable ensemble Kalman methods
`gauss_flows`	Rotation-based Gaussianization and normalizing flows
`xtremax`	Extreme-value and point-process models; computes the L5 extremes
`optax_bayes`	The Bayesian learning rule as optax gradient transformations

Each library’s adapter (*.adapters.pipekit) makes it composable into pipekit pipelines; the core never imports pipekit. See the filterax example.

Package	What it does	Read
`somax`	Ocean and atmosphere model zoo	—
`plumax`	Atmospheric plume transport	plumax · 1, package boundaries
`radtrax`	Differentiable radiative transfer (formerly rtmx)	radtrax · 3

The forward models each ship a thin adapters/pipekit.py that exposes the operator surface (ForwardModel, ObservationOperator, …) without touching the algorithmic core.

8Key code patterns¶

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

Define an operator

How an operator like PercentileClip is written.

Compose pipelines

Sequential, the pipe syntax, and DAGs with Graph.

Catalogs

Build, query and match a geocatalog.

Patcher

Geometry × sampler × window × aggregation.

Data assimilation cycle

Protocol-satisfying parts composed into a DACycle.

Swap in an emulator

Replace a classical forward model with a neural one.

Evaluation pipeline

Multi-track evaluation as a content-addressed artifact.

Benchmark contract

A pre-registerable, content-addressed benchmark.

Reproducibility artifact

Freeze a pipeline for a regulatory cold rerun.

9Open questions¶

None of these block work; all deserve thought before committing to a decision.

10A plan¶

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

Table 4:Phases at a glance.

Phase	Scope	Status	Time
1	GeoNumerics: gaussx (v0.0.6), spectraldiffx, finitevolx	gaussx done; the rest drafts	ongoing
2	GeoML: geonnax, kernellib, optax_bayes, pyrox, filterax, vardax, gauss_flows, xtremax	in progress	each at its own pace
3	GeoModels: somax, plumax, radtrax	some done	project-driven
4	Pipeline framework	current focus	~3–4 months
5	Evaluation, benchmarks and cycle discipline	next	~2 months
6	Deferred items, each with a trigger	when projects demand	—
7	Polish and community	open-ended	—

10.1Phases 4 and 5, month by month¶

Month 1

Month 2

Month 3

Month 4

Month 5

Month 6

pipekit 0.0.2: framework core, fourteen groups and state primitives (shipped; Report 2)
pipekit-array 0.0.2: Array API operators (shipped, phase A of four; Report 3)
the pipekit-ecosystem monorepo with uv workspaces

Gate: A1 config round-trip (Table 6).

10.2Deferred until a project needs it¶

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

Table 5:Phase 6: what is deferred, and what triggers it.

Item	Trigger
pipekit-jax’s operator layer (operators as PyTrees; 0.0.2 ships only a model wrapper and a diffrax forward model, Report 5)	the first project needing differentiable retrievals at scale or JIT-compiled inference
Deployment adapters (`pipekit.deploy.fused`, `.litserve`, `.modal`)	the first deployment beyond hand-shimming
Slice Protocol refactor (`ProfileSlice`, `TrajectorySlice`)	the first project needing first-class in-situ catalogs (likely the BGC benchmark)
More in-situ backends (Argo, ICOADS, SOCAT, …)	one per quarter, project-driven

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, not when its packages exist (Figure 12). Every pass condition is a number or a closed form, and several reuse the walkthrough’s closed-form stages as oracles (Table 6).

Figure 12:The framework ships first; every month closes on its gates.

Phases 1–3, one per lower layer, run alongside and meet the chain twice: the DA adapters land in month 3 and the GeoML models are trained in month 4. Phases 6 and 7 wait for projects to pull them.

Table 6:Gates: what must pass before each month closes.

Gate	Workstream	Passes when
A1 Config round-trip	A	For every shipped operator not marked `forbid_in_yaml`, `loads(dumps(op))` has the same config hash and returns the same output on a fixed input; once the topology serialiser lands, the same holds for a `Graph` (as Report 9’s A1)
B1 Backend parity	B	The InMemory and DuckDB geocatalog backends return the identical slice set for the same bbox and time query
B2 State rows	B	A cycle’s ensemble history, written as one file per run and member and indexed with `build_state_catalog`, loads back with `load_states` with no run dropped and equal to the history to 0 (Report 13, gates A2 and B1)
C1 Patch identity	C	An identity model through patch → stitch with Hann-window overlap-add reproduces the input to 1e-6 in float64, away from scene edges
D1 Kriging oracle	D	`KrigingOperator`’s mean and variance match (8) and (9) to 1e-6 in float64
D2 Kalman oracle	D	With linear 𝓜 and H, no inflation and no localisation, the square-root ensemble analysis mean matches the Kalman update computed from the same ensemble’s sample mean and covariance, to 1e-6 in float64
E1 Registry round-trip	E	Loading a stored operator returns one with the same content hash and bitwise-equal outputs
F1 Metric identities	F	CRPS (12) with one member equals absolute error exactly; RMSE of a field against itself is 0
G1 End to end	all	Stages 1–5 of the walkthrough run from one YAML file. A cold rerun of the frozen inference artifact reproduces the analysis history’s hash, and rerunning the frozen benchmark (its task, estimators, scorers and pinned TCCON references) on that history reproduces the report’s hash

10.4Risks¶

The schedule slips. Six months for Phases 4–5 is optimistic for one person; 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: wrap jitted functions in pipekit-jax’s JaxModelOp until the full backend lands, and jump phases; the sequence is a default, not a contract.
The Slice refactor ripples late. The longer GeoSlice stays concrete, the more code binds to it. Fallback: 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 maintenance. Fallback: one default backend per use case; 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. Gaps that matter become workstreams with proposals; the rest go to open questions or upcoming work. Every number on a page comes from a notebook that asserts it. Pseudocode follows one convention: a shapes header names the dimensions, the equation sits in a comment above the line that computes it, shapes go on the right, and every call into the stack carries a status marker (markers).

11.2The pages¶

Start with the capability index and the overview examples; 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–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–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, 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 lives only in the JAX packages. The plan is a default sequence that projects will reshape: pick a page from Where to read next and come back with questions, since no part of the stack requires the rest.

13Glossary¶

carrier: The data structure an operator acts on: NumPy or JAX arrays, a GeoTensor, or an xarray object. pipekit operators are carrier-agnostic.
GeoTensor: georeader’s raster carrier: an array with its geotransform and CRS. The carrier for L0–L2.
GeoSlice: geocatalog’s wire format for one spatiotemporal slice of an observation catalog.
state profile: geocatalog’s reserved columns for model states (run time, ensemble member, producer hash and parent state) and the loader and queries in geocatalog.states that use them. It replaces the planned statecatalog package and its StateSlice.
matchup: Pairing observations from different sources that coincide in space and time; the operation where the satellite and in-situ streams converge.
oracle: The classical method (step 2) used as the trusted reference to validate emulators and amortized predictors.