zarr_indexing.chunk_resolution

zarr_indexing.chunk_resolution

Chunk resolution — mapping transforms to chunk-level I/O.

Given an IndexTransform (which coordinates a request reads) and one grid per storage dimension (how storage is divided into chunks), chunk resolution answers:

For each chunk, which storage coordinates does this transform touch,
and where do those values land in the request?

The public result is a lazy, reusable ChunkPlan whose rows are ChunkProjections. Each identifies a chunk and pairs a chunk-local transform with a transform back to the request's cells, over one shared zero-origin cell domain, without assuming NumPy selectors, a codec pipeline, or a scheduler.

The plan is computed in factored form, the GridPartition. Restricting a transform to a chunk box distributes over output dimensions whenever each output map reads its own input axis — every basic and orthogonal selection — so each axis is resolved once against its grid into a table:

• StridedSet — a ConstantMap or DimensionMap axis: one row per touched chunk, holding the chunk-local start, the extent, the request position of the first cell, and whether the row covers its chunk exactly once.
• IndexedSet — an orthogonal ArrayMap axis: its coordinates grouped by chunk in CSR form, with the request positions they fill.
• JointSet — one connected index-array component. Arrays sharing input axes are grouped together; independent components remain separate tables.

A projection is one row of each table combined. Building the tables costs the sum of the touched chunks per axis rather than their product, rows are materialized only on request, and a consumer may read the tables directly instead. Two output maps that read one input axis through a DimensionMap (a diagonal, which no selection produces) have no factored form and are rejected with ValueError; a correlated index array varying over an axis a DimensionMap also reads is rejected with NotImplementedError.

ChunkCoverage

ChunkCoverage = Literal['full', 'partial', 'unknown']

ChunkPlan dataclass

A reusable, lazy partition of an index transform over a chunk grid.

Construct plans with plan_chunks. The plan's factored form, a GridPartition, is built on first use and memoized; iterating the plan or projections() materializes fresh ChunkProjection rows from it.

Examples:

Row 1 of a (3, 4) array with (2, 2) chunks crosses two chunks, and the plan can be walked again after it is exhausted:

>>> from zarr_indexing import IndexTransform
>>> from zarr_indexing.grid import dimension_grids_from_chunks
>>> grids = dimension_grids_from_chunks((2, 2), shape=(3, 4))
>>> plan = plan_chunks(IndexTransform.from_shape((3, 4))[1, :], grids)
>>> [p.chunk_coords for p in plan]
[(0, 0), (0, 1)]
>>> [p.chunk_coords for p in plan.projections()]
[(0, 0), (0, 1)]

Source code in src/zarr_indexing/chunk_resolution.py

@dataclass(frozen=True, slots=True)
class ChunkPlan:
    """A reusable, lazy partition of an index transform over a chunk grid.

    Construct plans with `plan_chunks`. The plan's factored form, a
    `GridPartition`, is built on first use and memoized; iterating the plan
    or `projections()` materializes fresh `ChunkProjection` rows from it.

    Examples
    --------
    Row 1 of a `(3, 4)` array with `(2, 2)` chunks crosses two chunks, and
    the plan can be walked again after it is exhausted:

    >>> from zarr_indexing import IndexTransform
    >>> from zarr_indexing.grid import dimension_grids_from_chunks
    >>> grids = dimension_grids_from_chunks((2, 2), shape=(3, 4))
    >>> plan = plan_chunks(IndexTransform.from_shape((3, 4))[1, :], grids)
    >>> [p.chunk_coords for p in plan]
    [(0, 0), (0, 1)]
    >>> [p.chunk_coords for p in plan.projections()]
    [(0, 0), (0, 1)]
    """

    transform: IndexTransform
    """The composed request this plan partitions."""

    dimension_grids: tuple[DimensionGridLike, ...]
    """One grid per storage dimension, defining the chunk layout the plan walks."""

    # Memoized factored form; derived state, excluded from identity (see
    # `IndexDomain._shape` for the same pattern).
    _partition: GridPartition | None = field(default=None, init=False, repr=False, compare=False)

    def partition(self) -> GridPartition:
        """The plan in factored, columnar form: one table per axis plus connected index-array tables.

        Built once per plan and memoized. Raises `ValueError` if two output
        maps read one input axis through a `DimensionMap` (a diagonal, which
        no selection produces): its tables are per output dimension, and a
        diagonal needs a strided set spanning several.
        """
        cached = self._partition
        if cached is None:
            cached = _partition_transform(self.transform, self.dimension_grids)
            object.__setattr__(self, "_partition", cached)
        return cached

    def projections(self) -> Iterator[ChunkProjection]:
        """Return a fresh iterator over the chunks touched by this plan."""
        return iter(self.partition())

    def __iter__(self) -> Iterator[ChunkProjection]:
        """Equivalent to `projections()`."""
        return self.projections()

dimension_grids instance-attribute

dimension_grids: tuple[DimensionGridLike, ...]

One grid per storage dimension, defining the chunk layout the plan walks.

transform instance-attribute

transform: IndexTransform

The composed request this plan partitions.

__init__

__init__(
    transform: IndexTransform,
    dimension_grids: tuple[DimensionGridLike, ...],
) -> None

__iter__

__iter__() -> Iterator[ChunkProjection]

Equivalent to projections().

Source code in src/zarr_indexing/chunk_resolution.py

def __iter__(self) -> Iterator[ChunkProjection]:
    """Equivalent to `projections()`."""
    return self.projections()

partition

partition() -> GridPartition

The plan in factored, columnar form: one table per axis plus connected index-array tables.

Built once per plan and memoized. Raises ValueError if two output maps read one input axis through a DimensionMap (a diagonal, which no selection produces): its tables are per output dimension, and a diagonal needs a strided set spanning several.

Source code in src/zarr_indexing/chunk_resolution.py

def partition(self) -> GridPartition:
    """The plan in factored, columnar form: one table per axis plus connected index-array tables.

    Built once per plan and memoized. Raises `ValueError` if two output
    maps read one input axis through a `DimensionMap` (a diagonal, which
    no selection produces): its tables are per output dimension, and a
    diagonal needs a strided set spanning several.
    """
    cached = self._partition
    if cached is None:
        cached = _partition_transform(self.transform, self.dimension_grids)
        object.__setattr__(self, "_partition", cached)
    return cached

projections

