Skip to article frontmatterSkip to article content
Site not loading correctly?

This may be due to an incorrect BASE_URL configuration. See the MyST Documentation for reference.

Overview examples

Code sketches for the patterns in the GeoStack overview

Each section links back to the part of the GeoStack overview it illustrates, and every sketch follows the series’ pseudocode convention (how the pages are written).

1Methane attribution, end to end

The pipeline behind Figure 11: catalog the scenes, retrieve, gap-fill, assimilate, index, evaluate and register. Every stage is an operator, so the whole thing composes, serialises and reruns. The problem each stage solves, with its equations, is in the walkthrough.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
import jax
import pipekit as pk                                         # today
import pipekit_cycle as pc                                   # today
import pipekit_train as pt                                   # today
import pipekit_jax as pj                                     # today
import pipekit_evaluate as pe                                # today
import pipekit_experiment as px                              # today

import geocatalog as gc  # today (ships as geotoolz-catalog)
import geopatcher as gp  # today (ships as geotoolz-patcher)
import geotoolz as gz                                        # today
import geocatalog.states as gcs                              # proposed (B)

from filterax.adapters.pipekit import EnKFAnalysis           # not pinned
from plumax.adapters.pipekit import PlumeForward, ColumnObs  # not pinned

Imports shared by every stage.

1 · Catalog
2 · Retrieve
3 · Gap-fill
4 · Assimilate
5 · Index
6 · Evaluate
7 · Register
1
2
3
4
5
6
7
8
9
10
11
# Shapes: one GeoSlice per scene or station record
emit_cat = gc.open_catalog("s3://imeo/emit.parquet")      # today
tccon_cat = gc.open_catalog("s3://imeo/tccon.parquet")    # today
cams_cat = gc.open_catalog("s3://cams/ch4_reanalysis.parquet")  # today
aircraft_cat = gc.open_catalog("s3://methaneair/2024.parquet")  # today

# Split TCCON once, by station, before anything is fitted: no site is in
# both halves, so held-out scores never see a training label.
split = pt.splitters.GroupSplit(group="site", test_fraction=0.3,
                                seed=cfg.seed)            # proposed (E)
tccon_train, tccon_test = split(tccon_cat)                # disjoint sites

Open the satellite and in-situ catalogs (stage 1).

tccon_train supplies the matchup pairs a learned retrieval trains on. tccon_test is the only TCCON data used to validate or evaluate.

2Tier examples

gaussx kernel matrix

Tier 1 numerics (briefs): kernels and solvers are lineax operators, so the Gram matrix never has to be materialised.

1
2
3
4
5
6
7
8
9
# Shapes: M training points, d input dimensions
from gaussx.kernels import Matern52                 # not pinned
from gaussx.linalg import cholesky_solve            # not pinned

k = Matern52(length_scale=10.0, variance=1.0)
# K = [k(xᵢ, xⱼ)]ᵢⱼ
K = k.gram(x)  # (M, d) → (M, M) operator, never dense
# m = K (K + σ² I)⁻¹ y
mean = K @ cholesky_solve(K + sigma2 * I, y_train)  # (M,) → (M,)

A Matérn-5/2 GP posterior mean with gaussx.

filterax adapter as a DA analysis step

Tier 2 algorithms plug in through adapters; this is the analysis step of the methane pipeline’s assimilate tab, with its localisation radius. It satisfies pipekit-cycle’s AnalysisStep protocol structurally: no base class needed. pipekit-cycle already ships EnsembleInflator and CovarianceLocalizer, so the adapter can take those objects instead of string options.

1
2
3
4
5
6
7
8
9
10
# Shapes: N state, N_e members
from filterax.adapters.pipekit import EnKFAnalysis  # not pinned

# Pᵇ ← ρ ∘ (λ Pᵇ), λ = 1.05, ρ a Gaspari–Cohn taper of 500 km
analysis = EnKFAnalysis(                            # (N_e, N) → (N_e, N)
    inflation=1.05,
    localization="gaspari-cohn",
    localization_radius_km=500.0,
)
assert isinstance(analysis, pc.AnalysisStep)        # the protocol: today

An EnKF analysis step from filterax.

3Key code patterns

Define an operator

A simplified sketch of how geotoolz.PercentileClip is written upstream.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
from typing import ClassVar

import numpy as np
from pipekit import Operator                           # today


