zarr_indexing.boundary

zarr_indexing.boundary

The positional (NumPy) selection dialect, lowered onto the transform algebra.

The transform algebra uses literal domain coordinates: an index is a point in the view's own coordinate system, which after view = arr[10:50] runs from 10 to 49, and a negative index is a negative coordinate rather than an offset from the end (TensorStore's convention — see zarr_indexing.transform).

NumPy uses positions: index 0 always means the first element of the object being indexed, and -1 means the last. This module translates between the two. It validates a selection against the view's shape with NumPy semantics, then shifts every coordinate by the domain's origin so the transform layer sees literal coordinates.

Note

zarr.Array currently carries its own copy of this normalization, tuned to a different boundary contract (Array.lazy[...] deliberately exposes the literal dialect, so a view's coordinates keep their meaning across composition). This module is the generic, zarr-free version used by LazyArray; consolidating zarr's copy onto it is left to a follow-up.

SelectionMode module-attribute

SelectionMode = Literal["basic", "orthogonal", "vectorized"]

normalize_positional_selection

normalize_positional_selection(
    selection: Any, domain: IndexDomain, mode: SelectionMode
) -> Any

Translate a positional (NumPy-dialect) selection into literal coordinates.

Positions are zero-based offsets into the current view; negatives wrap from the end. The returned selection addresses the same cells in the literal coordinate system domain uses, ready for zarr_indexing.transform.selection_to_transform.

Parameters:

• selection (Any) –

  A NumPy-style selection: integers, slices, Ellipsis, integer arrays or lists, or boolean arrays.

• domain (IndexDomain) –

  The domain of the view being indexed. Its shape defines the positional bounds and its origin the coordinate shift.

• mode (SelectionMode) –

  Which selection dialect the entries follow: "basic" (integers and slices), "orthogonal" (per-axis arrays, outer product), or "vectorized" (correlated coordinate arrays or a single mask).

Returns:

• tuple[Any, ...] –

  The selection with every coordinate expressed in literal domain coordinates.

Raises:

• IndexError –

  If a boolean scalar is used as an index, a boolean mask does not match the shape of the axes it covers, an index is out of bounds, or too many indices are supplied.

Source code in src/zarr_indexing/boundary.py

def normalize_positional_selection(
    selection: Any,
    domain: IndexDomain,
    mode: SelectionMode,
) -> Any:
    """Translate a positional (NumPy-dialect) selection into literal coordinates.

    Positions are zero-based offsets into the current view; negatives wrap
    from the end. The returned selection addresses the same cells in the
    literal coordinate system `domain` uses, ready for
    `zarr_indexing.transform.selection_to_transform`.

    Parameters
    ----------
    selection
        A NumPy-style selection: integers, slices, `Ellipsis`, integer arrays or
        lists, or boolean arrays.
    domain
        The domain of the view being indexed. Its shape defines the positional
        bounds and its origin the coordinate shift.
    mode
        Which selection dialect the entries follow: `"basic"` (integers and
        slices), `"orthogonal"` (per-axis arrays, outer product), or
        `"vectorized"` (correlated coordinate arrays or a single mask).

    Returns
    -------
    tuple[Any, ...]
        The selection with every coordinate expressed in literal domain
        coordinates.

    Raises
    ------
    IndexError
        If a boolean scalar is used as an index, a boolean mask does not match
        the shape of the axes it covers, an index is out of bounds, or too many
        indices are supplied.
    """
    entries = selection if isinstance(selection, tuple) else (selection,)
    shape = domain.shape
    origin = domain.inclusive_min
    ndim = domain.ndim

    if mode in ("orthogonal", "vectorized"):
        validate_advanced_selection(selection, domain, mode)
    else:
        for sel in entries:
            if is_bool_scalar(sel):
                raise IndexError(
                    "boolean scalars are not valid indices; use a boolean array "
                    "matching the shape of the axes it selects"
                )

    n_ellipsis = sum(1 for sel in entries if sel is Ellipsis)
    if n_ellipsis > 1:
        raise IndexError("an index can only have a single ellipsis ('...')")
    consumed = sum(_axes_consumed(sel, mode) for sel in entries if sel is not Ellipsis)
    if consumed > ndim:
        raise IndexError(
            f"too many indices for array: array has {ndim} dimensions, but {consumed} were indexed"
        )

    result: list[Any] = []
    axis = 0
    for sel in entries:
        if sel is Ellipsis:
            # Passed through rather than expanded: every mode of
            # `selection_to_transform` expands an ellipsis (and pads short
            # selections) to whole-axis slices itself, and `vectorized` mode
            # rejects an explicit slice, so expanding here would turn a legal
            # partial coordinate selection into an error.
            result.append(Ellipsis)
            axis += ndim - consumed
            continue
        if sel is None:
            # newaxis: no axis of the view is consumed, and there is no
            # coordinate to shift. The transform layer decides whether the mode
            # accepts it.
            result.append(None)
            continue

        arr = _as_index_array(sel)
        if arr is not None and arr.dtype == np.bool_:
            n_axes = _axes_consumed(sel, mode)
            expected = shape[axis : axis + n_axes]
            if arr.shape != tuple(expected):
                raise IndexError(
                    f"boolean index has shape {arr.shape} but the axes it "
                    f"covers have shape {tuple(expected)}"
                )
            for offset, positions in enumerate(np.nonzero(arr)):
                result.append(positions.astype(np.intp) + origin[axis + offset])
            axis += n_axes
            continue

        if axis >= ndim:
            raise IndexError(
                f"too many indices for array: array has {ndim} dimensions, "
                f"but {consumed} were indexed"
            )
        size = shape[axis]
        if arr is not None:
            result.append(_normalize_int_array(arr, size, axis) + origin[axis])
        elif isinstance(sel, slice):
            start, stop, step = _normalize_slice(sel, size, axis)
            result.append(slice(start + origin[axis], stop + origin[axis], step))
        elif (scalar := as_scalar_index(sel)) is not None:
            result.append(_normalize_int(scalar, size, axis) + origin[axis])
        else:
            raise IndexError(f"unsupported selection type: {type(sel)!r}")
        axis += 1

    # Axes the selection did not mention are left to `selection_to_transform`,
    # which pads them with whole-axis slices in every mode.
    return tuple(result)

split_scalar_axes

split_scalar_axes(
    selection: Any, domain: IndexDomain, mode: SelectionMode
) -> tuple[tuple[Any, ...] | None, Any]

Peel scalar integer indices out of a fancy selection.

A scalar integer drops its axis, and neither the orthogonal nor the vectorized path of the transform algebra models that — both widen a scalar into a length-1 index array, which keeps the axis — so the scalars are split off here and applied as a separate basic step first.

Applying them first is this package's rule, not NumPy's. NumPy groups a scalar with the advanced indices for the purpose of placing the broadcast result, so the two disagree when a scalar and an index array are separated: a[0, ..., [1, 2]] has shape (2, 3) for a (2, 3, 4) array, where a[0][..., [1, 2]] has shape (3, 2). The earlier claim here that they always agree rested on a[0, [1, 2], :], where the indices are adjacent and they happen to. Scalar-first is the documented dialect (see the lazy_array module docstring) — the divergence is deliberate, and this note exists so that the correct end is not "fixed" later.

Parameters:

• selection (Any) –

  A positional orthogonal or vectorized selection.

• domain (IndexDomain) –

  The domain of the view being indexed.

• mode (SelectionMode) –

  "orthogonal" or "vectorized"; controls how many axes each entry covers, which decides where the scalars sit.

Returns:

• tuple[tuple[Any, ...] | None, Any] –

  (basic_selection, remaining_selection). basic_selection is a full-rank basic selection in literal domain coordinates that drops the scalar axes, or None when the selection has no scalar entries (in which case remaining_selection is selection unchanged).

Raises:

• IndexError –

  If a boolean scalar is used as an index, an index is out of bounds, or too many indices are supplied.

Source code in src/zarr_indexing/boundary.py

def split_scalar_axes(
    selection: Any,
    domain: IndexDomain,
    mode: SelectionMode,
) -> tuple[tuple[Any, ...] | None, Any]:
    """Peel scalar integer indices out of a fancy selection.

    A scalar integer drops its axis, and neither the orthogonal nor the
    vectorized path of the transform algebra models that — both widen a scalar
    into a length-1 index array, which keeps the axis — so the scalars are split
    off here and applied as a separate basic step first.

    Applying them *first* is this package's rule, not NumPy's. NumPy groups a
    scalar with the advanced indices for the purpose of placing the broadcast
    result, so the two disagree when a scalar and an index array are separated:
    `a[0, ..., [1, 2]]` has shape `(2, 3)` for a `(2, 3, 4)` array, where
    `a[0][..., [1, 2]]` has shape `(3, 2)`. The earlier claim here that they
    always agree rested on `a[0, [1, 2], :]`, where the indices are adjacent and
    they happen to. Scalar-first is the documented dialect (see the `lazy_array`
    module docstring) — the divergence is deliberate, and this note exists so
    that the correct end is not "fixed" later.

    Parameters
    ----------
    selection
        A positional orthogonal or vectorized selection.
    domain
        The domain of the view being indexed.
    mode
        `"orthogonal"` or `"vectorized"`; controls how many axes each entry
        covers, which decides where the scalars sit.

    Returns
    -------
    tuple[tuple[Any, ...] | None, Any]
        `(basic_selection, remaining_selection)`. `basic_selection` is a
        full-rank basic selection in **literal** domain coordinates that drops
        the scalar axes, or `None` when the selection has no scalar entries (in
        which case `remaining_selection` is `selection` unchanged).

    Raises
    ------
    IndexError
        If a boolean scalar is used as an index, an index is out of bounds, or
        too many indices are supplied.
    """
    entries = selection if isinstance(selection, tuple) else (selection,)
    axes = _expanded_axis_walk(entries, domain.ndim, mode)

    scalar_axes: dict[int, int] = {}
    remaining: list[Any] = []
    for sel, axis in zip(entries, axes, strict=True):
        scalar = as_scalar_index(sel)
        if scalar is not None:
            scalar_axes[axis] = _normalize_int(scalar, domain.shape[axis], axis)
        else:
            remaining.append(sel)

    if len(scalar_axes) == 0:
        return None, selection

    basic: list[Any] = []
    for axis in range(domain.ndim):
        lo = domain.inclusive_min[axis]
        if axis in scalar_axes:
            basic.append(lo + scalar_axes[axis])
        else:
            basic.append(slice(lo, domain.exclusive_max[axis]))
    return tuple(basic), tuple(remaining)

validate_advanced_selection

validate_advanced_selection(
    selection: Any,
    domain: IndexDomain,
    mode: Literal["orthogonal", "vectorized"],
) -> None

Validate advanced-index selector dtypes and boolean mask extents.

This is the validation shared by positional callers such as LazyArray and direct IndexTransform.oindex / .vindex callers. It deliberately does not normalize coordinates: direct transforms use literal coordinates, whereas positional callers shift and wrap them separately.

Source code in src/zarr_indexing/boundary.py

def validate_advanced_selection(
    selection: Any,
    domain: IndexDomain,
    mode: Literal["orthogonal", "vectorized"],
) -> None:
    """Validate advanced-index selector dtypes and boolean mask extents.

    This is the validation shared by positional callers such as `LazyArray`
    and direct `IndexTransform.oindex` / `.vindex` callers. It deliberately
    does not normalize coordinates: direct transforms use literal coordinates,
    whereas positional callers shift and wrap them separately.
    """
    entries: tuple[Any, ...] = selection if isinstance(selection, tuple) else (selection,)
    axes = _expanded_axis_walk(entries, domain.ndim, mode)

    for sel, axis in zip(entries, axes, strict=True):
        arr = _as_index_array(sel)
        if arr is None:
            continue
        if arr.dtype == np.bool_:
            n_axes = _axes_consumed(sel, mode)
            expected = domain.shape[axis : axis + n_axes]
            if arr.shape != tuple(expected):
                extent = (
                    f"dimension {expected[0]}"
                    if len(expected) == 1
                    else f"dimensions {tuple(expected)}"
                )
                raise IndexError(
                    f"boolean index has shape {arr.shape} but {extent} has shape {tuple(expected)}"
                )
        elif arr.dtype.kind not in "iu":
            raise IndexError(
                f"arrays used as indices must be of integer or boolean type; got dtype {arr.dtype}"
            )