projections() -> Iterator[ChunkProjection]

Return a fresh iterator over the chunks touched by this plan.

Source code in src/zarr_indexing/chunk_resolution.py

def projections(self) -> Iterator[ChunkProjection]:
    """Return a fresh iterator over the chunks touched by this plan."""
    return iter(self.partition())

ChunkProjection dataclass

One source-independent projection of a request through a chunk.

Both transforms share a synthetic input domain. chunk_transform maps that domain to chunk-local storage coordinates; cell_transform maps it to the original request domain.

Attributes:

• chunk_coords (tuple[int, ...]) –

  Coordinates of the selected cell in the caller's grid.

• chunk_domain (IndexDomain) –

  Bounds of that grid cell in global storage coordinates.

• chunk_transform (IndexTransform) –

  Mapping from the shared synthetic domain to chunk-local storage.

• cell_transform (IndexTransform) –

  Mapping from the shared synthetic domain to request coordinates.

• coverage (ChunkCoverage) –

  Whether the request is proven to cover the whole grid cell exactly once. Fancy selections are conservatively "unknown".

Examples:

Row 1 of a (3, 4) array with (2, 2) chunks touches only part of the first chunk, whose domain spans rows [0, 2) and columns [0, 2):

>>> from zarr_indexing import IndexTransform
>>> from zarr_indexing.grid import dimension_grids_from_chunks
>>> grids = dimension_grids_from_chunks((2, 2), shape=(3, 4))
>>> plan = plan_chunks(IndexTransform.from_shape((3, 4))[1, :], grids)
>>> first = next(iter(plan))
>>> first.chunk_coords
(0, 0)
>>> first.chunk_domain.shape
(2, 2)
>>> first.coverage
'partial'

Source code in src/zarr_indexing/chunk_resolution.py

@dataclass(frozen=True, slots=True)
class ChunkProjection:
    """One source-independent projection of a request through a chunk.

    Both transforms share a synthetic input domain. ``chunk_transform`` maps
    that domain to chunk-local storage coordinates; ``cell_transform`` maps it
    to the original request domain.

    Attributes
    ----------
    chunk_coords
        Coordinates of the selected cell in the caller's grid.
    chunk_domain
        Bounds of that grid cell in global storage coordinates.
    chunk_transform
        Mapping from the shared synthetic domain to chunk-local storage.
    cell_transform
        Mapping from the shared synthetic domain to request coordinates.
    coverage
        Whether the request is proven to cover the whole grid cell exactly
        once. Fancy selections are conservatively ``"unknown"``.

    Examples
    --------
    Row 1 of a `(3, 4)` array with `(2, 2)` chunks touches only part of the
    first chunk, whose domain spans rows `[0, 2)` and columns `[0, 2)`:

    >>> from zarr_indexing import IndexTransform
    >>> from zarr_indexing.grid import dimension_grids_from_chunks
    >>> grids = dimension_grids_from_chunks((2, 2), shape=(3, 4))
    >>> plan = plan_chunks(IndexTransform.from_shape((3, 4))[1, :], grids)
    >>> first = next(iter(plan))
    >>> first.chunk_coords
    (0, 0)
    >>> first.chunk_domain.shape
    (2, 2)
    >>> first.coverage
    'partial'
    """

    chunk_coords: tuple[int, ...]
    chunk_domain: IndexDomain
    chunk_transform: IndexTransform
    cell_transform: IndexTransform
    coverage: ChunkCoverage

    def __post_init__(self) -> None:
        if self.chunk_transform.domain != self.cell_transform.domain:
            raise ValueError(
                "chunk_transform and cell_transform must share an input domain; "
                f"got {self.chunk_transform.domain!r} and {self.cell_transform.domain!r}"
            )

cell_transform instance-attribute

cell_transform: IndexTransform

chunk_coords instance-attribute

chunk_coords: tuple[int, ...]

chunk_domain instance-attribute

chunk_domain: IndexDomain

chunk_transform instance-attribute

chunk_transform: IndexTransform

coverage instance-attribute

coverage: ChunkCoverage

__init__

__init__(
    chunk_coords: tuple[int, ...],
    chunk_domain: IndexDomain,
    chunk_transform: IndexTransform,
    cell_transform: IndexTransform,
    coverage: ChunkCoverage,
) -> None

__post_init__

__post_init__() -> None

Source code in src/zarr_indexing/chunk_resolution.py

def __post_init__(self) -> None:
    if self.chunk_transform.domain != self.cell_transform.domain:
        raise ValueError(
            "chunk_transform and cell_transform must share an input domain; "
            f"got {self.chunk_transform.domain!r} and {self.cell_transform.domain!r}"
        )

GridPartition dataclass

A plan in factored form: per-axis tables whose product is the chunk walk.

sets holds one StridedSet or IndexedSet per output dimension the transform reads independently, in output-dimension order; joint_sets holds the connected index-array components. A projection is one row of each table, so the partition has prod(row_shape) rows, walked in row-major order over row_shape (the component tables after sets). Rows are materialized into ChunkProjection objects only on request; a vectorized consumer can read the tables directly.

Take one from ChunkPlan.partition.

Examples:

arr[1:6:2, 5:] on a (7, 9) array with (3, 4) chunks touches two chunks along each axis, so the partition has four rows:

>>> from zarr_indexing import IndexTransform, plan_chunks
>>> from zarr_indexing.grid import dimension_grids_from_chunks
>>> grids = dimension_grids_from_chunks((3, 4), shape=(7, 9))
>>> partition = plan_chunks(IndexTransform.from_shape((7, 9))[1:6:2, 5:], grids).partition()
>>> partition.row_shape, len(partition)
((2, 2), 4)
>>> partition.chunk_coords().tolist()
[[0, 1], [0, 2], [1, 1], [1, 2]]
>>> [projection.chunk_transform.selection_repr for projection in partition][3]
'{ [0, 4) step 2, [0, 1) }'

Source code in src/zarr_indexing/chunk_resolution.py

