Skip to content

zarr_indexing.transform

zarr_indexing.transform

Index transforms — composable, lazy coordinate mappings.

An IndexTransform pairs an input domain (the coordinates a user sees) with a tuple of output maps (the storage coordinates those inputs map to). One output map per storage dimension. See output_map.py for the three output map types.

Key operations:

  • Indexing (transform[2:8], .oindex[idx], .vindex[idx]) — produces a new transform with a narrower input domain and adjusted output maps. No I/O occurs. This is how lazy slicing works.

  • intersect(output_domain) — restrict to storage coordinates within a region. This is chunk resolution: "which of my coordinates fall in this chunk?"

  • translate(shift) — shift all output coordinates. This makes coordinates chunk-local: "express my coordinates relative to the chunk origin."

  • compose(outer, inner) — chain two transforms. See composition.py.

The transform is the atomic unit that connects user-facing indexing to chunk-level I/O. Every Array holds a transform (identity by default). Array.lazy[...] composes a new transform lazily. Reading resolves the transform against the chunk grid via intersect + translate.

IndexTransform dataclass

A composable mapping from input coordinates to storage coordinates.

An IndexTransform has:

  • domain: an IndexDomain describing the valid input coordinates (the user-facing shape, possibly with non-zero origin).
  • output: a tuple of output maps (one per storage dimension), each describing which storage coordinates the inputs touch.

For a freshly opened array, the transform is the identity: input coordinate i maps to storage coordinate i. Indexing operations compose new transforms without I/O.

Source code in packages/zarr-indexing/src/zarr_indexing/transform.py 
 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