class PercentileClip(Operator):
    """Clip a GeoTensor to [p_min, p_max] percentiles per band."""

    forbid_in_yaml: ClassVar[bool] = False
    p_min: float = 2.0
    p_max: float = 98.0
    axis: tuple = (-2, -1)

    def _apply(self, gt):
        # Shapes: B bands, H × W pixels
        # q_lo = Q_{p_min}(x_b), q_hi = Q_{p_max}(x_b) for each band b
        q = dict(axis=self.axis, keepdims=True)
        lo = np.quantile(gt.data, self.p_min / 100, **q)  # → (B, 1, 1)
        hi = np.quantile(gt.data, self.p_max / 100, **q)  # → (B, 1, 1)
        # x_clip = min(max(x, q_lo), q_hi)
        return gt.with_data(np.clip(gt.data, lo, hi))  # (B, H, W) → (B, H, W)

An operator is a class with fields and an _apply method.

ConfigMixin, inherited from Operator, derives get_config() from the fields. The operator is picklable, serialisable to YAML, registrable in the sandboxed loader and composable with |.

Compose pipelines

Sequential
Pipe syntax
Graph
1
2
3
4
5
# Shapes: B bands, H × W pixels; every step maps (B, H, W) → (B, H, W)
preprocess = pk.Sequential([               # today
    gz.ToFloat32(),                        # today
    gz.PercentileClip(p_min=2, p_max=98),  # today
])

A linear pipeline.

Catalogs

Observation catalogs index L0–L2 (lifecycle); matchup is where the two streams converge (two streams).

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# Shapes: one GeoSlice per file; M matched pairs
cat = gc.build_raster_catalog(                                   # proposed (B)
    "s3://imeo/sentinel2/2024/*.tif",
    backend="duckdb",
    out_uri="s3://imeo/sentinel2_index_2024.parquet",
)
cat = gc.open_catalog("s3://imeo/sentinel2_index_2024.parquet")  # today

# rows with footprint ∩ aoi ≠ ∅ and t₀ ≤ t ≤ t₁
rows = gc.query(cat, bbox=aoi, time=(t0, t1))                    # today*

# Wⱼ = {k : |tⱼ − tₖ| ≤ 3 h, d(sⱼ, sₖ) ≤ 25 km}
# k⋆(j) = argmin over Wⱼ of (d / 25 km)² + (|Δt| / 3 h)²
matched = gc.queries.matchup(  # → M ≤ |tccon| pairs  proposed (B)
    primary=tccon_train,  # training sites only
    secondary=cat,  # the dense satellite stream
    time_tolerance=pd.Timedelta(hours=3),
    space_tolerance_km=25.0,
    select="nearest",  # one k per j; ties: earlier time, then index
)

Build, open, query and match a geocatalog.

gc.intersect (today) already pairs overlapping rows across two catalogs; matchup adds the time and distance tolerances on top of it.

Patcher

The names below are geopatcher’s since 0.9.0, which dropped the Spatial prefixes (Report 7); Report 16, stage 2 runs the same pipeline with a matched filter in place of the model.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
# Shapes: C channels, H × W scene; 256² patches, stride 224
from geopatcher import SpatialPatcher, spatial     # today
from geopatcher.integrations.pipekit import (     # today
    ApplyToChips,
    GridSampler,
    MergePatches,
)

patcher = SpatialPatcher(                          # today
    geometry=spatial.geometry.Rectangular(size=(256, 256)),
    sampler=spatial.sampler.RegularStride(step=(224, 224)),
    window=spatial.window.Hann(),
    aggregation=spatial.aggregation.OverlapAdd(),
)
field = gp.RasterField(scene_reader)               # GeoTensor field  today

# y(s) = ∑ₚ wₚ(s) f(xₚ)(s) / ∑ₚ wₚ(s), wₚ the Hann window of p
inference = pk.Sequential([                        # today
    GridSampler(patcher=patcher),           # → P × (C, 256, 256)  today
    ApplyToChips(operator=preprocess | gz.ModelOp(model)),  # today
    MergePatches(aggregation=spatial.aggregation.OverlapAdd(),
                 domain=field.domain),      # P patches → (K, H, W)  today
])
result = inference(field)

Four-axis patching around a model, for any carrier.

GridSampler and MergePatches are marked forbid_in_yaml, so this pipeline runs but does not serialise.

Data assimilation cycle

The standalone, TROPOMI version of the methane pipeline’s assimilate tab: every ingredient satisfies a pipekit-cycle protocol, so any part that fits can be swapped in. The emulator and evaluation examples below build on it.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
# Shapes: N state, M observations, N_e = 40 members, T = 24 steps
tropomi_cat = gc.open_catalog("s3://imeo/tropomi.parquet")  # today
tropomi_source = pk.Sequential([                  # today
    gc.CatalogTimeQuery(tropomi_cat, window="1h"),  # proposed (B)
    pk.ThreadMap(gz.ReadNetCDF(variable="ch4")),  # → S × (M,)  today*
])
key = jax.random.PRNGKey(cfg.seed)  # one seed per run, from the config