@dataclass(frozen=True, slots=True)
class GridPartition:
    """A plan in factored form: per-axis tables whose product is the chunk walk.

    `sets` holds one `StridedSet` or `IndexedSet` per output dimension the
    transform reads independently, in output-dimension order; `joint_sets`
    holds the connected index-array components. A projection is one row of each
    table, so the partition has ``prod(row_shape)`` rows, walked
    in row-major order over `row_shape` (the component tables after `sets`). Rows are materialized into
    `ChunkProjection` objects only on request; a vectorized consumer can read
    the tables directly.

    Take one from `ChunkPlan.partition`.

    Examples
    --------
    `arr[1:6:2, 5:]` on a `(7, 9)` array with `(3, 4)` chunks touches two
    chunks along each axis, so the partition has four rows:

    >>> from zarr_indexing import IndexTransform, plan_chunks
    >>> from zarr_indexing.grid import dimension_grids_from_chunks
    >>> grids = dimension_grids_from_chunks((3, 4), shape=(7, 9))
    >>> partition = plan_chunks(IndexTransform.from_shape((7, 9))[1:6:2, 5:], grids).partition()
    >>> partition.row_shape, len(partition)
    ((2, 2), 4)
    >>> partition.chunk_coords().tolist()
    [[0, 1], [0, 2], [1, 1], [1, 2]]
    >>> [projection.chunk_transform.selection_repr for projection in partition][3]
    '{ [0, 4) step 2, [0, 1) }'
    """

    transform: IndexTransform
    """The transform this partition factors."""

    dimension_grids: tuple[DimensionGridLike, ...]
    """One grid per storage dimension."""

    sets: tuple[StridedSet | IndexedSet, ...]
    """Independent per-axis tables, in output-dimension order."""

    joint_sets: tuple[JointSet, ...]
    """Connected index-array components, ordered by their first output dimension."""

    row_shape: tuple[int, ...]
    """Rows per table: `sets` followed by `joint_sets`, in row-major iteration order."""

    def __len__(self) -> int:
        return math.prod(self.row_shape)

    def chunk_coords(self) -> np.ndarray[Any, np.dtype[np.intp]]:
        """Chunk coordinates of every row, shape ``(len(self), output rank)``, without materializing rows."""
        count = len(self)
        out = np.empty((count, self.transform.output_rank), dtype=np.intp)
        if count == 0 or not self.row_shape:
            return out
        indices = np.unravel_index(np.arange(count, dtype=np.intp), self.row_shape)
        for axis, rows in zip(self.sets, indices, strict=False):
            out[:, axis.output_dimension] = axis.chunk[rows]
        for joint, rows in zip(self.joint_sets, indices[len(self.sets) :], strict=True):
            out[:, list(joint.output_dimensions)] = joint.chunk[rows]
        return out

    def __iter__(self) -> Iterator[ChunkProjection]:
        """Materialize every row in order."""
        if math.prod(self.row_shape) == 0:
            return
        if not self.joint_sets:
            yield from self._iter_factorized()
        else:
            yield from self._iter_correlated()

    # -- factorized assembly ------------------------------------------------

    def _pieces(self, axis: StridedSet | IndexedSet, row: int) -> _AxisPiece:
        domain = self.transform.domain
        if isinstance(axis, StridedSet):
            origin = (
                0 if axis.input_dimension is None else domain.inclusive_min[axis.input_dimension]
            )
            return _strided_piece(axis, row, origin)
        return _indexed_piece(axis, row, domain.ndim, domain.inclusive_min[axis.input_dimension])

    def _unbound(self) -> tuple[list[OutputIndexMap | None], bool]:
        """Cell maps for request axes no output reads, and whether they permit full coverage.

        A whole-chunk cover must biject onto the chunk, so an unread axis can
        only be a singleton.
        """
        domain = self.transform.domain
        bound = {axis.input_dimension for axis in self.sets if axis.input_dimension is not None}
        cell_maps: list[OutputIndexMap | None] = [None] * domain.ndim
        unbound_ok = True
        for axis in range(domain.ndim):
            if axis not in bound:
                cell_maps[axis] = DimensionMap(
                    input_dimension=axis, offset=domain.inclusive_min[axis]
                )
                unbound_ok = unbound_ok and domain.shape[axis] <= 1
        return cell_maps, unbound_ok

    def _iter_factorized(self) -> Iterator[ChunkProjection]:
        pieces_per_set = [
            [self._pieces(axis, row) for row in range(len(axis))] for axis in self.sets
        ]
        base_cell_maps, unbound_ok = self._unbound()
        for combo in itertools.product(*pieces_per_set):
            yield self._factorized_projection(list(combo), base_cell_maps, unbound_ok)

    def _factorized_projection(
        self,
        pieces: list[_AxisPiece],
        base_cell_maps: list[OutputIndexMap | None],
        unbound_ok: bool,
    ) -> ChunkProjection:
        domain = self.transform.domain
        shape = list(domain.shape)
        cell_maps = list(base_cell_maps)
        chunk_maps: list[OutputIndexMap] = []
        chunk_coords: list[int] = []
        chunk_min: list[int] = []
        chunk_max: list[int] = []
        has_array = False
        full = unbound_ok
        for c, c_start, c_extent, k, extent, chunk_map, cell_map, piece_full in pieces:
            chunk_coords.append(c)
            chunk_min.append(c_start)
            chunk_max.append(c_start + c_extent)
            chunk_maps.append(chunk_map)
            if isinstance(chunk_map, ArrayMap):
                has_array = True
            if k is not None:
                shape[k] = extent
                cell_maps[k] = cell_map
            full = full and piece_full
        synthetic = IndexDomain._unchecked((0,) * domain.ndim, tuple(shape))  # pyright: ignore[reportPrivateUsage]
        if has_array:
            coverage: ChunkCoverage = "unknown"
        elif full:
            coverage = "full"
        else:
            coverage = "partial"
        return ChunkProjection(
            chunk_coords=tuple(chunk_coords),
            chunk_domain=IndexDomain._unchecked(tuple(chunk_min), tuple(chunk_max)),  # pyright: ignore[reportPrivateUsage]
            chunk_transform=IndexTransform._unchecked(synthetic, tuple(chunk_maps)),  # pyright: ignore[reportPrivateUsage]
            cell_transform=IndexTransform._unchecked(  # pyright: ignore[reportPrivateUsage]
                synthetic, tuple(cast("list[OutputIndexMap]", cell_maps))
            ),
            coverage=coverage,
        )

    # -- correlated assembly ------------------------------------------------

    def _slots(self) -> tuple[tuple[int | None, ...], dict[int, int]]:
        """One synthetic axis per nonconstant component, followed by residual axes."""
        component_slots: list[int | None] = []
        bound: set[int] = set()
        n_lead = 0
        for joint in self.joint_sets:
            bound.update(joint.broadcast_axes)
            component_slots.append(n_lead if joint.broadcast_axes else None)
            n_lead += bool(joint.broadcast_axes)
        residual_axes = [axis for axis in range(self.transform.input_rank) if axis not in bound]
        return tuple(component_slots), {
            axis: n_lead + slot for slot, axis in enumerate(residual_axes)
        }

    def _iter_correlated(self) -> Iterator[ChunkProjection]:
        component_slots, slot_of = self._slots()
        lo_all = self.transform.domain.inclusive_min
        pieces_per_set = [
            [
                _strided_piece(
                    cast("StridedSet", axis),
                    row,
                    0 if axis.input_dimension is None else lo_all[axis.input_dimension],
                    None if axis.input_dimension is None else slot_of[axis.input_dimension],
                )
                for row in range(len(axis))
            ]
            for axis in self.sets
        ]
        for residual in itertools.product(*pieces_per_set):
            for rows in itertools.product(*(range(len(joint)) for joint in self.joint_sets)):
                yield self._correlated_projection(list(residual), rows, component_slots, slot_of)

    def _correlated_projection(
        self,
        residual: list[_AxisPiece],
        rows: tuple[int, ...],
        component_slots: tuple[int | None, ...],
        slot_of: dict[int, int],
    ) -> ChunkProjection:
        transform = self.transform
        domain = transform.domain
        lo_all = domain.inclusive_min
        output_rank = transform.output_rank
        runs = tuple(joint.run(row) for joint, row in zip(self.joint_sets, rows, strict=True))
        points_shape = tuple(
            run.stop - run.start
            for run, slot in zip(runs, component_slots, strict=True)
            if slot is not None
        )
        n_lead = len(points_shape)
        rank = n_lead + len(slot_of)
        chunk_coords = [0] * output_rank
        chunk_min = [0] * output_rank
        chunk_max = [0] * output_rank
        chunk_maps: list[OutputIndexMap | None] = [None] * output_rank
        cell_maps: list[OutputIndexMap | None] = [None] * domain.ndim
        # Unread input axes are independent repetitions, not part of an index
        # array component. Retain their extent without expanding any arrays.
        extents = [domain.shape[axis] for axis in slot_of]
        for axis, slot in slot_of.items():
            cell_maps[axis] = DimensionMap(input_dimension=slot, offset=lo_all[axis])
        for out_dim, (c, c_start, c_extent, k, extent, chunk_map, cell_map, _full) in zip(
            (axis.output_dimension for axis in self.sets), residual, strict=True
        ):
            chunk_coords[out_dim] = c
            chunk_min[out_dim] = c_start
            chunk_max[out_dim] = c_start + c_extent
            chunk_maps[out_dim] = chunk_map
            if k is not None:
                slot = slot_of[k]
                extents[slot - n_lead] = extent
                cell_maps[k] = DimensionMap(
                    input_dimension=slot, offset=cast("DimensionMap", cell_map).offset
                )
        for joint, row, run, component_slot in zip(
            self.joint_sets, rows, runs, component_slots, strict=True
        ):
            component_shape = [1] * rank
            if component_slot is not None:
                component_shape[component_slot] = run.stop - run.start
            for column, out_dim in enumerate(joint.output_dimensions):
                c_start = int(joint.chunk_start[row, column])
                chunk_coords[out_dim] = int(joint.chunk[row, column])
                chunk_min[out_dim] = c_start
                chunk_max[out_dim] = c_start + int(joint.chunk_extent[row, column])
                chunk_maps[out_dim] = ArrayMap(
                    index_array=joint.index[run, column].reshape(component_shape),
                    offset=joint.offsets[column] - c_start,
                    stride=joint.strides[column],
                )
            for column, axis in enumerate(joint.broadcast_axes):
                cell_maps[axis] = ArrayMap(
                    index_array=joint.block_coordinates[run, column].reshape(component_shape),
                    offset=lo_all[axis],
                )
        synthetic = IndexDomain._unchecked((0,) * rank, points_shape + tuple(extents))  # pyright: ignore[reportPrivateUsage]
        return ChunkProjection(
            chunk_coords=tuple(chunk_coords),
            chunk_domain=IndexDomain._unchecked(tuple(chunk_min), tuple(chunk_max)),  # pyright: ignore[reportPrivateUsage]
            chunk_transform=IndexTransform._unchecked(  # pyright: ignore[reportPrivateUsage]
                synthetic, tuple(cast("list[OutputIndexMap]", chunk_maps))
            ),
            cell_transform=IndexTransform._unchecked(  # pyright: ignore[reportPrivateUsage]
                synthetic, tuple(cast("list[OutputIndexMap]", cell_maps))
            ),
            coverage="unknown",
        )

