Sister libraries on pipekit

The layers above pipekit¶

The proposal drew a strict stack: geotoolz and xrtoolz on top, delegating their array maths down to pipekit-array, which sits on pipekit core. What shipped has the same three carriers and different wiring (Figure 1). geotoolz sits directly on pipekit core. xrtoolz sits on a small package of its own, xrcore, which subclasses pipekit’s Operator. pipekit-array also sits on pipekit core, but neither domain library imports it. A fourth sibling, pipekit-jax, has shipped too. Table 1 sets each proposed layer against what it became. The stack-wide layering and its ownership rules are in Report 16.

The proposal’s sizes were estimates. The shipped line counts include docstrings and exclude tests.

Part 1 · pipekit-array: arrays through the Array API¶

What the Array API gives us¶

The Python Array API standard specifies a common array interface. A conforming array implements __array_namespace__(), which returns a module-like namespace with mean, sum, where, reshape and the rest. One function then works on any conforming backend:

1
2
3
4
5
from pipekit_array import array_namespace

def channel_mean(x):                          # x: (..., C)
    xp = array_namespace(x)                   # numpy, jax.numpy, cupy, …
    return xp.mean(x, axis=-1, keepdims=True)  # x̄ = (1/C) ∑ᶜ xᶜ  →  (..., 1)

The proposal’s conformance table was out of date in three places. It is corrected below from the NumPy and array-api-compat documentation (Table 2):

pipekit-array’s dispatch handles that wrinkle correctly. array_namespace tries array-api-compat first, which is shipped as the optional [compat] extra and needed for PyTorch, then falls back to x.__array_namespace__(). Either way it fails with the same message when no input conforms or the inputs mix backends.

What ships¶

pipekit-array is planned in four phases (its docs/design/boundaries.md, §4). Phase A has shipped:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
import numpy as np
from pipekit import Signature
from pipekit_array import ConcatenateAlong, MeanScalar, StackAlong, array_namespace

x = np.arange(24.0).reshape(2, 3, 4)          # (band, y, x)
array_namespace(x).__name__                   # → 'numpy'
MeanScalar()(x)                               # → 11.5            full mean, scalar
MeanScalar(axis=(-2, -1))(x)                  # → [ 5.5 17.5]     per band, (band,)
StackAlong(axis=0)([x[0], x[1]]).shape        # → (2, 3, 4)
ConcatenateAlong(axis=0)([x, x]).shape        # → (4, 3, 4)

pipe = StackAlong(axis=0) | MeanScalar(axis=(-2, -1))
pipe([x[0], x[1]])                            # → [ 5.5 17.5]

sig = Signature(dims=(("band", 2), ("y", 3), ("x", 4)), dtype="float64")
MeanScalar(axis=(-2, -1)).compute_output_signature(sig)
# → Signature((band=2) float64)               the reduced dims are dropped

MeanScalar is the first operator on pipekit’s Signature that does real shape inference: it drops the reduced dimensions, which is the payoff Report 2 said core leaves to the sister libraries. xrtoolz’s operators (in geo, interpolate, metrics and xreinx) infer shapes too, but on xrcore’s separate Signature (Section ??). Table 3 lists the whole planned catalogue by phase; only phase A ships.

The trade-offs¶

These held up as written, and they are why pipekit-array is worth its own package.

What works:

• NumPy code is bog-standard.

• JAX array maths works, and so do jax.jit, jax.vmap and jax.grad through a pipeline of JAX operators: a tracer is not a Node, so dispatch runs _apply (Report 5). Gradients with respect to parameters held by operators do not reach them (Section ??).

• CuPy arrays stay on the GPU through the pipeline.

• PyTorch keeps gradient tracking, so a pipeline can be differentiated inside an autograd context.

• The same operator runs on NumPy or JAX unchanged; the input’s namespace picks the backend. Inputs that mix backends in one call are rejected.

What doesn’t:

• NaN handling differs between backends at the edges, so operators that use NaN as a fill value need per-backend tests.

• In-place operations break on JAX, whose arrays are immutable. pipekit-array operators are therefore pure.

• Operations missing from the standard, such as some FFTs and advanced linear algebra, need backend-specific fallbacks. This is pipekit-array’s own open question Q1.

• Dask is lazy: a pipeline needs an explicit .compute() at the end (open question Q5).

Dependency surface¶

Shipped as proposed: dependencies = ["pipekit"], with one extra per backend (numpy, jax, torch, cupy, dask), plus a compat extra the proposal did not have (array-api-compat, needed for PyTorch). With no backend installed, array_namespace raises a TypeError naming the extras to install, rather than an ImportError.