@dataclass(frozen=True, slots=True)
class IndexTransform:
    """A composable mapping from input coordinates to storage coordinates.

    An `IndexTransform` has:

    - `domain`: an `IndexDomain` describing the valid input coordinates
      (the user-facing shape, possibly with non-zero origin).
    - `output`: a tuple of output maps (one per storage dimension), each
      describing which storage coordinates the inputs touch.

    For a freshly opened array, the transform is the identity: input
    coordinate `i` maps to storage coordinate `i`. Indexing operations
    compose new transforms without I/O.
    """

    domain: IndexDomain
    output: tuple[OutputIndexMap, ...]

    def __post_init__(self) -> None:
        for i, m in enumerate(self.output):
            if isinstance(m, DimensionMap):
                if m.input_dimension < 0 or m.input_dimension >= self.domain.ndim:
                    raise ValueError(
                        f"output[{i}].input_dimension = {m.input_dimension} "
                        f"is out of range for input rank {self.domain.ndim}"
                    )
            elif isinstance(m, ArrayMap) and m.index_array.ndim > self.domain.ndim:
                # ArrayMap index arrays produced by indexing and chunk resolution
                # are normalized to the full input rank (an axis the array varies
                # over is full-sized, every other axis a singleton). A rank
                # *exceeding* the domain is always a bug. A rank *below* it is
                # tolerated: TensorStore-format JSON (external input) may supply a
                # lower-rank index array that broadcasts against the input domain,
                # and `_array_map_dependency_axes` treats any missing leading axes
                # as singleton dependencies.
                raise ValueError(
                    f"output[{i}].index_array has {m.index_array.ndim} dims "
                    f"but input domain has {self.domain.ndim} dims"
                )

    @property
    def input_rank(self) -> int:
        return self.domain.ndim

    @property
    def output_rank(self) -> int:
        return len(self.output)

    @classmethod
    def identity(cls, domain: IndexDomain) -> IndexTransform:
        output = tuple(DimensionMap(input_dimension=i) for i in range(domain.ndim))
        return cls(domain=domain, output=output)

    @classmethod
    def from_shape(cls, shape: tuple[int, ...]) -> IndexTransform:
        return cls.identity(IndexDomain.from_shape(shape))

    @property
    def selection_repr(self) -> str:
        """Compact domain string, e.g. `'{ [2, 8), [0, 10) }'`.

        Follows TensorStore's IndexDomain notation: each dimension shown
        as `[inclusive_min, exclusive_max)` with stride annotation if not 1.
        Constant (integer-indexed) dimensions show as a single value.
        Array-indexed dimensions show the set of selected coordinates.
        """
        parts: list[str] = []
        for m in self.output:
            if isinstance(m, ConstantMap):
                parts.append(str(m.offset))
            elif isinstance(m, DimensionMap):
                d = m.input_dimension
                lo = self.domain.inclusive_min[d]
                hi = self.domain.exclusive_max[d]
                start = m.offset + m.stride * lo
                stop = m.offset + m.stride * hi
                if m.stride == 1:
                    parts.append(f"[{start}, {stop})")
                else:
                    parts.append(f"[{start}, {stop}) step {m.stride}")
            else:
                # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap)
                storage = m.offset + m.stride * m.index_array
                n = int(storage.size)  # .size, not len(): index_array may be 0-d
                if n <= 5:
                    vals = ", ".join(str(int(v)) for v in storage.ravel())
                    parts.append("{" + vals + "}")
                else:
                    parts.append("{" + f"array({n})" + "}")
        return "{ " + ", ".join(parts) + " }"

    def __repr__(self) -> str:
        maps: list[str] = []
        for i, m in enumerate(self.output):
            if isinstance(m, ConstantMap):
                maps.append(f"out[{i}] = {m.offset}")
            elif isinstance(m, DimensionMap):
                maps.append(f"out[{i}] = {m.offset} + {m.stride} * in[{m.input_dimension}]")
            else:
                # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap)
                maps.append(f"out[{i}] = {m.offset} + {m.stride} * arr{m.index_array.shape}[in]")
        maps_str = ", ".join(maps)
        return f"IndexTransform(domain={self.domain}, {maps_str})"

    def intersect(
        self, output_domain: IndexDomain
    ) -> (
        tuple[
            IndexTransform,
            dict[int, np.ndarray[Any, np.dtype[np.intp]]]
            | np.ndarray[Any, np.dtype[np.intp]]
            | None,
        ]
        | None
    ):
        """Restrict this transform to storage coordinates within output_domain.

        Returns `(restricted_transform, out_indices)` or None if empty.

        `out_indices` carries the surviving output positions: `None` when all
        positions survive (ConstantMap/DimensionMap only), a single integer array
        for one ArrayMap (or correlated/vectorized ArrayMaps), or a dict keyed by
        output dimension for >= 2 orthogonal ArrayMaps (an outer product).
        """
        return _intersect(self, output_domain)

    def translate(self, shift: tuple[int, ...]) -> IndexTransform:
        """Shift all output coordinates by `shift`."""
        if len(shift) != self.output_rank:
            raise ValueError(f"shift must have length {self.output_rank}, got {len(shift)}")
        new_output: list[OutputIndexMap] = []
        for m, s in zip(self.output, shift, strict=True):
            if isinstance(m, ConstantMap):
                new_output.append(ConstantMap(offset=m.offset + s))
            elif isinstance(m, DimensionMap):
                new_output.append(
                    DimensionMap(
                        input_dimension=m.input_dimension,
                        offset=m.offset + s,
                        stride=m.stride,
                    )
                )
            else:
                # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap)
                new_output.append(
                    ArrayMap(
                        index_array=m.index_array,
                        offset=m.offset + s,
                        stride=m.stride,
                        input_dimension=m.input_dimension,
                    )
                )
        return IndexTransform(domain=self.domain, output=tuple(new_output))

    def __getitem__(self, selection: Any) -> IndexTransform:
        return _apply_basic_indexing(self, selection)

    def translate_domain_by(self, shift: tuple[int, ...]) -> IndexTransform:
        """Shift the *input* domain by `shift`, preserving which cells are addressed.

        TensorStore's `translate_by`: the domain moves, and every output map is
        re-offset so that new coordinate `c` addresses the cell that `c - shift`
        addressed before. ArrayMaps are indexed positionally over the domain, so
        their index arrays are unchanged.
        """
        if len(shift) != self.input_rank:
            raise ValueError(f"shift must have length {self.input_rank}, got {len(shift)}")
        new_domain = self.domain.translate(shift)
        new_output: list[OutputIndexMap] = []
        for m in self.output:
            if isinstance(m, DimensionMap):
                s = shift[m.input_dimension]
                new_output.append(
                    DimensionMap(
                        input_dimension=m.input_dimension,
                        offset=m.offset - m.stride * s,
                        stride=m.stride,
                    )
                )
            else:
                # ConstantMap: no input dependence. ArrayMap: positional over
                # the domain, invariant under domain translation.
                new_output.append(m)
        return IndexTransform(domain=new_domain, output=tuple(new_output))

    def translate_domain_to(self, origins: tuple[int, ...]) -> IndexTransform:
        """Move the input domain so its per-dimension origins equal `origins`.

        TensorStore's `translate_to`; `translate_domain_to((0,) * rank)`
        re-zeros a view's coordinate system without changing which cells it
        addresses.
        """
        if len(origins) != self.input_rank:
            raise ValueError(f"origins must have length {self.input_rank}, got {len(origins)}")
        shift = tuple(o - m for o, m in zip(origins, self.domain.inclusive_min, strict=True))
        return self.translate_domain_by(shift)

    @property
    def oindex(self) -> _OIndexHelper:
        return _OIndexHelper(self)

    @property
    def vindex(self) -> _VIndexHelper:
        return _VIndexHelper(self)