dimension_grids instance-attribute

dimension_grids: tuple[DimensionGridLike, ...]

One grid per storage dimension.

joint_sets instance-attribute

joint_sets: tuple[JointSet, ...]

Connected index-array components, ordered by their first output dimension.

row_shape instance-attribute

row_shape: tuple[int, ...]

Rows per table: sets followed by joint_sets, in row-major iteration order.

sets instance-attribute

sets: tuple[StridedSet | IndexedSet, ...]

Independent per-axis tables, in output-dimension order.

transform instance-attribute

transform: IndexTransform

The transform this partition factors.

__init__

__init__(
    transform: IndexTransform,
    dimension_grids: tuple[DimensionGridLike, ...],
    sets: tuple[StridedSet | IndexedSet, ...],
    joint_sets: tuple[JointSet, ...],
    row_shape: tuple[int, ...],
) -> None

__iter__

__iter__() -> Iterator[ChunkProjection]

Materialize every row in order.

Source code in src/zarr_indexing/chunk_resolution.py

def __iter__(self) -> Iterator[ChunkProjection]:
    """Materialize every row in order."""
    if math.prod(self.row_shape) == 0:
        return
    if not self.joint_sets:
        yield from self._iter_factorized()
    else:
        yield from self._iter_correlated()

__len__

__len__() -> int

Source code in src/zarr_indexing/chunk_resolution.py

def __len__(self) -> int:
    return math.prod(self.row_shape)

chunk_coords

chunk_coords() -> ndarray[Any, dtype[intp]]