Migrating the NumPy operators¶

The proposal moved the NumPy-flavoured operators out of geotoolz and xrtoolz into pipekit-array, then re-exported them with carrier-specific defaults:

1
2
3
4
5
# geotoolz/qc.py (proposed)
from pipekit_array.qc import AssertValueRange as _AssertValueRange

class AssertValueRange(_AssertValueRange):
    """Read GeoTensor's fill_value as the default bounds when none are given."""

None of this has happened, because the pipekit-array side (phase C) does not exist yet. Neither does the geotoolz side: geotoolz has no value-range, NaN or valid-fraction assertions at all. Its assert_* functions in _src/contract.py are test-contract helpers, not operators.

Part 2 · geotoolz: GeoTensor operators¶

Scope¶

The proposal made geotoolz a thinner library focused on remote sensing over georeader.GeoTensor, with its old framework code (geotoolz.core) reduced to a shim that re-exports pipekit. The shipped geotoolz has no core at all: it depends on pipekit (pipekit>=0.0.2) and its top-level __init__ re-exports 16 pipekit names directly (Table 4). gz.Operator is pipekit.Operator.

geotoolz defines no class with the same name as a pipekit public name. It deliberately keeps augment.Compose, io.SinkOperator and io.SourceOperator out of its top level, so they cannot be confused with pipekit’s compose and Sink. Its operators accept either a GeoTensor or a plain np.ndarray and return the same kind, so the carrier is wider than the proposal assumed.

Layout¶

Table 5 maps geotoolz’s subpackages onto the proposed layout.

Three proposed subpackages moved into packages of their own, and two did not ship:

• patch became geopatcher (geotoolz-patcher). It depends on pipekit only through an optional [pipekit] extra, and geotoolz bridges to it through patch_ops. ExtractPatches and StitchPatches do not exist under those names.

• catalog became geocatalog (geotoolz-catalog), where GeoCatalog is a Protocol. CatalogPipeline does not exist anywhere.

• presets and readers moved downstream into geoproducts (Section ??).

• geocloud (geotoolz-cloud) is not the proposal’s cloud (masks). It is cloud-storage I/O: stores, filesystems, COG and HDF access, caching, credentials.

• core does not exist, and no qa assertions shipped (see below).

What geotoolz owns, and what has not shipped¶

Table 6 lists the named symbols that did not ship under their proposed name.

The ownership argument itself held: geotoolz keeps everything that needs GeoTensor’s crs, transform and bounds, plus the sensor knowledge and the remote-sensing domain operators. Two of the proposal’s reasons for keeping things there did not turn into code: geo-specific QC (AssertCRSEquals, AssertResolutionWithin) and per-sensor calibration tables in geotoolz.

Sensor readers¶

The proposal put eight sensor modules in geotoolz.readers, each bundling a reader, sensor-specific operators and zero-argument presets. There is no geotoolz.readers. Readers and presets live in geoproducts, downstream of geotoolz (its [operators] extra pulls geotoolz in), and two of the eight have shipped (Table 7):

The GOES and Himawari presets are NDVI, TrueColor, NaturalColor, DayCloudPhase, FireTemperature, MaskClouds, ParallaxCorrect, a synthetic or hybrid green band, and a Recipe. geoproducts also has a toy sensor for tests and a Carbon Mapper client. georeader keeps its own readers for Sentinel-2 SAFE, EMIT, EnMAP, PRISMA, PROBA-V, SPOT-VGT and Earth Engine.

ModelOp and SklearnOp¶

Both live in geotoolz.learn, not in pipekit-array.

• ModelOp wraps any callable or object with .predict. It calls np.asarray on the input, supports batch_size, and returns the model’s output without re-wrapping it as a GeoTensor.

• SklearnOp flattens pixels into samples, runs a scikit-learn estimator, and restores the grid. It has fit, partial_fit, transform, predict and save_state.

Both set forbid_in_yaml = True, because both hold a model object.

Part 3 · xrtoolz: xarray operators¶

Scope¶

The proposal made xr_toolz a focused xarray library, with its framework code reduced to a shim. It shipped as xrtoolz, and the framework layer became a small package of its own, xrcore (xrtoolz-core, 418 lines). xrcore’s Operator subclasses pipekit’s and adds one dispatch mode:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
import pipekit as pk
import xarray as xr
import xrtoolz

class Double(xrtoolz.Operator):
    def __init__(self, var: str = "u") -> None:
        self.var = var

    def _apply(self, ds: xr.Dataset) -> xr.Dataset:
        return (2 * ds[self.var]).rename(f"{self.var}2").to_dataset()

