Skip to content

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

  1. overview.md — the shared model: DataArray as the carrier, named dims as the type, lazy backend imports, Layer-0 functions + thin Operator wrappers.
  2. linalg/vision.md — gaussx-flavoured structured linear algebra with physical dims as row/column labels.
  3. prob/vision.md — distributions whose parameters are DataArrays.
  4. Inside each subpackage: architecture.md → api.md → decisions.md.

Cross-cutting principles

These hold for all three modules; see overview.md for the reasoning.

  1. DataArray-first. Layer-0 functions take and return DataArrays. Dataset is supported via small adapters (pack / unpack), and DataTree is inherited for free from xrtoolz.Operator's leaf-wise dispatch.
  2. 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.
  3. 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.
  4. Layer-0 functions + Layer-1 operators. Every public function has a thin Operator subclass with the same name (e.g. einsum → Einsum) so the same logic runs in Sequential / Graph. No logic is duplicated; the Operator calls the function in _apply.
  5. Backend deps are lazy and opt-in. Following the existing xrtoolz.inference pattern (D4): backends are imported inside functions, never at module load. Each module ships as a separate extra — install one at a time with pip install xrtoolz[einx], pip install xrtoolz[linalg], or pip install xrtoolz[prob], or combine them in a single command (pip install "xrtoolz[einx,linalg,prob]").
  6. 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 is pipekit_array.einx (geostack OP-O2), and pack_dataset / unpack_dataset live in xrops.utils.
  • linalg, prob — still draft / pre-implementation.

The docs describe the proposed API surface; concrete signatures may shift once each implementation lands.