Skip to content

pipekit-array

Array-API operators on top of pipekit-ops. One operator runs on NumPy, JAX, CuPy, PyTorch or dask because it dispatches on array_namespace(x). The design doc has the architecture, the ADRs and the v0.1 operator catalogue.

What ships today: namespace dispatch, MeanScalar, StackAlong and ConcatenateAlong, the named carrier DimArray, all the dim-level functions with their operator twins (ApplyAlong included), padded batches and BatchAlong, sample tables and Estimator, the einx bridge, KernelOp, the xarray adapter and the test kit. What is still to build from the v0.1 catalogue: the numeric QC family (AssertClose, AssertValueRange, AssertNoNaN, AssertValidFraction), the array-aware Histogram and the trained-model wrapper ModelOp.

array_namespace

pipekit_array._namespace

Array API namespace dispatch shim.

The single point of backend dispatch for every pipekit-array operator. array_namespace(*xs) returns the Array API namespace (numpy, jax.numpy, torch via array-api-compat, …) shared by all input arrays.

Tries the optional array-api-compat library first — without it, PyTorch tensors don't expose __array_namespace__ natively and torch backends won't dispatch. Falls back to native x.__array_namespace__() lookup if the compat shim isn't installed.

See docs/design/architecture.md §4 and ADR A1 for the rationale.

array_namespace(*xs)

Return the Array API namespace shared by all input arrays.

Parameters:

Name Type Description Default
*xs Any

One or more array-like objects. All must share the same Array API backend.

 ()

Returns:

Type Description
Any

The module-like namespace object (e.g. numpy, jax.numpy,
Any

array_api_compat.torch).

Raises:

Type Description
TypeError

If no input implements the Array API, or if the inputs span more than one backend. The error message is consistent whether or not array-api-compat is installed.
Source code in packages/pipekit-array/src/pipekit_array/_namespace.py 
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
def array_namespace(*xs: Any) -> Any:
    """Return the Array API namespace shared by all input arrays.

    Args:
        *xs: One or more array-like objects. All must share the same
            Array API backend.

    Returns:
        The module-like namespace object (e.g. `numpy`, `jax.numpy`,
        `array_api_compat.torch`).

    Raises:
        TypeError: If no input implements the Array API, or if the
            inputs span more than one backend. The error message is
            consistent whether or not `array-api-compat` is installed.
    """
    compat_ns = _compat_namespace()

    if compat_ns is not None:
        try:
            return compat_ns(*xs)
        except TypeError as exc:
            # array-api-compat's error messages differ from our
            # contract; re-raise with stable wording so tests and
            # users see a consistent shape.
            msg = str(exc)
            if "Multiple namespaces" in msg:
                names = sorted({type(x).__module__.split(".")[0] for x in xs})
                raise TypeError(_multi_backend_msg(names)) from exc
            raise TypeError(_NO_BACKEND_MSG) from exc

    namespaces: set[Any] = set()
    for x in xs:
        ns = getattr(x, "__array_namespace__", None)
        if ns is None:
            raise TypeError(
                f"Input of type {type(x).__name__!r} does not implement "
                f"the Array API (`__array_namespace__` missing). " + _NO_BACKEND_MSG
            )
        namespaces.add(ns())
    if len(namespaces) == 0:
        raise TypeError(_NO_BACKEND_MSG)
    if len(namespaces) > 1:
        names = sorted({ns.__name__ for ns in namespaces})
        raise TypeError(_multi_backend_msg(names))
    return next(iter(namespaces))

Reductions

pipekit_array.reductions

Reductions — MeanScalar.

Wraps xp.mean as a pipekit Operator. See docs/design/api/operators.md for the spec.

MeanScalar

Bases: Operator

Reduce an array to a scalar (or per-axis means) via xp.mean.

Defaults to a full reduction (axis=None) yielding a scalar of rank 0. Output dtype matches the input's floating dtype per the Array API mean contract.

Parameters:

Name Type Description Default
axis int | tuple[int, ...] | None

Axis or tuple of axes to reduce over. None reduces all axes.

 None
keepdims bool

If True, retains the reduced axes with size 1.

 False

Raises:

Type Description
ValueError

From compute_output_signature when axis is out of range for the input signature's rank (xp.mean would raise the backend-native equivalent at apply time).

Example::

out = MeanScalar()(image)               # scalar
out = MeanScalar(axis=(-2, -1))(image)  # per-band mean
out = MeanScalar(axis=0, keepdims=True)(image)  # shape (1, H, W)
Source code in packages/pipekit-array/src/pipekit_array/reductions.py 
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
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
class MeanScalar(Operator):
    """Reduce an array to a scalar (or per-axis means) via `xp.mean`.

    Defaults to a full reduction (`axis=None`) yielding a scalar of
    rank 0. Output dtype matches the input's floating dtype per the
    Array API `mean` contract.

    Args:
        axis: Axis or tuple of axes to reduce over. `None` reduces
            all axes.
        keepdims: If True, retains the reduced axes with size 1.

    Raises:
        ValueError: From `compute_output_signature` when ``axis`` is out
            of range for the input signature's rank (``xp.mean`` would
            raise the backend-native equivalent at apply time).

    Example::

        out = MeanScalar()(image)               # scalar
        out = MeanScalar(axis=(-2, -1))(image)  # per-band mean
        out = MeanScalar(axis=0, keepdims=True)(image)  # shape (1, H, W)
    """

    #: saved states name the module it had before the rename (pipekit-array 0.0.3)
    __former_modules__: ClassVar[tuple[str, ...]] = ("pipekit_array.reduce",)

    def __init__(
        self,
        axis: int | tuple[int, ...] | None = None,
        keepdims: bool = False,
    ) -> None:
        # a tuple of axes reloads from JSON as a list, which `xp.mean` rejects
        self.axis = tuple(axis) if isinstance(axis, list) else axis
        self.keepdims = keepdims

    def _apply(self, x: Any) -> Any:
        xp = array_namespace(x)
        return xp.mean(x, axis=self.axis, keepdims=self.keepdims)

    def get_config(self) -> dict[str, Any]:
        return {"axis": self.axis, "keepdims": self.keepdims}

    def compute_output_signature(self, input_signature: Any) -> Any:
        if input_signature is None:
            return None

        input_sig: Signature = input_signature
        n_dims = len(input_sig.dims)

        def _validate(axis: int) -> int:
            """Return positive axis index; raise on out-of-range."""
            if n_dims == 0 or axis < -n_dims or axis >= n_dims:
                raise ValueError(
                    f"MeanScalar.axis={axis!r} is out of range for input "
                    f"signature with {n_dims} dim(s); xp.mean would raise."
                )
            return axis % n_dims

        if self.axis is None:
            axes_to_reduce: tuple[int, ...] = tuple(range(n_dims))
        elif isinstance(self.axis, int):
            axes_to_reduce = (_validate(self.axis),)
        else:
            axes_to_reduce = tuple(_validate(a) for a in self.axis)

        if self.keepdims:
            new_dims = tuple(
                (name, 1) if i in axes_to_reduce else (name, size)
                for i, (name, size) in enumerate(input_sig.dims)
            )
        else:
            new_dims = tuple(
                (name, size)
                for i, (name, size) in enumerate(input_sig.dims)
                if i not in axes_to_reduce
            )
        return Signature(dims=new_dims, dtype=input_sig.dtype)

Combinators

pipekit_array.combinators

Combinators — StackAlong, ConcatenateAlong.

Thin wrappers over xp.stack and xp.concat (the Array API standardised names) so the operations compose into a pipeline. See docs/design/api/operators.md for the spec.

The planned ApplyToBands and array-shaped BatchedMap shipped as pipekit_array.ApplyAlong and pipekit_array.BatchAlong (ADRs A15, A19).

ConcatenateAlong

Bases: Operator

Concatenate a list-of-arrays along an existing axis via xp.concat.

xp.concat is the Array API standardised name; numpy 2.0+ exposes it as concat alongside the legacy concatenate alias. All input arrays must share shape (except along axis), dtype, and backend.

Parameters:

Name Type Description Default
axis int | None

Existing axis to concatenate along, for raw arrays. Default 0.

 None
dim str | None

For named carriers instead: the dim to join along (pipekit_array.concat). axis and dim are mutually exclusive.

 None

Raises:

Type Description
ValueError

From _apply when the input sequence is empty (also via array_namespace when the inputs mix backends).

Example::

out = ConcatenateAlong(axis=0)([arr_a, arr_b])  # shape (2*N, H, W)
Source code in packages/pipekit-array/src/pipekit_array/combinators.py 
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
class ConcatenateAlong(Operator):
    """Concatenate a list-of-arrays along an existing axis via `xp.concat`.

    `xp.concat` is the Array API standardised name; numpy 2.0+
    exposes it as `concat` alongside the legacy `concatenate` alias.
    All input arrays must share shape (except along `axis`), dtype,
    and backend.

    Args:
        axis: Existing axis to concatenate along, for raw arrays. Default 0.
        dim: For named carriers instead: the dim to join along
            (`pipekit_array.concat`). ``axis`` and ``dim`` are mutually
            exclusive.

    Raises:
        ValueError: From ``_apply`` when the input sequence is empty
            (also via `array_namespace` when the inputs mix backends).

    Example::

        out = ConcatenateAlong(axis=0)([arr_a, arr_b])  # shape (2*N, H, W)
    """

    def __init__(self, axis: int | None = None, dim: str | None = None) -> None:
        if axis is not None and dim is not None:
            raise ValueError(
                "ConcatenateAlong takes axis= (raw arrays) or dim=, not both."
            )
        self.axis = 0 if axis is None and dim is None else axis
        self.dim = dim

    def _apply(self, xs: Sequence[Any]) -> Any:
        if len(xs) == 0:
            raise ValueError("ConcatenateAlong requires at least one input array.")
        if self.dim is not None:
            return dims.concat(list(xs), self.dim)
        xp = array_namespace(*xs)
        return xp.concat(list(xs), axis=self.axis)

    def get_config(self) -> dict[str, Any]:
        if self.dim is not None:
            return {"dim": self.dim}
        return {"axis": self.axis}

StackAlong

Bases: Operator

Stack a list-of-arrays along a new axis via xp.stack.

All input arrays must share the same shape, dtype, and backend.

Parameters:

Name Type Description Default
axis int | None

Position of the new axis in the output, for raw arrays. Default 0.

 None
dim str | None

For named carriers instead: the new dim's name (pipekit_array.stack; it becomes the leading dim). axis and dim are mutually exclusive.

 None

Raises:

Type Description
ValueError

From _apply when the input sequence is empty (also via array_namespace when the inputs mix backends).

Example::

out = StackAlong(axis=0)([arr_a, arr_b, arr_c])  # shape (3, H, W)
Source code in packages/pipekit-array/src/pipekit_array/combinators.py 
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
49
50
51
52
53
54
55
56
57
58
59
60
class StackAlong(Operator):
    """Stack a list-of-arrays along a new axis via `xp.stack`.

    All input arrays must share the same shape, dtype, and backend.

    Args:
        axis: Position of the new axis in the output, for raw arrays.
            Default 0.
        dim: For named carriers instead: the new dim's name
            (`pipekit_array.stack`; it becomes the leading dim). ``axis``
            and ``dim`` are mutually exclusive.

    Raises:
        ValueError: From ``_apply`` when the input sequence is empty
            (also via `array_namespace` when the inputs mix backends).

    Example::

        out = StackAlong(axis=0)([arr_a, arr_b, arr_c])  # shape (3, H, W)
    """

    def __init__(self, axis: int | None = None, dim: str | None = None) -> None:
        if axis is not None and dim is not None:
            raise ValueError("StackAlong takes axis= (raw arrays) or dim=, not both.")
        self.axis = 0 if axis is None and dim is None else axis
        self.dim = dim

    def _apply(self, xs: Sequence[Any]) -> Any:
        if len(xs) == 0:
            raise ValueError("StackAlong requires at least one input array.")
        if self.dim is not None:
            return dims.stack(list(xs), self.dim)
        xp = array_namespace(*xs)
        return xp.stack(list(xs), axis=self.axis)

    def get_config(self) -> dict[str, Any]:
        if self.dim is not None:
            return {"dim": self.dim}
        return {"axis": self.axis}

Named arrays

pipekit_array.dimarray

DimArray — a raw array with named dims, and its reference adapter.

What a geo carrier becomes when a kernel consumes one of its spatial dims (per-column statistics consume row and keep col), and the plain named carrier for tests and for libraries with none of their own (PK-P7, decided 23 Oct 2026).

  • It holds a raw array (NumPy, JAX or dask), one name per axis, and the static labels of any dim (band names, level values).
  • It is a JAX pytree whose only leaf is the data: dims and labels are static, so a new set of values never recompiles. Under a bare jax.vmap (in_axes=0) the mapped dim is the first one: it is dropped inside and comes back as "batch" (vmapped_dims); to map a named dim, use pipekit_array.map_over.
  • It does not hold a georeference, coordinates beyond labels, or attrs, and does not do alignment, arithmetic or selection by value; DimArray.to_xarray crosses to xarray when an analysis needs those.

DimArrayAdapter is the reference pipekit_ops.NamedArrayAdapter: the test kit (pipekit_array.testing) compares every other adapter against it.

DimArray dataclass

A raw array with one name per axis and static labels on some dims.

Parameters:

Name Type Description Default
data Any

A NumPy, JAX or dask array.

 required
dims tuple[str, ...]

One unique name per axis of data.

 required
labels Mapping[str, tuple[Any, ...]]

Static labels per dim, e.g. {"band": ("B11", "B12")}; each has one entry per index of its dim. Dims left out have none. Values are stored as tuples.

 dict()

Raises:

Type Description
ValueError

The dims do not match the data's rank, repeat a name, or a label tuple's length differs from its dim's size.
Example

import numpy as np mu = DimArray(np.zeros((3, 2)), ("col", "band"), {"band": ["B11", "B12"]}) mu.dims, mu.shape, mu.labels["band"] (('col', 'band'), (3, 2), ('B11', 'B12'))

Source code in packages/pipekit-array/src/pipekit_array/dimarray.py 
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
@dataclasses.dataclass(frozen=True, eq=False)
class DimArray:
    """A raw array with one name per axis and static labels on some dims.

    Args:
        data: A NumPy, JAX or dask array.
        dims: One unique name per axis of ``data``.
        labels: Static labels per dim, e.g. ``{"band": ("B11", "B12")}``;
            each has one entry per index of its dim. Dims left out have
            none. Values are stored as tuples.

    Raises:
        ValueError: The dims do not match the data's rank, repeat a name,
            or a label tuple's length differs from its dim's size.

    Example:
        >>> import numpy as np
        >>> mu = DimArray(np.zeros((3, 2)), ("col", "band"), {"band": ["B11", "B12"]})
        >>> mu.dims, mu.shape, mu.labels["band"]
        (('col', 'band'), (3, 2), ('B11', 'B12'))
    """

    data: Any
    dims: tuple[str, ...]
    labels: Mapping[str, tuple[Any, ...]] = dataclasses.field(default_factory=dict)

    def __post_init__(self) -> None:
        dims = tuple(self.dims)
        shape = tuple(self.data.shape)
        if len(dims) != len(shape):
            raise ValueError(f"DimArray dims {dims} for data of shape {shape}.")
        if len(set(dims)) != len(dims):
            raise ValueError(f"DimArray dims must be unique, got {dims}.")
        labels: dict[str, tuple[Any, ...]] = {}
        for dim, values in dict(self.labels).items():
            if dim not in dims:
                raise ValueError(f"labels for {dim!r}, which is not in dims {dims}.")
            entries = tuple(values)
            size = shape[dims.index(dim)]
            if len(entries) != size:
                raise ValueError(
                    f"{len(entries)} labels for dim {dim!r} of size {size}."
                )
            labels[dim] = entries
        object.__setattr__(self, "dims", dims)
        object.__setattr__(self, "labels", labels)

    @property
    def shape(self) -> tuple[int, ...]:
        """The data's shape."""
        return tuple(self.data.shape)

    @property
    def dtype(self) -> Any:
        """The data's dtype."""
        return self.data.dtype

    def __array__(self, dtype: Any = None, copy: bool | None = None) -> Any:
        """The data as a NumPy array, so ``np.asarray(d)`` works (dims are dropped)."""
        import numpy as np

        if copy is None:
            return np.asarray(self.data, dtype=dtype)
        # copy=True always copies; copy=False raises if a copy is needed
        return np.array(self.data, dtype=dtype, copy=copy)

    def to_xarray(self) -> Any:
        """An ``xr.DataArray`` with the same data and dims; labels become coordinates.

        Raises:
            ImportError: xarray is not installed (``pipekit-array[xarray]``).
        """
        try:
            import xarray as xr
        except ImportError as exc:
            raise ImportError(_XARRAY_HINT) from exc
        coords = {dim: (dim, list(values)) for dim, values in self.labels.items()}
        return xr.DataArray(self.data, dims=self.dims, coords=coords)

    def __repr__(self) -> str:
        labelled = f", labels={sorted(self.labels)}" if self.labels else ""
        return f"DimArray[{' '.join(self.dims)}]{self.shape}{labelled}"

dtype property

The data's dtype.

shape property

The data's shape.

__array__(dtype=None, copy=None)

The data as a NumPy array, so np.asarray(d) works (dims are dropped).

Source code in packages/pipekit-array/src/pipekit_array/dimarray.py 
 95
 96
 97
 98
 99
100
101
102
def __array__(self, dtype: Any = None, copy: bool | None = None) -> Any:
    """The data as a NumPy array, so ``np.asarray(d)`` works (dims are dropped)."""
    import numpy as np

    if copy is None:
        return np.asarray(self.data, dtype=dtype)
    # copy=True always copies; copy=False raises if a copy is needed
    return np.array(self.data, dtype=dtype, copy=copy)

to_xarray()

An xr.DataArray with the same data and dims; labels become coordinates.

Raises:

Type Description
ImportError

xarray is not installed (pipekit-array[xarray]).
Source code in packages/pipekit-array/src/pipekit_array/dimarray.py 
104
105
106
107
108
109
110
111
112
113
114
115
def to_xarray(self) -> Any:
    """An ``xr.DataArray`` with the same data and dims; labels become coordinates.

    Raises:
        ImportError: xarray is not installed (``pipekit-array[xarray]``).
    """
    try:
        import xarray as xr
    except ImportError as exc:
        raise ImportError(_XARRAY_HINT) from exc
    coords = {dim: (dim, list(values)) for dim, values in self.labels.items()}
    return xr.DataArray(self.data, dims=self.dims, coords=coords)

DimArrayAdapter

The reference pipekit_ops.NamedArrayAdapter, for DimArray.

Source code in packages/pipekit-array/src/pipekit_array/dimarray.py 
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
class DimArrayAdapter:
    """The reference `pipekit_ops.NamedArrayAdapter`, for `DimArray`."""

    def dims(self, x: DimArray) -> tuple[str, ...]:
        """``x.dims``."""
        return x.dims

    def data(self, x: DimArray) -> Any:
        """``x.data``."""
        return x.data

    def labels(self, x: DimArray, dim: str) -> tuple[Any, ...] | None:
        """The labels of ``dim``, or None."""
        return x.labels.get(dim)

    def isel(self, x: DimArray, indexers: dict[str, Any]) -> DimArray:
        """Select by position; integers drop their dim, labels follow the data.

        Raises:
            KeyError: An indexer names a dim ``x`` does not have.
        """
        unknown = sorted(set(indexers) - set(x.dims))
        if unknown:
            raise KeyError(f"isel: unknown dim(s) {unknown}; dims are {x.dims}.")
        sizes = dict(zip(x.dims, x.data.shape, strict=True))
        normalized = {
            dim: _positions(index, dim, sizes[dim]) for dim, index in indexers.items()
        }
        # ints and slices at once; each position array on its own axis, so
        # two arrays select independently instead of pairing up (outer indexing)
        basic = tuple(
            index if isinstance(index, slice) or _is_scalar(index) else slice(None)
            for index in (normalized.get(dim, slice(None)) for dim in x.dims)
        )
        data = x.data[basic]
        kept_dims = [
            dim
            for dim, index in zip(x.dims, basic, strict=True)
            if not _is_scalar(index)
        ]
        for dim, index in normalized.items():
            if isinstance(index, slice) or _is_scalar(index):
                continue
            axis = kept_dims.index(dim)
            data = data[(slice(None),) * axis + (index,)]
        key = tuple(normalized.get(dim, slice(None)) for dim in x.dims)
        dims: list[str] = []
        labels: dict[str, tuple[Any, ...]] = {}
        for dim, index in zip(x.dims, key, strict=True):
            if _is_scalar(index):
                continue
            dims.append(dim)
            if dim in x.labels:
                kept = _index_labels(x.labels[dim], index)
                if kept is not None:
                    labels[dim] = kept
        return _build(data, tuple(dims), labels)

    def rewrap(self, x: DimArray, data: Any, dims: tuple[str, ...]) -> DimArray:
        """A DimArray around new data; a kept dim keeps its labels if its size did.

        Raises:
            ValueError: ``dims`` does not match ``data``'s rank or repeats.
        """
        labels: dict[str, tuple[Any, ...]] = {}
        shape = tuple(data.shape)
        if len(shape) == len(dims):
            for dim, size in zip(dims, shape, strict=True):
                values = x.labels.get(dim)
                if values is not None and len(values) == size:
                    labels[dim] = values
        return DimArray(data, tuple(dims), labels)

    def relabel(
        self, x: DimArray, labels: Mapping[str, tuple[Any, ...] | None]
    ) -> DimArray:
        """``x`` with the labels of the named dims set; None removes them.

        The optional `pipekit_ops.LabelledAdapter` method (ADR A22).

        Raises:
            ValueError: A dim is not in ``x.dims``, or a label tuple's
                length differs from its dim's size.
        """
        unknown = sorted(set(labels) - set(x.dims))
        if unknown:
            raise ValueError(f"relabel: unknown dim(s) {unknown}; dims are {x.dims}.")
        merged = dict(x.labels)
        for dim, values in labels.items():
            if values is None:
                merged.pop(dim, None)
            else:
                merged[dim] = tuple(values)
        # the constructor checks each tuple against its dim's size
        return DimArray(x.data, x.dims, merged)

    def check_compatible(self, *xs: DimArray) -> tuple[DimArray, ...]:
        """The arrays back, if every shared dim has one size and one set of labels.

        Raises:
            ValueError: A dim has two sizes, or two different label sets.
        """
        sizes: dict[str, int] = {}
        labels: dict[str, tuple[Any, ...]] = {}
        for i, x in enumerate(xs):
            for dim, size in zip(x.dims, x.shape, strict=True):
                if sizes.setdefault(dim, size) != size:
                    raise ValueError(
                        f"input {i}: dim {dim!r} has size {size}, not {sizes[dim]}."
                    )
                values = x.labels.get(dim)
                if values is None:
                    continue
                if labels.setdefault(dim, values) != values:
                    raise ValueError(
                        f"input {i}: dim {dim!r} has labels {values}, "
                        f"not {labels[dim]}."
                    )
        return xs

    def signature(self, x: DimArray) -> Signature:
        """Dims with their sizes, and the dtype's name."""
        dims = tuple(zip(x.dims, x.shape, strict=True))
        return Signature(dims=dims, dtype=str(x.dtype))

check_compatible(*xs)

The arrays back, if every shared dim has one size and one set of labels.

Raises:

Type Description
ValueError

A dim has two sizes, or two different label sets.
Source code in packages/pipekit-array/src/pipekit_array/dimarray.py 
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
def check_compatible(self, *xs: DimArray) -> tuple[DimArray, ...]:
    """The arrays back, if every shared dim has one size and one set of labels.

    Raises:
        ValueError: A dim has two sizes, or two different label sets.
    """
    sizes: dict[str, int] = {}
    labels: dict[str, tuple[Any, ...]] = {}
    for i, x in enumerate(xs):
        for dim, size in zip(x.dims, x.shape, strict=True):
            if sizes.setdefault(dim, size) != size:
                raise ValueError(
                    f"input {i}: dim {dim!r} has size {size}, not {sizes[dim]}."
                )
            values = x.labels.get(dim)
            if values is None:
                continue
            if labels.setdefault(dim, values) != values:
                raise ValueError(
                    f"input {i}: dim {dim!r} has labels {values}, "
                    f"not {labels[dim]}."
                )
    return xs

data(x)

x.data.

Source code in packages/pipekit-array/src/pipekit_array/dimarray.py 
203
204
205
def data(self, x: DimArray) -> Any:
    """``x.data``."""
    return x.data

dims(x)

x.dims.

Source code in packages/pipekit-array/src/pipekit_array/dimarray.py 
199
200
201
def dims(self, x: DimArray) -> tuple[str, ...]:
    """``x.dims``."""
    return x.dims

isel(x, indexers)

Select by position; integers drop their dim, labels follow the data.

Raises:

Type Description
KeyError

An indexer names a dim x does not have.
Source code in packages/pipekit-array/src/pipekit_array/dimarray.py 
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
def isel(self, x: DimArray, indexers: dict[str, Any]) -> DimArray:
    """Select by position; integers drop their dim, labels follow the data.

    Raises:
        KeyError: An indexer names a dim ``x`` does not have.
    """
    unknown = sorted(set(indexers) - set(x.dims))
    if unknown:
        raise KeyError(f"isel: unknown dim(s) {unknown}; dims are {x.dims}.")
    sizes = dict(zip(x.dims, x.data.shape, strict=True))
    normalized = {
        dim: _positions(index, dim, sizes[dim]) for dim, index in indexers.items()
    }
    # ints and slices at once; each position array on its own axis, so
    # two arrays select independently instead of pairing up (outer indexing)
    basic = tuple(
        index if isinstance(index, slice) or _is_scalar(index) else slice(None)
        for index in (normalized.get(dim, slice(None)) for dim in x.dims)
    )
    data = x.data[basic]
    kept_dims = [
        dim
        for dim, index in zip(x.dims, basic, strict=True)
        if not _is_scalar(index)
    ]
    for dim, index in normalized.items():
        if isinstance(index, slice) or _is_scalar(index):
            continue
        axis = kept_dims.index(dim)
        data = data[(slice(None),) * axis + (index,)]
    key = tuple(normalized.get(dim, slice(None)) for dim in x.dims)
    dims: list[str] = []
    labels: dict[str, tuple[Any, ...]] = {}
    for dim, index in zip(x.dims, key, strict=True):
        if _is_scalar(index):
            continue
        dims.append(dim)
        if dim in x.labels:
            kept = _index_labels(x.labels[dim], index)
            if kept is not None:
                labels[dim] = kept
    return _build(data, tuple(dims), labels)

labels(x, dim)

The labels of dim, or None.

Source code in packages/pipekit-array/src/pipekit_array/dimarray.py 
207
208
209
def labels(self, x: DimArray, dim: str) -> tuple[Any, ...] | None:
    """The labels of ``dim``, or None."""
    return x.labels.get(dim)

relabel(x, labels)

x with the labels of the named dims set; None removes them.

The optional pipekit_ops.LabelledAdapter method (ADR A22).

Raises:

Type Description
ValueError

A dim is not in x.dims, or a label tuple's length differs from its dim's size.
Source code in packages/pipekit-array/src/pipekit_array/dimarray.py 
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
def relabel(
    self, x: DimArray, labels: Mapping[str, tuple[Any, ...] | None]
) -> DimArray:
    """``x`` with the labels of the named dims set; None removes them.

    The optional `pipekit_ops.LabelledAdapter` method (ADR A22).

    Raises:
        ValueError: A dim is not in ``x.dims``, or a label tuple's
            length differs from its dim's size.
    """
    unknown = sorted(set(labels) - set(x.dims))
    if unknown:
        raise ValueError(f"relabel: unknown dim(s) {unknown}; dims are {x.dims}.")
    merged = dict(x.labels)
    for dim, values in labels.items():
        if values is None:
            merged.pop(dim, None)
        else:
            merged[dim] = tuple(values)
    # the constructor checks each tuple against its dim's size
    return DimArray(x.data, x.dims, merged)

rewrap(x, data, dims)

A DimArray around new data; a kept dim keeps its labels if its size did.

Raises:

Type Description
ValueError

dims does not match data's rank or repeats.
Source code in packages/pipekit-array/src/pipekit_array/dimarray.py 
254
255
256
257
258
259
260
261
262
263
264
265
266
267
def rewrap(self, x: DimArray, data: Any, dims: tuple[str, ...]) -> DimArray:
    """A DimArray around new data; a kept dim keeps its labels if its size did.

    Raises:
        ValueError: ``dims`` does not match ``data``'s rank or repeats.
    """
    labels: dict[str, tuple[Any, ...]] = {}
    shape = tuple(data.shape)
    if len(shape) == len(dims):
        for dim, size in zip(dims, shape, strict=True):
            values = x.labels.get(dim)
            if values is not None and len(values) == size:
                labels[dim] = values
    return DimArray(data, tuple(dims), labels)

signature(x)

Dims with their sizes, and the dtype's name.

Source code in packages/pipekit-array/src/pipekit_array/dimarray.py 
316
317
318
319
def signature(self, x: DimArray) -> Signature:
    """Dims with their sizes, and the dtype's name."""
    dims = tuple(zip(x.dims, x.shape, strict=True))
    return Signature(dims=dims, dtype=str(x.dtype))

is_mask(index)

A boolean array, or a non-empty list / tuple of bools.

Source code in packages/pipekit-array/src/pipekit_array/dimarray.py 
139
140
141
142
143
144
145
146
def is_mask(index: Any) -> bool:
    """A boolean array, or a non-empty list / tuple of bools."""
    dtype = getattr(index, "dtype", None)
    return (dtype is not None and dtype == bool) or (
        isinstance(index, (list, tuple))
        and len(index) > 0
        and all(isinstance(v, bool) for v in index)
    )

vmapped_dims(dims, shape)

Static dims matched to leaves jax.vmap added or removed an axis from.

vmap maps leaves but not static dims: inside, the data has lost its leading axis; outside, the results gained one. With the default in_axes=0 / out_axes=0 that axis is the first, so drop the first dims, or prepend "batch" ("batch_1", … if taken). A shape that already matches, or no shape (a tracing placeholder), leaves dims alone. Map over any other axis with map_over.

Source code in packages/pipekit-array/src/pipekit_array/dimarray.py 
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
def vmapped_dims(dims: tuple[str, ...], shape: Any) -> tuple[str, ...]:
    """Static dims matched to leaves ``jax.vmap`` added or removed an axis from.

    ``vmap`` maps leaves but not static dims: inside, the data has lost
    its leading axis; outside, the results gained one. With the default
    ``in_axes=0`` / ``out_axes=0`` that axis is the first, so drop the
    first dims, or prepend ``"batch"`` (``"batch_1"``, … if taken). A
    shape that already matches, or no shape (a tracing placeholder),
    leaves ``dims`` alone. Map over any other axis with `map_over`.
    """
    if shape is None or not hasattr(shape, "__len__") or len(shape) == len(dims):
        return dims
    if len(shape) < len(dims):
        return dims[len(dims) - len(shape) :]
    added: list[str] = []
    k = 0
    while len(added) < len(shape) - len(dims):
        name = "batch" if k == 0 else f"batch_{k}"
        k += 1
        if name not in dims:
            added.append(name)
    return (*added, *dims)

Dim-level functions

pipekit_array.dims

Dim-level functions: selection, reduction, stacking and mapping by dim name.

Each is written once, over the carrier's registered pipekit_ops.NamedArrayAdapter, and behaves the same on every registered carrier: a DimArray, an xr.DataArray ([xarray]), a geoarray carrier. The rule for joining this list (GA-Q3): a function never reads coordinate values or geometry. resample, groupby and rolling stay xarray's; reprojection and sampling at points stay in geoarray.

  • isel, sel — by position, by exact label; an integer or a single label drops its dim.
  • transpose — reorder dims by name.
  • reduce, argreduce — over named dims; reduce(..., where=valid) leaves masked entries out (a padded batch, nodata).
  • take_along — gather along a dim with an index carrier (argreduce then take_along is a max-NDVI composite).
  • map_over — a function on each item along a dim.
  • stack, concat — a new dim, or along an existing one.
  • stack_items — items that may each carry their own geometry (patches, per-frame grids) along a new leading dim; stack where that is the same thing.
  • combine — an element-wise function of carriers with the same dims.

Fast paths. An adapter may provide a method of the same name (the xarray adapter's reduce stays lazy on dask; geoarray's map_over is one jax.vmap over data and transform rows). It is called with the same arguments instead of the generic path; the test kit (pipekit_array.testing) proves the two agree.

The Operator forms live in pipekit_array.dim_operators.

adapter_of(x, what)

The adapter of a named carrier; a TypeError for an unnamed value.

Source code in packages/pipekit-array/src/pipekit_array/dims.py 
64
65
66
67
68
69
70
71
72
def adapter_of(x: Any, what: str) -> NamedArrayAdapter:
    """The adapter of a named carrier; a `TypeError` for an unnamed value."""
    adapter = get_adapter(x)
    if adapter is None:
        raise TypeError(
            f"{what} works on named carriers, got {type(x).__name__}; wrap a raw "
            "array as pipekit_array.DimArray(data, dims)."
        )
    return adapter

argreduce(x, fn, dim)

The index of a reduction along one dim ("argmax", "argmin").

Example