domain instance-attribute 

domain: IndexDomain

input_rank property 

input_rank: int

oindex property 

oindex: _OIndexHelper

output instance-attribute 

output: tuple[OutputIndexMap, ...]

output_rank property 

output_rank: int

selection_repr property 

selection_repr: str

Compact domain string, e.g. '{ [2, 8), [0, 10) }'.

Follows TensorStore's IndexDomain notation: each dimension shown as [inclusive_min, exclusive_max) with stride annotation if not 1. Constant (integer-indexed) dimensions show as a single value. Array-indexed dimensions show the set of selected coordinates.

vindex property 

vindex: _VIndexHelper

__getitem__ 

__getitem__(selection: Any) -> IndexTransform
Source code in packages/zarr-indexing/src/zarr_indexing/transform.py 
197
198
def __getitem__(self, selection: Any) -> IndexTransform:
    return _apply_basic_indexing(self, selection)

__init__ 

__init__(
    domain: IndexDomain, output: tuple[OutputIndexMap, ...]
) -> None

__post_init__ 

__post_init__() -> None
Source code in packages/zarr-indexing/src/zarr_indexing/transform.py 
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
def __post_init__(self) -> None:
    for i, m in enumerate(self.output):
        if isinstance(m, DimensionMap):
            if m.input_dimension < 0 or m.input_dimension >= self.domain.ndim:
                raise ValueError(
                    f"output[{i}].input_dimension = {m.input_dimension} "
                    f"is out of range for input rank {self.domain.ndim}"
                )
        elif isinstance(m, ArrayMap) and m.index_array.ndim > self.domain.ndim:
            # ArrayMap index arrays produced by indexing and chunk resolution
            # are normalized to the full input rank (an axis the array varies
            # over is full-sized, every other axis a singleton). A rank
            # *exceeding* the domain is always a bug. A rank *below* it is
            # tolerated: TensorStore-format JSON (external input) may supply a
            # lower-rank index array that broadcasts against the input domain,
            # and `_array_map_dependency_axes` treats any missing leading axes
            # as singleton dependencies.
            raise ValueError(
                f"output[{i}].index_array has {m.index_array.ndim} dims "
                f"but input domain has {self.domain.ndim} dims"
            )

__repr__ 