forward = PlumeForward(species="CH4", dt=3600.0)  # ForwardModel  not pinned
obs_op = ColumnObs(instrument="TROPOMI")  # ObservationOperator  not pinned
analysis = EnKFAnalysis(inflation=1.05)           # AnalysisStep  not pinned

da = pc.EnsembleDACycle(                          # today*
    forward_model=forward,
    obs_op=obs_op,
    analysis_step=analysis,
    obs_source=tropomi_source,                    # TROPOMI columns, each hour
    n_steps=24,
    n_members=40,
    save_history=True,
    key=key,
)

carry0 = pc.DAState(                              # today*
    background_state=x0,                          # (N,)
    ensemble_members=X0,                          # (N_e, N)
    t=t0,
    cycle_count=0,
)
final_state, history = da(x0, carry0)             # history: (T, N_e, N)

An ensemble DA cycle with explicit state carry.

Swap in a neural emulator

The same DA cycle with a different forward model: steps 3→4 of the six-step cycle.

Train the emulator
Swap it in
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# Shapes: S = 10,000 simulations, N state
# φ⋆ = argmin_φ 𝔼_θ ‖f(θ) − f_φ(θ)‖², f_φ the emulator
loop = pt.TrainingLoop(                                      # today*
    model_op=pj.JaxModelOp(ChemistryEmulatorNet()),          # today
    dataset=pt.SimulationDataset(  # S × (θ, f(θ))  today
        forward_model=PlumeForward(species="CH4", dt=3600.0),
        prior=AtmosphericPrior(distribution="climatology"),
        n_samples=10_000,
    ),
    loss=pt.MSE(),                                           # today
    optimizer_config={"name": "adam", "lr": 1e-3},
    n_epochs=100,
    backend="equinox",
    callbacks=[pt.LogToExperiment(tracker=mlflow_tracker)],  # today
)
emulator_op = loop.run()
emulator_hash = registry.store(                        # → content hash  today
    emulator_op, name="methane_emulator_v3",
    weights=emulator_op.serialize_weights(),           # → weights.bin  today
)

Train an emulator offline on simulations.

Evaluation pipeline

What exists today is the benchmark cube. Task, Estimator, Estimate and Scorer are protocols: you satisfy them with your own objects, structurally, and Benchmark fits each estimator, resolves its reference and scores it. These small classes belong upstream (plumax’s adapters); they are inline here only to show the shape. The multi-track helpers sketched in the methane pipeline should be built on the same cube.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
# Shapes: N state, N_e members
from dataclasses import dataclass


@dataclass
class EnsembleEstimate:              # satisfies pe.Estimate
    members: jax.Array               # (N_e, N)

    def point(self):                 # x̄ = (1/N_e) ∑ᵢ x⁽ⁱ⁾
        return self.members.mean(axis=0)               # (N_e, N) → (N,)

    def sample(self, n, *, key=None):  # n draws from the empirical ensemble
        if key is None:
            raise ValueError("sampling needs a key")
        i = jax.random.choice(key, len(self.members), (n,))  # with replacement
        return self.members[i]                         # (N_e, N) → (n, N)

    def log_prob(self, theta):
        raise NotImplementedError    # optional in the protocol


@dataclass
class CycleEstimator:                # satisfies pe.Estimator
    cycle: pc.EnsembleDACycle
    x0: jax.Array                    # (N,) initial state, frozen with it
    carry0: pc.DAState               # initial ensemble and time
    rung: int                        # 2 = classical, 4 = emulator-based
    complexity: int = 0

    def fit(self, task):             # a DA cycle has nothing to train
        return self

    def __call__(self, obs):         # the benchmark's obs drive the run
        cycle = self.cycle.replace(obs_source=pk.Const(obs))  # proposed (D)
        _, history = cycle(self.x0, self.carry0)       # (T, N_e, N)
        return EnsembleEstimate(history[-1])           # final analysis


classical = CycleEstimator(da, x0, carry0, rung=2)     # the oracle
emulated = CycleEstimator(da_emulated, x0, carry0, rung=4)

bench = pe.Benchmark(                                  # today
    task=plume_task,             # a pe.Task from plumax  not pinned
    estimators=[classical, emulated],
    scorers=[pe.metrics.CRPS()],   # 𝔼|X − y| − ½ 𝔼|X − X′|  proposed (F)
    reference=classical,         # explicit oracle (default: lowest rung 2)
)
report = bench.run()                                   # → BenchmarkReport