import numpy as np from pipekit_array import DimArray ndvi = DimArray(np.random.default_rng(0).random((12, 4)), ("time", "x")) argreduce(ndvi, "argmax", "time").dims ('x',)

Source code in packages/pipekit-array/src/pipekit_array/dims.py 
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
def argreduce(x: Any, fn: Reduction, dim: str) -> Any:
    """The index of a reduction along one dim (``"argmax"``, ``"argmin"``).

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> ndvi = DimArray(np.random.default_rng(0).random((12, 4)), ("time", "x"))
        >>> argreduce(ndvi, "argmax", "time").dims
        ('x',)
    """
    adapter = adapter_of(x, "argreduce")
    fast = _fast_path(adapter, "argreduce")
    if fast is not None:
        return fast(x, fn, dim)
    dims = adapter.dims(x)
    (axis,) = _axes(dims, (dim,), "argreduce")
    data = adapter.data(x)
    func = _resolve(fn, namespace_of(data), "argreduce")
    out = func(data, axis=axis)
    kept = tuple(d for d in dims if d != dim)
    return adapter.rewrap(x, out, kept)

check_index_labels(adapter, x, index_adapter, index, dims)

take_along's check: on each shared dim the index's labels are x's.

Raises:

Type Description
ValueError

Labels differ in value or order (the index is positional).
Source code in packages/pipekit-array/src/pipekit_array/dims.py 
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
def check_index_labels(
    adapter: NamedArrayAdapter,
    x: Any,
    index_adapter: NamedArrayAdapter,
    index: Any,
    dims: Sequence[str],
) -> None:
    """`take_along`'s check: on each shared dim the index's labels are x's.

    Raises:
        ValueError: Labels differ in value or order (the index is positional).
    """
    for d in dims:
        mine, theirs = adapter.labels(x, d), index_adapter.labels(index, d)
        if mine is not None and theirs is not None and tuple(mine) != tuple(theirs):
            raise ValueError(
                f"take_along: the index's {d!r} labels {tuple(theirs)} are not x's "
                f"{tuple(mine)}; the index is applied by position, so they must match."
            )

combine(fn, *xs)

fn of the data of carriers with the same dims; never aligns.

Inputs are put in the first one's dim order by name before fn sees them, and the result has those dims.

Example

import numpy as np from pipekit_array import DimArray nir = DimArray(np.full((2, 3), 0.4), ("y", "x")) red = DimArray(np.full((2, 3), 0.1), ("y", "x")) nd = combine(lambda a, b: (a - b) / (a + b), nir, red) nd.dims, bool(np.allclose(nd.data, 0.6)) (('y', 'x'), True)

Source code in packages/pipekit-array/src/pipekit_array/dims.py 
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
def combine(fn: Callable[..., Any], *xs: Any) -> Any:
    """``fn`` of the data of carriers with the same dims; never aligns.

    Inputs are put in the first one's dim order by name before ``fn`` sees
    them, and the result has those dims.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> nir = DimArray(np.full((2, 3), 0.4), ("y", "x"))
        >>> red = DimArray(np.full((2, 3), 0.1), ("y", "x"))
        >>> nd = combine(lambda a, b: (a - b) / (a + b), nir, red)
        >>> nd.dims, bool(np.allclose(nd.data, 0.6))
        (('y', 'x'), True)
    """
    if len(xs) == 0:
        raise ValueError("combine needs at least one carrier.")
    adapter = adapter_of(xs[0], "combine")
    fast = _fast_path(adapter, "combine")
    if fast is not None:
        return fast(fn, *xs)
    checked = adapter.check_compatible(*xs)
    require_one_backend([adapter.data(x) for x in checked], "combine")
    dims = adapter.dims(checked[0])
    datas = []
    for x in checked:
        other = adapter.dims(x)
        if sorted(other) != sorted(dims):
            raise ValueError(f"combine: dims {other} differ from {dims}.")
        data = adapter.data(x)
        if other != dims:
            data = namespace_of(data).permute_dims(data, _axes(other, dims, "combine"))
        datas.append(data)
    return adapter.rewrap(checked[0], fn(*datas), dims)

concat(xs, dim)

Join carriers along an existing dim; the other dims must agree.

The generic path drops dim's labels (the joined dim has a new size); an adapter's concat fast path may join them.

Source code in packages/pipekit-array/src/pipekit_array/dims.py 
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
def concat(xs: Sequence[Any], dim: str) -> Any:
    """Join carriers along an existing dim; the other dims must agree.

    The generic path drops ``dim``'s labels (the joined dim has a new size);
    an adapter's ``concat`` fast path may join them.
    """
    if len(xs) == 0:
        raise ValueError("concat needs at least one carrier.")
    adapter = adapter_of(xs[0], "concat")
    fast = _fast_path(adapter, "concat")
    if fast is not None:
        return fast(list(xs), dim)
    dims = _same_dims(adapter, xs, "concat")
    (axis,) = _axes(dims, (dim,), "concat")
    # the other dims must combine; dim itself may differ in size
    # the other dims must combine; an empty slice of dim checks them even
    # when an input has no entries along dim
    others = [adapter.isel(x, {dim: slice(0, 0)}) for x in xs]
    adapter.check_compatible(*others)
    datas = [adapter.data(x) for x in xs]
    require_one_backend(datas, "concat")
    xp = namespace_of(datas[0])
    return adapter.rewrap(xs[0], xp.concat(datas, axis=axis), dims)

isel(x, /, **indexers)

Select by position on named dims; an integer drops its dim.

x is positional-only, so a dim may be called x.

Parameters:

Name Type Description Default
x Any

A named carrier.

 required
**indexers Any

dim= an int, a slice or an integer index array.

 {}
Example

import numpy as np from pipekit_array import DimArray cube = DimArray(np.zeros((12, 2, 4)), ("time", "band", "x")) isel(cube, time=0).dims ('band', 'x')

Source code in packages/pipekit-array/src/pipekit_array/dims.py 
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
def isel(x: Any, /, **indexers: Any) -> Any:
    """Select by position on named dims; an integer drops its dim.

    ``x`` is positional-only, so a dim may be called ``x``.

    Args:
        x: A named carrier.
        **indexers: ``dim=`` an int, a slice or an integer index array.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> cube = DimArray(np.zeros((12, 2, 4)), ("time", "band", "x"))
        >>> isel(cube, time=0).dims
        ('band', 'x')
    """
    adapter = adapter_of(x, "isel")
    _axes(adapter.dims(x), list(indexers), "isel")
    return adapter.isel(x, dict(indexers))

map_over(x, fn, dim)

fn on each item along dim; the results stacked along dim, first.

Each item is isel(x, dim=i); fn returns a named carrier, and every result must have the same dims.

Example

import numpy as np from pipekit_array import DimArray cube = DimArray(np.ones((3, 2, 4)), ("time", "band", "x")) map_over(cube, lambda item: reduce(item, "sum", "band"), "time").dims ('time', 'x')

Source code in packages/pipekit-array/src/pipekit_array/dims.py 
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
def map_over(x: Any, fn: Callable[[Any], Any], dim: str) -> Any:
    """``fn`` on each item along ``dim``; the results stacked along ``dim``, first.

    Each item is ``isel(x, dim=i)``; ``fn`` returns a named carrier, and every
    result must have the same dims.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> cube = DimArray(np.ones((3, 2, 4)), ("time", "band", "x"))
        >>> map_over(cube, lambda item: reduce(item, "sum", "band"), "time").dims
        ('time', 'x')
    """
    adapter = adapter_of(x, "map_over")
    fast = _fast_path(adapter, "map_over")
    if fast is not None:
        return fast(x, fn, dim)
    dims = adapter.dims(x)
    (axis,) = _axes(dims, (dim,), "map_over")
    size = adapter.data(x).shape[axis]
    items = [fn(adapter.isel(x, {dim: i})) for i in range(size)]
    out = stack(items, dim)
    # each output position is the input's: keep the mapped dim's labels
    labels = adapter.labels(x, dim)
    if labels is not None and isinstance(out, DimArray):
        return DimArray(out.data, out.dims, {**out.labels, dim: labels})
    return out

namespace_of(data)

The array namespace of raw data, dask included (which has none of its own).

Source code in packages/pipekit-array/src/pipekit_array/dims.py 
119
120
121
122
123
124
125
126
127
def namespace_of(data: Any) -> Any:
    """The array namespace of raw data, dask included (which has none of its own)."""
    if type(data).__module__.startswith("dask."):
        if not _DASK:
            import dask.array as da

            _DASK.append(_DaskNamespace(da))
        return _DASK[0]
    return array_namespace(data)

reduce(x, fn, dim, *, where=None)

Reduce over one or more named dims.

Parameters:

Name Type Description Default
x Any

A named carrier.

 required
fn Reduction

An array-namespace name ("mean", "max", "nanmean") or a function taking axis= (and where=, when given).

 required
dim str | Sequence[str]

A dim name, or several.

 required
where Any | None