Chunk coordinates of every row, shape (len(self), output rank), without materializing rows.

Source code in src/zarr_indexing/chunk_resolution.py

def chunk_coords(self) -> np.ndarray[Any, np.dtype[np.intp]]:
    """Chunk coordinates of every row, shape ``(len(self), output rank)``, without materializing rows."""
    count = len(self)
    out = np.empty((count, self.transform.output_rank), dtype=np.intp)
    if count == 0 or not self.row_shape:
        return out
    indices = np.unravel_index(np.arange(count, dtype=np.intp), self.row_shape)
    for axis, rows in zip(self.sets, indices, strict=False):
        out[:, axis.output_dimension] = axis.chunk[rows]
    for joint, rows in zip(self.joint_sets, indices[len(self.sets) :], strict=True):
        out[:, list(joint.output_dimensions)] = joint.chunk[rows]
    return out

IndexedSet dataclass

One output dimension read through an orthogonal ArrayMap, one row per chunk.

The map's coordinates are grouped by chunk in CSR form: row i owns index[pointer[i]:pointer[i + 1]] (the index-array values, in request order) and positions[pointer[i]:pointer[i + 1]] (their positions along the request axis, ascending). local gives the same values as chunk-local storage coordinates. Columns are read-only NumPy arrays.

Examples:

>>> import numpy as np
>>> from zarr_indexing import IndexTransform, plan_chunks
>>> from zarr_indexing.grid import dimension_grids_from_chunks
>>> grids = dimension_grids_from_chunks((4,), shape=(10,))
>>> transform = IndexTransform.from_shape((10,)).oindex[np.array([9, 1, 2, 8])]
>>> (axis,) = plan_chunks(transform, grids).partition().sets
>>> axis.chunk.tolist(), axis.pointer.tolist()
([0, 2], [0, 2, 4])
>>> axis.local.tolist(), axis.positions.tolist()
([1, 2, 1, 0], [1, 2, 0, 3])

Source code in src/zarr_indexing/chunk_resolution.py

@dataclass(frozen=True, slots=True)
class IndexedSet:
    """One output dimension read through an orthogonal `ArrayMap`, one row per chunk.

    The map's coordinates are grouped by chunk in CSR form: row ``i`` owns
    ``index[pointer[i]:pointer[i + 1]]`` (the index-array values, in request
    order) and ``positions[pointer[i]:pointer[i + 1]]`` (their positions along
    the request axis, ascending). `local` gives the same values as chunk-local
    storage coordinates. Columns are read-only NumPy arrays.

    Examples
    --------
    >>> import numpy as np
    >>> from zarr_indexing import IndexTransform, plan_chunks
    >>> from zarr_indexing.grid import dimension_grids_from_chunks
    >>> grids = dimension_grids_from_chunks((4,), shape=(10,))
    >>> transform = IndexTransform.from_shape((10,)).oindex[np.array([9, 1, 2, 8])]
    >>> (axis,) = plan_chunks(transform, grids).partition().sets
    >>> axis.chunk.tolist(), axis.pointer.tolist()
    ([0, 2], [0, 2, 4])
    >>> axis.local.tolist(), axis.positions.tolist()
    ([1, 2, 1, 0], [1, 2, 0, 3])
    """

    output_dimension: int
    """The storage axis this table describes."""

    input_dimension: int
    """The request axis the index array varies over."""

    offset: int
    """The map's affine offset: storage is ``offset + stride * index``."""

    stride: int
    """The map's affine stride."""

    chunk: np.ndarray[Any, np.dtype[np.intp]]
    """Chunk index along the axis, one per row, ascending."""

    chunk_start: np.ndarray[Any, np.dtype[np.intp]]
    """Storage origin of each chunk."""

    chunk_extent: np.ndarray[Any, np.dtype[np.intp]]
    """Data extent of each chunk."""

    pointer: np.ndarray[Any, np.dtype[np.intp]]
    """CSR row pointer: row ``i`` owns entries ``pointer[i]`` to ``pointer[i + 1]``."""

    index: np.ndarray[Any, np.dtype[np.intp]]
    """Index-array values grouped by chunk."""

    positions: np.ndarray[Any, np.dtype[np.intp]]
    """Positions along the request axis, grouped by chunk, ascending within a row."""

    def __post_init__(self) -> None:
        _own_columns(self, "chunk", "chunk_start", "chunk_extent", "pointer", "index", "positions")

    def __len__(self) -> int:
        return int(self.chunk.size)

    @property
    def counts(self) -> np.ndarray[Any, np.dtype[np.intp]]:
        """Entries per row."""
        return np.diff(self.pointer)

    _local: np.ndarray[Any, np.dtype[np.intp]] | None = field(
        default=None, init=False, repr=False, compare=False
    )

    @property
    def local(self) -> np.ndarray[Any, np.dtype[np.intp]]:
        """Memoized, read-only chunk-local coordinates, grouped like `index`."""
        cached = self._local
        if cached is None:
            storage = checked_affine(self.offset, self.stride, self.index)
            cached = storage - np.repeat(self.chunk_start, self.counts)
            cached = _own(cached)
            object.__setattr__(self, "_local", cached)
        return cached

    def run(self, row: int) -> slice:
        """The slice of `index` / `positions` a row owns."""
        return slice(int(self.pointer[row]), int(self.pointer[row + 1]))

chunk instance-attribute

chunk: ndarray[Any, dtype[intp]]

Chunk index along the axis, one per row, ascending.

chunk_extent instance-attribute

chunk_extent: ndarray[Any, dtype[intp]]

Data extent of each chunk.

chunk_start instance-attribute

chunk_start: ndarray[Any, dtype[intp]]

Storage origin of each chunk.

counts property

counts: ndarray[Any, dtype[intp]]

Entries per row.

index instance-attribute

index: ndarray[Any, dtype[intp]]

Index-array values grouped by chunk.

input_dimension instance-attribute

input_dimension: int

The request axis the index array varies over.

local property

local: ndarray[Any, dtype[intp]]

Memoized, read-only chunk-local coordinates, grouped like index.

offset instance-attribute

offset: int

The map's affine offset: storage is offset + stride * index.

output_dimension instance-attribute

output_dimension: int

The storage axis this table describes.

pointer instance-attribute

pointer: ndarray[Any, dtype[intp]]

CSR row pointer: row i owns entries pointer[i] to pointer[i + 1].