__repr__() -> str
Source code in packages/zarr-indexing/src/zarr_indexing/transform.py 
134
135
136
137
138
139
140
141
142
143
144
145
def __repr__(self) -> str:
    maps: list[str] = []
    for i, m in enumerate(self.output):
        if isinstance(m, ConstantMap):
            maps.append(f"out[{i}] = {m.offset}")
        elif isinstance(m, DimensionMap):
            maps.append(f"out[{i}] = {m.offset} + {m.stride} * in[{m.input_dimension}]")
        else:
            # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap)
            maps.append(f"out[{i}] = {m.offset} + {m.stride} * arr{m.index_array.shape}[in]")
    maps_str = ", ".join(maps)
    return f"IndexTransform(domain={self.domain}, {maps_str})"

from_shape classmethod 

from_shape(shape: tuple[int, ...]) -> IndexTransform
Source code in packages/zarr-indexing/src/zarr_indexing/transform.py 
96
97
98
@classmethod
def from_shape(cls, shape: tuple[int, ...]) -> IndexTransform:
    return cls.identity(IndexDomain.from_shape(shape))

identity classmethod 

identity(domain: IndexDomain) -> IndexTransform
Source code in packages/zarr-indexing/src/zarr_indexing/transform.py 
91
92
93
94
@classmethod
def identity(cls, domain: IndexDomain) -> IndexTransform:
    output = tuple(DimensionMap(input_dimension=i) for i in range(domain.ndim))
    return cls(domain=domain, output=output)

intersect 

intersect(
    output_domain: IndexDomain,
) -> (
    tuple[
        IndexTransform,
        dict[int, ndarray[Any, dtype[intp]]]
        | ndarray[Any, dtype[intp]]
        | None,
    ]
    | None
)

Restrict this transform to storage coordinates within output_domain.

Returns (restricted_transform, out_indices) or None if empty.

out_indices carries the surviving output positions: None when all positions survive (ConstantMap/DimensionMap only), a single integer array for one ArrayMap (or correlated/vectorized ArrayMaps), or a dict keyed by output dimension for >= 2 orthogonal ArrayMaps (an outer product).

Source code in packages/zarr-indexing/src/zarr_indexing/transform.py 
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
def intersect(
    self, output_domain: IndexDomain
) -> (
    tuple[
        IndexTransform,
        dict[int, np.ndarray[Any, np.dtype[np.intp]]]
        | np.ndarray[Any, np.dtype[np.intp]]
        | None,
    ]
    | None
):
    """Restrict this transform to storage coordinates within output_domain.

    Returns `(restricted_transform, out_indices)` or None if empty.

    `out_indices` carries the surviving output positions: `None` when all
    positions survive (ConstantMap/DimensionMap only), a single integer array
    for one ArrayMap (or correlated/vectorized ArrayMaps), or a dict keyed by
    output dimension for >= 2 orthogonal ArrayMaps (an outer product).
    """
    return _intersect(self, output_domain)

translate 

translate(shift: tuple[int, ...]) -> IndexTransform

Shift all output coordinates by shift.

Source code in packages/zarr-indexing/src/zarr_indexing/transform.py 
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
def translate(self, shift: tuple[int, ...]) -> IndexTransform:
    """Shift all output coordinates by `shift`."""
    if len(shift) != self.output_rank:
        raise ValueError(f"shift must have length {self.output_rank}, got {len(shift)}")
    new_output: list[OutputIndexMap] = []
    for m, s in zip(self.output, shift, strict=True):
        if isinstance(m, ConstantMap):
            new_output.append(ConstantMap(offset=m.offset + s))
        elif isinstance(m, DimensionMap):
            new_output.append(
                DimensionMap(
                    input_dimension=m.input_dimension,
                    offset=m.offset + s,
                    stride=m.stride,
                )
            )
        else:
            # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap)
            new_output.append(
                ArrayMap(
                    index_array=m.index_array,
                    offset=m.offset + s,
                    stride=m.stride,
                    input_dimension=m.input_dimension,
                )
            )
    return IndexTransform(domain=self.domain, output=tuple(new_output))