Optional bool mask of the entries to keep, passed to fn as where= ("mean" then divides by the kept count). Either a named carrier, aligned to x by dim name (its dims must be among x's), or a plain array covering x's leading dims (where.shape == shape[:where.ndim]): a (sample,) mask on (sample, y, x) data, as Padded.valid is. Never right-broadcast. "max" / "min" need an initial for a mask: pass a function that supplies one.

 None

Raises:

Type Description
ValueError

A plain where does not match x's leading dims, or a named one has a dim x lacks.
TypeError

where on data other than NumPy or JAX (the reductions of dask, PyTorch and CuPy take no where=).
Example

import numpy as np from pipekit_array import DimArray cube = DimArray(np.ones((12, 2, 4)), ("time", "band", "x")) reduce(cube, "mean", "time").dims ('band', 'x') batch = DimArray(np.array([[1.0], [3.0], [9.0]]), ("sample", "x")) reduce(batch, "mean", "sample", where=np.array([True, True, False])).data array([2.])

Source code in packages/pipekit-array/src/pipekit_array/dims.py 
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
def reduce(
    x: Any, fn: Reduction, dim: str | Sequence[str], *, where: Any | None = None
) -> Any:
    """Reduce over one or more named dims.

    Args:
        x: A named carrier.
        fn: An array-namespace name (``"mean"``, ``"max"``, ``"nanmean"``) or a
            function taking ``axis=`` (and ``where=``, when given).
        dim: A dim name, or several.
        where: Optional bool mask of the entries to keep, passed to ``fn``
            as ``where=`` (``"mean"`` then divides by the kept count). Either
            a named carrier, aligned to ``x`` by dim name (its dims must be
            among ``x``'s), or a plain array covering ``x``'s *leading* dims
            (``where.shape == shape[:where.ndim]``): a ``(sample,)`` mask on
            ``(sample, y, x)`` data, as `Padded.valid` is. Never
            right-broadcast. ``"max"`` / ``"min"`` need an ``initial`` for
            a mask: pass a function that supplies one.

    Raises:
        ValueError: A plain ``where`` does not match ``x``'s leading dims,
            or a named one has a dim ``x`` lacks.
        TypeError: ``where`` on data other than NumPy or JAX (the
            reductions of dask, PyTorch and CuPy take no ``where=``).

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> cube = DimArray(np.ones((12, 2, 4)), ("time", "band", "x"))
        >>> reduce(cube, "mean", "time").dims
        ('band', 'x')
        >>> batch = DimArray(np.array([[1.0], [3.0], [9.0]]), ("sample", "x"))
        >>> reduce(batch, "mean", "sample", where=np.array([True, True, False])).data
        array([2.])
    """
    adapter = adapter_of(x, "reduce")
    names = (dim,) if isinstance(dim, str) else tuple(dim)
    fast = _fast_path(adapter, "reduce")
    dims = adapter.dims(x)
    data = adapter.data(x)
    if where is not None:
        _require_where_backend(data)
    mask = None if where is None else _aligned_mask(where, dims, data, "reduce")
    if fast is not None:
        return fast(x, fn, names) if mask is None else fast(x, fn, names, where=mask)
    axes = _axes(dims, names, "reduce")
    func = _resolve(fn, namespace_of(data), "reduce")
    out = func(data, axis=axes) if mask is None else func(data, axis=axes, where=mask)
    kept = tuple(d for d in dims if d not in names)
    return adapter.rewrap(x, out, kept)

require_one_backend(datas, what)

Raise BackendError when raw arrays sit on more than one backend.

Inputs never move between backends silently: a NumPy array next to a JAX one would otherwise be converted depending on input order.

Source code in packages/pipekit-array/src/pipekit_array/dims.py 
130
131
132
133
134
135
136
137
138
139
140
141
def require_one_backend(datas: Sequence[Any], what: str) -> None:
    """Raise `BackendError` when raw arrays sit on more than one backend.

    Inputs never move between backends silently: a NumPy array next to a
    JAX one would otherwise be converted depending on input order.
    """
    backends = {backend_name(d) for d in datas}
    if len(backends) > 1:
        raise BackendError(
            f"{what}: inputs on {sorted(backends)}; call to_jax() on the NumPy "
            "carriers first (inputs never move between backends silently)."
        )

sel(x, /, **labels)

Select by exact label on named dims: no ranges, no nearest.

x is positional-only, so a dim may be called x.

A single label drops its dim; a list of labels keeps it.

Raises:

Type Description
KeyError

A dim has no labels, or a label is not among them.
Example

import numpy as np from pipekit_array import DimArray scene = DimArray(np.zeros((2, 4)), ("band", "x"), {"band": ("B4", "B8")}) sel(scene, band="B8").dims ('x',)

Source code in packages/pipekit-array/src/pipekit_array/dims.py 
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
def sel(x: Any, /, **labels: Any) -> Any:
    """Select by exact label on named dims: no ranges, no nearest.

    ``x`` is positional-only, so a dim may be called ``x``.

    A single label drops its dim; a list of labels keeps it.

    Raises:
        KeyError: A dim has no labels, or a label is not among them.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> scene = DimArray(np.zeros((2, 4)), ("band", "x"), {"band": ("B4", "B8")})
        >>> sel(scene, band="B8").dims
        ('x',)
    """
    adapter = adapter_of(x, "sel")
    _axes(adapter.dims(x), list(labels), "sel")
    indexers: dict[str, Any] = {}
    for dim, wanted in labels.items():
        values = adapter.labels(x, dim)
        if values is None:
            raise KeyError(f"sel: dim {dim!r} has no labels; use isel.")
        position = {value: i for i, value in enumerate(values)}
        many = _many_labels(wanted, position)
        if many and not isinstance(wanted, (list, tuple)):
            wanted = list(wanted.tolist())
        for label in wanted if many else (wanted,):
            if label not in position:
                raise KeyError(f"sel: {label!r} is not a label of {dim!r}: {values}.")
        if many:
            indexers[dim] = [position[label] for label in wanted]
        else:
            indexers[dim] = position[wanted]
    return adapter.isel(x, indexers)

stack(xs, dim)

A new leading dim dim from carriers with the same dims; never aligns.

Source code in packages/pipekit-array/src/pipekit_array/dims.py 
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
def stack(xs: Sequence[Any], dim: str) -> Any:
    """A new leading dim ``dim`` from carriers with the same dims; never aligns."""
    if len(xs) == 0:
        raise ValueError("stack needs at least one carrier.")
    adapter = adapter_of(xs[0], "stack")
    fast = _fast_path(adapter, "stack")
    if fast is not None:
        return fast(list(xs), dim)
    checked = adapter.check_compatible(*xs)
    dims = _same_dims(adapter, checked, "stack")
    if dim in dims:
        raise ValueError(f"stack: {dim!r} is already a dim of {dims}; use concat.")
    datas = [adapter.data(x) for x in checked]
    require_one_backend(datas, "stack")
    xp = namespace_of(datas[0])
    return adapter.rewrap(checked[0], xp.stack(datas), (dim, *dims))

stack_items(xs, dim)

Items along a new leading dim dim, each keeping its own geometry.

The inverse of the split map_over makes: where stack requires one shared grid (geoarray's stack never aligns two), stack_items lets each item carry its own per-item leaves (a patch's transform row, a frame's grid) and makes dim a batch dim. A carrier family with no per-item geometry (DimArray, xr.DataArray) has nothing to keep apart, so this is stack; geoarray's adapter provides the stack_items fast path.

Raises:

Type Description
ValueError

xs is empty, or the items' dims differ.
Example

import numpy as np from pipekit_array import DimArray items = [DimArray(np.ones((2, 3)), ("y", "x")) for _ in range(4)] stack_items(items, "sample").dims ('sample', 'y', 'x')

Source code in packages/pipekit-array/src/pipekit_array/dims.py 
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
def stack_items(xs: Sequence[Any], dim: str) -> Any:
    """Items along a new leading dim ``dim``, each keeping its own geometry.

    The inverse of the split `map_over` makes: where `stack` requires one
    shared grid (geoarray's `stack` never aligns two), ``stack_items`` lets
    each item carry its own per-item leaves (a patch's transform row, a
    frame's grid) and makes ``dim`` a batch dim. A carrier family with no
    per-item geometry (`DimArray`, ``xr.DataArray``) has nothing to keep
    apart, so this is `stack`; geoarray's adapter provides the
    ``stack_items`` fast path.

    Raises:
        ValueError: ``xs`` is empty, or the items' dims differ.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> items = [DimArray(np.ones((2, 3)), ("y", "x")) for _ in range(4)]
        >>> stack_items(items, "sample").dims
        ('sample', 'y', 'x')
    """
    if len(xs) == 0:
        raise ValueError("stack_items needs at least one carrier.")
    adapter = adapter_of(xs[0], "stack_items")
    fast = _fast_path(adapter, "stack_items")
    if fast is not None:
        return fast(list(xs), dim)
    return stack(xs, dim)

take_along(x, index, dim)

Gather one entry of dim per position, with an index carrier.

index has x's dims without dim (an argreduce result, in any dim order); the result has those dims too.

Example

import numpy as np from pipekit_array import DimArray cube = DimArray(np.arange(24.0).reshape(3, 2, 4), ("time", "band", "x")) best = DimArray(np.array([0, 2, 1, 0]), ("x",)) take_along(cube, best, "time").dims ('band', 'x')

Source code in packages/pipekit-array/src/pipekit_array/dims.py 
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
def take_along(x: Any, index: Any, dim: str) -> Any:
    """Gather one entry of ``dim`` per position, with an index carrier.

    ``index`` has ``x``'s dims without ``dim`` (an `argreduce` result, in any
    dim order); the result has those dims too.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> cube = DimArray(np.arange(24.0).reshape(3, 2, 4), ("time", "band", "x"))
        >>> best = DimArray(np.array([0, 2, 1, 0]), ("x",))
        >>> take_along(cube, best, "time").dims
        ('band', 'x')
    """
    adapter = adapter_of(x, "take_along")
    fast = _fast_path(adapter, "take_along")
    if fast is not None:
        return fast(x, index, dim)
    index_adapter = adapter_of(index, "take_along")
    dims = adapter.dims(x)
    (axis,) = _axes(dims, (dim,), "take_along")
    rest = tuple(d for d in dims if d != dim)
    index_dims = index_adapter.dims(index)
    if not set(index_dims) <= set(rest):
        raise ValueError(f"take_along: index dims {index_dims} are not among {rest}.")
    data = adapter.data(x)
    xp = namespace_of(data)
    idx = index_adapter.data(index)
    # a JAX index next to NumPy data would be converted (or fail under jit)
    require_one_backend([data, idx], "take_along")
    # index → x's dim order with dim as a size-1 axis; missing dims broadcast
    check_index_labels(adapter, x, index_adapter, index, index_dims)
    sizes = dict(zip(dims, data.shape, strict=True))
    for d, n in zip(index_dims, idx.shape, strict=True):
        if n != sizes[d]:
            raise ValueError(
                f"take_along: index dim {d!r} has size {n}, but {d!r} has size "
                f"{sizes[d]}; the index is never broadcast."
            )
    present = [d for d in rest if d in index_dims]
    idx = xp.permute_dims(idx, tuple(index_dims.index(d) for d in present))
    full = tuple(1 if d == dim or d not in index_dims else None for d in dims)
    for position, size in enumerate(full):
        if size == 1:
            idx = xp.expand_dims(idx, axis=position)
    shape = tuple(1 if d == dim else n for d, n in zip(dims, data.shape, strict=True))
    idx = xp.broadcast_to(idx, shape)
    gathered = xp.take_along_axis(data, idx, axis=axis)
    out = xp.squeeze(gathered, axis=axis)
    return adapter.rewrap(x, out, rest)

transpose(x, order)

Reorder dims by name; order names every dim once.

Source code in packages/pipekit-array/src/pipekit_array/dims.py 
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
def transpose(x: Any, order: Sequence[str]) -> Any:
    """Reorder dims by name; ``order`` names every dim once."""
    adapter = adapter_of(x, "transpose")
    fast = _fast_path(adapter, "transpose")
    if fast is not None:
        return fast(x, tuple(order))
    dims = adapter.dims(x)
    if sorted(order) != sorted(dims):
        raise ValueError(
            f"transpose: order {tuple(order)} is not a permutation of {dims}."
        )
    data = adapter.data(x)
    xp = namespace_of(data)
    moved = xp.permute_dims(data, _axes(dims, order, "transpose"))
    return adapter.rewrap(x, moved, tuple(order))

pipekit_array.dim_shapes

Dim-level functions that change the dims: rename, add, drop, broadcast, merge, split.

Phase 1 of the dim-functions plan (docs/design/api/dim_functions.md, ADRs A11 to A14 and A22). Like the functions of pipekit_array.dims, each runs on any carrier with a registered pipekit_ops.NamedArrayAdapter, works on positions only (labels follow the data, their values are never read), and calls an adapter method of the same name, when there is one, instead of the generic path.

  • rename — new dim names; data, order and labels unchanged.
  • expand_dims, squeeze — add or drop length-1 dims.
  • broadcast_like — add another carrier's missing dims, in front.
  • merge_dims, split_dims — y x into pixel and back; labels become tuples and are split back.

Labels on new or renamed dims (A22). rewrap keeps labels by dim name, so it cannot label a dim whose name or labels are new. rename, broadcast_like, merge_dims and split_dims compute those labels and set them with the adapter's optional relabel (pipekit_ops.LabelledAdapter). An adapter without it runs them when no labels would be set, and raises TypeError when some would: labels are never dropped silently.

The Operator forms live in pipekit_array.dim_operators.

broadcast_like(x, other, /, *, exclude=())

Add other's missing dims to x, in front, in other's order.

y_{a…, d…} = x_{d…}, for every index a of the added dims.

Shared dims must agree exactly in size and labels (xarray outer-joins instead). x's own dims keep their order, so the result's order is xarray's only when the shared dims are in other's order. The added dims carry other's labels.

Parameters:

Name Type Description Default
x Any

A named carrier.

 required
other Any

A carrier of the same family; only its dims, sizes and labels are read.

 required
exclude Sequence[str]

Dims of other neither added nor checked.

 ()

Raises:

Type Description
TypeError

The adapter has no relabel and an added dim has labels.
ValueError

A shared dim differs in size or labels.
Example

import numpy as np from pipekit_array import DimArray cube = DimArray(np.zeros((12, 2, 4)), ("time", "band", "x")) median = DimArray(np.ones((2, 4)), ("band", "x")) broadcast_like(median, cube).shape (12, 2, 4)

Source code in packages/pipekit-array/src/pipekit_array/dim_shapes.py 
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
def broadcast_like(x: Any, other: Any, /, *, exclude: Sequence[str] = ()) -> Any:
    """Add ``other``'s missing dims to ``x``, in front, in ``other``'s order.

    y_{a…, d…} = x_{d…}, for every index a of the added dims.

    Shared dims must agree exactly in size and labels (xarray outer-joins
    instead). ``x``'s own dims keep their order, so the result's order is
    xarray's only when the shared dims are in ``other``'s order. The added
    dims carry ``other``'s labels.

    Args:
        x: A named carrier.
        other: A carrier of the same family; only its dims, sizes and
            labels are read.
        exclude: Dims of ``other`` neither added nor checked.

    Raises:
        TypeError: The adapter has no ``relabel`` and an added dim has
            labels.
        ValueError: A shared dim differs in size or labels.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> cube = DimArray(np.zeros((12, 2, 4)), ("time", "band", "x"))
        >>> median = DimArray(np.ones((2, 4)), ("band", "x"))
        >>> broadcast_like(median, cube).shape
        (12, 2, 4)
    """
    adapter = adapter_of(x, "broadcast_like")
    other_adapter = adapter_of(other, "broadcast_like")
    excluded = tuple(exclude)
    fast = _fast_path(adapter, "broadcast_like")
    if fast is not None:
        return fast(x, other, exclude=excluded)
    dims = adapter.dims(x)
    other_dims = other_adapter.dims(other)
    shared = [d for d in dims if d in other_dims and d not in excluded]
    added = tuple(d for d in other_dims if d not in dims and d not in excluded)
    # only the shared dims must combine: an empty slice of every other dim
    # checks them without touching the data
    x_part = adapter.isel(x, {d: slice(0, 0) for d in dims if d not in shared})
    other_part = other_adapter.isel(
        other, {d: slice(0, 0) for d in other_dims if d not in shared}
    )
    adapter.check_compatible(x_part, other_part)
    data = adapter.data(x)
    xp = namespace_of(data)
    sizes = dict(zip(other_dims, other_adapter.data(other).shape, strict=True))
    expanded = data
    for _ in added:
        expanded = xp.expand_dims(expanded, axis=0)
    target = (*(sizes[d] for d in added), *data.shape)
    broadcast = xp.broadcast_to(expanded, target)
    out = adapter.rewrap(x, broadcast, (*added, *dims))
    labels = {d: other_adapter.labels(other, d) for d in added}
    return set_labels(adapter, out, labels, "broadcast_like")

expand_dims(x, dim, /, *, axis=0)

Insert length-1 dims, consecutively, at position axis of the result.

y_{…, 0, …} = x_{…}: the new dims have size 1 and no labels.

A negative axis counts from the end, so axis=-1 appends.

Parameters:

Name Type Description Default
x Any

A named carrier.

 required
dim str | Sequence[str]

A new dim name, or several.

 required
axis int

Where the first new dim goes in the result.

 0

Raises:

Type Description
ValueError

A new name is already a dim, a name repeats, or axis is out of range.

Examples:

>>> import numpy as np
>>> from pipekit_array import DimArray
>>> scene = DimArray(np.zeros((2, 4)), ("band", "x"))
>>> expand_dims(scene, "member").dims
('member', 'band', 'x')
>>> expand_dims(scene, ("member", "run"), axis=-1).shape
(2, 4, 1, 1)
Source code in packages/pipekit-array/src/pipekit_array/dim_shapes.py 
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
def expand_dims(x: Any, dim: str | Sequence[str], /, *, axis: int = 0) -> Any:
    """Insert length-1 dims, consecutively, at position ``axis`` of the result.

    y_{…, 0, …} = x_{…}: the new dims have size 1 and no labels.

    A negative ``axis`` counts from the end, so ``axis=-1`` appends.

    Args:
        x: A named carrier.
        dim: A new dim name, or several.
        axis: Where the first new dim goes in the result.

    Raises:
        ValueError: A new name is already a dim, a name repeats, or
            ``axis`` is out of range.

    Examples:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> scene = DimArray(np.zeros((2, 4)), ("band", "x"))
        >>> expand_dims(scene, "member").dims
        ('member', 'band', 'x')
        >>> expand_dims(scene, ("member", "run"), axis=-1).shape
        (2, 4, 1, 1)
    """
    adapter = adapter_of(x, "expand_dims")
    names = (dim,) if isinstance(dim, str) else tuple(dim)
    dims = adapter.dims(x)
    taken = [n for n in names if n in dims]
    if taken:
        raise ValueError(f"expand_dims: {taken} are already dims of {dims}.")
    if len(set(names)) != len(names):
        raise ValueError(f"expand_dims: the new dims {names} repeat a name.")
    # a negative axis counts from the end of the result: -1 appends
    start = axis if axis >= 0 else axis + len(dims) + 1
    if not 0 <= start <= len(dims):
        raise ValueError(f"expand_dims: axis {axis} is out of range for dims {dims}.")
    fast = _fast_path(adapter, "expand_dims")
    if fast is not None:
        return fast(x, names, axis=start)
    data = adapter.data(x)
    xp = namespace_of(data)
    expanded = data
    for offset in range(len(names)):
        expanded = xp.expand_dims(expanded, axis=start + offset)
    new_dims = (*dims[:start], *names, *dims[start:])
    return adapter.rewrap(x, expanded, new_dims)

merge_dims(x, /, **dims)

Merge the listed dims into one, placed last, in C order (xarray's stack).

y_{…, p} = x_{…, i, j}, with p = i·X + j for pixel=("y", "x").

The merged dim's labels are tuples of the source labels, with positions for an unlabelled source; it has none when no source has labels. split_dims undoes it.

Parameters:

Name Type Description Default
x Any

A named carrier.

 required
**dims Sequence[str]

new=(source, …) for each merged dim; several merges go last in the order given.

 {}

Raises:

Type Description
KeyError

A source is not a dim of x.
ValueError

A merge lists no source, a source is merged twice, or the result repeats a dim name.
TypeError

A merge's sources are one string, or the adapter has no relabel and the merged dim has labels.

Examples:

>>> import numpy as np
>>> from pipekit_array import DimArray
>>> bands = {"band": ("B4", "B8")}
>>> cube = DimArray(np.zeros((2, 3, 4)), ("band", "y", "x"), bands)
>>> pixels = merge_dims(cube, pixel=("y", "x"))
>>> pixels.dims, pixels.shape
(('band', 'pixel'), (2, 12))
>>> merge_dims(cube, pb=("band", "y")).labels["pb"][:2]
(('B4', 0), ('B4', 1))
Source code in packages/pipekit-array/src/pipekit_array/dim_shapes.py 
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
def merge_dims(x: Any, /, **dims: Sequence[str]) -> Any:
    """Merge the listed dims into one, placed last, in C order (xarray's ``stack``).

    y_{…, p} = x_{…, i, j}, with p = i·X + j for ``pixel=("y", "x")``.

    The merged dim's labels are tuples of the source labels, with
    positions for an unlabelled source; it has none when no source has
    labels. `split_dims` undoes it.

    Args:
        x: A named carrier.
        **dims: ``new=(source, …)`` for each merged dim; several merges go
            last in the order given.

    Raises:
        KeyError: A source is not a dim of ``x``.
        ValueError: A merge lists no source, a source is merged twice, or
            the result repeats a dim name.
        TypeError: A merge's sources are one string, or the adapter has no
            ``relabel`` and the merged dim has labels.

    Examples:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> bands = {"band": ("B4", "B8")}
        >>> cube = DimArray(np.zeros((2, 3, 4)), ("band", "y", "x"), bands)
        >>> pixels = merge_dims(cube, pixel=("y", "x"))
        >>> pixels.dims, pixels.shape
        (('band', 'pixel'), (2, 12))
        >>> merge_dims(cube, pb=("band", "y")).labels["pb"][:2]
        (('B4', 0), ('B4', 1))
    """
    adapter = adapter_of(x, "merge_dims")
    current = adapter.dims(x)
    for name, group in dims.items():
        if isinstance(group, str):
            raise TypeError(
                f"merge_dims: {name}= takes a sequence of dim names, got {group!r}."
            )
        if len(group) == 0:
            raise ValueError(f"merge_dims: {name}= lists no dim to merge.")
    groups = {name: tuple(group) for name, group in dims.items()}
    sources = [s for group in groups.values() for s in group]
    _axes(current, sources, "merge_dims")
    if len(set(sources)) != len(sources):
        raise ValueError(f"merge_dims: a dim is merged twice in {groups}.")
    kept = tuple(d for d in current if d not in sources)
    new_dims = (*kept, *groups)
    if len(set(new_dims)) != len(new_dims):
        raise ValueError(f"merge_dims: the new dims {new_dims} repeat a name.")
    fast = _fast_path(adapter, "merge_dims")
    if fast is not None:
        return fast(x, groups)
    data = adapter.data(x)
    xp = namespace_of(data)
    sizes = dict(zip(current, data.shape, strict=True))
    moved = xp.permute_dims(data, _axes(current, (*kept, *sources), "merge_dims"))
    # each merged dim has the product of its sources' sizes
    merged_sizes = [math.prod(sizes[s] for s in group) for group in groups.values()]
    merged = xp.reshape(moved, (*(sizes[d] for d in kept), *merged_sizes))
    out = adapter.rewrap(x, merged, new_dims)
    labels = {
        name: merged_labels(
            [adapter.labels(x, s) for s in group], [sizes[s] for s in group]
        )
        for name, group in groups.items()
    }
    return set_labels(adapter, out, labels, "merge_dims")

merged_labels(sources, sizes)

The labels of a merged dim: tuples of the source labels, in C order.

An unlabelled source contributes its positions; when no source has labels, neither has the merged dim.

Example

merged_labels([("B1", "B2"), None], [2, 3])[:4] (('B1', 0), ('B1', 1), ('B1', 2), ('B2', 0)) merged_labels([None, None], [2, 3]) is None True

Source code in packages/pipekit-array/src/pipekit_array/dim_shapes.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def merged_labels(sources: Sequence[Labels], sizes: Sequence[int]) -> Labels:
    """The labels of a merged dim: tuples of the source labels, in C order.

    An unlabelled source contributes its positions; when no source has
    labels, neither has the merged dim.

    Example:
        >>> merged_labels([("B1", "B2"), None], [2, 3])[:4]
        (('B1', 0), ('B1', 1), ('B1', 2), ('B2', 0))
        >>> merged_labels([None, None], [2, 3]) is None
        True
    """
    if all(values is None for values in sources):
        return None
    parts = [
        tuple(range(size)) if values is None else tuple(values)
        for values, size in zip(sources, sizes, strict=True)
    ]
    return tuple(itertools.product(*parts))

rename(x, /, **names)

New dim names; data, order and labels unchanged.

y = x, with each dim d called names[d].

x is positional-only, so a dim may be called x.

Parameters:

Name Type Description Default
x Any

A named carrier.

 required
**names str

old=new for each dim to rename.

 {}

Raises:

Type Description
KeyError

An old name is not a dim of x.
ValueError

The new names repeat a dim.
TypeError

The adapter has neither a rename fast path nor relabel, and a renamed dim has labels.

Examples:

>>> import numpy as np
>>> from pipekit_array import DimArray
>>> scene = DimArray(np.zeros((2, 4)), ("band", "x"), {"band": ("B4", "B8")})
>>> channels = rename(scene, band="channel")
>>> channels.dims, channels.labels["channel"]
(('channel', 'x'), ('B4', 'B8'))
>>> rename(channels, channel="band").labels == scene.labels
True
Source code in packages/pipekit-array/src/pipekit_array/dim_shapes.py 
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
def rename(x: Any, /, **names: str) -> Any:
    """New dim names; data, order and labels unchanged.

    y = x, with each dim ``d`` called ``names[d]``.

    ``x`` is positional-only, so a dim may be called ``x``.

    Args:
        x: A named carrier.
        **names: ``old=new`` for each dim to rename.

    Raises:
        KeyError: An old name is not a dim of ``x``.
        ValueError: The new names repeat a dim.
        TypeError: The adapter has neither a ``rename`` fast path nor
            ``relabel``, and a renamed dim has labels.

    Examples:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> scene = DimArray(np.zeros((2, 4)), ("band", "x"), {"band": ("B4", "B8")})
        >>> channels = rename(scene, band="channel")
        >>> channels.dims, channels.labels["channel"]
        (('channel', 'x'), ('B4', 'B8'))
        >>> rename(channels, channel="band").labels == scene.labels
        True
    """
    adapter = adapter_of(x, "rename")
    dims = adapter.dims(x)
    _axes(dims, list(names), "rename")
    new_dims = tuple(names.get(d, d) for d in dims)
    if len(set(new_dims)) != len(new_dims):
        raise ValueError(f"rename: the new dims {new_dims} repeat a name.")
    fast = _fast_path(adapter, "rename")
    if fast is not None:
        return fast(x, dict(names))
    out = adapter.rewrap(x, adapter.data(x), new_dims)
    labels = {new: adapter.labels(x, old) for old, new in names.items()}
    return set_labels(adapter, out, labels, "rename")

set_labels(adapter, out, labels, what)

out with the labels of the named dims set through relabel.

Called after rewrap with every new or renamed dim, None included: rewrap may have matched a stale label to a reused name.

Raises:

Type Description
TypeError

The adapter has no relabel and a dim's labels must be set or removed.
Source code in packages/pipekit-array/src/pipekit_array/dim_shapes.py 
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
def set_labels(
    adapter: NamedArrayAdapter, out: Any, labels: Mapping[str, Labels], what: str
) -> Any:
    """``out`` with the labels of the named dims set through ``relabel``.

    Called after ``rewrap`` with every new or renamed dim, None included:
    ``rewrap`` may have matched a stale label to a reused name.

    Raises:
        TypeError: The adapter has no ``relabel`` and a dim's labels must
            be set or removed.
    """
    if not labels:
        return out
    labelled = _labelled(adapter)
    if labelled is not None:
        return labelled.relabel(out, dict(labels))
    for dim, values in labels.items():
        if values is not None or adapter.labels(out, dim) is not None:
            raise TypeError(
                f"{what}: {type(adapter).__name__} has no relabel, so the labels "
                f"of dim {dim!r} cannot be set; give the adapter relabel "
                f"(pipekit_ops.LabelledAdapter) or a {what} fast path."
            )
    return out

split_dims(x, /, **dims)

Split one dim into several, in place, in C order (xarray's unstack).

y_{…, i, j} = x_{…, p}, with p = i·X + j for pixel={"y": Y, "x": X}.

The sizes are given, not read from the labels, so they are static under jit (A14). Tuple labels from merge_dims are split back.

Parameters:

Name Type Description Default
x Any

A named carrier.

 required
**dims Mapping[str, int]

dim={part: size, …} for each dim to split; the sizes must multiply to the dim's length. A part may reuse the dim's name (time={"time": 4, "window": 3}).

 {}

Raises:

Type Description
KeyError

A dim to split is not a dim of x.
ValueError

The sizes do not multiply to the dim's length, a split lists no part, or the result repeats a dim name.
TypeError

The adapter has no relabel and a part has labels.

Examples:

>>> import numpy as np
>>> from pipekit_array import DimArray
>>> pixels = DimArray(np.zeros((2, 12)), ("band", "pixel"))
>>> split_dims(pixels, pixel={"y": 3, "x": 4}).dims
('band', 'y', 'x')
>>> cube = DimArray(np.zeros((3, 4)), ("y", "x"), {"x": tuple("abcd")})
>>> back = split_dims(merge_dims(cube, p=("y", "x")), p={"y": 3, "x": 4})
>>> back.labels
{'x': ('a', 'b', 'c', 'd')}
Source code in packages/pipekit-array/src/pipekit_array/dim_shapes.py 
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
def split_dims(x: Any, /, **dims: Mapping[str, int]) -> Any:
    """Split one dim into several, in place, in C order (xarray's ``unstack``).

    y_{…, i, j} = x_{…, p}, with p = i·X + j for ``pixel={"y": Y, "x": X}``.

    The sizes are given, not read from the labels, so they are static
    under ``jit`` (A14). Tuple labels from `merge_dims` are split back.

    Args:
        x: A named carrier.
        **dims: ``dim={part: size, …}`` for each dim to split; the sizes
            must multiply to the dim's length. A part may reuse the
            dim's name (``time={"time": 4, "window": 3}``).

    Raises:
        KeyError: A dim to split is not a dim of ``x``.
        ValueError: The sizes do not multiply to the dim's length, a split
            lists no part, or the result repeats a dim name.
        TypeError: The adapter has no ``relabel`` and a part has labels.

    Examples:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> pixels = DimArray(np.zeros((2, 12)), ("band", "pixel"))
        >>> split_dims(pixels, pixel={"y": 3, "x": 4}).dims
        ('band', 'y', 'x')
        >>> cube = DimArray(np.zeros((3, 4)), ("y", "x"), {"x": tuple("abcd")})
        >>> back = split_dims(merge_dims(cube, p=("y", "x")), p={"y": 3, "x": 4})
        >>> back.labels
        {'x': ('a', 'b', 'c', 'd')}
    """
    adapter = adapter_of(x, "split_dims")
    current = adapter.dims(x)
    _axes(current, list(dims), "split_dims")
    parts = {name: dict(part) for name, part in dims.items()}
    data = adapter.data(x)
    sizes = dict(zip(current, data.shape, strict=True))
    for name, part in parts.items():
        if not part:
            raise ValueError(f"split_dims: {name}= lists no part.")
        product = math.prod(part.values())
        if product != sizes[name]:
            raise ValueError(
                f"split_dims: the sizes {part} multiply to {product}, but "
                f"{name!r} has size {sizes[name]}."
            )
    new_dims = tuple(p for d in current for p in parts.get(d, {d: 0}))
    if len(set(new_dims)) != len(new_dims):
        raise ValueError(f"split_dims: the new dims {new_dims} repeat a name.")
    fast = _fast_path(adapter, "split_dims")
    if fast is not None:
        return fast(x, parts)
    xp = namespace_of(data)
    shape = tuple(
        n for d in current for n in (parts[d].values() if d in parts else (sizes[d],))
    )
    out = adapter.rewrap(x, xp.reshape(data, shape), new_dims)
    labels: dict[str, Labels] = {}
    for name, part in parts.items():
        recovered = split_labels(adapter.labels(x, name), tuple(part.values()))
        labels.update(zip(part, recovered, strict=True))
    return set_labels(adapter, out, labels, "split_dims")

split_labels(labels, sizes)

Each part's labels, recovered from a merged dim's tuple labels.

The inverse of merged_labels. A part whose labels are exactly its positions 0 … n-1 comes back unlabelled, as it was merged. Labels that are not tuples of one entry per part, or not a C-order grid, give no labels.

Example

split_labels((("B1", 0), ("B1", 1), ("B2", 0), ("B2", 1)), [2, 2]) (('B1', 'B2'), None)

Source code in packages/pipekit-array/src/pipekit_array/dim_shapes.py 
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
def split_labels(labels: Labels, sizes: Sequence[int]) -> tuple[Labels, ...]:
    """Each part's labels, recovered from a merged dim's tuple labels.

    The inverse of `merged_labels`. A part whose labels are exactly its
    positions ``0 … n-1`` comes back unlabelled, as it was merged. Labels
    that are not tuples of one entry per part, or not a C-order grid,
    give no labels.

    Example:
        >>> split_labels((("B1", 0), ("B1", 1), ("B2", 0), ("B2", 1)), [2, 2])
        (('B1', 'B2'), None)
    """
    unknown: tuple[Labels, ...] = (None,) * len(sizes)
    if labels is None:
        return unknown
    if not all(isinstance(v, tuple) and len(v) == len(sizes) for v in labels):
        return unknown
    parts = []
    for k, size in enumerate(sizes):
        # entry i of part k sits at flat position i·∏_{j>k} n_j
        stride = math.prod(sizes[k + 1 :])
        parts.append(tuple(labels[i * stride][k] for i in range(size)))
    if tuple(itertools.product(*parts)) != tuple(labels):
        return unknown
    return tuple(None if _are_positions(part) else part for part in parts)

squeeze(x, dim)

Drop named length-1 dims.

y_{…} = x_{…, 0, …}.

dim is required: squeezing every length-1 dim would also drop a batch of one (A12).

Parameters:

Name Type Description Default
x Any

A named carrier.

 required
dim str | Sequence[str]

A dim name, or several; each must have size 1.

 required

Raises:

Type Description
KeyError

A name is not a dim of x.
ValueError

A named dim's size is not 1.
Example

import numpy as np from pipekit_array import DimArray squeeze(DimArray(np.zeros((1, 2, 4)), ("time", "band", "x")), "time").dims ('band', 'x')

Source code in packages/pipekit-array/src/pipekit_array/dim_shapes.py 
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
def squeeze(x: Any, dim: str | Sequence[str], /) -> Any:
    """Drop named length-1 dims.

    y_{…} = x_{…, 0, …}.

    ``dim`` is required: squeezing every length-1 dim would also drop a
    batch of one (A12).

    Args:
        x: A named carrier.
        dim: A dim name, or several; each must have size 1.

    Raises:
        KeyError: A name is not a dim of ``x``.
        ValueError: A named dim's size is not 1.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> squeeze(DimArray(np.zeros((1, 2, 4)), ("time", "band", "x")), "time").dims
        ('band', 'x')
    """
    adapter = adapter_of(x, "squeeze")
    names = (dim,) if isinstance(dim, str) else tuple(dim)
    dims = adapter.dims(x)
    axes = _axes(dims, names, "squeeze")
    shape = adapter.data(x).shape
    wrong = [
        f"{n!r} has size {shape[a]}"
        for n, a in zip(names, axes, strict=True)
        if shape[a] != 1
    ]
    if wrong:
        raise ValueError(f"squeeze: only length-1 dims drop; {', '.join(wrong)}.")
    fast = _fast_path(adapter, "squeeze")
    if fast is not None:
        return fast(x, names)
    return adapter.isel(x, dict.fromkeys(names, 0))

pipekit_array.dim_shifts

Dim-level functions that move data by position: pad, roll, shift, diff, …

Phase 2 of the dim-functions plan (docs/design/api/dim_functions.md, ADRs A11, A12 and A16). Like the functions of pipekit_array.dims, each runs on any carrier with a registered pipekit_ops.NamedArrayAdapter and calls an adapter method of the same name, when there is one, instead of the generic path.

  • pad — widen named dims, with a constant or from the data's edges.
  • roll, shift — move data along named dims, cyclically or with a fill.
  • diff — the n-th forward difference along one dim.
  • cumreduce — a cumulative function along one dim ("cumulative_sum").
  • where — keep entries where a mask is True, aligned by dim name.

Index-only. None reads a label. Labels stay on their positions (xarray's roll_coords=False), and a dim whose size changes loses its labels, which rewrap does by itself.

Dtype preserved. A fill of NaN needs a floating dtype: pad, shift and where raise TypeError on integer or bool data unless a fill value is given, where xarray would promote to float (A12).

The Operator forms live in pipekit_array.dim_operators.

cumreduce(x, fn, dim)

A cumulative function along dim; labels and length unchanged.

c_t = ∑_{τ ≤ t} x_τ for "cumulative_sum".

No implicit skipna (xarray's cumsum skips NaN by default): ask for "nancumsum".

Parameters:

Name Type Description Default
x Any

A named carrier.

 required
fn Reduction

An array-namespace name ("cumulative_sum", "cumulative_prod", "nancumsum") or a function taking axis= that keeps the length.

 required
dim str

The dim to accumulate along.

 required

Raises:

Type Description
KeyError

dim is not a dim of x.
ValueError

The namespace has no fn, or fn changed the length of dim.

Examples:

>>> import numpy as np
>>> from pipekit_array import DimArray
>>> change = DimArray(np.array([1.0, 2.0, 3.0]), ("time",))
>>> cumreduce(change, "cumulative_sum", "time").data
array([1., 3., 6.])
>>> gaps = DimArray(np.array([1.0, np.nan, 3.0]), ("time",))
>>> cumreduce(gaps, "nancumsum", "time").data
array([1., 1., 4.])
Source code in packages/pipekit-array/src/pipekit_array/dim_shifts.py 
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
def cumreduce(x: Any, fn: Reduction, dim: str) -> Any:
    """A cumulative function along ``dim``; labels and length unchanged.

    c_t = ∑_{τ ≤ t} x_τ for ``"cumulative_sum"``.

    No implicit ``skipna`` (xarray's ``cumsum`` skips NaN by default): ask
    for ``"nancumsum"``.

    Args:
        x: A named carrier.
        fn: An array-namespace name (``"cumulative_sum"``,
            ``"cumulative_prod"``, ``"nancumsum"``) or a function taking
            ``axis=`` that keeps the length.
        dim: The dim to accumulate along.

    Raises:
        KeyError: ``dim`` is not a dim of ``x``.
        ValueError: The namespace has no ``fn``, or ``fn`` changed the
            length of ``dim``.

    Examples:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> change = DimArray(np.array([1.0, 2.0, 3.0]), ("time",))
        >>> cumreduce(change, "cumulative_sum", "time").data
        array([1., 3., 6.])
        >>> gaps = DimArray(np.array([1.0, np.nan, 3.0]), ("time",))
        >>> cumreduce(gaps, "nancumsum", "time").data
        array([1., 1., 4.])
    """
    adapter = adapter_of(x, "cumreduce")
    dims = adapter.dims(x)
    (axis,) = _axes(dims, (dim,), "cumreduce")
    fast = _fast_path(adapter, "cumreduce")
    if fast is not None:
        return fast(x, fn, dim)
    data = adapter.data(x)
    func = _resolve(fn, namespace_of(data), "cumreduce")
    # c_t = ∑_{τ ≤ t} x_τ (or fn's own accumulation)
    out = func(data, axis=axis)
    if tuple(out.shape) != tuple(data.shape):
        raise ValueError(
            f"cumreduce: fn changed the shape {tuple(data.shape)} to "
            f"{tuple(out.shape)}; a cumulative function keeps the length of {dim!r}."
        )
    return adapter.rewrap(x, out, dims)

diff(x, dim, /, n=1, *, label='upper')

The n-th forward difference along dim (xarray's diff).

Δy_t = y_{t+1} - y_t, applied n times: N entries become N - n.

The result's labels on dim are the minuend's ("upper", positions n … N - 1) or the subtrahend's ("lower", positions 0 … N - n - 1).

Parameters:

Name Type Description Default
x Any

A named carrier.

 required
dim str

The dim to difference along.

 required
n int

The order, n ≥ 0; static under jit.

 1
label str

"upper" or "lower".

 'upper'

Raises:

Type Description
KeyError

dim is not a dim of x.
ValueError

n is negative, or label is not "upper" or "lower".

Examples:

>>> import numpy as np
>>> from pipekit_array import DimArray
>>> series = DimArray(np.array([1.0, 4.0, 9.0, 16.0]), ("time",))
>>> diff(series, "time").data
array([3., 5., 7.])
>>> labelled = DimArray(series.data, ("time",), {"time": tuple("abcd")})
>>> second = diff(labelled, "time", 2, label="lower")
>>> second.data, second.labels["time"]
(array([2., 2.]), ('a', 'b'))
Source code in packages/pipekit-array/src/pipekit_array/dim_shifts.py 
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
def diff(x: Any, dim: str, /, n: int = 1, *, label: str = "upper") -> Any:
    """The n-th forward difference along ``dim`` (xarray's ``diff``).

    Δy_t = y_{t+1} - y_t, applied n times: N entries become N - n.

    The result's labels on ``dim`` are the minuend's (``"upper"``, positions
    n … N - 1) or the subtrahend's (``"lower"``, positions 0 … N - n - 1).

    Args:
        x: A named carrier.
        dim: The dim to difference along.
        n: The order, n ≥ 0; static under ``jit``.
        label: ``"upper"`` or ``"lower"``.

    Raises:
        KeyError: ``dim`` is not a dim of ``x``.
        ValueError: ``n`` is negative, or ``label`` is not ``"upper"`` or
            ``"lower"``.

    Examples:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> series = DimArray(np.array([1.0, 4.0, 9.0, 16.0]), ("time",))
        >>> diff(series, "time").data
        array([3., 5., 7.])
        >>> labelled = DimArray(series.data, ("time",), {"time": tuple("abcd")})
        >>> second = diff(labelled, "time", 2, label="lower")
        >>> second.data, second.labels["time"]
        (array([2., 2.]), ('a', 'b'))
    """
    adapter = adapter_of(x, "diff")
    dims = adapter.dims(x)
    _axes(dims, (dim,), "diff")
    if n < 0:
        raise ValueError(f"diff: the order n must be ≥ 0, got {n}.")
    if label not in ("upper", "lower"):
        raise ValueError(f"diff: label must be 'upper' or 'lower', got {label!r}.")
    fast = _fast_path(adapter, "diff")
    if fast is not None:
        return fast(x, dim, n=n, label=label)
    out = x
    for _ in range(n):
        upper = adapter.isel(out, {dim: slice(1, None)})
        lower = adapter.isel(out, {dim: slice(None, -1)})
        xp = namespace_of(adapter.data(upper))
        # Δy_t = y_{t+1} - y_t
        difference = xp.subtract(adapter.data(upper), adapter.data(lower))
        # the labelled slice carries its labels onto the difference
        labelled = upper if label == "upper" else lower
        out = adapter.rewrap(labelled, difference, dims)
    return out

pad(x, /, *, mode='constant', constant_values=None, **pad_width)

Widen named dims at either end (xarray's pad, four modes).

y_{…, t, …} = x_{…, t - b, …} for b ≤ t < b + N; the new entries follow mode.

A padded dim loses its labels (xarray pads its coordinate with NaN).

Parameters:

Name Type Description Default
x Any

A named carrier.

 required
mode str

"constant" (constant_values), "edge" (repeat the end entry), "reflect" (mirror, end excluded) or "wrap" (the other end's entries).

 'constant'
constant_values float | None

The fill for mode="constant"; None is NaN.

 None
**pad_width int | tuple[int, int]

dim= a width b for both ends, or a pair (before, after).

 {}

Raises:

Type Description
KeyError

A padded dim is not a dim of x.
ValueError

mode is not one of the four, a width is negative, or constant_values is given with another mode.
TypeError

mode="constant" on integer or bool data without constant_values.

Examples:

>>> import numpy as np
>>> from pipekit_array import DimArray
>>> series = DimArray(np.arange(3.0), ("time",))
>>> pad(series, time=(2, 0)).data
array([nan, nan,  0.,  1.,  2.])
>>> pad(series, mode="wrap", time=1).data
array([2., 0., 1., 2., 0.])
>>> counts = DimArray(np.arange(3), ("time",))
>>> pad(counts, time=1)
Traceback (most recent call last):
...
TypeError: pad: int64 data has no NaN to fill with; pass constant_values=
Source code in packages/pipekit-array/src/pipekit_array/dim_shifts.py 
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
def pad(
    x: Any,
    /,
    *,
    mode: str = "constant",
    constant_values: float | None = None,
    **pad_width: int | tuple[int, int],
) -> Any:
    """Widen named dims at either end (xarray's ``pad``, four modes).

    y_{…, t, …} = x_{…, t - b, …} for b ≤ t < b + N; the new entries follow
    ``mode``.

    A padded dim loses its labels (xarray pads its coordinate with NaN).

    Args:
        x: A named carrier.
        mode: ``"constant"`` (``constant_values``), ``"edge"`` (repeat the
            end entry), ``"reflect"`` (mirror, end excluded) or ``"wrap"``
            (the other end's entries).
        constant_values: The fill for ``mode="constant"``; None is NaN.
        **pad_width: ``dim=`` a width b for both ends, or a pair
            ``(before, after)``.

    Raises:
        KeyError: A padded dim is not a dim of ``x``.
        ValueError: ``mode`` is not one of the four, a width is negative,
            or ``constant_values`` is given with another mode.
        TypeError: ``mode="constant"`` on integer or bool data without
            ``constant_values``.

    Examples:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> series = DimArray(np.arange(3.0), ("time",))
        >>> pad(series, time=(2, 0)).data
        array([nan, nan,  0.,  1.,  2.])
        >>> pad(series, mode="wrap", time=1).data
        array([2., 0., 1., 2., 0.])
        >>> counts = DimArray(np.arange(3), ("time",))
        >>> pad(counts, time=1)  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        TypeError: pad: int64 data has no NaN to fill with; pass constant_values=
    """
    adapter = adapter_of(x, "pad")
    dims = adapter.dims(x)
    _axes(dims, list(pad_width), "pad")
    if mode not in PAD_MODES:
        raise ValueError(f"pad: mode {mode!r} is not one of {PAD_MODES}.")
    if mode != "constant" and constant_values is not None:
        raise ValueError(f"pad: constant_values is for mode='constant', not {mode!r}.")
    widths = {d: _widths(d, w) for d, w in pad_width.items()}
    data = adapter.data(x)
    xp = namespace_of(data)
    if (
        mode == "constant"
        and constant_values is None
        and not _holds_nan(xp, data.dtype)
    ):
        raise _no_nan_fill("pad", "constant_values", data.dtype)
    fast = _fast_path(adapter, "pad")
    if fast is not None:
        return fast(x, widths, mode=mode, constant_values=constant_values)
    full_widths = tuple(widths.get(d, (0, 0)) for d in dims)
    if mode == "constant":
        fill = float("nan") if constant_values is None else constant_values
        padded = xp.pad(data, full_widths, mode=mode, constant_values=fill)
    else:
        padded = xp.pad(data, full_widths, mode=mode)
    return adapter.rewrap(x, padded, dims)

roll(x, /, **shifts)

Shift data cyclically along named dims; labels stay (xarray's roll).

y_{…, t, …} = x_{…, (t - k) mod N, …}.

Parameters:

Name Type Description Default
x Any

A named carrier.

 required
**shifts int

dim=k; a positive k moves data to higher positions. k may be traced under jit.

 {}

Raises:

Type Description
KeyError

A rolled dim is not a dim of x.

Examples:

>>> import numpy as np
>>> from pipekit_array import DimArray
>>> series = DimArray(np.arange(4.0), ("time",), {"time": tuple("abcd")})
>>> rolled = roll(series, time=1)
>>> rolled.data, rolled.labels["time"]
(array([3., 0., 1., 2.]), ('a', 'b', 'c', 'd'))
>>> bool(np.array_equal(roll(rolled, time=-1).data, series.data))
True
Source code in packages/pipekit-array/src/pipekit_array/dim_shifts.py 
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
def roll(x: Any, /, **shifts: int) -> Any:
    """Shift data cyclically along named dims; labels stay (xarray's ``roll``).

    y_{…, t, …} = x_{…, (t - k) mod N, …}.

    Args:
        x: A named carrier.
        **shifts: ``dim=k``; a positive k moves data to higher positions.
            k may be traced under ``jit``.

    Raises:
        KeyError: A rolled dim is not a dim of ``x``.

    Examples:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> series = DimArray(np.arange(4.0), ("time",), {"time": tuple("abcd")})
        >>> rolled = roll(series, time=1)
        >>> rolled.data, rolled.labels["time"]
        (array([3., 0., 1., 2.]), ('a', 'b', 'c', 'd'))
        >>> bool(np.array_equal(roll(rolled, time=-1).data, series.data))
        True
    """
    adapter = adapter_of(x, "roll")
    dims = adapter.dims(x)
    axes = _axes(dims, list(shifts), "roll")
    fast = _fast_path(adapter, "roll")
    if fast is not None:
        return fast(x, dict(shifts))
    data = adapter.data(x)
    xp = namespace_of(data)
    # y_t = x_{(t - k) mod N}
    rolled = xp.roll(data, shift=tuple(shifts.values()), axis=axes)
    return adapter.rewrap(x, rolled, dims)

shift(x, /, *, fill_value=None, **shifts)

Shift data along named dims, filling the vacated entries (xarray's shift).

y_{…, t, …} = x_{…, t - k, …} if 0 ≤ t - k < N, else the fill.

Labels stay on their positions. A shift of at least the dim's length fills it entirely.

Parameters:

Name Type Description Default
x Any

A named carrier.

 required
fill_value float | None

The value of the vacated entries; None is NaN.

 None
**shifts int

dim=k, a Python int (static under jit); a positive k moves data to higher positions.

 {}

Raises:

Type Description
KeyError

A shifted dim is not a dim of x.
TypeError

Integer or bool data without fill_value.

Examples:

>>> import numpy as np
>>> from pipekit_array import DimArray
>>> series = DimArray(np.arange(4.0), ("time",))
>>> shift(series, time=1).data
array([nan,  0.,  1.,  2.])
>>> shift(DimArray(np.arange(4), ("time",)), time=-2, fill_value=-1).data
array([ 2,  3, -1, -1])
Source code in packages/pipekit-array/src/pipekit_array/dim_shifts.py 
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
def shift(x: Any, /, *, fill_value: float | None = None, **shifts: int) -> Any:
    """Shift data along named dims, filling the vacated entries (xarray's ``shift``).

    y_{…, t, …} = x_{…, t - k, …} if 0 ≤ t - k < N, else the fill.

    Labels stay on their positions. A shift of at least the dim's length
    fills it entirely.

    Args:
        x: A named carrier.
        fill_value: The value of the vacated entries; None is NaN.
        **shifts: ``dim=k``, a Python int (static under ``jit``); a
            positive k moves data to higher positions.

    Raises:
        KeyError: A shifted dim is not a dim of ``x``.
        TypeError: Integer or bool data without ``fill_value``.

    Examples:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> series = DimArray(np.arange(4.0), ("time",))
        >>> shift(series, time=1).data
        array([nan,  0.,  1.,  2.])
        >>> shift(DimArray(np.arange(4), ("time",)), time=-2, fill_value=-1).data
        array([ 2,  3, -1, -1])
    """
    adapter = adapter_of(x, "shift")
    dims = adapter.dims(x)
    axes = _axes(dims, list(shifts), "shift")
    data = adapter.data(x)
    xp = namespace_of(data)
    if fill_value is None and not _holds_nan(xp, data.dtype):
        raise _no_nan_fill("shift", "fill_value", data.dtype)
    fast = _fast_path(adapter, "shift")
    if fast is not None:
        return fast(x, dict(shifts), fill_value=fill_value)
    fill = float("nan") if fill_value is None else fill_value
    shifted = data
    for axis, offset in zip(axes, shifts.values(), strict=True):
        size = shifted.shape[axis]
        vacated = min(abs(offset), size)
        if vacated == 0:
            continue
        # the entries that stay: the first N - |k| for k > 0, the last for k < 0
        kept = slice(0, size - vacated) if offset > 0 else slice(vacated, size)
        index = tuple(kept if a == axis else slice(None) for a in range(data.ndim))
        filled_shape = tuple(
            vacated if a == axis else n for a, n in enumerate(shifted.shape)
        )
        filled = xp.full(filled_shape, fill, dtype=data.dtype)
        parts = [filled, shifted[index]] if offset > 0 else [shifted[index], filled]
        shifted = xp.concat(parts, axis=axis)
    return adapter.rewrap(x, shifted, dims)

where(x, cond, /, other=None)

Keep x where cond is True, other elsewhere (xarray's where).

y = x where v, else the fill; True keeps (A16).

cond aligns by dim name, as reduce(where=) does: a named mask whose dims are among x's, so a time y x cloud mask applies to a time band y x cube; or a plain array covering x's leading dims. It is never right-broadcast.

Parameters:

Name Type Description Default
x Any

A named carrier.

 required
cond Any

A bool mask v, named or plain.

 required
other float | None

The fill where cond is False; None is NaN.

 None

Raises:

Type Description
ValueError

A named cond has a dim x lacks or a size that differs, or a plain one does not match x's leading dims.
TypeError

Integer or bool data without other.

Examples:

>>> import numpy as np
>>> from pipekit_array import DimArray
>>> cube = DimArray(np.ones((2, 3, 2)), ("time", "band", "x"))
>>> valid = DimArray(np.array([[True, False], [True, True]]), ("time", "x"))
>>> clear = where(cube, valid)
>>> clear.dims, clear.data[0, :, 1]
(('time', 'band', 'x'), array([nan, nan, nan]))
>>> where(cube, valid, 0.0).data[0, 0]
array([1., 0.])
Source code in packages/pipekit-array/src/pipekit_array/dim_shifts.py 
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
def where(x: Any, cond: Any, /, other: float | None = None) -> Any:
    """Keep ``x`` where ``cond`` is True, ``other`` elsewhere (xarray's ``where``).

    y = x where v, else the fill; True keeps (A16).

    ``cond`` aligns by dim name, as ``reduce(where=)`` does: a named mask
    whose dims are among ``x``'s, so a ``time y x`` cloud mask applies to a
    ``time band y x`` cube; or a plain array covering ``x``'s leading dims.
    It is never right-broadcast.

    Args:
        x: A named carrier.
        cond: A bool mask v, named or plain.
        other: The fill where ``cond`` is False; None is NaN.

    Raises:
        ValueError: A named ``cond`` has a dim ``x`` lacks or a size that
            differs, or a plain one does not match ``x``'s leading dims.
        TypeError: Integer or bool data without ``other``.

    Examples:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> cube = DimArray(np.ones((2, 3, 2)), ("time", "band", "x"))
        >>> valid = DimArray(np.array([[True, False], [True, True]]), ("time", "x"))
        >>> clear = where(cube, valid)
        >>> clear.dims, clear.data[0, :, 1]
        (('time', 'band', 'x'), array([nan, nan, nan]))
        >>> where(cube, valid, 0.0).data[0, 0]
        array([1., 0.])
    """
    adapter = adapter_of(x, "where")
    dims = adapter.dims(x)
    data = adapter.data(x)
    mask = _aligned_mask(cond, dims, data, "where")
    xp = namespace_of(data)
    if other is None and not _holds_nan(xp, data.dtype):
        raise _no_nan_fill("where", "other", data.dtype)
    fast = _fast_path(adapter, "where")
    if fast is not None:
        return fast(x, mask, other)
    fill = xp.asarray(float("nan") if other is None else other, dtype=data.dtype)
    # y = x if v else the fill
    kept = xp.where(mask, data, fill)
    return adapter.rewrap(x, kept, dims)

pipekit_array.dim_windows

Dim-level functions that cut a dim into windows: sliding and non-overlapping.

Phase 3 of the dim-functions plan (docs/design/api/dim_functions.md, ADR A15). Like the functions of pipekit_array.dims, each runs on any carrier with a registered pipekit_ops.NamedArrayAdapter, works on positions only (labels follow the data, their values are never read), and calls an adapter method of the same name, when there is one, instead of the generic path.

  • window — sliding windows along one dim, as a new last dim; full windows only, each labelled by its last entry.
  • coarsen — non-overlapping blocks reduced by a function: split_dims then reduce; the coarsened dim loses its labels.

The Operator forms live in pipekit_array.dim_operators; the operators that apply another operator along a dim (ApplyAlong) or in chunks of a dim (pipekit_array.batching.BatchAlong) live next to MapOver and batch.

coarsen(x, fn, /, *, boundary='exact', **windows)

Reduce non-overlapping blocks of k entries along each named dim.

z_{…, i', …} = fn_{0 ≤ r < k} x_{…, k·i' + r, …}, i' < ⌊N/k⌋.

A composition: split_dims each dim into (dim, block) then reduce over the blocks, so coarsen(x, "mean", time=3) is reduce(split_dims(x, time={"time": 4, "window": 3}), "mean", "window"). A coarsened dim loses its labels: xarray's coord_func reads coordinate values, which this layer never does (A15).

Parameters:

Name Type Description Default
x Any

A named carrier.

 required
fn Reduction

An array-namespace name ("mean", "nanmean", "max") or a function taking axis=.

 required
boundary str

"exact" requires k to divide N; "trim" drops the last N mod k entries.

 'exact'
**windows int

dim=k for each dim to coarsen, with k ≥ 1.

 {}

Raises:

Type Description
KeyError

A dim is not a dim of x.
ValueError

No dim is given, k is not positive, boundary is not one of BOUNDARIES, or under "exact" k does not divide N.

Examples:

>>> import numpy as np
>>> from pipekit_array import DimArray
>>> ndvi = DimArray(np.arange(24.0).reshape(2, 12), ("y", "x"))
>>> coarse = coarsen(ndvi, "mean", x=4)
>>> coarse.dims, coarse.data.tolist()
(('y', 'x'), [[1.5, 5.5, 9.5], [13.5, 17.5, 21.5]])
>>> coarsen(ndvi, "max", boundary="trim", x=5, y=2).data.tolist()
[[16.0, 21.0]]
>>> coarsen(ndvi, "mean", x=5)
Traceback (most recent call last):
...
ValueError: coarsen: 5 does not divide 'x' of size 12; ...
Source code in packages/pipekit-array/src/pipekit_array/dim_windows.py 
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
def coarsen(
    x: Any, fn: Reduction, /, *, boundary: str = "exact", **windows: int
) -> Any:
    """Reduce non-overlapping blocks of k entries along each named dim.

    z_{…, i', …} = fn_{0 ≤ r < k} x_{…, k·i' + r, …},  i' < ⌊N/k⌋.

    A composition: `split_dims` each dim into ``(dim, block)`` then
    `reduce` over the blocks, so ``coarsen(x, "mean", time=3)`` is
    ``reduce(split_dims(x, time={"time": 4, "window": 3}), "mean",
    "window")``. A coarsened dim loses its labels: xarray's ``coord_func``
    reads coordinate values, which this layer never does (A15).

    Args:
        x: A named carrier.
        fn: An array-namespace name (``"mean"``, ``"nanmean"``, ``"max"``)
            or a function taking ``axis=``.
        boundary: ``"exact"`` requires k to divide N; ``"trim"`` drops
            the last N mod k entries.
        **windows: ``dim=k`` for each dim to coarsen, with k ≥ 1.

    Raises:
        KeyError: A dim is not a dim of ``x``.
        ValueError: No dim is given, k is not positive, ``boundary`` is
            not one of `BOUNDARIES`, or under ``"exact"`` k does not
            divide N.

    Examples:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> ndvi = DimArray(np.arange(24.0).reshape(2, 12), ("y", "x"))
        >>> coarse = coarsen(ndvi, "mean", x=4)
        >>> coarse.dims, coarse.data.tolist()
        (('y', 'x'), [[1.5, 5.5, 9.5], [13.5, 17.5, 21.5]])
        >>> coarsen(ndvi, "max", boundary="trim", x=5, y=2).data.tolist()
        [[16.0, 21.0]]
        >>> coarsen(ndvi, "mean", x=5)  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        ValueError: coarsen: 5 does not divide 'x' of size 12; ...
    """
    adapter = adapter_of(x, "coarsen")
    if not windows:
        raise ValueError("coarsen: give at least one dim=k.")
    if boundary not in BOUNDARIES:
        raise ValueError(
            f"coarsen: boundary must be one of {BOUNDARIES}, got {boundary!r}."
        )
    dims = adapter.dims(x)
    axes = _axes(dims, list(windows), "coarsen")
    data = adapter.data(x)
    lengths = {dim: data.shape[axis] for dim, axis in zip(windows, axes, strict=True)}
    for dim, block in windows.items():
        if block < 1:
            raise ValueError(f"coarsen: {dim}={block} must be positive.")
        if boundary == "exact" and lengths[dim] % block != 0:
            raise ValueError(
                f"coarsen: {block} does not divide {dim!r} of size {lengths[dim]}; "
                "pass boundary='trim' or pad first."
            )
    fast = _fast_path(adapter, "coarsen")
    if fast is not None:
        return fast(x, fn, dict(windows), boundary=boundary)
    # trim: keep the first k·⌊N/k⌋ entries
    tails = {
        dim: slice(0, lengths[dim] - lengths[dim] % block)
        for dim, block in windows.items()
        if lengths[dim] % block != 0
    }
    trimmed = adapter.isel(x, tails) if tails else x
    blocks: dict[str, str] = {}
    for dim in windows:
        blocks[dim] = _block_dim(dim, (*dims, *blocks.values()))
    # x_{…, i', r, …} = x_{…, k·i' + r, …}
    parts = {
        dim: {dim: lengths[dim] // block, blocks[dim]: block}
        for dim, block in windows.items()
    }
    split = split_dims(trimmed, **parts)
    # z_{…, i', …} = fn over r of x_{…, i', r, …}
    return _dims.reduce(split, fn, tuple(blocks.values()))

window(x, /, *, window_dim='window', stride=1, **sizes)

Sliding windows along one dim; the window becomes a new last dim.

w_{…, n, …, r} = x_{…, n·s + r, …}, 0 ≤ r < S, 0 ≤ n < ⌊(N - S)/s⌋ + 1.

Full windows only: there is no NaN fill, so float32 stays float32 and integers work (A15). Window n carries the label of its last entry, position n·s + S - 1, as xarray's rolling(...).construct(window_dim).isel(dim=slice(S - 1, None, s)) does. For xarray's length-N result, pad the dim by S - 1 first.

Algorithm
  1. Count the full windows, n = ⌊(N - S)/s⌋ + 1.
  2. Take the S strided slices of the data, x_{r::s}, each n long.
  3. Stack them on a new last axis and rewrap on the slice of end positions, so the windowed dim keeps the labels of those.

The generic path copies the data S times; the xarray adapter's fast path is a view.

Parameters:

Name Type Description Default
x Any

A named carrier.

 required
window_dim str

The new last dim, indexing the entries r of a window.

 'window'
stride int

The step s ≥ 1 between window starts.

 1
**sizes int

dim=S for the one dim to window, with 1 ≤ S ≤ N.

 {}

Raises:

Type Description
KeyError

The dim is not a dim of x.
ValueError

sizes names no dim or several, S or s is not positive, S exceeds N, or window_dim is already a dim.

Examples:

>>> import numpy as np
>>> from pipekit_array import DimArray
>>> ndvi = DimArray(np.arange(12.0), ("time",), {"time": tuple(range(12))})
>>> windows = window(ndvi, time=3, stride=2)
>>> windows.dims, windows.shape
(('time', 'window'), (5, 3))
>>> windows.labels["time"]
(2, 4, 6, 8, 10)
>>> window(DimArray(np.arange(4), ("time",)), time=2).data.tolist()
[[0, 1], [1, 2], [2, 3]]
Source code in packages/pipekit-array/src/pipekit_array/dim_windows.py 
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
def window(
    x: Any, /, *, window_dim: str = "window", stride: int = 1, **sizes: int
) -> Any:
    """Sliding windows along one dim; the window becomes a new last dim.

    w_{…, n, …, r} = x_{…, n·s + r, …},  0 ≤ r < S,  0 ≤ n < ⌊(N - S)/s⌋ + 1.

    Full windows only: there is no NaN fill, so float32 stays float32 and
    integers work (A15). Window ``n`` carries the label of its last entry,
    position n·s + S - 1, as xarray's
    ``rolling(...).construct(window_dim).isel(dim=slice(S - 1, None, s))``
    does. For xarray's length-N result, `pad` the dim by S - 1 first.

    Algorithm:
        1. Count the full windows, n = ⌊(N - S)/s⌋ + 1.
        2. Take the S strided slices of the data, x_{r::s}, each n long.
        3. Stack them on a new last axis and rewrap on the slice of end
           positions, so the windowed dim keeps the labels of those.

    The generic path copies the data S times; the xarray adapter's fast
    path is a view.

    Args:
        x: A named carrier.
        window_dim: The new last dim, indexing the entries r of a window.
        stride: The step s ≥ 1 between window starts.
        **sizes: ``dim=S`` for the one dim to window, with 1 ≤ S ≤ N.

    Raises:
        KeyError: The dim is not a dim of ``x``.
        ValueError: ``sizes`` names no dim or several, S or s is not
            positive, S exceeds N, or ``window_dim`` is already a dim.

    Examples:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> ndvi = DimArray(np.arange(12.0), ("time",), {"time": tuple(range(12))})
        >>> windows = window(ndvi, time=3, stride=2)
        >>> windows.dims, windows.shape
        (('time', 'window'), (5, 3))
        >>> windows.labels["time"]
        (2, 4, 6, 8, 10)
        >>> window(DimArray(np.arange(4), ("time",)), time=2).data.tolist()
        [[0, 1], [1, 2], [2, 3]]
    """
    adapter = adapter_of(x, "window")
    if len(sizes) != 1:
        raise ValueError(f"window: give one dim=size, got {sizes}.")
    ((dim, size),) = sizes.items()
    dims = adapter.dims(x)
    (axis,) = _axes(dims, (dim,), "window")
    if window_dim in dims:
        raise ValueError(f"window: {window_dim!r} is already a dim of {dims}.")
    if size < 1 or stride < 1:
        raise ValueError(
            f"window: size and stride must be positive, got {dim}={size}, "
            f"stride={stride}."
        )
    data = adapter.data(x)
    length = data.shape[axis]
    if size > length:
        raise ValueError(
            f"window: a window of {size} does not fit {dim!r} of size {length}."
        )
    fast = _fast_path(adapter, "window")
    if fast is not None:
        return fast(x, dim, size, window_dim=window_dim, stride=stride)
    xp = namespace_of(data)
    # n = ⌊(N - S)/s⌋ + 1
    count = (length - size) // stride + 1
    # the last start is (n - 1)·s
    last_start = (count - 1) * stride
    leading = (slice(None),) * axis
    # x_{r::s}, n entries each, for 0 ≤ r < S
    strided = [
        data[(*leading, slice(r, r + last_start + 1, stride))] for r in range(size)
    ]
    # w_{…, n, …, r} = x_{…, n·s + r, …}
    stacked = xp.stack(strided, axis=-1)
    # window n carries the label of position n·s + S - 1
    ends = adapter.isel(x, {dim: slice(size - 1, size + last_start, stride)})
    out = adapter.rewrap(ends, stacked, (*dims, window_dim))
    return set_labels(adapter, out, {window_dim: None}, "window")

pipekit_array.dim_algebra

Dim-level statistics and algebra: quantiles over dims, and contraction.

Phase 4 of the dim-functions plan (docs/design/api/dim_functions.md, ADRs A17, A18 and A22). Like the functions of pipekit_array.dims, each runs on any carrier with a registered pipekit_ops.NamedArrayAdapter, works on positions only, and calls an adapter method of the same name, when there is one, instead of the generic path.

  • quantile — quantiles over named dims; a sequence q adds a leading quantile dim labelled with q.
  • dot — the sum of a product over named dims, as xr.dot; a shared dim not summed is a batch dim and appears once.

Both set the labels of their new dims through the adapter's optional relabel (A22), as pipekit_array.dim_shapes does.

The Operator forms live in pipekit_array.dim_operators.

check_method(method, what)

Raise ValueError unless method is one of METHODS.

Source code in packages/pipekit-array/src/pipekit_array/dim_algebra.py 
53
54
55
56
57
58
59
def check_method(method: str, what: str) -> None:
    """Raise `ValueError` unless ``method`` is one of `METHODS`."""
    if method not in METHODS:
        raise ValueError(
            f"{what}: method {method!r} is not supported; use one of {METHODS} "
            "(the six JAX computes)."
        )

dot(*xs, dim)

The sum of the product of carriers over named dims (xr.dot).

s_{o…} = ∑{d ∈ dim} ∏ₖ x⁽ᵏ⁾, where o… are the dims not in dim.

The output dims are every input dim not in dim, in order of first appearance. A dim shared by several inputs and not in dim is a batch dim: it appears once, so dot is one named-index einsum, never tensordot (A17). Shared dims must agree exactly in size and labels; dims from inputs after the first carry those inputs' labels.

Algorithm
  1. Give each distinct dim a letter, in order of first appearance.
  2. Write one subscript per input from its dims, and the output's from the dims not in dim.
  3. xp.einsum; for a namespace without it, broadcast every input to the union of the dims, multiply, and xp.sum over dim (exact, but holds the full product).
  4. rewrap on the first input, then relabel the dims taken from later inputs.

Parameters:

Name Type Description Default
*xs Any

Named carriers of one family and one backend.

 ()
dim str | Sequence[str]

A dim name, or several, to sum over; required (A17): xarray would sum every shared dim, so a dim shared by accident would vanish.

 required

Raises:

Type Description
ValueError

No carrier is given, a shared dim differs in size or labels, or there are more dims than letters.
KeyError

A name in dim is a dim of no input.
BackendError

The inputs sit on more than one backend.
TypeError

The adapter has no relabel and a dim from a later input has labels.

Examples:

>>> import numpy as np
>>> from pipekit_array import DimArray
>>> x = DimArray(np.ones((12, 13)), ("time", "band"))
>>> w = DimArray(np.ones((12, 13, 3)), ("time", "band", "component"))
>>> s = dot(x, w, dim="band")
>>> s.dims, bool(np.allclose(s.data, 13.0))
(('time', 'component'), True)
>>> pixels = DimArray(np.ones((2, 13, 6)), ("time", "band", "pixel"))
>>> labelled = {"component": ("pc1", "pc2", "pc3")}
>>> weights = DimArray(np.ones((13, 3)), ("band", "component"), labelled)
>>> scores = dot(pixels, weights, dim="band")
>>> scores.dims, scores.labels["component"]
(('time', 'pixel', 'component'), ('pc1', 'pc2', 'pc3'))
Source code in packages/pipekit-array/src/pipekit_array/dim_algebra.py 
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
def dot(*xs: Any, dim: str | Sequence[str]) -> Any:
    """The sum of the product of carriers over named dims (``xr.dot``).

    s_{o…} = ∑_{d ∈ dim} ∏ₖ x⁽ᵏ⁾_{…}, where o… are the dims not in ``dim``.

    The output dims are every input dim not in ``dim``, in order of first
    appearance. A dim shared by several inputs and not in ``dim`` is a
    batch dim: it appears once, so ``dot`` is one named-index ``einsum``,
    never ``tensordot`` (A17). Shared dims must agree exactly in size and
    labels; dims from inputs after the first carry those inputs' labels.

    Algorithm:
        1. Give each distinct dim a letter, in order of first appearance.
        2. Write one subscript per input from its dims, and the output's
           from the dims not in ``dim``.
        3. ``xp.einsum``; for a namespace without it, broadcast every
           input to the union of the dims, multiply, and ``xp.sum`` over
           ``dim`` (exact, but holds the full product).
        4. ``rewrap`` on the first input, then ``relabel`` the dims taken
           from later inputs.

    Args:
        *xs: Named carriers of one family and one backend.
        dim: A dim name, or several, to sum over; required (A17): xarray
            would sum every shared dim, so a dim shared by accident would
            vanish.

    Raises:
        ValueError: No carrier is given, a shared dim differs in size or
            labels, or there are more dims than letters.
        KeyError: A name in ``dim`` is a dim of no input.
        BackendError: The inputs sit on more than one backend.
        TypeError: The adapter has no ``relabel`` and a dim from a later
            input has labels.

    Examples:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> x = DimArray(np.ones((12, 13)), ("time", "band"))
        >>> w = DimArray(np.ones((12, 13, 3)), ("time", "band", "component"))
        >>> s = dot(x, w, dim="band")
        >>> s.dims, bool(np.allclose(s.data, 13.0))
        (('time', 'component'), True)
        >>> pixels = DimArray(np.ones((2, 13, 6)), ("time", "band", "pixel"))
        >>> labelled = {"component": ("pc1", "pc2", "pc3")}
        >>> weights = DimArray(np.ones((13, 3)), ("band", "component"), labelled)
        >>> scores = dot(pixels, weights, dim="band")
        >>> scores.dims, scores.labels["component"]
        (('time', 'pixel', 'component'), ('pc1', 'pc2', 'pc3'))
    """
    if not xs:
        raise ValueError("dot needs at least one carrier.")
    adapters = [adapter_of(x, "dot") for x in xs]
    adapter = adapters[0]
    names = (dim,) if isinstance(dim, str) else tuple(dim)
    each_dims = [a.dims(x) for a, x in zip(adapters, xs, strict=True)]
    union = tuple(dict.fromkeys(d for dims in each_dims for d in dims))
    _axes(union, names, "dot")
    if len(union) > len(string.ascii_letters):
        raise ValueError(
            f"dot: {len(union)} distinct dims, but einsum has only "
            f"{len(string.ascii_letters)} letters."
        )
    datas = [a.data(x) for a, x in zip(adapters, xs, strict=True)]
    require_one_backend(datas, "dot")
    fast = _fast_path(adapter, "dot")
    if fast is not None:
        return fast(list(xs), names)
    # only the shared dims must combine: an empty slice of every other dim
    # checks them without touching the data
    counts = Counter(d for dims in each_dims for d in dims)
    parts = [
        a.isel(x, {d: slice(0, 0) for d in dims if counts[d] == 1})
        for a, x, dims in zip(adapters, xs, each_dims, strict=True)
    ]
    adapter.check_compatible(*parts)
    out_dims = tuple(d for d in union if d not in names)
    xp = namespace_of(datas[0])
    einsum = getattr(xp, "einsum", None)
    if einsum is None:
        out = _product_sum(xp, datas, each_dims, union, names)
    else:
        letters = dict(zip(union, string.ascii_letters, strict=False))
        inputs = ",".join("".join(letters[d] for d in dims) for dims in each_dims)
        output = "".join(letters[d] for d in out_dims)
        # s_{o…} = ∑_{d ∈ dim} ∏ₖ x⁽ᵏ⁾, e.g. "ab,abc->ac" for time band, time
        # band component over band
        out = einsum(f"{inputs}->{output}", *datas)
    result = adapter.rewrap(xs[0], out, out_dims)
    # dims the first input lacks take the labels of the first input with them
    labels: dict[str, Labels] = {}
    for a, x, dims in zip(adapters[1:], xs[1:], each_dims[1:], strict=True):
        for d in dims:
            if d in out_dims and d not in each_dims[0] and d not in labels:
                labels[d] = a.labels(x, d)
    return set_labels(adapter, result, labels, "dot")

probabilities(q, what)

q as one float, or a tuple of floats; each must lie in [0, 1].

Raises:

Type Description
ValueError

A probability is outside [0, 1].
Source code in packages/pipekit-array/src/pipekit_array/dim_algebra.py 
62
63
64
65
66
67
68
69
70
71
72
def probabilities(q: float | Sequence[float], what: str) -> float | tuple[float, ...]:
    """``q`` as one float, or a tuple of floats; each must lie in [0, 1].

    Raises:
        ValueError: A probability is outside [0, 1].
    """
    values = float(q) if isinstance(q, (int, float)) else tuple(float(v) for v in q)
    outside = [v for v in _as_tuple(values) if not 0.0 <= v <= 1.0]
    if outside:
        raise ValueError(f"{what}: q must lie in [0, 1], got {outside}.")
    return values

quantile(x, q, dim, /, *, method='linear', skipna=False)

Quantiles over named dims (xarray's quantile).

Q_{k…} = quantile q_k of {x_{…, d…} : d over dim}.

A sequence q adds a leading dim quantile, labelled with q; a float drops it. A float input keeps its dtype. NaN propagates unless skipna (no implicit skipna, A12).

Parameters:

Name Type Description Default
x Any

A named carrier.

 required
q float | Sequence[float]

A probability in [0, 1], or a sequence of them; static under jit.

 required
dim str | Sequence[str]

A dim name, or several, to reduce over.

 required
method str

One of METHODS, the six JAX supports.

 'linear'
skipna bool

Ignore NaN (nanquantile).

 False

Raises:

Type Description
KeyError

A name is not a dim of x.
ValueError

method is not one of METHODS, a probability is outside [0, 1], or a sequence q meets a kept dim named quantile.
TypeError

The adapter has no relabel and q is a sequence.

Examples:

>>> import numpy as np
>>> from pipekit_array import DimArray
>>> cube = DimArray(np.arange(24.0).reshape(4, 2, 3), ("time", "band", "x"))
>>> spread = quantile(cube, [0.1, 0.9], "time")
>>> spread.dims, spread.labels["quantile"]
(('quantile', 'band', 'x'), (0.1, 0.9))
>>> quantile(cube, 0.5, ("time", "x"), method="lower").dims
('band',)
Source code in packages/pipekit-array/src/pipekit_array/dim_algebra.py 
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
def quantile(
    x: Any,
    q: float | Sequence[float],
    dim: str | Sequence[str],
    /,
    *,
    method: str = "linear",
    skipna: bool = False,
) -> Any:
    """Quantiles over named dims (xarray's ``quantile``).

    Q_{k…} = quantile q_k of {x_{…, d…} : d over ``dim``}.

    A sequence ``q`` adds a leading dim ``quantile``, labelled with ``q``;
    a float drops it. A float input keeps its dtype. NaN propagates unless
    ``skipna`` (no implicit ``skipna``, A12).

    Args:
        x: A named carrier.
        q: A probability in [0, 1], or a sequence of them; static under
            ``jit``.
        dim: A dim name, or several, to reduce over.
        method: One of `METHODS`, the six JAX supports.
        skipna: Ignore NaN (``nanquantile``).

    Raises:
        KeyError: A name is not a dim of ``x``.
        ValueError: ``method`` is not one of `METHODS`, a probability is
            outside [0, 1], or a sequence ``q`` meets a kept dim named
            ``quantile``.
        TypeError: The adapter has no ``relabel`` and ``q`` is a sequence.

    Examples:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> cube = DimArray(np.arange(24.0).reshape(4, 2, 3), ("time", "band", "x"))
        >>> spread = quantile(cube, [0.1, 0.9], "time")
        >>> spread.dims, spread.labels["quantile"]
        (('quantile', 'band', 'x'), (0.1, 0.9))
        >>> quantile(cube, 0.5, ("time", "x"), method="lower").dims
        ('band',)
    """
    adapter = adapter_of(x, "quantile")
    names = (dim,) if isinstance(dim, str) else tuple(dim)
    dims = adapter.dims(x)
    axes = _axes(dims, names, "quantile")
    check_method(method, "quantile")
    values = probabilities(q, "quantile")
    many = isinstance(values, tuple)
    kept = tuple(d for d in dims if d not in names)
    if many and "quantile" in kept:
        raise ValueError(
            f"quantile: 'quantile' is already a dim of {dims}; rename it first."
        )
    fast = _fast_path(adapter, "quantile")
    if fast is not None:
        return fast(x, values, names, method=method, skipna=skipna)
    import numpy as np

    data = adapter.data(x)
    if backend_name(data) == "dask":
        # dask's quantile needs every reduced axis in one chunk
        data = data.rechunk(dict.fromkeys(axes, -1))
    xp = namespace_of(data)
    func = _resolve("nanquantile" if skipna else "quantile", xp, "quantile")
    floating = np.issubdtype(data.dtype, np.floating)
    # q in the data's float dtype: NumPy promotes float32 to float64 for a
    # float64 sequence q; dask needs a NumPy q
    q_array = np.asarray(values, dtype=data.dtype if floating else None)
    # one q as a NumPy scalar: dask takes a 0-d array for a sequence
    q_value = q_array if many else q_array[()]
    out = func(data, q_value, axis=axes, method=method)
    result = adapter.rewrap(x, out, ("quantile", *kept) if many else kept)
    if not isinstance(values, tuple):
        return result
    labels: dict[str, Labels] = {"quantile": values}
    return set_labels(adapter, result, labels, "quantile")

pipekit_array.dim_operators

The dim-level functions as Operators, for pipelines and config.

Each wraps one function of pipekit_array.dims, pipekit_array.dim_shapes or pipekit_array.dim_shifts and computes its output Signature without running. Operators that take only names and numbers (ISel, Sel, Transpose, Reduce, ArgReduce, TakeAlong, Rename, ExpandDims, Squeeze, BroadcastLike, MergeDims, SplitDims, Pad, Roll, Shift, Diff, CumReduce, Where) round-trip through dumps / loads; those that hold a user function (MapOver, Combine) set forbid_in_yaml, like pipekit_ops.Lambda.

Reduce and ArgReduce take the reduction by its array-namespace name ("mean", "argmax") so they serialise; the function forms in pipekit_array.dims also take a callable.

Phase 3 (pipekit_array.dim_windows): Window and Coarsen (by the reduction's name) round-trip; ApplyAlong holds an Operator, not a function, so it round-trips whenever that operator does.

Phase 4 (pipekit_array.dim_algebra): Quantile and Dot round-trip; q is stored as a float or an ordered list.

ApplyAlong

Bases: Operator

inner on each item along dim; the results stacked along dim, first.

map_over with an Operator instead of a function, so it serialises whenever inner does: inner is stored as a nested config, and ApplyAlong itself does not set forbid_in_yaml (a flagged inner is found by the serialiser's walk). Use MapOver for a plain function. Replaces the planned ApplyToBands (A15).

Parameters:

Name Type Description Default
inner Operator

The operator applied to each item, isel(x, dim=i); it returns named carriers with the same dims.

 required
dim str

The dim to map over.

 required
Example

import numpy as np from pipekit_array import DimArray cube = DimArray(np.ones((3, 2, 4)), ("time", "band", "x")) ApplyAlong(Reduce("sum", dim="band"), dim="time")(cube).dims ('time', 'x')

Source code in packages/pipekit-array/src/pipekit_array/dim_operators.py 
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
class ApplyAlong(Operator):
    """``inner`` on each item along ``dim``; the results stacked along ``dim``, first.

    `map_over` with an `Operator` instead of a function, so it serialises
    whenever ``inner`` does: ``inner`` is stored as a nested config, and
    `ApplyAlong` itself does not set ``forbid_in_yaml`` (a flagged
    ``inner`` is found by the serialiser's walk). Use `MapOver` for a
    plain function. Replaces the planned ``ApplyToBands`` (A15).

    Args:
        inner: The operator applied to each item, ``isel(x, dim=i)``; it
            returns named carriers with the same dims.
        dim: The dim to map over.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> cube = DimArray(np.ones((3, 2, 4)), ("time", "band", "x"))
        >>> ApplyAlong(Reduce("sum", dim="band"), dim="time")(cube).dims
        ('time', 'x')
    """

    def __init__(self, inner: Operator, dim: str) -> None:
        require_operator(inner, "ApplyAlong", "inner")
        self.inner = inner
        self.dim = dim

    def _apply(self, x: Any) -> Any:
        return _dims.map_over(x, self.inner, self.dim)

    def get_config(self) -> dict[str, Any]:
        return {"inner": nested_config(self.inner), "dim": self.dim}

    def compute_output_signature(self, input_signature: Any) -> Any:
        sig = _signature(input_signature)
        if sig is None:
            return None
        item = self.inner.compute_output_signature(sig.drop_dims((self.dim,)))
        if not isinstance(item, Signature):
            return None
        return Signature(
            dims=((self.dim, dict(sig.dims)[self.dim]), *item.dims), dtype=item.dtype
        )

ArgReduce

Bases: Operator

The index of a reduction along one dim ("argmax", "argmin").

The output signature's dtype is None: the integer width depends on the backend.

Source code in packages/pipekit-array/src/pipekit_array/dim_operators.py 
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
class ArgReduce(Operator):
    """The index of a reduction along one dim (``"argmax"``, ``"argmin"``).

    The output signature's dtype is ``None``: the integer width depends on
    the backend.
    """

    def __init__(self, fn: str, dim: str) -> None:
        if not isinstance(fn, str):
            raise TypeError("ArgReduce takes the reduction's name, e.g. 'argmax'.")
        self.fn = fn
        self.dim = dim

    def _apply(self, x: Any) -> Any:
        return _dims.argreduce(x, self.fn, self.dim)

    def compute_output_signature(self, input_signature: Any) -> Any:
        sig = _signature(input_signature)
        if sig is None:
            return None
        return Signature(dims=sig.drop_dims((self.dim,)).dims, dtype=None)

BroadcastLike

Bases: Operator

Add other's missing dims to x, in front: op(x, other).

compute_output_signature takes the pair (x's, other's): from x's alone the added dims are unknown, so it returns None.

Parameters:

Name Type Description Default
exclude Sequence[str]

Dims of other neither added nor checked.

 ()
Source code in packages/pipekit-array/src/pipekit_array/dim_operators.py 
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
class BroadcastLike(Operator):
    """Add ``other``'s missing dims to ``x``, in front: ``op(x, other)``.

    `compute_output_signature` takes the pair ``(x's, other's)``: from
    ``x``'s alone the added dims are unknown, so it returns None.

    Args:
        exclude: Dims of ``other`` neither added nor checked.
    """

    def __init__(self, exclude: Sequence[str] = ()) -> None:
        self.exclude = tuple(exclude)

    def _apply(self, x: Any, other: Any) -> Any:
        return _dim_shapes.broadcast_like(x, other, exclude=self.exclude)

    def get_config(self) -> dict[str, Any]:
        return {"exclude": list(self.exclude)}

    def compute_output_signature(self, input_signature: Any) -> Any:
        if not isinstance(input_signature, tuple) or len(input_signature) != 2:
            return None
        sig, other = (_signature(s) for s in input_signature)
        if sig is None or other is None:
            return None
        names = {d for d, _ in sig.dims}
        added = tuple(
            (d, n) for d, n in other.dims if d not in names and d not in self.exclude
        )
        return Signature(dims=(*added, *sig.dims), dtype=sig.dtype)

Coarsen

Bases: Operator

Reduce non-overlapping blocks along named dims by a named reduction.

Parameters:

Name Type Description Default
fn str

The reduction's name in the array namespace: "mean", "nanmean", "max", …

 required
boundary str

"exact" or "trim".

 'exact'
**windows int

dim=k for each dim to coarsen.

 {}
Example

import numpy as np from pipekit_array import DimArray Coarsen("mean", y=2, x=2)(DimArray(np.ones((4, 6)), ("y", "x"))).shape (2, 3)

Source code in packages/pipekit-array/src/pipekit_array/dim_operators.py 
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
class Coarsen(Operator):
    """Reduce non-overlapping blocks along named dims by a named reduction.

    Args:
        fn: The reduction's name in the array namespace: ``"mean"``,
            ``"nanmean"``, ``"max"``, …
        boundary: ``"exact"`` or ``"trim"``.
        **windows: ``dim=k`` for each dim to coarsen.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> Coarsen("mean", y=2, x=2)(DimArray(np.ones((4, 6)), ("y", "x"))).shape
        (2, 3)
    """

    def __init__(self, fn: str, boundary: str = "exact", **windows: int) -> None:
        if not isinstance(fn, str):
            raise TypeError(
                "Coarsen takes the reduction's name (e.g. 'mean') so it serialises; "
                "for a function, use pipekit_array.coarsen(x, fn, ...)."
            )
        self.fn = fn
        self.boundary = boundary
        self.windows = dict(windows)

    def _apply(self, x: Any) -> Any:
        return _dim_windows.coarsen(x, self.fn, boundary=self.boundary, **self.windows)

    def get_config(self) -> dict[str, Any]:
        return {"fn": self.fn, "boundary": self.boundary, **self.windows}

    def compute_output_signature(self, input_signature: Any) -> Any:
        sig = _signature(input_signature)
        if sig is None:
            return None
        sizes = dict(sig.dims)
        coarse: dict[str, int | None] = {}
        for d, k in self.windows.items():
            length = sizes[d]
            # ⌊N/k⌋ blocks; "exact" raises at run time unless k divides N
            coarse[d] = None if length is None else length // k
        dims = sig.replace_dims(coarse).dims
        return Signature(dims=dims, dtype=_reduced_dtype(self.fn, sig.dtype))

Combine

Bases: Operator

An element-wise function of carriers with the same dims: op(a, b, …).

Holds a user function, so its config is a debug repr only.

Source code in packages/pipekit-array/src/pipekit_array/dim_operators.py 
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
class Combine(Operator):
    """An element-wise function of carriers with the same dims: ``op(a, b, …)``.

    Holds a user function, so its config is a debug repr only.
    """

    forbid_in_yaml: ClassVar[bool] = True

    def __init__(self, fn: Callable[..., Any]) -> None:
        self.fn = fn

    def _apply(self, *xs: Any) -> Any:
        return _dims.combine(self.fn, *xs)

    def compute_output_signature(self, input_signature: Any) -> Any:
        sig = _signature(input_signature)
        if sig is None:
            return None
        # same dims; fn may change the dtype
        return Signature(dims=sig.dims, dtype=None)

CumReduce

Bases: Operator

A cumulative function along dim, by its array-namespace name.

Parameters:

Name Type Description Default
fn str

"cumulative_sum", "cumulative_prod", "nancumsum", …

 required
dim str

The dim to accumulate along.

 required
Example

import numpy as np from pipekit_array import DimArray CumReduce("cumulative_sum", "time")(DimArray(np.ones(3), ("time",))).data array([1., 2., 3.])

Source code in packages/pipekit-array/src/pipekit_array/dim_operators.py 
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
class CumReduce(Operator):
    """A cumulative function along ``dim``, by its array-namespace name.

    Args:
        fn: ``"cumulative_sum"``, ``"cumulative_prod"``, ``"nancumsum"``, …
        dim: The dim to accumulate along.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> CumReduce("cumulative_sum", "time")(DimArray(np.ones(3), ("time",))).data
        array([1., 2., 3.])
    """

    def __init__(self, fn: str, dim: str) -> None:
        if not isinstance(fn, str):
            raise TypeError(
                "CumReduce takes the function's name (e.g. 'cumulative_sum') so it "
                "serialises; for a function, use pipekit_array.cumreduce(x, fn, dim)."
            )
        self.fn = fn
        self.dim = dim

    def _apply(self, x: Any) -> Any:
        return _dim_shifts.cumreduce(x, self.fn, self.dim)

    def get_config(self) -> dict[str, Any]:
        return {"fn": self.fn, "dim": self.dim}

    def compute_output_signature(self, input_signature: Any) -> Any:
        sig = _signature(input_signature)
        if sig is None:
            return None
        keeps = (
            self.fn in _CUMULATIVE_FLOAT_PRESERVING
            and sig.dtype is not None
            and "float" in sig.dtype
        )
        return Signature(dims=sig.dims, dtype=sig.dtype if keeps else None)

Diff

Bases: Operator

The n-th forward difference along dim: N entries become N - n.

Parameters:

Name Type Description Default
dim str

The dim to difference along.

 required
n int

The order, n ≥ 0.

 1
label str

The labels kept on dim: "upper" (the minuend's) or "lower" (the subtrahend's).

 'upper'
Example

import numpy as np from pipekit_array import DimArray Diff("time")(DimArray(np.zeros((12, 4)), ("time", "x"))).shape (11, 4)

Source code in packages/pipekit-array/src/pipekit_array/dim_operators.py 
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
class Diff(Operator):
    """The n-th forward difference along ``dim``: N entries become N - n.

    Args:
        dim: The dim to difference along.
        n: The order, n ≥ 0.
        label: The labels kept on ``dim``: ``"upper"`` (the minuend's) or
            ``"lower"`` (the subtrahend's).

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> Diff("time")(DimArray(np.zeros((12, 4)), ("time", "x"))).shape
        (11, 4)
    """

    def __init__(self, dim: str, n: int = 1, label: str = "upper") -> None:
        self.dim = dim
        self.n = n
        self.label = label

    def _apply(self, x: Any) -> Any:
        return _dim_shifts.diff(x, self.dim, self.n, label=self.label)

    def get_config(self) -> dict[str, Any]:
        return {"dim": self.dim, "n": self.n, "label": self.label}

    def compute_output_signature(self, input_signature: Any) -> Any:
        sig = _signature(input_signature)
        if sig is None:
            return None
        size = dict(sig.dims)[self.dim]
        return sig.replace_dims(
            {self.dim: None if size is None else max(size - self.n, 0)}
        )

Dot

Bases: Operator

The sum of the product of carriers over named dims: op(*xs).

compute_output_signature takes the tuple of the inputs' signatures: from one signature the other inputs' dims are unknown, so it returns None.

Parameters:

Name Type Description Default
dim str | Sequence[str]

A dim name, or several, to sum over.

 required
Example

import numpy as np from pipekit_array import DimArray x = DimArray(np.ones((12, 13)), ("time", "band")) w = DimArray(np.ones((13, 3)), ("band", "component")) Dot(dim="band")(x, w).dims ('time', 'component')

Source code in packages/pipekit-array/src/pipekit_array/dim_operators.py 
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
class Dot(Operator):
    """The sum of the product of carriers over named dims: ``op(*xs)``.

    `compute_output_signature` takes the tuple of the inputs' signatures:
    from one signature the other inputs' dims are unknown, so it returns
    None.

    Args:
        dim: A dim name, or several, to sum over.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> x = DimArray(np.ones((12, 13)), ("time", "band"))
        >>> w = DimArray(np.ones((13, 3)), ("band", "component"))
        >>> Dot(dim="band")(x, w).dims
        ('time', 'component')
    """

    def __init__(self, dim: str | Sequence[str]) -> None:
        self.dim = _names(dim)

    def _apply(self, *xs: Any) -> Any:
        return _dim_algebra.dot(*xs, dim=self.dim)

    def get_config(self) -> dict[str, Any]:
        return {"dim": _names_config(self.dim)}

    def compute_output_signature(self, input_signature: Any) -> Any:
        if not isinstance(input_signature, tuple) or not input_signature:
            return None
        sigs = [s for s in map(_signature, input_signature) if s is not None]
        if len(sigs) != len(input_signature):
            return None
        names = (self.dim,) if isinstance(self.dim, str) else self.dim
        sizes: dict[str, int | None] = {}
        for sig in sigs:
            for d, n in sig.dims:
                # the first known size; check_compatible makes them agree
                if sizes.get(d) is None:
                    sizes[d] = n
        dims = tuple((d, n) for d, n in sizes.items() if d not in names)
        dtypes = {sig.dtype for sig in sigs}
        return Signature(dims=dims, dtype=dtypes.pop() if len(dtypes) == 1 else None)

ExpandDims

Bases: Operator

Insert length-1 dims at position axis of the result.

Parameters:

Name Type Description Default
dim str | Sequence[str]

A new dim name, or several.

 required
axis int

Where the first new dim goes; negative counts from the end.

 0
Example

import numpy as np from pipekit_array import DimArray ExpandDims("member")(DimArray(np.zeros((2, 4)), ("band", "x"))).dims ('member', 'band', 'x')

Source code in packages/pipekit-array/src/pipekit_array/dim_operators.py 
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
class ExpandDims(Operator):
    """Insert length-1 dims at position ``axis`` of the result.

    Args:
        dim: A new dim name, or several.
        axis: Where the first new dim goes; negative counts from the end.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> ExpandDims("member")(DimArray(np.zeros((2, 4)), ("band", "x"))).dims
        ('member', 'band', 'x')
    """

    def __init__(self, dim: str | Sequence[str], axis: int = 0) -> None:
        self.dim = _names(dim)
        self.axis = axis

    def _apply(self, x: Any) -> Any:
        return _dim_shapes.expand_dims(x, self.dim, axis=self.axis)

    def get_config(self) -> dict[str, Any]:
        return {"dim": _names_config(self.dim), "axis": self.axis}

    def compute_output_signature(self, input_signature: Any) -> Any:
        sig = _signature(input_signature)
        if sig is None:
            return None
        names = (self.dim,) if isinstance(self.dim, str) else self.dim
        start = self.axis if self.axis >= 0 else self.axis + len(sig.dims) + 1
        added = tuple((n, 1) for n in names)
        dims = (*sig.dims[:start], *added, *sig.dims[start:])
        return Signature(dims=dims, dtype=sig.dtype)

ISel

Bases: Operator

Select by position along one dim; an integer drops the dim.

Parameters:

Name Type Description Default
dim str

The dim to index.

 required
index int | slice

An int, or a slice (stored in config as start/stop/step).

 required
Example

import numpy as np from pipekit_array import DimArray ISel(dim="time", index=0)(DimArray(np.zeros((12, 4)), ("time", "x"))).dims ('x',)

Source code in packages/pipekit-array/src/pipekit_array/dim_operators.py 
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
class ISel(Operator):
    """Select by position along one dim; an integer drops the dim.

    Args:
        dim: The dim to index.
        index: An int, or a slice (stored in config as start/stop/step).

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> ISel(dim="time", index=0)(DimArray(np.zeros((12, 4)), ("time", "x"))).dims
        ('x',)
    """

    def __init__(self, dim: str, index: int | slice) -> None:
        self.dim = dim
        index = _slice_from_config(index)
        if type(index).__module__ == "numpy":
            # NumPy positions or a mask (array or scalar) as int / list, so
            # the config is JSON; numpy is optional, so no isinstance
            index = index.tolist()
        self.index = index

    def _apply(self, x: Any) -> Any:
        return _dims.isel(x, **{self.dim: self.index})

    def get_config(self) -> dict[str, Any]:
        return {"dim": self.dim, "index": _slice_config(self.index)}

    def compute_output_signature(self, input_signature: Any) -> Any:
        sig = _signature(input_signature)
        if sig is None:
            return None
        if isinstance(self.index, slice):
            size = dict(sig.dims)[self.dim]
            new = None if size is None else len(range(*self.index.indices(size)))
            return sig.replace_dims({self.dim: new})
        if getattr(self.index, "ndim", None) == 0 or isinstance(self.index, int):
            return sig.drop_dims((self.dim,))
        # a list or 1-D array of positions keeps the dim, one entry each; a
        # boolean mask keeps one entry per True
        return sig.replace_dims({self.dim: _selected_count(self.index)})

MapOver

Bases: Operator

fn on each item along dim; the results stacked along dim, first.

Holds a user function, so its config is a debug repr only.

Source code in packages/pipekit-array/src/pipekit_array/dim_operators.py 
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
class MapOver(Operator):
    """``fn`` on each item along ``dim``; the results stacked along ``dim``, first.

    Holds a user function, so its config is a debug repr only.
    """

    forbid_in_yaml: ClassVar[bool] = True

    def __init__(self, fn: Callable[[Any], Any], dim: str) -> None:
        self.fn = fn
        self.dim = dim

    def _apply(self, x: Any) -> Any:
        return _dims.map_over(x, self.fn, self.dim)

    def compute_output_signature(self, input_signature: Any) -> Any:
        # fn is opaque: its output dims are unknown without running it
        return None

MergeDims

Bases: Operator

Merge the listed dims into one, placed last, in C order.

Its config lists the merges in order (JSON objects sort their keys).

Parameters:

Name Type Description Default
**dims Sequence[str]

new=(source, …) for each merged dim.

 {}
Example

import numpy as np from pipekit_array import DimArray cube = DimArray(np.zeros((2, 3, 4)), ("band", "y", "x")) MergeDims(pixel=("y", "x"))(cube).shape (2, 12)

Source code in packages/pipekit-array/src/pipekit_array/dim_operators.py 
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
class MergeDims(Operator):
    """Merge the listed dims into one, placed last, in C order.

    Its config lists the merges in order (JSON objects sort their keys).

    Args:
        **dims: ``new=(source, …)`` for each merged dim.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> cube = DimArray(np.zeros((2, 3, 4)), ("band", "y", "x"))
        >>> MergeDims(pixel=("y", "x"))(cube).shape
        (2, 12)
    """

    def __init__(self, **dims: Sequence[str]) -> None:
        self.dims = {name: _merge_group(name, group) for name, group in dims.items()}

    def _apply(self, x: Any) -> Any:
        return _dim_shapes.merge_dims(x, **self.dims)

    def get_config(self) -> dict[str, Any]:
        return {"dims": [[name, list(group)] for name, group in self.dims.items()]}

    @classmethod
    def from_config(cls, config: dict[str, Any]) -> Self:
        return cls(**{name: tuple(group) for name, group in config["dims"]})

    def compute_output_signature(self, input_signature: Any) -> Any:
        sig = _signature(input_signature)
        if sig is None:
            return None
        sizes = dict(sig.dims)
        sources = [s for group in self.dims.values() for s in group]
        kept = sig.drop_dims(tuple(sources)).dims
        merged = tuple(
            (name, _product([sizes[s] for s in group]))
            for name, group in self.dims.items()
        )
        return Signature(dims=(*kept, *merged), dtype=sig.dtype)

Pad

Bases: Operator

Widen named dims at either end; a padded dim loses its labels.

Parameters:

Name Type Description Default
mode str

"constant", "edge", "reflect" or "wrap".

 'constant'
constant_values float | None

The fill for mode="constant"; None is NaN.

 None
**pad_width int | Sequence[int]

dim= a width for both ends, or (before, after).

 {}
Example

import numpy as np from pipekit_array import DimArray Pad(time=(2, 0))(DimArray(np.zeros((12, 4)), ("time", "x"))).shape (14, 4)

Source code in packages/pipekit-array/src/pipekit_array/dim_operators.py 
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
class Pad(Operator):
    """Widen named dims at either end; a padded dim loses its labels.

    Args:
        mode: ``"constant"``, ``"edge"``, ``"reflect"`` or ``"wrap"``.
        constant_values: The fill for ``mode="constant"``; None is NaN.
        **pad_width: ``dim=`` a width for both ends, or ``(before, after)``.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> Pad(time=(2, 0))(DimArray(np.zeros((12, 4)), ("time", "x"))).shape
        (14, 4)
    """

    def __init__(
        self,
        mode: str = "constant",
        constant_values: float | None = None,
        **pad_width: int | Sequence[int],
    ) -> None:
        self.mode = mode
        self.constant_values = constant_values
        # an int or a pair, as (before, after); JSON gives the pair as a list
        self.pad_width = {
            d: (w, w) if isinstance(w, int) else (w[0], w[1])
            for d, w in pad_width.items()
        }

    def _apply(self, x: Any) -> Any:
        return _dim_shifts.pad(
            x, mode=self.mode, constant_values=self.constant_values, **self.pad_width
        )

    def get_config(self) -> dict[str, Any]:
        widths = {d: list(w) for d, w in self.pad_width.items()}
        return {"mode": self.mode, "constant_values": self.constant_values, **widths}

    def compute_output_signature(self, input_signature: Any) -> Any:
        sig = _signature(input_signature)
        if sig is None:
            return None
        sizes = dict(sig.dims)
        grown: dict[str, int | None] = {}
        for d, (before, after) in self.pad_width.items():
            size = sizes[d]
            grown[d] = None if size is None else size + before + after
        return sig.replace_dims(grown)

Quantile

Bases: Operator

Quantiles over named dims; a sequence q adds a leading quantile dim.

q is stored in config as a float or an ordered list.

Parameters:

Name Type Description Default
q float | Sequence[float]

A probability in [0, 1], or a sequence of them.

 required
dim str | Sequence[str]

A dim name, or several, to reduce over.

 required
method str

One of pipekit_array.dim_algebra.METHODS.

 'linear'
skipna bool

Ignore NaN.

 False

Raises:

Type Description
ValueError

method is not supported, or a probability is outside [0, 1].
Example

import numpy as np from pipekit_array import DimArray cube = DimArray(np.ones((12, 2, 4)), ("time", "band", "x")) Quantile([0.1, 0.9], dim="time")(cube).dims ('quantile', 'band', 'x')

Source code in packages/pipekit-array/src/pipekit_array/dim_operators.py 
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
class Quantile(Operator):
    """Quantiles over named dims; a sequence ``q`` adds a leading ``quantile`` dim.

    ``q`` is stored in config as a float or an ordered list.

    Args:
        q: A probability in [0, 1], or a sequence of them.
        dim: A dim name, or several, to reduce over.
        method: One of `pipekit_array.dim_algebra.METHODS`.
        skipna: Ignore NaN.

    Raises:
        ValueError: ``method`` is not supported, or a probability is
            outside [0, 1].

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> cube = DimArray(np.ones((12, 2, 4)), ("time", "band", "x"))
        >>> Quantile([0.1, 0.9], dim="time")(cube).dims
        ('quantile', 'band', 'x')
    """

    def __init__(
        self,
        q: float | Sequence[float],
        dim: str | Sequence[str],
        method: str = "linear",
        skipna: bool = False,
    ) -> None:
        _dim_algebra.check_method(method, "Quantile")
        self.q = _dim_algebra.probabilities(q, "Quantile")
        self.dim = _names(dim)
        self.method = method
        self.skipna = skipna

    def _apply(self, x: Any) -> Any:
        return _dim_algebra.quantile(
            x, self.q, self.dim, method=self.method, skipna=self.skipna
        )

    def get_config(self) -> dict[str, Any]:
        q = list(self.q) if isinstance(self.q, tuple) else self.q
        return {
            "q": q,
            "dim": _names_config(self.dim),
            "method": self.method,
            "skipna": self.skipna,
        }

    def compute_output_signature(self, input_signature: Any) -> Any:
        sig = _signature(input_signature)
        if sig is None:
            return None
        names = (self.dim,) if isinstance(self.dim, str) else self.dim
        kept = sig.drop_dims(names).dims
        # a float input keeps its dtype; an integer one promotes by backend
        dtype = sig.dtype if sig.dtype is not None and "float" in sig.dtype else None
        if isinstance(self.q, tuple):
            dims = (("quantile", len(self.q)), *kept)
            return Signature(dims=dims, dtype=dtype)
        return Signature(dims=kept, dtype=dtype)

Reduce

Bases: Operator

Reduce over named dims with an array-namespace reduction.

Parameters:

Name Type Description Default
fn str

The reduction's name in the array namespace: "mean", "max", "nanmedian", …

 required
dim str | Sequence[str]

A dim name, or several.

 required
Example

import numpy as np from pipekit_array import DimArray cube = DimArray(np.ones((12, 2, 4)), ("time", "band", "x")) Reduce("median", dim="time")(cube).dims ('band', 'x')

Source code in packages/pipekit-array/src/pipekit_array/dim_operators.py 
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
class Reduce(Operator):
    """Reduce over named dims with an array-namespace reduction.

    Args:
        fn: The reduction's name in the array namespace: ``"mean"``, ``"max"``,
            ``"nanmedian"``, …
        dim: A dim name, or several.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> cube = DimArray(np.ones((12, 2, 4)), ("time", "band", "x"))
        >>> Reduce("median", dim="time")(cube).dims
        ('band', 'x')
    """

    def __init__(self, fn: str, dim: str | Sequence[str]) -> None:
        if not isinstance(fn, str):
            raise TypeError(
                "Reduce takes the reduction's name (e.g. 'mean') so it serialises; "
                "for a function, use pipekit_array.reduce(x, fn, dim)."
            )
        self.fn = fn
        self.dim = dim if isinstance(dim, str) else tuple(dim)

    def _apply(self, x: Any) -> Any:
        return _dims.reduce(x, self.fn, self.dim)

    def get_config(self) -> dict[str, Any]:
        dim = self.dim if isinstance(self.dim, str) else list(self.dim)
        return {"fn": self.fn, "dim": dim}

    def compute_output_signature(self, input_signature: Any) -> Any:
        sig = _signature(input_signature)
        if sig is None:
            return None
        names = (self.dim,) if isinstance(self.dim, str) else self.dim
        dropped = sig.drop_dims(tuple(names))
        return Signature(dims=dropped.dims, dtype=_reduced_dtype(self.fn, sig.dtype))

Rename

Bases: Operator

New dim names; data, order and labels unchanged.

Parameters:

Name Type Description Default
**names str

old=new for each dim to rename.

 {}
Example

import numpy as np from pipekit_array import DimArray Rename(band="channel")(DimArray(np.zeros((2, 4)), ("band", "x"))).dims ('channel', 'x')

Source code in packages/pipekit-array/src/pipekit_array/dim_operators.py 
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
class Rename(Operator):
    """New dim names; data, order and labels unchanged.

    Args:
        **names: ``old=new`` for each dim to rename.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> Rename(band="channel")(DimArray(np.zeros((2, 4)), ("band", "x"))).dims
        ('channel', 'x')
    """

    def __init__(self, **names: str) -> None:
        self.names = dict(names)

    def _apply(self, x: Any) -> Any:
        return _dim_shapes.rename(x, **self.names)

    def get_config(self) -> dict[str, Any]:
        return dict(self.names)

    def compute_output_signature(self, input_signature: Any) -> Any:
        sig = _signature(input_signature)
        if sig is None:
            return None
        dims = tuple((self.names.get(d, d), n) for d, n in sig.dims)
        return Signature(dims=dims, dtype=sig.dtype)

Roll

Bases: Operator

Shift data cyclically along named dims; labels stay.

Parameters:

Name Type Description Default
**shifts int

dim=k; a positive k moves data to higher positions.

 {}
Example

import numpy as np from pipekit_array import DimArray Roll(time=1)(DimArray(np.arange(3.0), ("time",))).data array([2., 0., 1.])

Source code in packages/pipekit-array/src/pipekit_array/dim_operators.py 
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
class Roll(Operator):
    """Shift data cyclically along named dims; labels stay.

    Args:
        **shifts: ``dim=k``; a positive k moves data to higher positions.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> Roll(time=1)(DimArray(np.arange(3.0), ("time",))).data
        array([2., 0., 1.])
    """

    def __init__(self, **shifts: int) -> None:
        self.shifts = dict(shifts)

    def _apply(self, x: Any) -> Any:
        return _dim_shifts.roll(x, **self.shifts)

    def get_config(self) -> dict[str, Any]:
        return dict(self.shifts)

    def compute_output_signature(self, input_signature: Any) -> Any:
        return _signature(input_signature)

Sel

Bases: Operator

Select by exact label along one dim; one label drops the dim.

Parameters:

Name Type Description Default
dim str

The dim to select on; it must have labels.

 required
label Any

One label, or a list of labels (the dim is kept). A tuple is one label when the dim has tuple-valued labels, else several.

 required
Source code in packages/pipekit-array/src/pipekit_array/dim_operators.py 
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
class Sel(Operator):
    """Select by exact label along one dim; one label drops the dim.

    Args:
        dim: The dim to select on; it must have labels.
        label: One label, or a list of labels (the dim is kept). A tuple is
            one label when the dim has tuple-valued labels, else several.
    """

    def __init__(self, dim: str, label: Any) -> None:
        self.dim = dim
        # a tuple that is itself a label (tuple-valued labels) stays one label;
        # a config from JSON spells it {"__tuple__": [...]}
        label = _decode_tuples(label)
        if getattr(label, "ndim", 0) == 1:
            # a 1-D array of labels is a list of labels: the dim is kept, and
            # the config stays JSON
            label = list(label.tolist())
        self.label = label

    def _apply(self, x: Any) -> Any:
        return _dims.sel(x, **{self.dim: self.label})

    def get_config(self) -> dict[str, Any]:
        # JSON has no tuples: keep "one tuple label" apart from "a list of labels"
        return {"dim": self.dim, "label": _encode_tuples(self.label)}

    def compute_output_signature(self, input_signature: Any) -> Any:
        sig = _signature(input_signature)
        if sig is None:
            return None
        if isinstance(self.label, tuple):
            return None  # one tuple-valued label or several: unknown without labels
        if isinstance(self.label, list):
            return sig.replace_dims({self.dim: len(self.label)})
        return sig.drop_dims((self.dim,))

Shift

Bases: Operator

Shift data along named dims, filling the vacated entries; labels stay.

Parameters:

Name Type Description Default
fill_value float | None

The value of the vacated entries; None is NaN.

 None
**shifts int

dim=k; a positive k moves data to higher positions.

 {}
Example

import numpy as np from pipekit_array import DimArray Shift(time=1)(DimArray(np.arange(3.0), ("time",))).data array([nan, 0., 1.])

Source code in packages/pipekit-array/src/pipekit_array/dim_operators.py 
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
class Shift(Operator):
    """Shift data along named dims, filling the vacated entries; labels stay.

    Args:
        fill_value: The value of the vacated entries; None is NaN.
        **shifts: ``dim=k``; a positive k moves data to higher positions.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> Shift(time=1)(DimArray(np.arange(3.0), ("time",))).data
        array([nan,  0.,  1.])
    """

    def __init__(self, fill_value: float | None = None, **shifts: int) -> None:
        self.fill_value = fill_value
        self.shifts = dict(shifts)

    def _apply(self, x: Any) -> Any:
        return _dim_shifts.shift(x, fill_value=self.fill_value, **self.shifts)

    def get_config(self) -> dict[str, Any]:
        return {"fill_value": self.fill_value, **self.shifts}

    def compute_output_signature(self, input_signature: Any) -> Any:
        return _signature(input_signature)

SplitDims

Bases: Operator

Split one dim into several, in place, in C order.

Its config lists each split's parts in order (JSON objects sort their keys).

Parameters:

Name Type Description Default
**dims Mapping[str, int]

dim={part: size, …} for each dim to split.

 {}
Example

import numpy as np from pipekit_array import DimArray pixels = DimArray(np.zeros((2, 12)), ("band", "pixel")) SplitDims(pixel={"y": 3, "x": 4})(pixels).dims ('band', 'y', 'x')

Source code in packages/pipekit-array/src/pipekit_array/dim_operators.py 
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
class SplitDims(Operator):
    """Split one dim into several, in place, in C order.

    Its config lists each split's parts in order (JSON objects sort their
    keys).

    Args:
        **dims: ``dim={part: size, …}`` for each dim to split.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> pixels = DimArray(np.zeros((2, 12)), ("band", "pixel"))
        >>> SplitDims(pixel={"y": 3, "x": 4})(pixels).dims
        ('band', 'y', 'x')
    """

    def __init__(self, **dims: Mapping[str, int]) -> None:
        self.dims = {name: dict(part) for name, part in dims.items()}

    def _apply(self, x: Any) -> Any:
        return _dim_shapes.split_dims(x, **self.dims)

    def get_config(self) -> dict[str, Any]:
        return {
            "dims": [
                [name, [[p, n] for p, n in part.items()]]
                for name, part in self.dims.items()
            ]
        }

    @classmethod
    def from_config(cls, config: dict[str, Any]) -> Self:
        return cls(**{name: dict(part) for name, part in config["dims"]})

    def compute_output_signature(self, input_signature: Any) -> Any:
        sig = _signature(input_signature)
        if sig is None:
            return None
        dims: list[tuple[str, int | None]] = []
        for d, n in sig.dims:
            dims.extend(self.dims[d].items() if d in self.dims else [(d, n)])
        return Signature(dims=tuple(dims), dtype=sig.dtype)

Squeeze

Bases: Operator

Drop named length-1 dims.

Parameters:

Name Type Description Default
dim str | Sequence[str]

A dim name, or several; each must have size 1.

 required
Source code in packages/pipekit-array/src/pipekit_array/dim_operators.py 
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
class Squeeze(Operator):
    """Drop named length-1 dims.

    Args:
        dim: A dim name, or several; each must have size 1.
    """

    def __init__(self, dim: str | Sequence[str]) -> None:
        self.dim = _names(dim)

    def _apply(self, x: Any) -> Any:
        return _dim_shapes.squeeze(x, self.dim)

    def get_config(self) -> dict[str, Any]:
        return {"dim": _names_config(self.dim)}

    def compute_output_signature(self, input_signature: Any) -> Any:
        sig = _signature(input_signature)
        if sig is None:
            return None
        names = (self.dim,) if isinstance(self.dim, str) else self.dim
        return sig.drop_dims(names)

TakeAlong

Bases: Operator

Gather one entry of dim per position: op(x, index).

index is a named carrier with x's dims without dim, such as an ArgReduce result.

Source code in packages/pipekit-array/src/pipekit_array/dim_operators.py 
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
class TakeAlong(Operator):
    """Gather one entry of ``dim`` per position: ``op(x, index)``.

    ``index`` is a named carrier with ``x``'s dims without ``dim``, such as
    an `ArgReduce` result.
    """

    def __init__(self, dim: str) -> None:
        self.dim = dim

    def _apply(self, x: Any, index: Any) -> Any:
        return _dims.take_along(x, index, self.dim)

    def compute_output_signature(self, input_signature: Any) -> Any:
        sig = _signature(input_signature)
        if sig is None:
            return None
        return sig.drop_dims((self.dim,))

Transpose

Bases: Operator

Reorder dims by name.

Parameters:

Name Type Description Default
order Sequence[str]

Every dim, once, in the new order.

 required
Source code in packages/pipekit-array/src/pipekit_array/dim_operators.py 
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
class Transpose(Operator):
    """Reorder dims by name.

    Args:
        order: Every dim, once, in the new order.
    """

    def __init__(self, order: Sequence[str]) -> None:
        self.order = tuple(order)

    def _apply(self, x: Any) -> Any:
        return _dims.transpose(x, self.order)

    def get_config(self) -> dict[str, Any]:
        return {"order": list(self.order)}

    def compute_output_signature(self, input_signature: Any) -> Any:
        sig = _signature(input_signature)
        if sig is None:
            return None
        sizes = dict(sig.dims)
        return Signature(dims=tuple((d, sizes[d]) for d in self.order), dtype=sig.dtype)

Where

Bases: Operator

Keep x where cond is True, other elsewhere: op(x, cond).

compute_output_signature takes x's signature, or the pair (x's, cond's): the output has x's dims and dtype either way.

Parameters:

Name Type Description Default
other float | None

The fill where cond is False; None is NaN.

 None
Example

import numpy as np from pipekit_array import DimArray valid = DimArray(np.array([True, False]), ("x",)) Where()(DimArray(np.ones(2), ("x",)), valid).data array([ 1., nan])

Source code in packages/pipekit-array/src/pipekit_array/dim_operators.py 
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
class Where(Operator):
    """Keep ``x`` where ``cond`` is True, ``other`` elsewhere: ``op(x, cond)``.

    `compute_output_signature` takes ``x``'s signature, or the pair
    ``(x's, cond's)``: the output has ``x``'s dims and dtype either way.

    Args:
        other: The fill where ``cond`` is False; None is NaN.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> valid = DimArray(np.array([True, False]), ("x",))
        >>> Where()(DimArray(np.ones(2), ("x",)), valid).data
        array([ 1., nan])
    """

    def __init__(self, other: float | None = None) -> None:
        self.other = other

    def _apply(self, x: Any, cond: Any) -> Any:
        return _dim_shifts.where(x, cond, self.other)

    def get_config(self) -> dict[str, Any]:
        return {"other": self.other}

    def compute_output_signature(self, input_signature: Any) -> Any:
        if isinstance(input_signature, tuple) and len(input_signature) == 2:
            input_signature = input_signature[0]
        return _signature(input_signature)

Window

Bases: Operator

Sliding windows along one dim, as a new last dim; full windows only.

Parameters:

Name Type Description Default
window_dim str

The new last dim.

 'window'
stride int

The step between window starts.

 1
**sizes int

dim=S for the one dim to window.

 {}
Example

import numpy as np from pipekit_array import DimArray Window(stride=2, time=3)(DimArray(np.zeros((12, 4)), ("time", "x"))).shape (5, 4, 3)

Source code in packages/pipekit-array/src/pipekit_array/dim_operators.py 
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
class Window(Operator):
    """Sliding windows along one dim, as a new last dim; full windows only.

    Args:
        window_dim: The new last dim.
        stride: The step between window starts.
        **sizes: ``dim=S`` for the one dim to window.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> Window(stride=2, time=3)(DimArray(np.zeros((12, 4)), ("time", "x"))).shape
        (5, 4, 3)
    """

    def __init__(
        self, window_dim: str = "window", stride: int = 1, **sizes: int
    ) -> None:
        if len(sizes) != 1:
            raise ValueError(f"Window: give one dim=size, got {sizes}.")
        self.window_dim = window_dim
        self.stride = stride
        self.sizes = dict(sizes)

    def _apply(self, x: Any) -> Any:
        return _dim_windows.window(
            x, window_dim=self.window_dim, stride=self.stride, **self.sizes
        )

    def get_config(self) -> dict[str, Any]:
        return {"window_dim": self.window_dim, "stride": self.stride, **self.sizes}

    def compute_output_signature(self, input_signature: Any) -> Any:
        sig = _signature(input_signature)
        if sig is None:
            return None
        ((dim, size),) = self.sizes.items()
        length = dict(sig.dims)[dim]
        # n = ⌊(N - S)/s⌋ + 1
        count = None if length is None else (length - size) // self.stride + 1
        windowed = sig.replace_dims({dim: count})
        return Signature(
            dims=(*windowed.dims, (self.window_dim, size)), dtype=sig.dtype
        )

Sample tables

to_rows(x, sample_dims) turns any named carrier into a (n_samples, n_features) table, with rows over the sample dims and columns over the rest. from_rows(rows, layout) turns a table back into a carrier like the input: with the original features, with a new feature dim (feature_dim="component") or, for one value per sample, with the sample dims only. Both run on the data's own backend (geostack OP-O3).

pipekit_array.rows

Sample tables by dim name: to_rows and from_rows.

A per-sample model (scikit-learn, a pixel classifier) wants a 2-D table: one row per sample, one column per feature. to_rows turns any named carrier into one, choosing the samples by dim name, and from_rows turns a table of rows back into a carrier like the input, on the data's own backend (NumPy, JAX under jit, dask):

x (time, band, y, x) --to_rows(sample_dims=("time", "y", "x"))-->
    rows (time·y·x, band), layout
rows (n, k) --from_rows(..., feature_dim="component")-->
    (time, component, y, x)

Samples are the C-order product of the sample dims, taken in the carrier's dim order; features likewise over every other dim. Like the rest of the dim-level functions, these never read coordinate values (GA-Q3): labels and geometry come back through the carrier's adapter, which keeps those of the dims that survive with their size.

RowLayout dataclass

How to_rows flattened a carrier, so from_rows can rebuild one.

Parameters:

Name Type Description Default
dims tuple[str, ...]

The carrier's dims, in its order.

 required
sizes tuple[int, ...]

The size of each dim, in the same order.

 required
sample_dims tuple[str, ...]

The dims that became rows, in the carrier's order.

 required
feature_dims tuple[str, ...]

The dims that became columns, in the carrier's order.

 required
template Any

The carrier to rewrap like (left out of equality and repr).

 None
Source code in packages/pipekit-array/src/pipekit_array/rows.py 
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
@dataclasses.dataclass(frozen=True)
class RowLayout:
    """How `to_rows` flattened a carrier, so `from_rows` can rebuild one.

    Args:
        dims: The carrier's dims, in its order.
        sizes: The size of each dim, in the same order.
        sample_dims: The dims that became rows, in the carrier's order.
        feature_dims: The dims that became columns, in the carrier's order.
        template: The carrier to rewrap like (left out of equality and repr).
    """

    dims: tuple[str, ...]
    sizes: tuple[int, ...]
    sample_dims: tuple[str, ...]
    feature_dims: tuple[str, ...]
    template: Any = dataclasses.field(default=None, compare=False, repr=False)

    @property
    def n_samples(self) -> int:
        """The number of rows: the product of the sample dims' sizes."""
        return math.prod(self._size(d) for d in self.sample_dims)

    @property
    def n_features(self) -> int:
        """The number of columns: the product of the feature dims' sizes."""
        return math.prod(self._size(d) for d in self.feature_dims)

    def _size(self, dim: str) -> int:
        return self.sizes[self.dims.index(dim)]

n_features property

The number of columns: the product of the feature dims' sizes.

n_samples property

The number of rows: the product of the sample dims' sizes.

from_rows(rows, layout, *, feature_dim=None)

A carrier like layout.template from a table of rows.

Parameters:

Name Type Description Default
rows Any

(n_samples, n_features) gives back the original feature dims; (n_samples, k) with feature_dim gives one new dim of size k in their place (PCA components, class probabilities); (n_samples,) gives the sample dims only (a label or a score per sample).

 required
layout RowLayout

From to_rows.

 required
feature_dim str | None

The name of the new feature dim, for a table whose columns are not the original features.

 None

Returns:

Type Description
Any

A carrier with layout's dims in their original order, the new
Any

feature dim (if any) where the first feature dim was, rewrapped by
Any

the template's adapter.

Raises:

Type Description
ValueError

The rows do not match the layout, or a new feature dim is needed and feature_dim is missing or already a dim.
Example

import numpy as np from pipekit_array import DimArray scene = DimArray(np.arange(24.0).reshape(2, 3, 4), ("band", "y", "x")) rows, layout = to_rows(scene, sample_dims=("y", "x")) from_rows(rows[:, :1] * 0, layout, feature_dim="score").dims ('score', 'y', 'x')

Source code in packages/pipekit-array/src/pipekit_array/rows.py 
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
def from_rows(rows: Any, layout: RowLayout, *, feature_dim: str | None = None) -> Any:
    """A carrier like ``layout.template`` from a table of rows.

    Args:
        rows: ``(n_samples, n_features)`` gives back the original feature
            dims; ``(n_samples, k)`` with ``feature_dim`` gives one new dim
            of size ``k`` in their place (PCA components, class
            probabilities); ``(n_samples,)`` gives the sample dims only
            (a label or a score per sample).
        layout: From `to_rows`.
        feature_dim: The name of the new feature dim, for a table whose
            columns are not the original features.

    Returns:
        A carrier with ``layout``'s dims in their original order, the new
        feature dim (if any) where the first feature dim was, rewrapped by
        the template's adapter.

    Raises:
        ValueError: The rows do not match the layout, or a new feature dim
            is needed and ``feature_dim`` is missing or already a dim.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> scene = DimArray(np.arange(24.0).reshape(2, 3, 4), ("band", "y", "x"))
        >>> rows, layout = to_rows(scene, sample_dims=("y", "x"))
        >>> from_rows(rows[:, :1] * 0, layout, feature_dim="score").dims
        ('score', 'y', 'x')
    """
    template = layout.template
    if template is None:
        raise ValueError("from_rows needs a layout from to_rows (it has no template).")
    adapter = adapter_of(template, "from_rows")
    shape = tuple(int(n) for n in rows.shape)
    if len(shape) not in (1, 2) or shape[0] != layout.n_samples:
        raise ValueError(
            f"from_rows: rows of shape {shape} do not fit {layout.n_samples} samples "
            f"over {layout.sample_dims}."
        )
    sample_sizes = tuple(layout._size(d) for d in layout.sample_dims)
    if len(shape) == 1:
        block_dims = layout.sample_dims
        block_shape = sample_sizes
        order = layout.sample_dims
    elif feature_dim is None:
        if shape[1] != layout.n_features:
            raise ValueError(
                f"from_rows: {shape[1]} columns are not the {layout.n_features} "
                f"features over {layout.feature_dims}; pass feature_dim= to name "
                "the new columns."
            )
        block_dims = (*layout.sample_dims, *layout.feature_dims)
        block_shape = (*sample_sizes, *(layout._size(d) for d in layout.feature_dims))
        order = layout.dims
    else:
        if feature_dim in layout.dims:
            raise ValueError(
                f"from_rows: feature_dim {feature_dim!r} is already a dim of "
                f"{layout.dims}."
            )
        block_dims = (*layout.sample_dims, feature_dim)
        block_shape = (*sample_sizes, shape[1])
        order = _with_new_feature_dim(layout, feature_dim)
    xp = namespace_of(rows)
    block = xp.reshape(rows, block_shape)
    axes = tuple(block_dims.index(d) for d in order)
    if axes != tuple(range(len(axes))):
        block = xp.permute_dims(block, axes)
    return adapter.rewrap(template, block, tuple(order))

to_rows(x, sample_dims)

x as a (n_samples, n_features) table, and how to undo it.

Parameters:

Name Type Description Default
x Any

A named carrier.

 required
sample_dims Sequence[str]

The dims whose entries are samples (rows); every other dim of x is a feature (column). At least one.

 required

Returns:

Type Description
Any

The table, on x's backend, and the RowLayout for from_rows.
RowLayout

With no feature dims the table has one column.

Raises:

Type Description
TypeError

x is not a named carrier.
ValueError

No sample dims, a repeated one, or one x lacks.
Example

import numpy as np from pipekit_array import DimArray scene = DimArray(np.arange(24.0).reshape(2, 3, 4), ("band", "y", "x")) rows, layout = to_rows(scene, sample_dims=("y", "x")) rows.shape, layout.feature_dims ((12, 2), ('band',))

Source code in packages/pipekit-array/src/pipekit_array/rows.py 
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
def to_rows(x: Any, sample_dims: Sequence[str]) -> tuple[Any, RowLayout]:
    """``x`` as a ``(n_samples, n_features)`` table, and how to undo it.

    Args:
        x: A named carrier.
        sample_dims: The dims whose entries are samples (rows); every other
            dim of ``x`` is a feature (column). At least one.

    Returns:
        The table, on ``x``'s backend, and the `RowLayout` for `from_rows`.
        With no feature dims the table has one column.

    Raises:
        TypeError: ``x`` is not a named carrier.
        ValueError: No sample dims, a repeated one, or one ``x`` lacks.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> scene = DimArray(np.arange(24.0).reshape(2, 3, 4), ("band", "y", "x"))
        >>> rows, layout = to_rows(scene, sample_dims=("y", "x"))
        >>> rows.shape, layout.feature_dims
        ((12, 2), ('band',))
    """
    adapter = adapter_of(x, "to_rows")
    dims = tuple(adapter.dims(x))
    wanted = tuple(sample_dims)
    if not wanted:
        raise ValueError("to_rows needs at least one sample dim.")
    if len(set(wanted)) != len(wanted):
        raise ValueError(f"to_rows: sample dims {wanted} repeat a name.")
    missing = [d for d in wanted if d not in dims]
    if missing:
        raise ValueError(f"to_rows: sample dims {missing} are not in {dims}.")
    samples = tuple(d for d in dims if d in wanted)
    features = tuple(d for d in dims if d not in wanted)
    sizes = tuple(int(n) for n in adapter.data(x).shape)
    layout = RowLayout(dims, sizes, samples, features, template=x)
    data = adapter.data(transpose(x, (*samples, *features)))
    xp = namespace_of(data)
    rows = xp.reshape(data, (layout.n_samples, layout.n_features))
    return rows, layout

scikit-learn estimators

Estimator fits and applies a scikit-learn-style estimator on any named carrier, through to_rows and from_rows ([sklearn], imported lazily; geostack OP-O3).

pipekit_array.estimator

Estimator: one scikit-learn wrapper for every named carrier (geostack OP-O3).

A scikit-learn-style estimator learns from a (n_samples, n_features) table. Estimator chooses the samples by dim name, flattens the carrier with pipekit_array.to_rows, fits or applies the estimator, and rebuilds a carrier like the input with pipekit_array.from_rows:

cube (time, band, y, x) ── Estimator(PCA(3), sample_dims=("time", "y", "x"),
                                     feature_dim="component") ──>
    (time, component, y, x)

It satisfies pipekit_ops.FittableTransformer: fit returns the operator with estimator_ and layout_ set, and transform (and predict, predict_proba, inverse_transform) only read them. Fitting and applying run on the host: scikit-learn computes on NumPy, so a dask input is computed and a JAX input converted, eagerly; inside a JAX trace fit raises.

NaN policies. nan_policy decides what a non-finite entry costs:

  • "propagate" (default): the rows go to the estimator as they are;
  • "raise": any non-finite entry raises, naming the policy;
  • "mask_samples": rows with a non-finite entry are left out of fitting and applying, and come back as NaN in the output;
  • "mask_features": columns non-finite in every fitted row are dropped at fit time and at transform time, and come back as NaN from inverse_transform;
  • "mask": "mask_features" then "mask_samples" on the columns kept, for a land-masked field that also has gaps.

The scikit-learn import is lazy: install the [sklearn] extra.

Estimator

Bases: Operator

A scikit-learn-style estimator fitted and applied on a named carrier.

Parameters:

Name Type Description Default
estimator Any

Any object with scikit-learn's fit and the methods you call (transform, predict, …). It is cloned on fit; the original is never fitted.

 required
sample_dims tuple[str, ...] | list[str]

The dims whose entries are samples; every other dim is a feature.

 required
feature_dim str | None

The name of the output dim when the estimator's output columns are not the input's features (PCA components, class probabilities). Left None, outputs must have the input's feature count.

 None
nan_policy NanPolicy

"propagate", "raise", "mask_samples", "mask_features" or "mask" (module docstring).

 'propagate'
method str

What calling the operator runs once fitted: "transform", "predict" or "predict_proba".

 'transform'

Raises:

Type Description
ValueError

An unknown nan_policy or method, or no sample dims.
Example

import numpy as np from sklearn.decomposition import PCA from pipekit_array import DimArray rng = np.random.default_rng(0) cube = DimArray(rng.normal(size=(5, 4, 3)), ("band", "y", "x")) pca = Estimator(PCA(2), sample_dims=("y", "x"), feature_dim="component") pca.fit(cube).transform(cube).dims ('component', 'y', 'x')

Source code in packages/pipekit-array/src/pipekit_array/estimator.py 
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
class Estimator(Operator):
    """A scikit-learn-style estimator fitted and applied on a named carrier.

    Args:
        estimator: Any object with scikit-learn's ``fit`` and the methods
            you call (``transform``, ``predict``, …). It is cloned on
            ``fit``; the original is never fitted.
        sample_dims: The dims whose entries are samples; every other dim
            is a feature.
        feature_dim: The name of the output dim when the estimator's output
            columns are not the input's features (PCA components, class
            probabilities). Left ``None``, outputs must have the input's
            feature count.
        nan_policy: ``"propagate"``, ``"raise"``, ``"mask_samples"``,
            ``"mask_features"`` or ``"mask"`` (module docstring).
        method: What calling the operator runs once fitted:
            ``"transform"``, ``"predict"`` or ``"predict_proba"``.

    Raises:
        ValueError: An unknown ``nan_policy`` or ``method``, or no sample
            dims.

    Example:
        >>> import numpy as np
        >>> from sklearn.decomposition import PCA
        >>> from pipekit_array import DimArray
        >>> rng = np.random.default_rng(0)
        >>> cube = DimArray(rng.normal(size=(5, 4, 3)), ("band", "y", "x"))
        >>> pca = Estimator(PCA(2), sample_dims=("y", "x"), feature_dim="component")
        >>> pca.fit(cube).transform(cube).dims
        ('component', 'y', 'x')
    """

    forbid_in_yaml: ClassVar[bool] = True
    backends: ClassVar[tuple[str, ...]] = ("numpy", "dask")
    __config_exclude__: ClassVar[tuple[str, ...]] = (
        "estimator_",
        "layout_",
        "kept_features_",
    )

    def __init__(
        self,
        estimator: Any,
        *,
        sample_dims: tuple[str, ...] | list[str],
        feature_dim: str | None = None,
        nan_policy: NanPolicy = "propagate",
        method: str = "transform",
    ) -> None:
        if nan_policy not in _POLICIES:
            raise ValueError(
                f"nan_policy must be one of {_POLICIES}; got {nan_policy!r}"
            )
        if method not in _METHODS:
            raise ValueError(f"method must be one of {_METHODS}; got {method!r}")
        if not tuple(sample_dims):
            raise ValueError("Estimator needs at least one sample dim.")
        self.estimator = estimator
        self.sample_dims = tuple(sample_dims)
        self.feature_dim = feature_dim
        self.nan_policy = nan_policy
        self.method = method
        self.estimator_: Any = None
        self.layout_: RowLayout | None = None
        self.kept_features_: np.ndarray | None = None

    # ------------------------------------------------------------------ fit

    def fit(self, x: Any, y: Any = None) -> Self:
        """Learn the estimator from ``x`` (and ``y``, for a supervised one).

        Args:
            x: A named carrier with the sample dims.
            y: Targets: a carrier whose dims are the sample dims (plus at
                most one target dim), or ``None``.

        Returns:
            This operator, fitted.

        Raises:
            TypeError: ``x`` is a JAX tracer (fit eagerly, outside ``jit``).
            ValueError: ``nan_policy="raise"`` and ``x`` has a non-finite
                entry.
        """
        from sklearn.base import clone

        rows, layout = to_rows(x, self.sample_dims)
        table = _host(rows, "fit")
        self._check(table)
        if self.nan_policy in ("mask_features", "mask"):
            self.kept_features_ = np.asarray(np.isfinite(table).any(axis=0))
            table = table[:, self.kept_features_]
        keep = self._rows(table)
        targets = None if y is None else _targets(y, self.sample_dims)
        estimator = clone(self.estimator)
        if targets is None:
            estimator.fit(table[keep])
        else:
            estimator.fit(table[keep], targets[keep])
        self.estimator_ = estimator
        # the layout without its template: the training data is not kept alive
        self.layout_ = dataclasses.replace(layout, template=None)
        return self

    # ------------------------------------------------------------- apply

    def transform(self, x: Any) -> Any:
        """The estimator's ``transform`` of ``x``, as a carrier like ``x``."""
        return self._apply_method("transform", x)

    def predict(self, x: Any) -> Any:
        """One prediction per sample: a carrier with the sample dims only."""
        return self._apply_method("predict", x)

    def predict_proba(self, x: Any) -> Any:
        """Class probabilities per sample, along ``feature_dim`` (``"class"``)."""
        return self._apply_method("predict_proba", x)

    def inverse_transform(self, z: Any) -> Any:
        """Map ``transform``'s output back to the fitted features.

        Args:
            z: A carrier with the sample dims and the output dim (``feature_dim``).

        Returns:
            A carrier with the fitted feature dims and sizes in place of the
            output dim; dropped features (``"mask_features"``) are ``NaN``.
        """
        layout = self._fitted()
        rows, z_layout = to_rows(z, self.sample_dims)
        table = _host(rows, "inverse_transform")
        keep = _finite_rows(table) if self._masks_samples else _all(table)
        out_cols = self._restore_features(
            self.estimator_.inverse_transform(table[keep])
        )
        full = _fill_rows(out_cols, keep, layout.n_features)
        target = RowLayout(
            dims=layout.dims,
            sizes=tuple(
                _size(z_layout, d) if d in layout.sample_dims else _size(layout, d)
                for d in layout.dims
            ),
            sample_dims=layout.sample_dims,
            feature_dims=layout.feature_dims,
            template=z,
        )
        return from_rows(full, target)

    def _apply(self, x: Any) -> Any:
        return getattr(self, self.method)(x)

    # ----------------------------------------------------------- helpers

    @property
    def _masks_samples(self) -> bool:
        return self.nan_policy in ("mask_samples", "mask")

    def _fitted(self) -> RowLayout:
        if self.estimator_ is None or self.layout_ is None:
            raise ValueError("Estimator is not fitted: call fit(x) first.")
        return self.layout_

    def _apply_method(self, name: str, x: Any) -> Any:
        layout = self._fitted()
        rows, x_layout = to_rows(x, self.sample_dims)
        _require_same_features(x_layout, layout)
        table = _host(rows, name)
        self._check(table)
        if self.kept_features_ is not None:
            table = table[:, self.kept_features_]
        keep = self._rows(table)
        out = np.asarray(getattr(self.estimator_, name)(table[keep]))
        if out.ndim == 1:
            return from_rows(_fill_rows(out, keep, None), x_layout)
        full = _fill_rows(out, keep, out.shape[1])
        if self.feature_dim is None and out.shape[1] == x_layout.n_features:
            return from_rows(full, x_layout)
        dim = self.feature_dim or ("class" if name == "predict_proba" else None)
        if dim is None:
            raise ValueError(
                f"Estimator.{name} returned {out.shape[1]} columns for "
                f"{x_layout.n_features} features; pass feature_dim= to name them."
            )
        return from_rows(full, x_layout, feature_dim=dim)

    def _check(self, table: np.ndarray) -> None:
        """Raise on a non-finite entry under ``nan_policy="raise"``."""
        if self.nan_policy == "raise" and not np.isfinite(table).all():
            raise ValueError(
                "Estimator: the input has non-finite entries (nan_policy='raise'); "
                "use 'mask_samples', 'mask_features' or 'mask' to leave them out."
            )

    def _rows(self, table: np.ndarray) -> np.ndarray:
        """The rows to use: the finite ones when the policy masks samples."""
        return _finite_rows(table) if self._masks_samples else _all(table)

    def _restore_features(self, cols: np.ndarray) -> np.ndarray:
        """Re-insert features dropped at fit time as ``NaN`` columns."""
        if self.kept_features_ is None:
            return cols
        out = np.full((cols.shape[0], self.kept_features_.size), np.nan)
        out[:, self.kept_features_] = cols
        return out

fit(x, y=None)

Learn the estimator from x (and y, for a supervised one).

Parameters:

Name Type Description Default
x Any

A named carrier with the sample dims.

 required
y Any

Targets: a carrier whose dims are the sample dims (plus at most one target dim), or None.

 None

Returns:

Type Description
Self

This operator, fitted.

Raises:

Type Description
TypeError

x is a JAX tracer (fit eagerly, outside jit).
ValueError

nan_policy="raise" and x has a non-finite entry.
Source code in packages/pipekit-array/src/pipekit_array/estimator.py 
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
def fit(self, x: Any, y: Any = None) -> Self:
    """Learn the estimator from ``x`` (and ``y``, for a supervised one).

    Args:
        x: A named carrier with the sample dims.
        y: Targets: a carrier whose dims are the sample dims (plus at
            most one target dim), or ``None``.

    Returns:
        This operator, fitted.

    Raises:
        TypeError: ``x`` is a JAX tracer (fit eagerly, outside ``jit``).
        ValueError: ``nan_policy="raise"`` and ``x`` has a non-finite
            entry.
    """
    from sklearn.base import clone

    rows, layout = to_rows(x, self.sample_dims)
    table = _host(rows, "fit")
    self._check(table)
    if self.nan_policy in ("mask_features", "mask"):
        self.kept_features_ = np.asarray(np.isfinite(table).any(axis=0))
        table = table[:, self.kept_features_]
    keep = self._rows(table)
    targets = None if y is None else _targets(y, self.sample_dims)
    estimator = clone(self.estimator)
    if targets is None:
        estimator.fit(table[keep])
    else:
        estimator.fit(table[keep], targets[keep])
    self.estimator_ = estimator
    # the layout without its template: the training data is not kept alive
    self.layout_ = dataclasses.replace(layout, template=None)
    return self

inverse_transform(z)

Map transform's output back to the fitted features.

Parameters:

Name Type Description Default
z Any

A carrier with the sample dims and the output dim (feature_dim).

 required

Returns:

Type Description
Any

A carrier with the fitted feature dims and sizes in place of the
Any

output dim; dropped features ("mask_features") are NaN.
Source code in packages/pipekit-array/src/pipekit_array/estimator.py 
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
def inverse_transform(self, z: Any) -> Any:
    """Map ``transform``'s output back to the fitted features.

    Args:
        z: A carrier with the sample dims and the output dim (``feature_dim``).

    Returns:
        A carrier with the fitted feature dims and sizes in place of the
        output dim; dropped features (``"mask_features"``) are ``NaN``.
    """
    layout = self._fitted()
    rows, z_layout = to_rows(z, self.sample_dims)
    table = _host(rows, "inverse_transform")
    keep = _finite_rows(table) if self._masks_samples else _all(table)
    out_cols = self._restore_features(
        self.estimator_.inverse_transform(table[keep])
    )
    full = _fill_rows(out_cols, keep, layout.n_features)
    target = RowLayout(
        dims=layout.dims,
        sizes=tuple(
            _size(z_layout, d) if d in layout.sample_dims else _size(layout, d)
            for d in layout.dims
        ),
        sample_dims=layout.sample_dims,
        feature_dims=layout.feature_dims,
        template=z,
    )
    return from_rows(full, target)

predict(x)

One prediction per sample: a carrier with the sample dims only.

Source code in packages/pipekit-array/src/pipekit_array/estimator.py 
170
171
172
def predict(self, x: Any) -> Any:
    """One prediction per sample: a carrier with the sample dims only."""
    return self._apply_method("predict", x)

predict_proba(x)

Class probabilities per sample, along feature_dim ("class").

Source code in packages/pipekit-array/src/pipekit_array/estimator.py 
174
175
176
def predict_proba(self, x: Any) -> Any:
    """Class probabilities per sample, along ``feature_dim`` (``"class"``)."""
    return self._apply_method("predict_proba", x)

transform(x)

The estimator's transform of x, as a carrier like x.

Source code in packages/pipekit-array/src/pipekit_array/estimator.py 
166
167
168
def transform(self, x: Any) -> Any:
    """The estimator's ``transform`` of ``x``, as a carrier like ``x``."""
    return self._apply_method("transform", x)

einx

pipekit_array.einx writes tensor algebra as einx patterns whose axis names are the carrier's dim names: rearrange, reduce, einsum, repeat and the Einx operator. It runs on NumPy or JAX and rewraps through the first input's adapter, so the dims that survive keep their labels, coordinates and grid. A dask input raises BackendError naming .compute(), because einx has no dask backend (geostack OP-O2).

pipekit_array.einx

A named-dims einx bridge on any carrier (geostack OP-O2).

einx <https://github.com/fferflo/einx>_ writes tensor algebra as patterns ("band y x -> y x band"). Here the pattern's axis names are the carrier's dim names: each input is matched to its slot by name, einx runs on the data's own backend (NumPy, JAX under jit), and the result is rewrapped through the first input's adapter with the output slot's dims, so it keeps the labels, coordinates and (for geoarray) the grid of the dims that survive:

rearrange("band y x -> y x band", scene)            # transpose by name
reduce("time band y x -> band y x", cube, op="mean") # over time
einsum("time band, band k -> time k", cube, w)       # contract band
repeat("y x -> member y x", field, member=8)         # broadcast

einsum, reduce and repeat match input slots to dims by name, in any order. rearrange also splits and merges dims ("(a b)"), so its input slot is read positionally: it lists the input's dims in order, a group standing for one dim that is split. A merged output group (a b) becomes one dim named a_b.

dask is not an einx backend: a dask-backed input raises pipekit_ops.BackendError naming .compute() rather than computing it silently. einx is imported lazily: install the [einx] extra.

Einx

Bases: Operator

One einx operation as an operator: Einx("reduce", pattern, op="mean").

Parameters:

Name Type Description Default
operation str

"rearrange", "reduce", "einsum" or "repeat".

 required
pattern str

The pattern, with dim names as axis names.

 required
op str | None

The reduction, for "reduce" (a name, so it serialises).

 None
sizes dict[str, int] | None

Axis sizes forwarded to einx.

 None
Example

import numpy as np from pipekit_array import DimArray cube = DimArray(np.ones((2, 3, 4)), ("time", "y", "x")) Einx("reduce", "time y x -> y x", op="mean")(cube).dims ('y', 'x')

Source code in packages/pipekit-array/src/pipekit_array/einx.py 
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
class Einx(Operator):
    """One einx operation as an operator: ``Einx("reduce", pattern, op="mean")``.

    Args:
        operation: ``"rearrange"``, ``"reduce"``, ``"einsum"`` or ``"repeat"``.
        pattern: The pattern, with dim names as axis names.
        op: The reduction, for ``"reduce"`` (a name, so it serialises).
        sizes: Axis sizes forwarded to einx.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> cube = DimArray(np.ones((2, 3, 4)), ("time", "y", "x"))
        >>> Einx("reduce", "time y x -> y x", op="mean")(cube).dims
        ('y', 'x')
    """

    backends: ClassVar[tuple[str, ...]] = ("numpy", "jax")

    def __init__(
        self,
        operation: str,
        pattern: str,
        *,
        op: str | None = None,
        sizes: dict[str, int] | None = None,
    ) -> None:
        if operation not in _FUNCTIONS:
            raise ValueError(f"operation must be one of {tuple(_FUNCTIONS)}")
        if (operation == "reduce") != (op is not None):
            raise ValueError("op is given exactly when operation is 'reduce'")
        _parse(pattern)  # fail at construction, not at the first call
        self.operation = operation
        self.pattern = pattern
        self.op = op
        self.sizes = dict(sizes or {})

    def _apply(self, *xs: Any) -> Any:
        fn = _FUNCTIONS[self.operation]
        if self.operation == "reduce":
            return fn(self.pattern, *xs, op=self.op, **self.sizes)
        return fn(self.pattern, *xs, **self.sizes)

PatternError

Bases: ValueError

A pattern this bridge cannot read, or one that does not fit its inputs.

Source code in packages/pipekit-array/src/pipekit_array/einx.py 
47
48
class PatternError(ValueError):
    """A pattern this bridge cannot read, or one that does not fit its inputs."""

einsum(pattern, *xs, **sizes)

Contract the dims named on the input slots and left off the output.

Parameters:

Name Type Description Default
pattern str

"time band, band k -> time k".

 required
*xs Any

One named carrier per input slot; carriers of one family are checked to combine first (geoarray: the same grid).

 ()
**sizes int

Forwarded to einx.

 {}

Returns:

Type Description
Any

A carrier with the output slot's dims, rewrapped by the first
Any

input's adapter.
Example

import numpy as np from pipekit_array import DimArray a = DimArray(np.ones((2, 3)), ("time", "band")) w = DimArray(np.ones((3, 4)), ("band", "k")) einsum("time band, band k -> time k", a, w).shape (2, 4)

Source code in packages/pipekit-array/src/pipekit_array/einx.py 
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
def einsum(pattern: str, *xs: Any, **sizes: int) -> Any:
    """Contract the dims named on the input slots and left off the output.

    Args:
        pattern: ``"time band, band k -> time k"``.
        *xs: One named carrier per input slot; carriers of one family are
            checked to combine first (geoarray: the same grid).
        **sizes: Forwarded to einx.

    Returns:
        A carrier with the output slot's dims, rewrapped by the first
        input's adapter.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> a = DimArray(np.ones((2, 3)), ("time", "band"))
        >>> w = DimArray(np.ones((3, 4)), ("band", "k"))
        >>> einsum("time band, band k -> time k", a, w).shape
        (2, 4)
    """
    einx = _einx()
    parsed = _parse(pattern)
    prepared = _by_name(xs, parsed, "einsum")
    datas = _datas(prepared, "einsum")
    return _wrap(prepared[0], einx.dot(pattern, *datas, **sizes), parsed)

rearrange(pattern, x, **sizes)

Transpose, split or merge dims by name.

Parameters:

Name Type Description Default
pattern str

One input slot and one output slot. The input slot lists x's dims in order; (a b) stands for one dim split into a and b (give their sizes in sizes).

 required
x Any

A named carrier.

 required
**sizes int

Sizes of split dims einx cannot infer (a=2).

 {}

Returns:

Type Description
Any

A carrier with the output slot's dims (a merged (a b) is
Any

a_b), rewrapped by x's adapter.

Raises:

Type Description
PatternError

The input slot does not describe x's dims.
Example

import numpy as np from pipekit_array import DimArray scene = DimArray(np.zeros((3, 4, 5)), ("band", "y", "x")) rearrange("band y x -> y x band", scene).dims ('y', 'x', 'band')

Source code in packages/pipekit-array/src/pipekit_array/einx.py 
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
def rearrange(pattern: str, x: Any, **sizes: int) -> Any:
    """Transpose, split or merge dims by name.

    Args:
        pattern: One input slot and one output slot. The input slot lists
            ``x``'s dims in order; ``(a b)`` stands for one dim split into
            ``a`` and ``b`` (give their sizes in ``sizes``).
        x: A named carrier.
        **sizes: Sizes of split dims einx cannot infer (``a=2``).

    Returns:
        A carrier with the output slot's dims (a merged ``(a b)`` is
        ``a_b``), rewrapped by ``x``'s adapter.

    Raises:
        PatternError: The input slot does not describe ``x``'s dims.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> scene = DimArray(np.zeros((3, 4, 5)), ("band", "y", "x"))
        >>> rearrange("band y x -> y x band", scene).dims
        ('y', 'x', 'band')
    """
    einx = _einx()
    parsed = _parse(pattern)
    if len(parsed.inputs) != 1:
        raise PatternError("rearrange takes one input slot.")
    dims = tuple(adapter_of(x, "rearrange").dims(x))
    slot = parsed.inputs[0]
    names = [el for el in slot if isinstance(el, str)]
    if len(slot) != len(dims) or any(
        isinstance(el, str) and el != d for el, d in zip(slot, dims, strict=True)
    ):
        if not any(isinstance(el, tuple) for el in slot) and set(names) == set(dims):
            x = transpose(x, tuple(names))  # a flat slot in another order: by name
        else:
            raise PatternError(
                f"rearrange: the input slot {slot} does not describe dims {dims} "
                "in order"
            )
    (data,) = _datas([x], "rearrange")
    return _wrap(x, einx.id(pattern, data, **sizes), parsed)

reduce(pattern, x, *, op, **sizes)

Reduce over the dims on the input slot that the output slot leaves out.

Parameters:

Name Type Description Default
pattern str

"time band y x -> band y x".

 required
x Any

A named carrier.

 required
op str | Callable[..., Any]

"sum", "mean", "min", "max", "prod", "std", "var", "any", "all", or a callable taking axis= (applied on the data's backend).

 required
**sizes int

Forwarded to einx.

 {}

Returns:

Type Description
Any

A carrier with the output slot's dims.
Example

import numpy as np from pipekit_array import DimArray cube = DimArray(np.ones((2, 3, 4)), ("time", "y", "x")) reduce("time y x -> y x", cube, op="sum").data[0, 0] np.float64(2.0)

Source code in packages/pipekit-array/src/pipekit_array/einx.py 
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
def reduce(pattern: str, x: Any, *, op: str | Callable[..., Any], **sizes: int) -> Any:
    """Reduce over the dims on the input slot that the output slot leaves out.

    Args:
        pattern: ``"time band y x -> band y x"``.
        x: A named carrier.
        op: ``"sum"``, ``"mean"``, ``"min"``, ``"max"``, ``"prod"``,
            ``"std"``, ``"var"``, ``"any"``, ``"all"``, or a callable taking
            ``axis=`` (applied on the data's backend).
        **sizes: Forwarded to einx.

    Returns:
        A carrier with the output slot's dims.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> cube = DimArray(np.ones((2, 3, 4)), ("time", "y", "x"))
        >>> reduce("time y x -> y x", cube, op="sum").data[0, 0]
        np.float64(2.0)
    """
    einx = _einx()
    parsed = _parse(pattern)
    (x,) = _by_name([x], parsed, "reduce")
    (data,) = _datas([x], "reduce")
    if isinstance(op, str):
        if op not in _REDUCERS:
            raise PatternError(f"reduce: op must be one of {_REDUCERS} or a callable")
        result = getattr(einx, op)(pattern, data, **sizes)
    else:
        result = _backend(data).adapt_numpylike_reduce(op)(pattern, data, **sizes)
    return _wrap(x, result, parsed)

repeat(pattern, x, **sizes)

Broadcast x along new dims (sizes in sizes).

Example

import numpy as np from pipekit_array import DimArray field = DimArray(np.zeros((4, 5)), ("y", "x")) repeat("y x -> member y x", field, member=8).shape (8, 4, 5)

Source code in packages/pipekit-array/src/pipekit_array/einx.py 
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
def repeat(pattern: str, x: Any, **sizes: int) -> Any:
    """Broadcast ``x`` along new dims (sizes in ``sizes``).

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> field = DimArray(np.zeros((4, 5)), ("y", "x"))
        >>> repeat("y x -> member y x", field, member=8).shape
        (8, 4, 5)
    """
    einx = _einx()
    parsed = _parse(pattern)
    (x,) = _by_name([x], parsed, "repeat")
    (data,) = _datas([x], "repeat")
    return _wrap(x, einx.id(pattern, data, **sizes), parsed)

Batches

batch groups carriers into fixed-size Padded batches (value, valid) along a new dim. A short last batch is padded with copies of its last item (pad="last"), yielded short ("none") or dropped ("drop"). stack_items is the stacking it uses; leave the padding out of a statistic with reduce(..., where=padded.valid) (ADR A10). BatchAlong(inner, size, dim, pad) is the operator form for one carrier: it cuts dim into chunks, runs inner on each and joins the results (ADR A19).

pipekit_array.batching

Fixed-size batches of named carriers, with a mask of the real entries.

batch groups a stream of carriers into Padded batches of one size, so a jitted step compiles once for every batch, the short last one included:

items ──batch(size=4, dim="sample", pad="last")──▶ Padded
  i₀ i₁ i₂ i₃ i₄ i₅                                 value  (sample, *item_dims)
                                                    valid  (sample,) bool
  batch 1: i₀ i₁ i₂ i₃    valid = T T T T
  batch 2: i₄ i₅ i₅ i₅    valid = T T F F   ← copies of the last real item

Padding copies the last real item: its values are finite and in distribution, so a masked loss and its gradient stay finite and a normalisation over the batch sees real statistics. Reduce with reduce(out, "mean", "sample", where=padded.valid) to leave the padding out.

BatchAlong is the operator form for one carrier: it cuts a dim into chunks, runs an operator on each and joins the results, with the same pad policies (A19).

BatchAlong

Bases: Operator

inner on chunks of size along dim of one carrier, joined back.

y = concat([inner(x_{0:m}), inner(x_{m:2m}), …], dim), m = size.

Not pipekit_ops.BatchedMap, which takes an iterable and calls its operator on lists of items: BatchAlong splits one named carrier along a dim (A19, which supersedes A5). The use case is a GPU-bound model on a cube that does not fit in device memory at once.

inner must keep dim, one output entry per input entry (a per-item operator), for the "last" padding to be trimmed and for dim's labels to come back; the result then equals inner(x).

Algorithm
  1. Cut dim into chunks of size; the last may be short.
  2. Apply pad (ADR A10) to a short last chunk: "none" keeps it short; "last" pads it to size with copies of its last item, so a jitted inner compiles once, and trims the padding from its output; "drop" drops it.
  3. Join the outputs along dim and set dim's labels to the input's (an adapter's generic concat drops them).

Parameters:

Name Type Description Default
inner Operator

The operator applied to each chunk.

 required
size int

Entries per chunk; must be positive.

 required
dim str

The dim to cut.

 required
pad PadPolicy

"none" (default), "last" or "drop".

 'none'

Raises:

Type Description
TypeError

inner is not an Operator.
ValueError

size is not positive or pad is not a policy.
Example

import numpy as np from pipekit_array import DimArray, Reduce cube = DimArray(np.ones((10, 2, 4)), ("sample", "band", "x")) per_chunk = BatchAlong(Reduce("sum", "band"), size=4, dim="sample") per_chunk(cube).dims, per_chunk(cube).shape (('sample', 'x'), (10, 4)) BatchAlong(Reduce("sum", "band"), 4, "sample", pad="drop")(cube).shape (8, 4)

Source code in packages/pipekit-array/src/pipekit_array/batching.py 
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
class BatchAlong(Operator):
    """``inner`` on chunks of ``size`` along ``dim`` of one carrier, joined back.

    y = concat([inner(x_{0:m}), inner(x_{m:2m}), …], dim),  m = ``size``.

    Not `pipekit_ops.BatchedMap`, which takes an iterable and calls its
    operator on lists of items: `BatchAlong` splits *one* named carrier
    along a dim (A19, which supersedes A5). The use case is a GPU-bound
    model on a cube that does not fit in device memory at once.

    ``inner`` must keep ``dim``, one output entry per input entry (a
    per-item operator), for the ``"last"`` padding to be trimmed and for
    ``dim``'s labels to come back; the result then equals ``inner(x)``.

    Algorithm:
        1. Cut ``dim`` into chunks of ``size``; the last may be short.
        2. Apply ``pad`` (ADR A10) to a short last chunk: ``"none"`` keeps
           it short; ``"last"`` pads it to ``size`` with copies of its
           last item, so a jitted ``inner`` compiles once, and trims the
           padding from its output; ``"drop"`` drops it.
        3. Join the outputs along ``dim`` and set ``dim``'s labels to the
           input's (an adapter's generic ``concat`` drops them).

    Args:
        inner: The operator applied to each chunk.
        size: Entries per chunk; must be positive.
        dim: The dim to cut.
        pad: ``"none"`` (default), ``"last"`` or ``"drop"``.

    Raises:
        TypeError: ``inner`` is not an `Operator`.
        ValueError: ``size`` is not positive or ``pad`` is not a policy.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray, Reduce
        >>> cube = DimArray(np.ones((10, 2, 4)), ("sample", "band", "x"))
        >>> per_chunk = BatchAlong(Reduce("sum", "band"), size=4, dim="sample")
        >>> per_chunk(cube).dims, per_chunk(cube).shape
        (('sample', 'x'), (10, 4))
        >>> BatchAlong(Reduce("sum", "band"), 4, "sample", pad="drop")(cube).shape
        (8, 4)
    """

    def __init__(
        self, inner: Operator, size: int, dim: str, pad: PadPolicy = "none"
    ) -> None:
        require_operator(inner, "BatchAlong", "inner")
        if size <= 0:
            raise ValueError(f"BatchAlong: size must be positive, got {size}.")
        if pad not in PAD_POLICIES:
            raise ValueError(
                f"BatchAlong: pad must be one of {PAD_POLICIES}, got {pad!r}."
            )
        self.inner = inner
        self.size = size
        self.dim = dim
        self.pad = pad

    def _apply(self, x: Any) -> Any:
        adapter = adapter_of(x, "BatchAlong")
        (axis,) = _axes(adapter.dims(x), (self.dim,), "BatchAlong")
        length = adapter.data(x).shape[axis]
        outputs = []
        covered = 0
        for start in range(0, length, self.size):
            stop = min(start + self.size, length)
            missing = self.size - (stop - start)
            if missing and self.pad == "drop":
                break
            covered = stop
            if missing and self.pad == "last":
                # copies of the chunk's last item fill it to size
                positions = [*range(start, stop), *([stop - 1] * missing)]
                padded = self.inner(adapter.isel(x, {self.dim: positions}))
                outputs.append(isel(padded, **{self.dim: slice(0, stop - start)}))
                continue
            outputs.append(self.inner(adapter.isel(x, {self.dim: slice(start, stop)})))
        if not outputs:
            raise ValueError(
                f"BatchAlong: pad='drop' drops every entry of {self.dim!r} "
                f"(size {length} < {self.size})."
            )
        out = concat(outputs, self.dim)
        labels = adapter.labels(x, self.dim)
        if labels is None:
            return out
        out_adapter = adapter_of(out, "BatchAlong")
        (out_axis,) = _axes(out_adapter.dims(out), (self.dim,), "BatchAlong")
        if out_adapter.data(out).shape[out_axis] != covered:
            return out
        # a per-item inner: entry i of the output is entry i of the input
        return set_labels(out_adapter, out, {self.dim: labels[:covered]}, "BatchAlong")

    def get_config(self) -> dict[str, Any]:
        return {
            "inner": nested_config(self.inner),
            "size": self.size,
            "dim": self.dim,
            "pad": self.pad,
        }

    def compute_output_signature(self, input_signature: Any) -> Any:
        if not isinstance(input_signature, Signature):
            return None
        length = dict(input_signature.dims)[self.dim]
        if length is None:
            return None
        chunk = input_signature.replace_dims({self.dim: min(self.size, length)})
        out = self.inner.compute_output_signature(chunk)
        if not isinstance(out, Signature) or self.dim not in dict(out.dims):
            return None
        # a per-item inner: as many entries as the input keeps
        kept = length - length % self.size if self.pad == "drop" else length
        return out.replace_dims({self.dim: kept})

Padded

Bases: NamedTuple

A batch: the items as one carrier and a mask of the real ones.

A NamedTuple, so JAX treats it as a pytree: a jitted step taking a Padded does not retrace when the mask changes. A full batch is a Padded whose valid is all True.

Attributes:

Name Type Description
value Any

The items stacked along the batch dim, first (stack_items): (n, *item_dims).
valid Any

(n,) bool, True for real items, in the value's array namespace.
Example

import numpy as np from pipekit_array import DimArray, reduce items = [DimArray(np.full(3, float(i)), ("x",)) for i in range(3)] padded = next(batch(items, size=4, dim="sample")) padded.valid.tolist() [True, True, True, False] reduce(padded.value, "mean", "sample", where=padded.valid).data array([1., 1., 1.])

Source code in packages/pipekit-array/src/pipekit_array/batching.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
class Padded(NamedTuple):
    """A batch: the items as one carrier and a mask of the real ones.

    A `NamedTuple`, so JAX treats it as a pytree: a jitted step taking a
    `Padded` does not retrace when the mask changes. A full batch is a
    `Padded` whose ``valid`` is all True.

    Attributes:
        value: The items stacked along the batch dim, first
            (`stack_items`): ``(n, *item_dims)``.
        valid: ``(n,)`` bool, True for real items, in the value's array
            namespace.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray, reduce
        >>> items = [DimArray(np.full(3, float(i)), ("x",)) for i in range(3)]
        >>> padded = next(batch(items, size=4, dim="sample"))
        >>> padded.valid.tolist()
        [True, True, True, False]
        >>> reduce(padded.value, "mean", "sample", where=padded.valid).data
        array([1., 1., 1.])
    """

    value: Any
    valid: Any

batch(items, *, size, dim, pad='last')

Group named carriers into Padded batches of size along dim.

Parameters:

Name Type Description Default
items Iterable[Any]

Named carriers with the same dims (each registered with pipekit's adapter registry), consumed lazily.

 required
size int

Items per batch; must be positive.

 required
dim str

The new leading batch dim ("sample", "patch", "member"); must not be a dim of the items.

 required
pad PadPolicy

What to do with a short last batch:

  • "last" (default): pad it to size with copies of its last real item, valid False for the copies, so every batch has one shape;
  • "none": yield it short, all valid;
  • "drop": drop it (drop_remainder in tf.data / Grain).
 'last'

Yields:

Type Description
Padded

One Padded per size items, plus the last short batch per
Padded

pad. Each value is stack_items of its items.

Raises:

Type Description
ValueError

size is not positive, pad is not a policy, or the items' dims differ (from stack_items).
Example

import numpy as np from pipekit_array import DimArray items = [DimArray(np.ones((2, 2)), ("y", "x")) for _ in range(5)] def sizes(pad): ... return [len(b.valid) for b in batch(items, size=2, dim="s", pad=pad)] sizes("last"), sizes("none"), sizes("drop") ([2, 2, 2], [2, 2, 1], [2, 2])

Source code in packages/pipekit-array/src/pipekit_array/batching.py 
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
def batch(
    items: Iterable[Any], *, size: int, dim: str, pad: PadPolicy = "last"
) -> Iterator[Padded]:
    """Group named carriers into `Padded` batches of ``size`` along ``dim``.

    Args:
        items: Named carriers with the same dims (each registered with
            pipekit's adapter registry), consumed lazily.
        size: Items per batch; must be positive.
        dim: The new leading batch dim (``"sample"``, ``"patch"``,
            ``"member"``); must not be a dim of the items.
        pad: What to do with a short last batch:

            - ``"last"`` (default): pad it to ``size`` with copies of its
              last real item, ``valid`` False for the copies, so every
              batch has one shape;
            - ``"none"``: yield it short, all valid;
            - ``"drop"``: drop it (``drop_remainder`` in tf.data / Grain).

    Yields:
        One `Padded` per ``size`` items, plus the last short batch per
        ``pad``. Each ``value`` is `stack_items` of its items.

    Raises:
        ValueError: ``size`` is not positive, ``pad`` is not a policy, or
            the items' dims differ (from `stack_items`).

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> items = [DimArray(np.ones((2, 2)), ("y", "x")) for _ in range(5)]
        >>> def sizes(pad):
        ...     return [len(b.valid) for b in batch(items, size=2, dim="s", pad=pad)]
        >>> sizes("last"), sizes("none"), sizes("drop")
        ([2, 2, 2], [2, 2, 1], [2, 2])
    """
    if size <= 0:
        raise ValueError(f"batch: size must be positive, got {size}.")
    if pad not in PAD_POLICIES:
        raise ValueError(f"batch: pad must be one of {PAD_POLICIES}, got {pad!r}.")
    group: list[Any] = []
    for item in items:
        group.append(item)
        if len(group) == size:
            yield _padded(group, dim, 0)
            group = []
    if group and pad != "drop":
        yield _padded(group, dim, size - len(group) if pad == "last" else 0)

Kernels

pipekit_array.kernel

KernelSpec and KernelOp — one L1 kernel, written once, on every carrier.

A shapes-only kernel declares the axes it consumes and produces with a generalized-ufunc signature ("(c)->()", "(r,c),(c),(c,d)->(r)"), the convention NumPy (np.vectorize), JAX (jnp.vectorize) and xarray (apply_ufunc core dims) already share. KernelOp wraps the kernel once and runs it:

raw arrays        → the kernel, called as is
named carriers    → core dims moved last by name, the kernel on the
                    data, the result rewrapped by the adapter
                    (NumPy eagerly, JAX traced under jit)
an adapter with   → that adapter's own path (the xarray adapter:
``apply_kernel``    one ``apply_ufunc``, lazy on dask)

Kernel conventions:

  1. Consumed axes last. Core dims arrive as the trailing axes.
  2. Broadcast over leading axes, or set vectorize=True and pay for np.vectorize / jnp.vectorize.
  3. Non-carrier inputs are keyword parameters (a background mean, a threshold), passed whole and never mapped. Geometry derived from the carrier (geometry=) is passed the same way; geometry that varies per pixel has the carrier's dims and is a positional input instead.
  4. Write against the array namespace (xp.linalg.solve) to run on every backend. A kernel that needs a JAX-only library (lineax) declares backends=("jax",) (spike finding F3).

Rules from the spike (geotoolz#516):

  • Outputs keep the first input's dim order (F2): a gufunc puts core dims last, so each output is transposed back, with dims new to the output last.
  • geometry= needs one geometry per call (F4): on a carrier whose adapter reports batch_dims (per-item transforms), KernelOp raises and names map_over.
  • An output-only core dim declares its size (F5): KernelSpec.sizes maps it to an input core dim, so a lazy backend can plan the output.

Carriers of different kinds combine by dim name and size (a geo carrier "band row col" with a DimArray "col band"); grid checks (check_compatible) run between carriers of the same kind.

KernelOp

Bases: Operator

An L1 kernel wrapped once, run on raw arrays and on every named carrier.

Parameters:

Name Type Description Default
kernel Callable[..., Any]

The L1 function: arrays (core dims last) and keyword params in, arrays out.

 required
spec KernelSpec | str

A KernelSpec, or its signature string.

 required
core dict[str, str]

Each core letter of the spec → the carrier dim it stands for, e.g. {"c": "band"}; the same kernel serves "band" and "variable".

 required
geometry dict[str, Callable[[Any], Any]] | None

Kernel keyword → a function of the first carrier returning arrays ({"spacing": spacing_of}), computed per call and passed whole.

 None
params Mapping[str, Any] | None

Keyword parameters passed to every call, as a mapping: the way to pass one named core, geometry or params.

 None
**kwargs Any

More keyword parameters passed to every call.

 {}

Raises:

Type Description
ValueError

core misses or invents a letter of the spec, or a parameter has the name of a derived-geometry keyword.
TypeError

geometry is not a mapping of callables (pass a kernel parameter named geometry through params=).

KernelOp holds a user function, so its config is a debug repr only (forbid_in_yaml, like pipekit_ops.Lambda); a named operator built on one by composition (geoops.NDVI) is what serialises.

Example

import numpy as np from pipekit_array import DimArray def nd(a, b, eps): ... return (a - b) / (a + b + eps) op = KernelOp(nd, "(),()->()", core={}, eps=1e-10) nir = DimArray(np.full((2, 3), 0.4), ("y", "x")) red = DimArray(np.full((2, 3), 0.1), ("y", "x")) op(nir, red).dims ('y', 'x')

Source code in packages/pipekit-array/src/pipekit_array/kernel.py 
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
class KernelOp(Operator):
    """An L1 kernel wrapped once, run on raw arrays and on every named carrier.

    Args:
        kernel: The L1 function: arrays (core dims last) and keyword params in,
            arrays out.
        spec: A `KernelSpec`, or its signature string.
        core: Each core letter of the spec → the carrier dim it stands for,
            e.g. ``{"c": "band"}``; the same kernel serves ``"band"`` and
            ``"variable"``.
        geometry: Kernel keyword → a function of the first carrier returning
            arrays (``{"spacing": spacing_of}``), computed per call and passed
            whole.
        params: Keyword parameters passed to every call, as a mapping: the
            way to pass one named ``core``, ``geometry`` or ``params``.
        **kwargs: More keyword parameters passed to every call.

    Raises:
        ValueError: ``core`` misses or invents a letter of the spec, or a
            parameter has the name of a derived-geometry keyword.
        TypeError: ``geometry`` is not a mapping of callables (pass a kernel
            parameter named ``geometry`` through ``params=``).

    ``KernelOp`` holds a user function, so its config is a debug repr only
    (``forbid_in_yaml``, like `pipekit_ops.Lambda`); a named operator built on
    one by composition (``geoops.NDVI``) is what serialises.

    Example:
        >>> import numpy as np
        >>> from pipekit_array import DimArray
        >>> def nd(a, b, eps):
        ...     return (a - b) / (a + b + eps)
        >>> op = KernelOp(nd, "(),()->()", core={}, eps=1e-10)
        >>> nir = DimArray(np.full((2, 3), 0.4), ("y", "x"))
        >>> red = DimArray(np.full((2, 3), 0.1), ("y", "x"))
        >>> op(nir, red).dims
        ('y', 'x')
    """

    forbid_in_yaml: ClassVar[bool] = True

    def __init__(
        self,
        kernel: Callable[..., Any],
        spec: KernelSpec | str,
        /,
        core: dict[str, str],
        geometry: dict[str, Callable[[Any], Any]] | None = None,
        params: Mapping[str, Any] | None = None,
        **kwargs: Any,
    ) -> None:
        self.kernel = kernel
        self.spec = spec if isinstance(spec, KernelSpec) else KernelSpec(spec)
        letters = set(self.spec.letters)
        if set(core) != letters:
            raise ValueError(
                f"core must map exactly the spec's letters {sorted(letters)}, "
                f"got {sorted(core)}."
            )
        self.core = dict(core)
        if geometry is not None and not (
            isinstance(geometry, Mapping) and all(map(callable, geometry.values()))
        ):
            raise TypeError(
                "geometry= maps kernel keywords to functions of the first carrier; "
                "pass a kernel parameter named 'geometry' as params={'geometry': ...}."
            )
        self.geometry = dict(geometry or {})
        self.params = {**dict(params or {}), **kwargs}
        clash = sorted(set(self.params) & set(self.geometry))
        if clash:
            raise ValueError(
                f"{clash} are both parameters and derived-geometry keywords; "
                "the kernel can take each name once."
            )

    def _apply(self, *carriers: Any) -> Any:
        if len(carriers) != len(self.spec.inputs):
            raise TypeError(
                f"{self.spec.signature!r} takes {len(self.spec.inputs)} input(s), "
                f"got {len(carriers)}."
            )
        derived = derive_geometry(self.geometry, carriers[0])
        kwargs = self.params | derived
        return run_kernel(self.kernel, self.spec, self.core, carriers, kwargs)

    def get_config(self) -> dict[str, Any]:
        module = getattr(self.kernel, "__module__", "?")
        name = getattr(self.kernel, "__qualname__", repr(self.kernel))
        return {
            "kernel": f"{module}:{name}",
            "spec": self.spec.signature,
            "core": self.core,
            "geometry": sorted(self.geometry),
            "params": sorted(self.params),
        }

    def compute_output_signature(self, input_signature: Any) -> Any:
        """The output signature(s) for a one-input kernel; ``None`` otherwise.

        A two-output kernel gives a tuple of signatures. With several inputs
        the other inputs' dims are unknown here, so the answer is ``None``.
        The dtype is the one ``spec.dtypes`` declares (``"same"``: the
        input's), else unknown: a kernel may return a mask or an index.
        """
        return kernel_output_signature(self.spec, self.core, input_signature)

compute_output_signature(input_signature)

The output signature(s) for a one-input kernel; None otherwise.

A two-output kernel gives a tuple of signatures. With several inputs the other inputs' dims are unknown here, so the answer is None. The dtype is the one spec.dtypes declares ("same": the input's), else unknown: a kernel may return a mask or an index.

Source code in packages/pipekit-array/src/pipekit_array/kernel.py 
572
573
574
575
576
577
578
579
580
def compute_output_signature(self, input_signature: Any) -> Any:
    """The output signature(s) for a one-input kernel; ``None`` otherwise.

    A two-output kernel gives a tuple of signatures. With several inputs
    the other inputs' dims are unknown here, so the answer is ``None``.
    The dtype is the one ``spec.dtypes`` declares (``"same"``: the
    input's), else unknown: a kernel may return a mask or an index.
    """
    return kernel_output_signature(self.spec, self.core, input_signature)

KernelSpec dataclass

A generalized-ufunc signature over abstract core-dim letters.

Parameters:

Name Type Description Default
signature str

E.g. "(),()->()" (element-wise), "(c)->()" (consumes one axis), "(r,c)->(c),(c,d)" (two outputs).

 required
backends tuple[str, ...]

The backends the kernel is written for, from "numpy", "jax", "dask", "xarray".

 ('numpy', 'jax', 'dask', 'xarray')
vectorize bool

The kernel handles only its core dims, and is mapped over the rest with np.vectorize / jnp.vectorize.

 False
sizes tuple[tuple[str, str], ...]

Output-only core letters and the input letter whose size each has, e.g. (("d", "c"),). Pairs, not a dict, so the spec stays hashable.

 ()
dtypes tuple[str, ...]

Each output's dtype, for output signatures and lazy backends: a dtype name ("bool") or "same" for the first input's. Empty (the default) leaves it unknown in signatures; dask then finds it by calling the kernel once.

 ()

Raises:

Type Description
ValueError

The signature does not parse, sizes names a letter that is not output-only or an input letter that does not exist, or dtypes is not one entry per output.
Example

spec = KernelSpec("(r,c)->(c),(c,d)", sizes=(("d", "c"),)) spec.inputs, spec.outputs ((('r', 'c'),), (('c',), ('c', 'd')))

Source code in packages/pipekit-array/src/pipekit_array/kernel.py 
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
@dataclasses.dataclass(frozen=True)
class KernelSpec:
    """A generalized-ufunc signature over abstract core-dim letters.

    Args:
        signature: E.g. ``"(),()->()"`` (element-wise), ``"(c)->()"``
            (consumes one axis), ``"(r,c)->(c),(c,d)"`` (two outputs).
        backends: The backends the kernel is written for, from ``"numpy"``,
            ``"jax"``, ``"dask"``, ``"xarray"``.
        vectorize: The kernel handles only its core dims, and is mapped over
            the rest with ``np.vectorize`` / ``jnp.vectorize``.
        sizes: Output-only core letters and the input letter whose size each
            has, e.g. ``(("d", "c"),)``. Pairs, not a dict, so the spec stays
            hashable.
        dtypes: Each output's dtype, for output signatures and lazy
            backends: a dtype name (``"bool"``) or ``"same"`` for the first
            input's. Empty (the default) leaves it unknown in signatures;
            dask then finds it by calling the kernel once.

    Raises:
        ValueError: The signature does not parse, ``sizes`` names a letter
            that is not output-only or an input letter that does not exist,
            or ``dtypes`` is not one entry per output.

    Example:
        >>> spec = KernelSpec("(r,c)->(c),(c,d)", sizes=(("d", "c"),))
        >>> spec.inputs, spec.outputs
        ((('r', 'c'),), (('c',), ('c', 'd')))
    """

    signature: str
    backends: tuple[str, ...] = ("numpy", "jax", "dask", "xarray")
    vectorize: bool = False
    sizes: tuple[tuple[str, str], ...] = ()
    dtypes: tuple[str, ...] = ()

    def __post_init__(self) -> None:
        if not _SIGNATURE.match(self.signature):
            raise ValueError(f"Not a gufunc signature: {self.signature!r}.")
        object.__setattr__(self, "backends", tuple(self.backends))
        object.__setattr__(self, "sizes", tuple(tuple(p) for p in self.sizes))
        object.__setattr__(self, "dtypes", tuple(map(str, self.dtypes)))
        if self.dtypes and len(self.dtypes) != len(self.outputs):
            raise ValueError(
                f"dtypes: {len(self.dtypes)} entries for {len(self.outputs)} output(s)."
            )
        in_letters = {c for group in self.inputs for c in group}
        out_only = {c for group in self.outputs for c in group} - in_letters
        declared = {out for out, _ in self.sizes}
        for out, source in self.sizes:
            if out not in out_only:
                raise ValueError(f"sizes: {out!r} is not an output-only letter.")
            if source not in in_letters:
                raise ValueError(f"sizes: {source!r} is not an input letter.")
        missing = sorted(out_only - declared)
        if missing:
            raise ValueError(
                f"Output-only core letter(s) {missing} need a size, e.g. "
                f"sizes=(({missing[0]!r}, '<input letter>'),)."
            )

    @functools.cached_property
    def inputs(self) -> tuple[tuple[str, ...], ...]:
        """Core letters of each input."""
        return _groups(self.signature.split("->")[0])

    @functools.cached_property
    def outputs(self) -> tuple[tuple[str, ...], ...]:
        """Core letters of each output."""
        return _groups(self.signature.split("->")[1])

    @functools.cached_property
    def letters(self) -> tuple[str, ...]:
        """Every core letter, in order of first appearance."""
        seen: dict[str, None] = {}
        for group in (*self.inputs, *self.outputs):
            for letter in group:
                seen.setdefault(letter, None)
        return tuple(seen)

inputs cached property

Core letters of each input.

letters cached property

Every core letter, in order of first appearance.

outputs cached property

Core letters of each output.

derive_geometry(geometry, first)

The kernel's derived-geometry keywords, computed from the first carrier.

Raises:

Type Description
ValueError

The carrier's adapter reports batch_dims (one geometry per item), but keyword inputs are passed whole (F4); the message names map_over.
Source code in packages/pipekit-array/src/pipekit_array/kernel.py 
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
def derive_geometry(
    geometry: dict[str, Callable[[Any], Any]], first: Any
) -> dict[str, Any]:
    """The kernel's derived-geometry keywords, computed from the first carrier.

    Raises:
        ValueError: The carrier's adapter reports ``batch_dims`` (one geometry
            per item), but keyword inputs are passed whole (F4); the message
            names ``map_over``.
    """
    if not geometry:
        return {}
    adapter = get_adapter(first)
    batch = getattr(adapter, "batch_dims", None)
    batched = batch(first) if batch is not None else ()
    if batched:
        raise ValueError(
            f"geometry= on a carrier batched over {batched} (one geometry per "
            "item), but keyword inputs are passed whole: use "
            f"map_over(x, op, {batched[0]!r}) so each item has its own."
        )
    return {name: fn(first) for name, fn in geometry.items()}

kernel_output_dims(spec, core, input_dims)

The dims of each output, from the inputs' dims, without running.

Loop dims are every input's dims minus its core dims, in order of first appearance. Output k has the loop dims plus its core dims, in the first input's order, with dims new to the output last (F2).

Source code in packages/pipekit-array/src/pipekit_array/kernel.py 
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
def kernel_output_dims(
    spec: KernelSpec, core: dict[str, str], input_dims: Sequence[tuple[str, ...]]
) -> tuple[tuple[str, ...], ...]:
    """The dims of each output, from the inputs' dims, without running.

    Loop dims are every input's dims minus its core dims, in order of first
    appearance. Output k has the loop dims plus its core dims, in the first
    input's order, with dims new to the output last (F2).
    """
    loop: list[str] = []
    for dims, group in zip(input_dims, spec.inputs, strict=True):
        cores = {core[c] for c in group}
        for d in dims:
            if d not in cores and d not in loop:
                loop.append(d)
    _check_core_not_loop(loop, [tuple(core[c] for c in g) for g in spec.inputs])
    first = tuple(input_dims[0]) if input_dims else ()
    outputs = []
    for group in spec.outputs:
        dims = (*loop, *(core[c] for c in group))
        ordered = [d for d in first if d in dims] + [d for d in dims if d not in first]
        outputs.append(tuple(ordered))
    return tuple(outputs)

kernel_output_signature(spec, core, input_signature)

KernelOp.compute_output_signature for any operator running spec.

Shared with pipekit_jax.JaxKernelOp, so a converted operator reports the same dims, sizes, dtypes and output count.

Source code in packages/pipekit-array/src/pipekit_array/kernel.py 
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
def kernel_output_signature(
    spec: KernelSpec, core: dict[str, str], input_signature: Any
) -> Any:
    """`KernelOp.compute_output_signature` for any operator running ``spec``.

    Shared with ``pipekit_jax.JaxKernelOp``, so a converted operator reports
    the same dims, sizes, dtypes and output count.
    """
    sig = signature_of(input_signature)
    if not isinstance(sig, Signature) or len(spec.inputs) != 1:
        return None
    sizes = dict(sig.dims)
    names = tuple(sizes)
    letter_size = {}
    for letter, dim in core.items():
        if dim in sizes:
            letter_size[letter] = sizes[dim]
    for out, source in spec.sizes:
        letter_size[out] = letter_size.get(source)
    dim_size = {core[k]: v for k, v in letter_size.items()} | {
        d: s for d, s in sizes.items()
    }
    dtypes = [sig.dtype if d == "same" else d for d in spec.dtypes] or [None] * len(
        spec.outputs
    )
    outs = tuple(
        Signature(dims=tuple((d, dim_size.get(d)) for d in dims), dtype=dtype)
        for dims, dtype in zip(
            kernel_output_dims(spec, core, [names]), dtypes, strict=True
        )
    )
    return outs if len(outs) > 1 else outs[0]

run_kernel(kernel, spec, core, carriers, kwargs)

Run kernel on raw arrays or named carriers; the dispatch behind KernelOp.

Raises:

Type Description
BackendError

Carriers on more than one backend, or a backend the spec does not list.
TypeError

A raw array next to named carriers.
ValueError

A core dim is missing from its input, or a dim has two sizes.
Source code in packages/pipekit-array/src/pipekit_array/kernel.py 
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
def run_kernel(
    kernel: Callable[..., Any],
    spec: KernelSpec,
    core: dict[str, str],
    carriers: Sequence[Any],
    kwargs: dict[str, Any],
) -> Any:
    """Run ``kernel`` on raw arrays or named carriers; the dispatch behind `KernelOp`.

    Raises:
        BackendError: Carriers on more than one backend, or a backend the
            spec does not list.
        TypeError: A raw array next to named carriers.
        ValueError: A core dim is missing from its input, or a dim has two
            sizes.
    """
    adapters = [get_adapter(x) for x in carriers]
    if all(a is None for a in adapters):
        backend = _require_backend({backend_name(x) for x in carriers}, spec)
        return _call(kernel, spec, kwargs, backend)(*carriers)
    for i, adapter in enumerate(adapters):
        if adapter is None:
            raise TypeError(
                f"input {i} is a raw array next to named carriers; wrap it as a "
                "pipekit_array.DimArray with its dims."
            )
    named: list[NamedArrayAdapter] = [a for a in adapters if a is not None]
    # a class attribute check: isinstance against a runtime Protocol costs ~10 µs
    own = [cast(_RunsKernels, a) for a in named if hasattr(type(a), "apply_kernel")]
    if own:
        if len({type(a) for a in named}) > 1:
            raise BackendError(
                f"{type(own[0]).__name__} carriers cannot mix with others in one "
                "kernel call; convert them to one family first."
            )
        return own[0].apply_kernel(kernel, spec, core, kwargs, *carriers)
    return _run_named(kernel, spec, core, list(carriers), named, kwargs)

xarray

pipekit_array.xarray

The xr.DataArray adapter (pipekit-array[xarray]).

Registered lazily: importing pipekit-array never imports xarray; the adapter loads the first time a DataArray reaches pipekit_ops.get_adapter (pipekit_ops.register_lazy_adapter). This module is the only place pipekit-array imports xarray.

  • Exact coordinates. XarrayAdapter.check_compatible aligns with join="exact": two cubes whose x differ by 1e-7 raise instead of inner-joining to an empty result.
  • Lazy on dask. KernelOp runs as one xr.apply_ufunc call, dask="parallelized", one task per chunk; the dim-level fast paths (reduce, argreduce, take_along, stack, concat, combine, transpose, rename, window) are native xarray calls, so nothing computes until .compute().
  • DataArray only (PK-P5): Dataset and DataTree dispatch stays in xrops.

from_xarray turns a DataArray into a DimArray, keeping the labels of dims that have a 1-D index coordinate.

XarrayAdapter

The pipekit_ops.NamedArrayAdapter for xr.DataArray, with fast paths.

Source code in packages/pipekit-array/src/pipekit_array/xarray.py 
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
class XarrayAdapter:
    """The `pipekit_ops.NamedArrayAdapter` for ``xr.DataArray``, with fast paths."""

    # -- the protocol ----------------------------------------------------

    def dims(self, x: xr.DataArray) -> tuple[str, ...]:
        """``x.dims``."""
        return tuple(str(d) for d in x.dims)

    def data(self, x: xr.DataArray) -> Any:
        """``x.data``: NumPy, dask or JAX, never computed."""
        return x.data

    def labels(self, x: xr.DataArray, dim: str) -> tuple[Any, ...] | None:
        """The values of ``dim``'s 1-D index coordinate, or None."""
        if dim in x.indexes and x.coords[dim].ndim == 1:
            values = x.coords[dim].values
            # tolist() turns datetime64[ns] into integers; keep the scalars
            if values.dtype.kind in "mM":
                return tuple(values)
            return tuple(values.tolist())
        return None

    def isel(self, x: xr.DataArray, indexers: dict[str, Any]) -> xr.DataArray:
        """``x.isel``; coordinates follow the data."""
        return x.isel(indexers)

    def rewrap(self, x: xr.DataArray, data: Any, dims: tuple[str, ...]) -> xr.DataArray:
        """A DataArray around new data, keeping coords and attrs of surviving dims.

        A coordinate is kept when all its dims survive with their sizes; a
        scalar one named after a new dim (left by ``isel``) is dropped, since
        xarray refuses a scalar variable named after a dim.
        """
        sizes = dict(zip(dims, data.shape, strict=True))
        coords = {
            name: coord
            for name, coord in x.coords.items()
            if all(sizes.get(d) == x.sizes[d] for d in coord.dims)
            and (name not in sizes or name in coord.dims)
        }
        out = xr.DataArray(data, dims=dims, coords=coords, attrs=x.attrs, name=x.name)
        # name=None lets xarray take a dask array's graph name; keep x's instead
        return out if out.name == x.name else out.rename(x.name)

    def check_compatible(self, *xs: xr.DataArray) -> tuple[xr.DataArray, ...]:
        """The inputs back if their shared coordinates are exactly equal.

        Raises:
            ValueError: Coordinates differ; xarray would otherwise inner-join.
        """
        try:
            xr.align(*xs, join="exact", copy=False)
        except ValueError as exc:
            raise ValueError(
                f"DataArray inputs are not aligned (join='exact'): {exc}"
            ) from exc
        return xs

    def signature(self, x: xr.DataArray) -> Signature:
        """Dims with their sizes, and the dtype's name."""
        dims = tuple((str(d), int(n)) for d, n in x.sizes.items())
        return Signature(dims=dims, dtype=str(x.dtype))

    def relabel(
        self, x: xr.DataArray, labels: Mapping[str, tuple[Any, ...] | None]
    ) -> xr.DataArray:
        """``assign_coords`` for a tuple, ``drop_vars`` for None (ADR A22).

        Tuple-valued labels (`pipekit_array.merge_dims`) become a 1-D object
        coordinate, one tuple per entry.

        Raises:
            ValueError: A dim is not in ``x.dims``, or a label tuple's
                length differs from its dim's size.
        """
        dims = self.dims(x)
        unknown = sorted(set(labels) - set(dims))
        if unknown:
            raise ValueError(f"relabel: unknown dim(s) {unknown}; dims are {dims}.")
        out = x.drop_vars([d for d in labels if d in x.coords])
        coords: dict[str, Any] = {}
        for dim, values in labels.items():
            if values is None:
                continue
            size = int(x.sizes[dim])
            if len(values) != size:
                raise ValueError(
                    f"{len(values)} labels for dim {dim!r} of size {size}."
                )
            coords[dim] = (dim, _label_array(values))
        return out.assign_coords(coords) if coords else out

    # -- dim-level fast paths --------------------------------------------

    def rename(self, x: xr.DataArray, names: dict[str, str]) -> xr.DataArray:
        """``x.rename``, which also renames the dim in non-index coordinates.

        The generic path (``rewrap`` then ``relabel``) would drop a
        coordinate such as ``quality(time, x)`` on a renamed dim (A22).
        """
        return x.rename(names)

    def transpose(self, x: xr.DataArray, order: tuple[str, ...]) -> xr.DataArray:
        """``x.transpose``."""
        return x.transpose(*order)

    def reduce(
        self,
        x: xr.DataArray,
        fn: str | Callable[..., Any],
        dims: tuple[str, ...],
        *,
        where: Any | None = None,
    ) -> xr.DataArray:
        """``x.reduce`` with the array-namespace function, lazy on dask.

        A name resolves to NumPy's function (``"nanmean"`` included), so NaN
        handling matches the generic path: DataArray's own ``mean`` would
        skip NaN by default. ``where`` arrives aligned to ``x``'s full shape
        (`pipekit_array.reduce` aligns it) and is passed to the function.
        """
        func = _numpy_function(fn, "reduce")
        if where is None:
            return x.reduce(func, dim=list(dims))
        return x.reduce(func, dim=list(dims), where=where)

    def argreduce(
        self, x: xr.DataArray, fn: str | Callable[..., Any], dim: str
    ) -> xr.DataArray:
        """``x.reduce`` with NumPy's ``argmax`` / ``argmin`` (or ``fn``)."""
        return x.reduce(_numpy_function(fn, "argreduce"), dim=dim)

    def take_along(self, x: xr.DataArray, index: Any, dim: str) -> xr.DataArray:
        """One entry of ``dim`` per position, lazily on dask, in ``x``'s dim order.

        An ``apply_ufunc`` over ``take_along_axis``, not a vectorized
        ``isel``: dask cannot index with a lazy indexer. Attributes are kept,
        auxiliary coordinates on ``dim`` (``quality(time, x)``) are gathered
        with the same index, and ``dim``'s own coordinate is dropped.
        """
        rest = [str(d) for d in x.dims if d != dim]
        idx = _as_dataarray(index)  # raises for an unnamed index
        index_dims = tuple(map(str, idx.dims))
        if not set(index_dims) <= set(rest):
            # the gather would broadcast over an unrelated index dim
            raise ValueError(
                f"take_along: index dims {index_dims} are not among {tuple(rest)}."
            )
        # the labels are checked on the index itself: idx drops its coordinates
        index_adapter = get_adapter(index)
        if index_adapter is not None:
            check_index_labels(self, x, index_adapter, index, index_adapter.dims(index))
        out = _gather(x, idx, dim)
        aux = {
            str(name): _gather(xr.DataArray(c.variable), idx, dim).variable
            for name, c in x.coords.items()
            if dim in c.dims and name != dim
        }
        if aux:
            out = out.assign_coords(aux)
        return out.transpose(*rest)

    def stack(self, xs: Sequence[xr.DataArray], dim: str) -> xr.DataArray:
        """A new leading dim; coordinates must be identical (never aligned)."""
        if dim in xs[0].dims:
            raise ValueError(
                f"stack: {dim!r} is already a dim of {tuple(xs[0].dims)}; use concat."
            )
        _require_same_dims(xs, "stack")
        checked = self.check_compatible(*xs)
        return xr.concat(
            checked, dim=dim, join="exact", coords="minimal", compat="identical"
        )

    def concat(self, xs: Sequence[xr.DataArray], dim: str) -> xr.DataArray:
        """Along an existing dim, joining its coordinate; others must be identical."""
        dims = _require_same_dims(xs, "concat")
        if dim not in dims:
            raise ValueError(f"concat: {dim!r} is not a dim of {dims}; use stack.")
        return xr.concat(
            list(xs), dim=dim, join="exact", coords="minimal", compat="identical"
        )

    def combine(self, fn: Callable[..., Any], *xs: xr.DataArray) -> xr.DataArray:
        """``apply_ufunc(fn, …, join="exact")``, lazy on dask."""
        checked = self.check_compatible(*xs)
        dims = tuple(checked[0].dims)
        ordered = [x.transpose(*dims) for x in checked]
        lazy = any(_is_dask(x) for x in ordered)
        return xr.apply_ufunc(
            fn,
            *ordered,
            dask="parallelized",
            join="exact",
            # dask needs the result dtype up front: read it from a one-element call
            output_dtypes=[_probe_dtype(fn, ordered)] if lazy else None,
        )

    def window(
        self, x: xr.DataArray, dim: str, size: int, *, window_dim: str, stride: int
    ) -> xr.DataArray:
        """``rolling(...).construct(window_dim)`` from the first full window: a view.

        The fill is zero, not NaN, so integers keep their dtype; the
        ``isel`` drops every filled window.
        """
        constructed = x.rolling({dim: size}).construct(
            window_dim, fill_value=np.zeros((), x.dtype)[()]
        )
        return constructed.isel({dim: slice(size - 1, None, stride)})

    # -- KernelOp ---------------------------------------------------------

    def apply_kernel(
        self,
        kernel: Callable[..., Any],
        spec: KernelSpec,
        core: dict[str, str],
        kwargs: dict[str, Any],
        *xs: xr.DataArray,
    ) -> Any:
        """A `KernelOp` call as one ``xr.apply_ufunc``: lazy, exact, in input order.

        Core dims move last; ``kwargs`` (params and derived geometry) are
        passed whole; output-only core dims take their sizes from
        ``spec.sizes``; outputs come back in the first input's dim order (F2).

        Raises:
            BackendError: The spec does not list ``"xarray"``.
        """
        if "xarray" not in spec.backends:
            raise BackendError(
                f"the kernel declares backends {spec.backends}, not 'xarray'."
            )
        checked = self.check_compatible(*xs)
        in_dims = [[core[c] for c in group] for group in spec.inputs]
        out_dims = [[core[c] for c in group] for group in spec.outputs]
        letter_size: dict[str, int] = {}
        for x, group in zip(checked, spec.inputs, strict=True):
            for letter in group:
                letter_size[letter] = int(x.sizes[core[letter]])
        output_sizes = {core[out]: letter_size[source] for out, source in spec.sizes}
        result = xr.apply_ufunc(
            kernel,
            *checked,
            input_core_dims=in_dims,
            output_core_dims=out_dims,
            kwargs=kwargs,
            vectorize=spec.vectorize,
            dask="parallelized",
            join="exact",
            output_dtypes=_output_dtypes(kernel, spec, kwargs, checked, letter_size),
            dask_gufunc_kwargs={"allow_rechunk": True, "output_sizes": output_sizes},
        )
        targets = kernel_output_dims(
            spec, core, [tuple(map(str, x.dims)) for x in checked]
        )
        if len(out_dims) > 1:
            return tuple(r.transpose(*t) for r, t in zip(result, targets, strict=True))
        return result.transpose(*targets[0])

apply_kernel(kernel, spec, core, kwargs, *xs)

A KernelOp call as one xr.apply_ufunc: lazy, exact, in input order.

Core dims move last; kwargs (params and derived geometry) are passed whole; output-only core dims take their sizes from spec.sizes; outputs come back in the first input's dim order (F2).

Raises:

Type Description
BackendError

The spec does not list "xarray".
Source code in packages/pipekit-array/src/pipekit_array/xarray.py 
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
def apply_kernel(
    self,
    kernel: Callable[..., Any],
    spec: KernelSpec,
    core: dict[str, str],
    kwargs: dict[str, Any],
    *xs: xr.DataArray,
) -> Any:
    """A `KernelOp` call as one ``xr.apply_ufunc``: lazy, exact, in input order.

    Core dims move last; ``kwargs`` (params and derived geometry) are
    passed whole; output-only core dims take their sizes from
    ``spec.sizes``; outputs come back in the first input's dim order (F2).

    Raises:
        BackendError: The spec does not list ``"xarray"``.
    """
    if "xarray" not in spec.backends:
        raise BackendError(
            f"the kernel declares backends {spec.backends}, not 'xarray'."
        )
    checked = self.check_compatible(*xs)
    in_dims = [[core[c] for c in group] for group in spec.inputs]
    out_dims = [[core[c] for c in group] for group in spec.outputs]
    letter_size: dict[str, int] = {}
    for x, group in zip(checked, spec.inputs, strict=True):
        for letter in group:
            letter_size[letter] = int(x.sizes[core[letter]])
    output_sizes = {core[out]: letter_size[source] for out, source in spec.sizes}
    result = xr.apply_ufunc(
        kernel,
        *checked,
        input_core_dims=in_dims,
        output_core_dims=out_dims,
        kwargs=kwargs,
        vectorize=spec.vectorize,
        dask="parallelized",
        join="exact",
        output_dtypes=_output_dtypes(kernel, spec, kwargs, checked, letter_size),
        dask_gufunc_kwargs={"allow_rechunk": True, "output_sizes": output_sizes},
    )
    targets = kernel_output_dims(
        spec, core, [tuple(map(str, x.dims)) for x in checked]
    )
    if len(out_dims) > 1:
        return tuple(r.transpose(*t) for r, t in zip(result, targets, strict=True))
    return result.transpose(*targets[0])

argreduce(x, fn, dim)

x.reduce with NumPy's argmax / argmin (or fn).

Source code in packages/pipekit-array/src/pipekit_array/xarray.py 
167
168
169
170
171
def argreduce(
    self, x: xr.DataArray, fn: str | Callable[..., Any], dim: str
) -> xr.DataArray:
    """``x.reduce`` with NumPy's ``argmax`` / ``argmin`` (or ``fn``)."""
    return x.reduce(_numpy_function(fn, "argreduce"), dim=dim)

check_compatible(*xs)

The inputs back if their shared coordinates are exactly equal.

Raises:

Type Description
ValueError

Coordinates differ; xarray would otherwise inner-join.
Source code in packages/pipekit-array/src/pipekit_array/xarray.py 
85
86
87
88
89
90
91
92
93
94
95
96
97
def check_compatible(self, *xs: xr.DataArray) -> tuple[xr.DataArray, ...]:
    """The inputs back if their shared coordinates are exactly equal.

    Raises:
        ValueError: Coordinates differ; xarray would otherwise inner-join.
    """
    try:
        xr.align(*xs, join="exact", copy=False)
    except ValueError as exc:
        raise ValueError(
            f"DataArray inputs are not aligned (join='exact'): {exc}"
        ) from exc
    return xs

combine(fn, *xs)

apply_ufunc(fn, …, join="exact"), lazy on dask.

Source code in packages/pipekit-array/src/pipekit_array/xarray.py 
224
225
226
227
228
229
230
231
232
233
234
235
236
237
def combine(self, fn: Callable[..., Any], *xs: xr.DataArray) -> xr.DataArray:
    """``apply_ufunc(fn, …, join="exact")``, lazy on dask."""
    checked = self.check_compatible(*xs)
    dims = tuple(checked[0].dims)
    ordered = [x.transpose(*dims) for x in checked]
    lazy = any(_is_dask(x) for x in ordered)
    return xr.apply_ufunc(
        fn,
        *ordered,
        dask="parallelized",
        join="exact",
        # dask needs the result dtype up front: read it from a one-element call
        output_dtypes=[_probe_dtype(fn, ordered)] if lazy else None,
    )

concat(xs, dim)

Along an existing dim, joining its coordinate; others must be identical.

Source code in packages/pipekit-array/src/pipekit_array/xarray.py 
215
216
217
218
219
220
221
222
def concat(self, xs: Sequence[xr.DataArray], dim: str) -> xr.DataArray:
    """Along an existing dim, joining its coordinate; others must be identical."""
    dims = _require_same_dims(xs, "concat")
    if dim not in dims:
        raise ValueError(f"concat: {dim!r} is not a dim of {dims}; use stack.")
    return xr.concat(
        list(xs), dim=dim, join="exact", coords="minimal", compat="identical"
    )

data(x)

x.data: NumPy, dask or JAX, never computed.

Source code in packages/pipekit-array/src/pipekit_array/xarray.py 
49
50
51
def data(self, x: xr.DataArray) -> Any:
    """``x.data``: NumPy, dask or JAX, never computed."""
    return x.data

dims(x)

x.dims.

Source code in packages/pipekit-array/src/pipekit_array/xarray.py 
45
46
47
def dims(self, x: xr.DataArray) -> tuple[str, ...]:
    """``x.dims``."""
    return tuple(str(d) for d in x.dims)

isel(x, indexers)

x.isel; coordinates follow the data.

Source code in packages/pipekit-array/src/pipekit_array/xarray.py 
63
64
65
def isel(self, x: xr.DataArray, indexers: dict[str, Any]) -> xr.DataArray:
    """``x.isel``; coordinates follow the data."""
    return x.isel(indexers)

labels(x, dim)

The values of dim's 1-D index coordinate, or None.

Source code in packages/pipekit-array/src/pipekit_array/xarray.py 
53
54
55
56
57
58
59
60
61
def labels(self, x: xr.DataArray, dim: str) -> tuple[Any, ...] | None:
    """The values of ``dim``'s 1-D index coordinate, or None."""
    if dim in x.indexes and x.coords[dim].ndim == 1:
        values = x.coords[dim].values
        # tolist() turns datetime64[ns] into integers; keep the scalars
        if values.dtype.kind in "mM":
            return tuple(values)
        return tuple(values.tolist())
    return None

reduce(x, fn, dims, *, where=None)

x.reduce with the array-namespace function, lazy on dask.

A name resolves to NumPy's function ("nanmean" included), so NaN handling matches the generic path: DataArray's own mean would skip NaN by default. where arrives aligned to x's full shape (pipekit_array.reduce aligns it) and is passed to the function.

Source code in packages/pipekit-array/src/pipekit_array/xarray.py 
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
def reduce(
    self,
    x: xr.DataArray,
    fn: str | Callable[..., Any],
    dims: tuple[str, ...],
    *,
    where: Any | None = None,
) -> xr.DataArray:
    """``x.reduce`` with the array-namespace function, lazy on dask.

    A name resolves to NumPy's function (``"nanmean"`` included), so NaN
    handling matches the generic path: DataArray's own ``mean`` would
    skip NaN by default. ``where`` arrives aligned to ``x``'s full shape
    (`pipekit_array.reduce` aligns it) and is passed to the function.
    """
    func = _numpy_function(fn, "reduce")
    if where is None:
        return x.reduce(func, dim=list(dims))
    return x.reduce(func, dim=list(dims), where=where)

relabel(x, labels)

assign_coords for a tuple, drop_vars for None (ADR A22).

Tuple-valued labels (pipekit_array.merge_dims) become a 1-D object coordinate, one tuple per entry.

Raises:

Type Description
ValueError

A dim is not in x.dims, or a label tuple's length differs from its dim's size.
Source code in packages/pipekit-array/src/pipekit_array/xarray.py 
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def relabel(
    self, x: xr.DataArray, labels: Mapping[str, tuple[Any, ...] | None]
) -> xr.DataArray:
    """``assign_coords`` for a tuple, ``drop_vars`` for None (ADR A22).

    Tuple-valued labels (`pipekit_array.merge_dims`) become a 1-D object
    coordinate, one tuple per entry.

    Raises:
        ValueError: A dim is not in ``x.dims``, or a label tuple's
            length differs from its dim's size.
    """
    dims = self.dims(x)
    unknown = sorted(set(labels) - set(dims))
    if unknown:
        raise ValueError(f"relabel: unknown dim(s) {unknown}; dims are {dims}.")
    out = x.drop_vars([d for d in labels if d in x.coords])
    coords: dict[str, Any] = {}
    for dim, values in labels.items():
        if values is None:
            continue
        size = int(x.sizes[dim])
        if len(values) != size:
            raise ValueError(
                f"{len(values)} labels for dim {dim!r} of size {size}."
            )
        coords[dim] = (dim, _label_array(values))
    return out.assign_coords(coords) if coords else out

rename(x, names)

x.rename, which also renames the dim in non-index coordinates.

The generic path (rewrap then relabel) would drop a coordinate such as quality(time, x) on a renamed dim (A22).

Source code in packages/pipekit-array/src/pipekit_array/xarray.py 
135
136
137
138
139
140
141
def rename(self, x: xr.DataArray, names: dict[str, str]) -> xr.DataArray:
    """``x.rename``, which also renames the dim in non-index coordinates.

    The generic path (``rewrap`` then ``relabel``) would drop a
    coordinate such as ``quality(time, x)`` on a renamed dim (A22).
    """
    return x.rename(names)

rewrap(x, data, dims)

A DataArray around new data, keeping coords and attrs of surviving dims.

A coordinate is kept when all its dims survive with their sizes; a scalar one named after a new dim (left by isel) is dropped, since xarray refuses a scalar variable named after a dim.

Source code in packages/pipekit-array/src/pipekit_array/xarray.py 
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
def rewrap(self, x: xr.DataArray, data: Any, dims: tuple[str, ...]) -> xr.DataArray:
    """A DataArray around new data, keeping coords and attrs of surviving dims.

    A coordinate is kept when all its dims survive with their sizes; a
    scalar one named after a new dim (left by ``isel``) is dropped, since
    xarray refuses a scalar variable named after a dim.
    """
    sizes = dict(zip(dims, data.shape, strict=True))
    coords = {
        name: coord
        for name, coord in x.coords.items()
        if all(sizes.get(d) == x.sizes[d] for d in coord.dims)
        and (name not in sizes or name in coord.dims)
    }
    out = xr.DataArray(data, dims=dims, coords=coords, attrs=x.attrs, name=x.name)
    # name=None lets xarray take a dask array's graph name; keep x's instead
    return out if out.name == x.name else out.rename(x.name)

signature(x)

Dims with their sizes, and the dtype's name.

Source code in packages/pipekit-array/src/pipekit_array/xarray.py 
 99
100
101
102
def signature(self, x: xr.DataArray) -> Signature:
    """Dims with their sizes, and the dtype's name."""
    dims = tuple((str(d), int(n)) for d, n in x.sizes.items())
    return Signature(dims=dims, dtype=str(x.dtype))

stack(xs, dim)

A new leading dim; coordinates must be identical (never aligned).

Source code in packages/pipekit-array/src/pipekit_array/xarray.py 
203
204
205
206
207
208
209
210
211
212
213
def stack(self, xs: Sequence[xr.DataArray], dim: str) -> xr.DataArray:
    """A new leading dim; coordinates must be identical (never aligned)."""
    if dim in xs[0].dims:
        raise ValueError(
            f"stack: {dim!r} is already a dim of {tuple(xs[0].dims)}; use concat."
        )
    _require_same_dims(xs, "stack")
    checked = self.check_compatible(*xs)
    return xr.concat(
        checked, dim=dim, join="exact", coords="minimal", compat="identical"
    )

take_along(x, index, dim)

One entry of dim per position, lazily on dask, in x's dim order.

An apply_ufunc over take_along_axis, not a vectorized isel: dask cannot index with a lazy indexer. Attributes are kept, auxiliary coordinates on dim (quality(time, x)) are gathered with the same index, and dim's own coordinate is dropped.

Source code in packages/pipekit-array/src/pipekit_array/xarray.py 
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
def take_along(self, x: xr.DataArray, index: Any, dim: str) -> xr.DataArray:
    """One entry of ``dim`` per position, lazily on dask, in ``x``'s dim order.

    An ``apply_ufunc`` over ``take_along_axis``, not a vectorized
    ``isel``: dask cannot index with a lazy indexer. Attributes are kept,
    auxiliary coordinates on ``dim`` (``quality(time, x)``) are gathered
    with the same index, and ``dim``'s own coordinate is dropped.
    """
    rest = [str(d) for d in x.dims if d != dim]
    idx = _as_dataarray(index)  # raises for an unnamed index
    index_dims = tuple(map(str, idx.dims))
    if not set(index_dims) <= set(rest):
        # the gather would broadcast over an unrelated index dim
        raise ValueError(
            f"take_along: index dims {index_dims} are not among {tuple(rest)}."
        )
    # the labels are checked on the index itself: idx drops its coordinates
    index_adapter = get_adapter(index)
    if index_adapter is not None:
        check_index_labels(self, x, index_adapter, index, index_adapter.dims(index))
    out = _gather(x, idx, dim)
    aux = {
        str(name): _gather(xr.DataArray(c.variable), idx, dim).variable
        for name, c in x.coords.items()
        if dim in c.dims and name != dim
    }
    if aux:
        out = out.assign_coords(aux)
    return out.transpose(*rest)

transpose(x, order)

x.transpose.

Source code in packages/pipekit-array/src/pipekit_array/xarray.py 
143
144
145
def transpose(self, x: xr.DataArray, order: tuple[str, ...]) -> xr.DataArray:
    """``x.transpose``."""
    return x.transpose(*order)

window(x, dim, size, *, window_dim, stride)

rolling(...).construct(window_dim) from the first full window: a view.

The fill is zero, not NaN, so integers keep their dtype; the isel drops every filled window.

Source code in packages/pipekit-array/src/pipekit_array/xarray.py 
239
240
241
242
243
244
245
246
247
248
249
250
def window(
    self, x: xr.DataArray, dim: str, size: int, *, window_dim: str, stride: int
) -> xr.DataArray:
    """``rolling(...).construct(window_dim)`` from the first full window: a view.

    The fill is zero, not NaN, so integers keep their dtype; the
    ``isel`` drops every filled window.
    """
    constructed = x.rolling({dim: size}).construct(
        window_dim, fill_value=np.zeros((), x.dtype)[()]
    )
    return constructed.isel({dim: slice(size - 1, None, stride)})

from_xarray(da)

A DimArray with the same data and dims; 1-D index coordinates become labels.

Source code in packages/pipekit-array/src/pipekit_array/xarray.py 
417
418
419
420
421
422
def from_xarray(da: xr.DataArray) -> DimArray:
    """A `DimArray` with the same data and dims; 1-D index coordinates become labels."""
    adapter = XarrayAdapter()
    dims = adapter.dims(da)
    labels = {d: adapter.labels(da, d) for d in dims}
    return DimArray(da.data, dims, {d: v for d, v in labels.items() if v is not None})

register()

Register XarrayAdapter for xr.DataArray (the lazy loader calls this).

Source code in packages/pipekit-array/src/pipekit_array/xarray.py 
425
426
427
def register() -> None:
    """Register `XarrayAdapter` for ``xr.DataArray`` (the lazy loader calls this)."""
    register_adapter(xr.DataArray, XarrayAdapter())

Test kit

The test kit ships in the wheel, for carrier libraries' own test suites (gates PK-K1 and PK-K4). pipekit's own tests run PK-K2, PK-K3 and, under -m slow, PK-K5.

pipekit_array.testing

Contract checks for named-array adapters and KernelOp (gates PK-K1, PK-K4).

A carrier library proves its adapter in its own test suite with one call, instead of copying tests::

from pipekit_array.testing import check_adapter, check_kernel_op

def test_adapter_contract():
    check_adapter(GeoArrayAdapter(), make_geoarray)

def test_ndvi_on_every_carrier():
    check_kernel_op(ndvi, [(nir_np, red_np), (nir_jax, red_jax), (nir_da, red_da)])
  • check_adapter(adapter, make_carrier) — every dim-level function (and the adapter's fast paths, which are what the functions call) on the carrier equals the generic path on a DimArray with the same data, dims and labels; the protocol methods behave. The functions that make new or renamed dims (pipekit_array.dim_shapes) must also give the same labels; those on the band labels run when the adapter has relabel (pipekit_ops.LabelledAdapter). So must the index-only transforms (pipekit_array.dim_shifts) and the windows (pipekit_array.dim_windows). So must quantile and dot (pipekit_array.dim_algebra), within rtol.
  • check_kernel_op(op, calls) — the same KernelOp on each carrier family gives equal data, equal dims, and the signature op.compute_output_signature predicts (PK-K1, PK-K4).

Both raise AssertionError naming the function and the carrier type. They need NumPy, and nothing from pytest.

check_adapter(adapter, make_carrier, *, rtol=1e-12, seed=0)

Assert that adapter and its fast paths agree with the reference adapter.

Parameters:

Name Type Description Default
adapter NamedArrayAdapter

The adapter under test; it must be the one registered for the carriers make_carrier builds.

 required
make_carrier MakeCarrier

(data, dims, labels) -> carrier, building a carrier in the DIMS layout (time band y x, band labels LABELS) from a float64 NumPy array of shape (3, 3, 4, 5). Put the data on the backend under test inside it (jnp.asarray, a dask array).

 required
rtol float

Tolerance for reductions; selections compare exactly.

 1e-12
seed int

For the random data.

 0

Raises:

Type Description
AssertionError

On the first mismatch, naming the function and the carrier type.
Source code in packages/pipekit-array/src/pipekit_array/testing.py 
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
def check_adapter(
    adapter: NamedArrayAdapter,
    make_carrier: MakeCarrier,
    *,
    rtol: float = 1e-12,
    seed: int = 0,
) -> None:
    """Assert that ``adapter`` and its fast paths agree with the reference adapter.

    Args:
        adapter: The adapter under test; it must be the one registered for
            the carriers ``make_carrier`` builds.
        make_carrier: ``(data, dims, labels) -> carrier``, building a carrier
            in the `DIMS` layout (``time band y x``, band labels `LABELS`)
            from a float64 NumPy array of shape (3, 3, 4, 5). Put the data on
            the backend under test inside it (``jnp.asarray``, a dask array).
        rtol: Tolerance for reductions; selections compare exactly.
        seed: For the random data.

    Raises:
        AssertionError: On the first mismatch, naming the function and the
            carrier type.
    """
    rng = np.random.default_rng(seed)
    data = rng.random((3, 3, 4, 5))
    carrier = make_carrier(data, DIMS, dict(LABELS))
    reference = DimArray(data, DIMS, dict(LABELS))
    name = type(carrier).__name__
    registered = get_adapter(carrier)
    # the same type: the functions dispatch through the registered instance
    assert type(registered) is type(adapter), (
        f"{name}: get_adapter returns {type(registered).__name__}, "
        f"not the adapter under test {type(adapter).__name__}"
    )
    _check_protocol(adapter, carrier, data, name)
    lazy = _dims.backend_name(adapter.data(carrier)) == "dask"
    labelled = getattr(adapter, "relabel", None) is not None
    shape_cases = _dim_shape_cases(labelled)
    shift_cases = [*_dim_shift_cases(), *_dim_window_cases()]
    algebra_cases = _dim_algebra_cases(labelled)
    shape_labels = {label for label, _ in [*shape_cases, *shift_cases, *algebra_cases]}
    for label, fn in [*_cases(), *shape_cases, *shift_cases, *algebra_cases]:
        if lazy and label.endswith("where)"):
            continue  # dask's reductions take no where= (reduce raises TypeError)
        expected = fn(reference)
        actual = fn(carrier)
        got = _dims_of(actual)
        want = _dims_of(expected)
        assert got == want, f"{label} on {name}: dims {got}, expected {want}"
        exact = not label.startswith(
            ("reduce", "map_over", "cumreduce", "coarsen", "quantile", "dot")
        )
        if label in shape_labels:
            _check_labels(actual, expected, label, name)
        a, e = values_of(actual), values_of(expected)
        if exact:
            np.testing.assert_array_equal(a, e, err_msg=f"{label} on {name}")
        else:
            np.testing.assert_allclose(a, e, rtol=rtol, err_msg=f"{label} on {name}")

check_kernel_op(op, calls, *, rtol=1e-06)

Assert that one KernelOp agrees across carrier families (PK-K1, PK-K4).

Parameters:

Name Type Description Default
op KernelOp

The operator under test.

 required
calls Sequence[Any]

The inputs of each call, one entry per carrier family: a carrier for a one-input kernel, or a tuple of inputs. The first entry is the reference.

 required
rtol float

Relative tolerance on the outputs' values.

 1e-06

Raises:

Type Description
AssertionError

Outputs differ in dims or values, or a one-input kernel's output dims differ from op.compute_output_signature.
Source code in packages/pipekit-array/src/pipekit_array/testing.py 
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
def check_kernel_op(op: KernelOp, calls: Sequence[Any], *, rtol: float = 1e-6) -> None:
    """Assert that one `KernelOp` agrees across carrier families (PK-K1, PK-K4).

    Args:
        op: The operator under test.
        calls: The inputs of each call, one entry per carrier family: a
            carrier for a one-input kernel, or a tuple of inputs. The first
            entry is the reference.
        rtol: Relative tolerance on the outputs' values.

    Raises:
        AssertionError: Outputs differ in dims or values, or a one-input
            kernel's output dims differ from ``op.compute_output_signature``.
    """
    if not calls:
        raise ValueError("check_kernel_op needs at least one call.")
    results = []
    for entry in calls:
        inputs = entry if isinstance(entry, tuple) else (entry,)
        out = op(*inputs)
        outs = out if isinstance(out, tuple) else (out,)
        results.append((inputs, outs))
    _, ref_outs = results[0]
    for inputs, outs in results:
        name = "/".join(type(x).__name__ for x in inputs)
        assert len(outs) == len(ref_outs), (
            f"{name}: {len(outs)} outputs, not {len(ref_outs)}"
        )
        for k, (o, r) in enumerate(zip(outs, ref_outs, strict=True)):
            got, want = _dims_of(o), _dims_of(r)
            assert got == want, f"output {k} on {name}: dims {got}, expected {want}"
            np.testing.assert_allclose(
                values_of(o), values_of(r), rtol=rtol, err_msg=f"output {k} on {name}"
            )
        if len(inputs) == 1:
            _check_signature(op, inputs[0], outs, name)

values_of(x)

A named carrier's data on the host, computing it if lazy.

Source code in packages/pipekit-array/src/pipekit_array/testing.py 
61
62
63
64
65
66
67
68
def values_of(x: Any) -> np.ndarray:
    """A named carrier's data on the host, computing it if lazy."""
    adapter = get_adapter(x)
    data = x if adapter is None else adapter.data(x)
    compute = getattr(data, "compute", None)
    if callable(compute) and type(data).__module__.startswith("dask."):
        data = compute()
    return np.asarray(data)