Score an emulator-based estimator against the classical oracle.

Benchmark contract

A benchmark is a frozen contract: the task as a level transition, the carriers, the split and the gap it reproduces, the references, the scorers and the baselines, hashed with pipekit’s stable_json and sha256_hex. Report 15 runs this contract’s synthetic twin; the classes are proposed there (workstreams A, B and D).

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
# Shapes: along-track obs binned to the grid (time, lat, lon), NaN off-track,
#         in; (time, lat, lon) maps out
contract = pe.BenchmarkContract(                                 # proposed (A)
    name="oceanbench_ssh_2025", version="1.0.0",
    task="L1->L2", operation="gap filling",   # sparse grid → dense grid
    variant="deterministic",
    input={"carrier": "sparse grid", "vars": ["ssh"], "units": "m"},
    output={"carrier": "grid", "dims": ["time", "lat", "lon"], "units": "m"},
    splits={                        # the one place the split lives
        "_target_": "geocatalog.splits.split_by",               # proposed (B)
        "column": "producer", "test": "c2",
        "gap": "mapping where one altimeter flew, from the others",
    },
    references={
        "reanalysis": "s3://glorys/2023/index.parquet",
        "analysis": "s3://duacs/2023/index.parquet",
        "observations": "s3://altimetry/2023/c2.parquet",
    },
    scorers=[...],
    baselines=[{"_target_": "xrtoolz.interpolate.optimal_interpolation",
                "length": 100e3, "time_scale": "30D"}],         # proposed (C)
    status="pre-registered",
)
digest = contract.content_hash()  # sha256_hex(stable_json(...))  proposed (A)
report = pe.Benchmark.from_contract(contract, estimators=entries).run()
report.contract_hash == digest                                  # proposed (A)

A pre-registerable, content-addressed benchmark contract.

Splits follow Report 15, stage 2; the withheld mission ("c2" here) is the curator’s choice, sketched. The task string "L1->L2" follows pipekit-evaluate’s level names, which call L2 gap-filled observations; the overview’s levels put gap filling at L2 → L3 (two names for L2).

Reproducibility artifact

pipekit-experiment already ships InferenceArtifact: a JSON snapshot of a pipeline’s YAML, the hashes of the registered models it loads, the lockfile and free-form metadata, with save and load. What it cannot do yet is rebuild and rerun the pipeline from its frozen inputs; that rerun is proposed. Two more gaps block a faithful cold rerun today: a Graph’s config is a debug representation, not a round-trippable topology, so pk.dumps of a pipeline containing one needs a topology serialiser (proposed, workstream A); and reload_models restores model configs but not the weights a JaxModelOp keeps in the registry’s weights.bin, so the sketch reloads them with from_registry. Gate G1 freezes the evaluation too, so the report can be regenerated, not just the analysis.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
from pathlib import Path

inference = px.InferenceArtifact(                      # today
    pipeline_yaml=pk.dumps(da_emulated),     # Graph topology  proposed (A)
    pinned_model_hashes={"forward_model": emulator_hash},
    deps_lock=Path("pixi.lock").read_text(),
    metadata={
        "obs_catalog_uri": "s3://imeo/tropomi/2024_Q3_frozen/",
        "obs_catalog_sha256": "ab12cd34...",
        "seed": cfg.seed,
        "doi": "10.xxxx/mars-2024Q3",
    },
)
inference.save("regulatory/mars_methane_2024Q3.json")  # today

evaluation = px.InferenceArtifact(                     # today
    pipeline_yaml=pk.dumps(bench),     # Benchmark → YAML  proposed (F)
    deps_lock=Path("pixi.lock").read_text(),
    metadata={
        "reference_catalog_uri": "s3://imeo/tccon/2024_Q3_frozen/",
        "reference_catalog_sha256": "9f3e...",
        "inference_artifact": "mars_methane_2024Q3.json",
    },
)
evaluation.save("regulatory/mars_methane_2024Q3_eval.json")  # today

# Years later: a cold rerun, no cloud accounts needed
art = px.InferenceArtifact.load("regulatory/mars_methane_2024Q3.json")  # today
# reload_models restores configs only; JAX weights come back separately
emulator = pj.JaxModelOp(ChemistryEmulatorNet()).from_registry(
    registry, art.pinned_model_hashes["forward_model"])  # today
history = px.rerun(art, {"forward_model": emulator})     # proposed (E)

Freeze the analysis and its evaluation for a cold rerun.

GeoStack
Overview
Examples
Benchmarks gallery