translate_domain_by 

translate_domain_by(
    shift: tuple[int, ...],
) -> IndexTransform

Shift the input domain by shift, preserving which cells are addressed.

TensorStore's translate_by: the domain moves, and every output map is re-offset so that new coordinate c addresses the cell that c - shift addressed before. ArrayMaps are indexed positionally over the domain, so their index arrays are unchanged.

Source code in packages/zarr-indexing/src/zarr_indexing/transform.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
def translate_domain_by(self, shift: tuple[int, ...]) -> IndexTransform:
    """Shift the *input* domain by `shift`, preserving which cells are addressed.

    TensorStore's `translate_by`: the domain moves, and every output map is
    re-offset so that new coordinate `c` addresses the cell that `c - shift`
    addressed before. ArrayMaps are indexed positionally over the domain, so
    their index arrays are unchanged.
    """
    if len(shift) != self.input_rank:
        raise ValueError(f"shift must have length {self.input_rank}, got {len(shift)}")
    new_domain = self.domain.translate(shift)
    new_output: list[OutputIndexMap] = []
    for m in self.output:
        if isinstance(m, DimensionMap):
            s = shift[m.input_dimension]
            new_output.append(
                DimensionMap(
                    input_dimension=m.input_dimension,
                    offset=m.offset - m.stride * s,
                    stride=m.stride,
                )
            )
        else:
            # ConstantMap: no input dependence. ArrayMap: positional over
            # the domain, invariant under domain translation.
            new_output.append(m)
    return IndexTransform(domain=new_domain, output=tuple(new_output))

translate_domain_to 

translate_domain_to(
    origins: tuple[int, ...],
) -> IndexTransform

Move the input domain so its per-dimension origins equal origins.

TensorStore's translate_to; translate_domain_to((0,) * rank) re-zeros a view's coordinate system without changing which cells it addresses.

Source code in packages/zarr-indexing/src/zarr_indexing/transform.py 
228
229
230
231
232
233
234
235
236
237
238
def translate_domain_to(self, origins: tuple[int, ...]) -> IndexTransform:
    """Move the input domain so its per-dimension origins equal `origins`.

    TensorStore's `translate_to`; `translate_domain_to((0,) * rank)`
    re-zeros a view's coordinate system without changing which cells it
    addresses.
    """
    if len(origins) != self.input_rank:
        raise ValueError(f"origins must have length {self.input_rank}, got {len(origins)}")
    shift = tuple(o - m for o, m in zip(origins, self.domain.inclusive_min, strict=True))
    return self.translate_domain_by(shift)

selection_to_transform 

selection_to_transform(
    selection: Any,
    transform: IndexTransform,
    mode: Literal["basic", "orthogonal", "vectorized"],
) -> IndexTransform

Convert a user selection into a composed IndexTransform.

Negative indices are treated as literal coordinates (TensorStore convention). The caller (Array layer) is responsible for converting numpy-style negative indices before calling this function.

Source code in packages/zarr-indexing/src/zarr_indexing/transform.py 
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
def selection_to_transform(
    selection: Any,
    transform: IndexTransform,
    mode: Literal["basic", "orthogonal", "vectorized"],
) -> IndexTransform:
    """Convert a user selection into a composed IndexTransform.

    Negative indices are treated as literal coordinates (TensorStore convention).
    The caller (Array layer) is responsible for converting numpy-style negative
    indices before calling this function.
    """
    if mode == "basic":
        _validate_basic_selection(selection)
        return transform[selection]
    elif mode == "orthogonal":
        _validate_array_selection(selection, transform.domain.shape, mode)
        return transform.oindex[selection]
    elif mode == "vectorized":
        _validate_array_selection(selection, transform.domain.shape, mode)
        return transform.vindex[selection]
    else:
        raise ValueError(f"Unknown mode: {mode!r}")