xrtoolz.Operator is pk.Operator            # → False: a subclass,
issubclass(xrtoolz.Operator, pk.Operator)  # → True   so it composes with pipekit
xrtoolz.Sequential is pk.Sequential        # → True   re-exported unchanged

ds = xr.Dataset({"u": ("x", [1.0, 2.0]), "v": ("x", [3.0, 4.0])})
tree = xr.DataTree.from_dict({"/a": ds, "/b": ds * 10})
Double("u")(tree)["b"].to_dataset()["u2"].values   # → [20. 40.]  every node

A Node argument still builds a graph and plain data still runs _apply. A DataTree argument runs _apply_tree, which by default maps _apply over every node with xr.map_over_datasets. That is the xr.DataTree support the proposal listed, built in at the base class.

xrtoolz re-exports Graph, Input, Node, Sequential and Tap from pipekit, all the same objects. It defines no Tap, Sequential, Graph, Identity, Fanout or Lambda of its own.

Layout¶

The proposal had sixteen subpackages. xrtoolz folds six of them into one geo subpackage, and moves data access out to xrreader (Table 8; the named symbols are in Table 9):

The xarray combinators¶

The proposal kept two combinators in xrtoolz because they merge with xarray, and both shipped:

1
2
3
4
5
6
7
from xrtoolz import ApplyToEach, Augment

sorted(Augment(Double("u"))(ds).data_vars)    # → ['u', 'u2', 'v']

each = ApplyToEach(Double(), kwarg="var", values=["u", "v"])
sorted(each(ds).data_vars)                    # → ['u', 'u2', 'v', 'v2']
# Double(var="u") runs, merges; then Double(var="v") runs, merges

• Augment(inner) runs inner(ds) and merges the result back in with ds.merge(result, compat="no_conflicts"): the “compute derived variables and append them” pattern behind chains like Augment(RelativeVorticity()) | Augment(KineticEnergy()) | Augment(OkuboWeiss()).

• ApplyToEach(prototype, *, kwarg, values) rebuilds the prototype once per value (cls(**{**config, kwarg: v})) and merges every result. It rejects composite prototypes at construction. Each instance runs on the accumulated dataset, not on the original input, so a later value can see an earlier one’s output.

The proposal asked for a docstring note that this merge differs from pipekit’s framework combinators. The note exists on pipekit’s side (combine.py, which still calls the library xr_toolz), not in xrtoolz. pipekit’s Tap replaced xrtoolz’s, as planned.

Inference¶

xrtoolz.inference ships ModelOp, SklearnModelOp and JaxModelOp, all xrcore operators. The proposal’s sample_dim argument became feature_dim: every dimension except feature_dim is treated as a sample dimension, and output_name and output_dim name the result. JaxModelOp imports JAX lazily and has a jit flag. xrsklearn separately ships SklearnOp and XarrayEstimator.

The five sibling packages¶

None of these were in the proposal.

• xrcore (xrtoolz-core): the Operator subclass and Signature above.

• xreinx (xrtoolz-einx): einx patterns whose axes are dimension names. It has the functions einsum, rearrange, reduce, repeat, matmul and pack_dataset, and an operator for each (Einsum, Rearrange, Reduce, …). Its Reduce is a different class from xrtoolz.geo.Reduce.

• xrgrad (xrtoolz-grad): finite-difference calculus on xarray through finitediffx (JAX). It has partial, gradient, divergence, curl and laplacian, with geometry= set to "cartesian", "rectilinear" or "spherical". It is plain functions, with no operators and no pipekit dependency.

• xrreader (xrtoolz-reader): data access, with no pipekit dependency. It has the sources CMEMSSource, CDSSource, AemetSource and LocalL2Source; request types (BBox, TimeRange, DepthRange, PressureLevels, Station); a CF variable registry with validate_dataset; and Level-2 openers for TROPOMI, EMIT and GHGSat CH₄.

• xrsklearn (xrtoolz-sklearn): XarrayEstimator, which stacks, delegates to scikit-learn and unstacks for a DataArray, Dataset or DataTree; .sklearn accessors; and SklearnOp.

The array maths was meant to delegate to pipekit-array (the reduction inside RMSE, for example). No xrtoolz package imports pipekit-array. RMSE reduces through NumPy kernels in metrics/_src/_pixel_kernels.py.

Part 4 · Where each symbol lives¶

The proposal flagged this as the most common confusion. Table 10 shows where each name lives at the pinned revisions.

Part 5 · Other libraries that could fit on top¶