positions instance-attribute

positions: ndarray[Any, dtype[intp]]

Positions along the request axis, grouped by chunk, ascending within a row.

stride instance-attribute

stride: int

The map's affine stride.

__init__

__init__(
    output_dimension: int,
    input_dimension: int,
    offset: int,
    stride: int,
    chunk: ndarray[Any, dtype[intp]],
    chunk_start: ndarray[Any, dtype[intp]],
    chunk_extent: ndarray[Any, dtype[intp]],
    pointer: ndarray[Any, dtype[intp]],
    index: ndarray[Any, dtype[intp]],
    positions: ndarray[Any, dtype[intp]],
) -> None

__len__

__len__() -> int

Source code in src/zarr_indexing/chunk_resolution.py

def __len__(self) -> int:
    return int(self.chunk.size)

__post_init__

__post_init__() -> None

Source code in src/zarr_indexing/chunk_resolution.py

def __post_init__(self) -> None:
    _own_columns(self, "chunk", "chunk_start", "chunk_extent", "pointer", "index", "positions")

run

run(row: int) -> slice

The slice of index / positions a row owns.

Source code in src/zarr_indexing/chunk_resolution.py

def run(self, row: int) -> slice:
    """The slice of `index` / `positions` a row owns."""
    return slice(int(self.pointer[row]), int(self.pointer[row + 1]))

JointSet dataclass

One connected component of index arrays, grouped by the chunk each point lands in.

Arrays connected by shared input axes, possibly transitively, constrain each other and are sorted into chunks together. Row i is one touched chunk, chunk[i] its coordinates on the output_dimensions, and CSR range pointer[i]:pointer[i + 1] its points: index holds their index-array values per output dimension, positions their flat positions in this component's input block, and block_coordinates those positions unravelled over the block. Columns are read-only NumPy arrays.

Examples:

>>> import numpy as np
>>> from zarr_indexing import IndexTransform, plan_chunks
>>> from zarr_indexing.grid import dimension_grids_from_chunks
>>> grids = dimension_grids_from_chunks((3, 4), shape=(7, 9))
>>> transform = IndexTransform.from_shape((7, 9)).vindex[
...     np.array([0, 6, 6, 1]), np.array([8, 0, 1, 8])
... ]
>>> (joint,) = plan_chunks(transform, grids).partition().joint_sets
>>> joint.chunk.tolist(), joint.pointer.tolist(), joint.positions.tolist()
([[0, 2], [2, 0]], [0, 2, 4], [0, 3, 1, 2])

Source code in src/zarr_indexing/chunk_resolution.py

@dataclass(frozen=True, slots=True)
class JointSet:
    """One connected component of index arrays, grouped by the chunk each point lands in.

    Arrays connected by shared input axes, possibly transitively, constrain
    each other and are sorted into chunks together.
    Row ``i`` is one touched chunk, `chunk[i]` its coordinates on the
    `output_dimensions`, and CSR range ``pointer[i]:pointer[i + 1]`` its
    points: `index` holds their index-array values per output dimension,
    `positions` their flat positions in this component's input block, and
    `block_coordinates` those positions unravelled over the block. Columns are
    read-only NumPy arrays.

    Examples
    --------
    >>> import numpy as np
    >>> from zarr_indexing import IndexTransform, plan_chunks
    >>> from zarr_indexing.grid import dimension_grids_from_chunks
    >>> grids = dimension_grids_from_chunks((3, 4), shape=(7, 9))
    >>> transform = IndexTransform.from_shape((7, 9)).vindex[
    ...     np.array([0, 6, 6, 1]), np.array([8, 0, 1, 8])
    ... ]
    >>> (joint,) = plan_chunks(transform, grids).partition().joint_sets
    >>> joint.chunk.tolist(), joint.pointer.tolist(), joint.positions.tolist()
    ([[0, 2], [2, 0]], [0, 2, 4], [0, 3, 1, 2])
    """

    output_dimensions: tuple[int, ...]
    """The storage axes read by correlated index arrays."""

    offsets: tuple[int, ...]
    """Affine offset of each array's map, aligned with `output_dimensions`."""

    strides: tuple[int, ...]
    """Affine stride of each array's map."""

    broadcast_axes: tuple[int, ...]
    """The request axes the arrays broadcast over."""

    broadcast_shape: tuple[int, ...]
    """The extent of those axes."""

    chunk: np.ndarray[Any, np.dtype[np.intp]]
    """Chunk coordinates on `output_dimensions`, shape ``(rows, k)``, lexicographic."""

    chunk_start: np.ndarray[Any, np.dtype[np.intp]]
    """Storage origin of each chunk on `output_dimensions`, shape ``(rows, k)``."""

    chunk_extent: np.ndarray[Any, np.dtype[np.intp]]
    """Data extent of each chunk on `output_dimensions`, shape ``(rows, k)``."""

    pointer: np.ndarray[Any, np.dtype[np.intp]]
    """CSR row pointer into `index`, `positions` and `block_coordinates`."""

    index: np.ndarray[Any, np.dtype[np.intp]]
    """Index-array values per point and output dimension, shape ``(points, k)``."""

    positions: np.ndarray[Any, np.dtype[np.intp]]
    """Flat block position of each point, ascending within a row."""

    _block_coordinates: np.ndarray[Any, np.dtype[np.intp]] | None = field(
        default=None, init=False, repr=False, compare=False
    )

    def __post_init__(self) -> None:
        _own_columns(self, "chunk", "chunk_start", "chunk_extent", "pointer", "index", "positions")

    def __len__(self) -> int:
        return int(self.chunk.shape[0])

    @property
    def counts(self) -> np.ndarray[Any, np.dtype[np.intp]]:
        """Points per row."""
        return np.diff(self.pointer)

    _local: np.ndarray[Any, np.dtype[np.intp]] | None = field(
        default=None, init=False, repr=False, compare=False
    )

    @property
    def local(self) -> np.ndarray[Any, np.dtype[np.intp]]:
        """Memoized, read-only chunk-local coordinates, shape ``(points, k)``."""
        cached = self._local
        if cached is None:
            storage = np.stack(
                [
                    checked_affine(offset, stride, self.index[:, column])
                    for column, (offset, stride) in enumerate(
                        zip(self.offsets, self.strides, strict=True)
                    )
                ],
                axis=1,
            )
            cached = storage - np.repeat(self.chunk_start, self.counts, axis=0)
            cached = _own(cached)
            object.__setattr__(self, "_local", cached)
        return cached

    @property
    def block_coordinates(self) -> np.ndarray[Any, np.dtype[np.intp]]:
        """Memoized `positions` unravelled over `broadcast_shape`, shape ``(points, len(broadcast_axes))``."""
        cached = self._block_coordinates
        if cached is None:
            if self.broadcast_shape:
                cached = np.stack(np.unravel_index(self.positions, self.broadcast_shape), axis=1)
            else:
                cached = np.empty((self.positions.size, 0), dtype=np.intp)
            cached = _own(np.asarray(cached, dtype=np.intp))
            object.__setattr__(self, "_block_coordinates", cached)
        return cached

    def run(self, row: int) -> slice:
        """The slice of the point arrays a row owns."""
        return slice(int(self.pointer[row]), int(self.pointer[row + 1]))

