xrtoolz bridge modules — design docs¶
Three helper subpackages that connect labeled xarray data (DataArray /
Dataset / DataTree) to the three Python toolchains xrtoolz pipelines
keep reaching for:
| Module | Bridges | Purpose |
|---|---|---|
xrtoolz.einx |
einx | Retired: the named-dims einx bridge is pipekit_array.einx (geostack OP-O2); its design pages were removed with xreinx. |
xrtoolz.linalg |
gaussx / lineax / matfree | Structured linear operators, solvers, logdet, sampling — all in DataArray-space, with physical dims as the row/column labels. |
xrtoolz.prob |
numpyro / pyrox | Probability/statistics on DataArrays: distributions parameterized by named tensors, sampling, log-probability, and posterior IO. |
The three modules are siblings of the existing domain packages
(xrtoolz.geo, xrtoolz.ocn, xrtoolz.atm, xrtoolz.rs,
xrtoolz.ice) and follow the same Layer-0 / Layer-1 split documented in
CLAUDE.md and docs/design/.
Structure¶
design_docs/xrtoolz/
├── README.md # This file — index + cross-cutting principles
├── overview.md # Shared motivation, contracts, and dependency model
├── linalg/ # gaussx + lineax bridge
│ ├── vision.md
│ ├── architecture.md
│ ├── api.md
│ └── decisions.md
└── prob/ # numpyro + pyrox bridge
├── vision.md
├── architecture.md
├── api.md
└── decisions.md
Reading order¶
- overview.md — the shared model: DataArray as the
carrier, named dims as the type, lazy backend imports, Layer-0
functions + thin
Operatorwrappers. - linalg/vision.md — gaussx-flavoured structured linear algebra with physical dims as row/column labels.
- prob/vision.md — distributions whose parameters are DataArrays.
- Inside each subpackage:
architecture.md→api.md→decisions.md.
Cross-cutting principles¶
These hold for all three modules; see overview.md for the reasoning.
- DataArray-first. Layer-0 functions take and return
DataArrays.Datasetis supported via small adapters (pack/unpack), andDataTreeis inherited for free fromxrtoolz.Operator's leaf-wise dispatch. - Named dims are the type. The dim names on the inputs determine which axis a backend operates on; sizes are inferred. No positional axis indices in the public API.
- Coordinate semantics. Each function documents which coords it preserves, which it drops, and how broadcasting interacts with labeled coords. The defaults match xarray's: aligned-by-coord on shared dims, broadcast on disjoint dims.
- Layer-0 functions + Layer-1 operators. Every public function has
a thin
Operatorsubclass with the same name (e.g.einsum→Einsum) so the same logic runs inSequential/Graph. No logic is duplicated; theOperatorcalls the function in_apply. - Backend deps are lazy and opt-in. Following the existing
xrtoolz.inferencepattern (D4): backends are imported inside functions, never at module load. Each module ships as a separate extra — install one at a time withpip install xrtoolz[einx],pip install xrtoolz[linalg], orpip install xrtoolz[prob], or combine them in a single command (pip install "xrtoolz[einx,linalg,prob]"). - No reimplementation. einx owns named-tensor semantics, gaussx owns linear-algebra dispatch, numpyro owns inference. These bridges only translate between labeled and array carriers.
Status¶
einx— retired. The named-dims einx bridge ispipekit_array.einx(geostack OP-O2), andpack_dataset/unpack_datasetlive inxrops.utils.linalg,prob— still draft / pre-implementation.
The docs describe the proposed API surface; concrete signatures may shift once each implementation lands.