The proposal asked, for each candidate, whether it deserves a library of its own.

• pipekit-array, for NumPy, JAX, CuPy and PyTorch: yes, and it exists. The Array API is the right abstraction, and one library covers all four backends.

• pipekit-numba: no, and none exists. Numba is a compiler, not a carrier. A numba-jitted operator is still NumPy-flavoured at the carrier level; at most, a few pipekit-array operators could get @njit inner kernels behind a [fast] extra. That extra does not exist.

• A JAX-traceable library: still deferred, and pipekit-jax is not it. Differentiating through a pipeline with respect to its operators’ parameters means operators that are PyTrees with static metadata and a pure _apply; dispatch already lets tracers through. The shipped pipekit-jax does something else:

  • JaxModelOp wraps an eqx.Module as an operator, with serialize_weights, with_weights and from_registry, so trained weights round-trip through a pipekit-experiment registry;

  • DiffraxForwardModel is a pipekit-cycle forward model on a diffrax solve;

  • to_diffrax_adjoint and truncated_scan interpret pipekit-cycle’s adjoint specs.

  A pipeline of JAX operators can already be jitted, vmapped and differentiated with respect to its inputs; what is missing is differentiating with respect to the operators’ own parameters, because no shipped operator is a PyTree (Report 5).

• pipekit-dask: out of scope, and none exists. Distributing work is the orchestrator’s job. Operators that pickle can go inside dask.bag.map or dask.delayed as they are.

• pipekit-cuda: covered by pipekit-array with CuPy. That is the [cupy] extra.

• pandas, polars, dicts: deferred, and none exists. Carrier = Any, so a DataFrame operator is just an Operator subclass. The methane, MARS and atmospheric-chemistry work has not needed a DataFrame library.

Part 6 · How to decide on a new sister library¶

The proposal’s three questions, applied to the shipped code:

1. Is there a clearly typed carrier? If not, the operators belong in an existing library or in core.

2. Are there at least eight operators that would consume it natively? If not, contribute them to the closest existing library.

3. Does the carrier already have a wide-enough standard? The Array API does for arrays, and xarray.Dataset is its own de facto standard. If the carrier has no coherent standard, you would be inventing one alongside the operators.

Table 11 applies them to each candidate.

Two points in the shipped code bear on question 2. The domain libraries grew packages of their own (geopatcher, geocatalog, xrreader, xrgrad), each with a clear job and its own dependencies, and that turned out to be how they stay focused. pipekit-array exists ahead of its operators. That is fine for a library planned in phases, but until phase C lands, the code does not justify the package.

Summary¶

Table 12 sets each layer’s proposed status against the pinned revisions.

The proposal’s takeaway was two new packages (pipekit, pipekit-array) and a refactor of two existing ones. The refactor happened: neither domain library carries framework code any more. The ecosystem grew beyond the proposal: pipekit’s workspace has seven packages, geotoolz’s five and xrtoolz’s six. The piece that has not landed is the middle layer the proposal was built around, a shared array library the domain libraries delegate to.

Open questions¶

Decisions the shipped code raises:

1. Is pipekit-array still the middle layer? Both domain libraries grew their own array code instead. Either pipekit-array’s phases B to D land and geotoolz and xrtoolz migrate onto them, or pipekit-array becomes an optional toolkit beside them and the layering in this report is retired.

2. What do pipekit-array’s BatchedMap and Histogram become? Core already ships both names. Options: give the array versions distinct names (BatchedArrayMap, say), make them subclasses that add axis-0 splitting, or drop them and extend the core classes.

3. One Signature or two? Resolved: one, pipekit’s (Report 16 §A).

4. One JaxModelOp name or two? Resolved: one model-operator family (Report 16 §B).

5. Where do the numeric QC asserts go? AssertValueRange, AssertNoNaN and AssertValidFraction were promised to pipekit-array, and pipekit’s qc.py still says they live there. They exist nowhere. Should they wait for phase C, or land in geotoolz and xrtoolz first?

6. Does geotoolz own geo QC? AssertCRSEquals, AssertResolutionWithin and AssertSchema were part of the case for geotoolz and did not ship.

7. Which sensors come next in geoproducts? Resolved for TROPOMI: the reader stays on the xarray side (Report 16 §E).

Upcoming work¶

Table 13 lists the follow-ups whose shape is clear.

References¶

Citations¶

To be added: the Python Array API standard, array-api-compat, NumPy’s Array API compatibility notes, xarray, and georeader, as entries in references.bib. Until then, the Array API standard, array-api-compat and NumPy’s compatibility notes are linked inline.