block_coordinates property

block_coordinates: ndarray[Any, dtype[intp]]

Memoized positions unravelled over broadcast_shape, shape (points, len(broadcast_axes)).

broadcast_axes instance-attribute

broadcast_axes: tuple[int, ...]

The request axes the arrays broadcast over.

broadcast_shape instance-attribute

broadcast_shape: tuple[int, ...]

The extent of those axes.

chunk instance-attribute

chunk: ndarray[Any, dtype[intp]]

Chunk coordinates on output_dimensions, shape (rows, k), lexicographic.

chunk_extent instance-attribute

chunk_extent: ndarray[Any, dtype[intp]]

Data extent of each chunk on output_dimensions, shape (rows, k).

chunk_start instance-attribute

chunk_start: ndarray[Any, dtype[intp]]

Storage origin of each chunk on output_dimensions, shape (rows, k).

counts property

counts: ndarray[Any, dtype[intp]]

Points per row.

index instance-attribute

index: ndarray[Any, dtype[intp]]

Index-array values per point and output dimension, shape (points, k).

local property

local: ndarray[Any, dtype[intp]]

Memoized, read-only chunk-local coordinates, shape (points, k).

offsets instance-attribute

offsets: tuple[int, ...]

Affine offset of each array's map, aligned with output_dimensions.

output_dimensions instance-attribute

output_dimensions: tuple[int, ...]

The storage axes read by correlated index arrays.

pointer instance-attribute

pointer: ndarray[Any, dtype[intp]]

CSR row pointer into index, positions and block_coordinates.

positions instance-attribute

positions: ndarray[Any, dtype[intp]]

Flat block position of each point, ascending within a row.

strides instance-attribute

strides: tuple[int, ...]

Affine stride of each array's map.

__init__

__init__(
    output_dimensions: tuple[int, ...],
    offsets: tuple[int, ...],
    strides: tuple[int, ...],
    broadcast_axes: tuple[int, ...],
    broadcast_shape: tuple[int, ...],
    chunk: ndarray[Any, dtype[intp]],
    chunk_start: ndarray[Any, dtype[intp]],
    chunk_extent: ndarray[Any, dtype[intp]],
    pointer: ndarray[Any, dtype[intp]],
    index: ndarray[Any, dtype[intp]],
    positions: ndarray[Any, dtype[intp]],
) -> None

__len__

__len__() -> int

Source code in src/zarr_indexing/chunk_resolution.py

def __len__(self) -> int:
    return int(self.chunk.shape[0])

__post_init__

__post_init__() -> None

Source code in src/zarr_indexing/chunk_resolution.py

def __post_init__(self) -> None:
    _own_columns(self, "chunk", "chunk_start", "chunk_extent", "pointer", "index", "positions")

run

run(row: int) -> slice

The slice of the point arrays a row owns.

Source code in src/zarr_indexing/chunk_resolution.py

def run(self, row: int) -> slice:
    """The slice of the point arrays a row owns."""
    return slice(int(self.pointer[row]), int(self.pointer[row + 1]))

StridedSet dataclass

One output dimension read through a ConstantMap or DimensionMap, one row per chunk.

Row i is the map restricted to chunk chunk[i] and re-based to chunk-local, zero-origin coordinates: the chunk-local map is DimensionMap(input_dimension, offset=local_start[i], stride=stride) over [0, extent[i]) (a ConstantMap(local_start[i]) for a constant), and its cells are request positions [origin[i], origin[i] + extent[i]) along the input axis, counted from the request domain's lower bound.

Columns are read-only NumPy arrays. extent and origin are measured along the request axis, whose bounds are arbitrary Python ints; they are intp unless a value does not fit, in which case they hold exact ints (dtype=object).

Examples:

>>> from zarr_indexing import IndexTransform, plan_chunks
>>> from zarr_indexing.grid import dimension_grids_from_chunks
>>> grids = dimension_grids_from_chunks((4,), shape=(10,))
>>> (axis,) = plan_chunks(IndexTransform.from_shape((10,))[1:9:2], grids).partition().sets
>>> axis.chunk.tolist(), axis.local_start.tolist(), axis.extent.tolist(), axis.origin.tolist()
([0, 1], [1, 1], [2, 2], [0, 2])

Source code in src/zarr_indexing/chunk_resolution.py

@dataclass(frozen=True, slots=True)
class StridedSet:
    """One output dimension read through a `ConstantMap` or `DimensionMap`, one row per chunk.

    Row ``i`` is the map restricted to chunk ``chunk[i]`` and re-based to
    chunk-local, zero-origin coordinates: the chunk-local map is
    `DimensionMap(input_dimension, offset=local_start[i], stride=stride)` over
    ``[0, extent[i])`` (a `ConstantMap(local_start[i])` for a constant), and
    its cells are request positions ``[origin[i], origin[i] + extent[i])``
    along the input axis, counted from the request domain's lower bound.

    Columns are read-only NumPy arrays. `extent` and `origin` are measured
    along the request axis, whose bounds are arbitrary Python ints; they are
    ``intp`` unless a value does not fit, in which case they hold exact ints
    (``dtype=object``).

    Examples
    --------
    >>> from zarr_indexing import IndexTransform, plan_chunks
    >>> from zarr_indexing.grid import dimension_grids_from_chunks
    >>> grids = dimension_grids_from_chunks((4,), shape=(10,))
    >>> (axis,) = plan_chunks(IndexTransform.from_shape((10,))[1:9:2], grids).partition().sets
    >>> axis.chunk.tolist(), axis.local_start.tolist(), axis.extent.tolist(), axis.origin.tolist()
    ([0, 1], [1, 1], [2, 2], [0, 2])
    """

    output_dimension: int
    """The storage axis this table describes."""

    input_dimension: int | None
    """The request axis the map reads, or `None` for a constant."""

    stride: int
    """Storage step per request cell; ``0`` for a constant."""

    chunk: np.ndarray[Any, np.dtype[np.intp]]
    """Chunk index along the axis, one per row, ascending."""

    chunk_start: np.ndarray[Any, np.dtype[np.intp]]
    """Storage origin of each chunk."""

    chunk_extent: np.ndarray[Any, np.dtype[np.intp]]
    """Data extent of each chunk (clipped at the array boundary)."""

    local_start: np.ndarray[Any, np.dtype[np.intp]]
    """Chunk-local storage coordinate of the row's first cell."""

    extent: np.ndarray[Any, np.dtype[np.intp]]
    """Cells the row selects along the request axis (``1`` for a constant)."""

    origin: np.ndarray[Any, np.dtype[np.intp]]
    """Position of the row's first cell along the request axis (``0`` for a constant)."""

    full: np.ndarray[Any, np.dtype[np.bool_]]
    """Whether the row covers its chunk's data extent exactly once (in either direction)."""

    def __post_init__(self) -> None:
        _own_columns(
            self, "chunk", "chunk_start", "chunk_extent", "local_start", "extent", "origin", "full"
        )

    def __len__(self) -> int:
        return int(self.chunk.size)

chunk instance-attribute

chunk: ndarray[Any, dtype[intp]]

Chunk index along the axis, one per row, ascending.

chunk_extent instance-attribute

chunk_extent: ndarray[Any, dtype[intp]]

Data extent of each chunk (clipped at the array boundary).

chunk_start instance-attribute

chunk_start: ndarray[Any, dtype[intp]]

Storage origin of each chunk.

extent instance-attribute

extent: ndarray[Any, dtype[intp]]

Cells the row selects along the request axis (1 for a constant).

full instance-attribute

full: ndarray[Any, dtype[bool_]]

Whether the row covers its chunk's data extent exactly once (in either direction).

input_dimension instance-attribute

input_dimension: int | None

The request axis the map reads, or None for a constant.

local_start instance-attribute

local_start: ndarray[Any, dtype[intp]]

Chunk-local storage coordinate of the row's first cell.

origin instance-attribute

origin: ndarray[Any, dtype[intp]]

Position of the row's first cell along the request axis (0 for a constant).

output_dimension instance-attribute

output_dimension: int

The storage axis this table describes.

stride instance-attribute

stride: int

Storage step per request cell; 0 for a constant.

__init__

__init__(
    output_dimension: int,
    input_dimension: int | None,
    stride: int,
    chunk: ndarray[Any, dtype[intp]],
    chunk_start: ndarray[Any, dtype[intp]],
    chunk_extent: ndarray[Any, dtype[intp]],
    local_start: ndarray[Any, dtype[intp]],
    extent: ndarray[Any, dtype[intp]],
    origin: ndarray[Any, dtype[intp]],
    full: ndarray[Any, dtype[bool_]],
) -> None

__len__

__len__() -> int

Source code in src/zarr_indexing/chunk_resolution.py

def __len__(self) -> int:
    return int(self.chunk.size)

__post_init__

__post_init__() -> None

Source code in src/zarr_indexing/chunk_resolution.py

def __post_init__(self) -> None:
    _own_columns(
        self, "chunk", "chunk_start", "chunk_extent", "local_start", "extent", "origin", "full"
    )

plan_chunks

plan_chunks(
    transform: IndexTransform,
    dimension_grids: Sequence[DimensionGridLike],
) -> ChunkPlan

Plan a transform against a caller-selected chunk grid.

Parameters:

• transform (IndexTransform) –

  Mapping from the request domain to storage coordinates.

• dimension_grids (Sequence[DimensionGridLike]) –

  One storage grid per transform output dimension.

Returns:

• ChunkPlan –

  A reusable plan whose projections are computed lazily; its partition() is the factored form they are derived from.

Examples:

Row 1 of a (3, 4) array with (2, 2) chunks touches the two chunks in the top grid row, each contributing a (2, 2) chunk domain:

>>> from zarr_indexing import IndexTransform
>>> from zarr_indexing.grid import dimension_grids_from_chunks
>>> grids = dimension_grids_from_chunks((2, 2), shape=(3, 4))
>>> plan = plan_chunks(IndexTransform.from_shape((3, 4))[1, :], grids)
>>> [p.chunk_coords for p in plan]
[(0, 0), (0, 1)]
>>> [p.chunk_domain.shape for p in plan]
[(2, 2), (2, 2)]

Source code in src/zarr_indexing/chunk_resolution.py

def plan_chunks(
    transform: IndexTransform,
    dimension_grids: Sequence[DimensionGridLike],
) -> ChunkPlan:
    """Plan a transform against a caller-selected chunk grid.

    Parameters
    ----------
    transform
        Mapping from the request domain to storage coordinates.
    dimension_grids
        One storage grid per transform output dimension.

    Returns
    -------
    ChunkPlan
        A reusable plan whose projections are computed lazily; its
        `partition()` is the factored form they are derived from.

    Examples
    --------
    Row 1 of a `(3, 4)` array with `(2, 2)` chunks touches the two chunks in
    the top grid row, each contributing a `(2, 2)` chunk domain:

    >>> from zarr_indexing import IndexTransform
    >>> from zarr_indexing.grid import dimension_grids_from_chunks
    >>> grids = dimension_grids_from_chunks((2, 2), shape=(3, 4))
    >>> plan = plan_chunks(IndexTransform.from_shape((3, 4))[1, :], grids)
    >>> [p.chunk_coords for p in plan]
    [(0, 0), (0, 1)]
    >>> [p.chunk_domain.shape for p in plan]
    [(2, 2), (2, 2)]
    """
    grids = tuple(dimension_grids)
    if len(grids) != transform.output_rank:
        raise ValueError(
            "dimension_grids must have one entry per transform output dimension; "
            f"got {len(grids)} grids for output rank {transform.output_rank}"
        )
    return ChunkPlan(transform=transform, dimension_grids=grids)