diff --git a/.github/workflows/check_changelogs.yml b/.github/workflows/check_changelogs.yml index 0033b43db2..d7a54fc2c4 100644 --- a/.github/workflows/check_changelogs.yml +++ b/.github/workflows/check_changelogs.yml @@ -29,3 +29,6 @@ jobs: - name: Check zarr-metadata changelog entries run: uv run --no-sync python ci/check_changelog_entries.py packages/zarr-metadata/changes + + - name: Check zarr-indexing changelog entries + run: uv run --no-sync python ci/check_changelog_entries.py packages/zarr-indexing/changes diff --git a/.github/workflows/zarr-indexing-release.yml b/.github/workflows/zarr-indexing-release.yml new file mode 100644 index 0000000000..7cfd571eae --- /dev/null +++ b/.github/workflows/zarr-indexing-release.yml @@ -0,0 +1,117 @@ +name: zarr-indexing release + +on: + workflow_dispatch: + push: + tags: + - 'zarr_indexing-v*' + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: false + +jobs: + build: + name: Build wheel and sdist + runs-on: ubuntu-latest + defaults: + run: + shell: bash + working-directory: packages/zarr-indexing + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + persist-credentials: false + fetch-depth: 0 # hatch-vcs needs full history + tags + + - name: Install Hatch + uses: pypa/hatch@257e27e51a6a5616ed08a39a408a21c35c9931bc + with: + version: '1.16.5' + + - name: Build + run: hatch build + + - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: zarr-indexing-dist + path: packages/zarr-indexing/dist + + test_artifacts: + name: Test built artifacts + needs: [build] + runs-on: ubuntu-latest + steps: + - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + name: zarr-indexing-dist + path: dist + + - name: Install uv + uses: astral-sh/setup-uv@11f9893b081a58869d3b5fccaea48c9e9e46f990 # v8.3.2 + with: + enable-cache: false + + - name: Set up Python + run: uv python install 3.12 + + - name: Install built wheel and run import smoke test + run: | + wheel=$(ls dist/*.whl) + uv run --with "${wheel}" --python 3.12 --no-project \ + python -c "import zarr_indexing; print('zarr_indexing', zarr_indexing.__version__)" + + upload_pypi: + name: Upload to PyPI + needs: [build, test_artifacts] + if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/zarr_indexing-v') + runs-on: ubuntu-latest + environment: + name: zarr-indexing-releases + url: https://pypi.org/p/zarr-indexing + permissions: + id-token: write # required for OIDC trusted publishing + attestations: write # required for artifact attestations + steps: + - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + name: zarr-indexing-dist + path: dist + + - name: Generate artifact attestation + uses: actions/attest@f7c74d28b9d84cb8768d0b8ca14a4bac6ef463e6 # v4.2.0 + with: + subject-path: dist/* + + - name: Publish package to PyPI + uses: pypa/gh-action-pypi-publish@cef221092ed1bacb1cc03d23a2d87d1d172e277b # v1.14.0 + + upload_testpypi: + name: Upload to TestPyPI + needs: [build, test_artifacts] + if: github.event_name == 'workflow_dispatch' + runs-on: ubuntu-latest + environment: + name: zarr-indexing-releases-test + url: https://test.pypi.org/p/zarr-indexing + permissions: + id-token: write + attestations: write + steps: + - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + name: zarr-indexing-dist + path: dist + + - name: Generate artifact attestation + uses: actions/attest@f7c74d28b9d84cb8768d0b8ca14a4bac6ef463e6 # v4.2.0 + with: + subject-path: dist/* + + - name: Publish package to TestPyPI + uses: pypa/gh-action-pypi-publish@cef221092ed1bacb1cc03d23a2d87d1d172e277b # v1.14.0 + with: + repository-url: https://test.pypi.org/legacy/ diff --git a/.github/workflows/zarr-indexing.yml b/.github/workflows/zarr-indexing.yml new file mode 100644 index 0000000000..2106b10916 --- /dev/null +++ b/.github/workflows/zarr-indexing.yml @@ -0,0 +1,123 @@ +name: zarr-indexing + +on: + push: + branches: [main] + paths: + - 'packages/zarr-indexing/**' + - '.github/workflows/zarr-indexing.yml' + pull_request: + paths: + - 'packages/zarr-indexing/**' + - '.github/workflows/zarr-indexing.yml' + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + test: + name: pytest py=${{ matrix.python-version }} + runs-on: ubuntu-latest + defaults: + run: + shell: bash + strategy: + fail-fast: false + matrix: + python-version: ['3.12', '3.13', '3.14'] + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + persist-credentials: false + - name: Install uv + uses: astral-sh/setup-uv@11f9893b081a58869d3b5fccaea48c9e9e46f990 # v8.3.2 + with: + enable-cache: true + - name: Set up Python ${{ matrix.python-version }} + run: uv python install ${{ matrix.python-version }} + # The transform tests exercise chunk resolution against zarr's ChunkGrid, + # so they run from the repo root against the root environment (which + # provides `zarr`) with this package as an editable overlay rather than in + # package isolation. + - name: Sync test dependency group + run: uv sync --group test --python ${{ matrix.python-version }} + - name: Run pytest + run: uv run --no-sync --group test --with-editable ./packages/zarr-indexing python -m pytest packages/zarr-indexing/tests + + ruff: + name: ruff + runs-on: ubuntu-latest + defaults: + run: + shell: bash + working-directory: packages/zarr-indexing + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + persist-credentials: false + - name: Install uv + uses: astral-sh/setup-uv@11f9893b081a58869d3b5fccaea48c9e9e46f990 # v8.3.2 + - name: Run ruff + run: uvx ruff check . + + pyright: + name: pyright + runs-on: ubuntu-latest + defaults: + run: + shell: bash + working-directory: packages/zarr-indexing + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + persist-credentials: false + - name: Install uv + uses: astral-sh/setup-uv@11f9893b081a58869d3b5fccaea48c9e9e46f990 # v8.3.2 + with: + enable-cache: true + - name: Set up Python + run: uv python install 3.12 + - name: Sync test dependency group + run: uv sync --group test --python 3.12 + - name: Run pyright + run: uv run --group test --with pyright pyright src + + docs: + name: docs + runs-on: ubuntu-latest + defaults: + run: + shell: bash + working-directory: packages/zarr-indexing + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + persist-credentials: false + - name: Install uv + uses: astral-sh/setup-uv@11f9893b081a58869d3b5fccaea48c9e9e46f990 # v8.3.2 + with: + enable-cache: true + - name: Install just + uses: extractions/setup-just@53165ef7e734c5c07cb06b3c8e7b647c5aa16db3 # v4 + - name: Build docs + # The strict mkdocs build lives in packages/zarr-indexing/justfile. + run: just docs-check + + zarr-indexing-complete: + name: zarr-indexing complete + needs: [test, ruff, pyright, docs] + if: always() + runs-on: ubuntu-latest + steps: + - name: Check failure + if: | + contains(needs.*.result, 'failure') || + contains(needs.*.result, 'cancelled') + run: exit 1 + - name: Success + run: echo Success! diff --git a/.readthedocs.yaml b/.readthedocs.yaml index 55b5d6fed0..dddf8449a4 100644 --- a/.readthedocs.yaml +++ b/.readthedocs.yaml @@ -6,15 +6,14 @@ build: python: "3.12" jobs: post_checkout: - # Cancel pull request builds whose changes are confined to the - # zarr-metadata package, which has its own Read the Docs project. Exit - # code 183 cancels the build and reports success to the Git provider. - # Scoped to PR builds ("external" versions) because origin/main is only - # a meaningful diff base there. Read the Docs strips shell quoting from - # commands, so the exclude pathspec must use the quote-free :! form, - # not ':(exclude)'. + # Cancel pull request builds whose changes are confined to the packages + # that have their own Read the Docs projects. Exit code 183 cancels the + # build and reports success to the Git provider. Scoped to PR builds + # ("external" versions) because origin/main is only a meaningful diff + # base there. Read the Docs strips shell quoting from commands, so the + # exclude pathspecs must use the quote-free :! form, not ':(exclude)'. - | - if [ "$READTHEDOCS_VERSION_TYPE" = "external" ] && git diff --quiet origin/main -- :!packages/zarr-metadata; + if [ "$READTHEDOCS_VERSION_TYPE" = "external" ] && git diff --quiet origin/main -- :!packages/zarr-metadata :!packages/zarr-indexing; then exit 183; fi diff --git a/mkdocs.yml b/mkdocs.yml index 87aaf23430..67431edd90 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -93,6 +93,7 @@ nav: - ' zarr.zeros': api/zarr/functions/zeros.md - ' zarr.zeros_like': api/zarr/functions/zeros_like.md - 'zarr-metadata ↪': https://zarr-metadata.readthedocs.io/ + - 'zarr-indexing ↪': https://zarr-indexing.readthedocs.io/ - release-notes.md - contributing.md hooks: diff --git a/packages/zarr-indexing/.readthedocs.yaml b/packages/zarr-indexing/.readthedocs.yaml new file mode 100644 index 0000000000..b8c7b76e2b --- /dev/null +++ b/packages/zarr-indexing/.readthedocs.yaml @@ -0,0 +1,30 @@ +# Read the Docs configuration for the zarr-indexing docs site, separate from +# the zarr-python site configured by the repo-root .readthedocs.yaml. The RTD +# project for zarr-indexing must set its configuration-file path to +# packages/zarr-indexing/.readthedocs.yaml. +version: 2 + +build: + os: ubuntu-22.04 + tools: + python: "3.12" + jobs: + post_checkout: + # Cancel pull request builds that do not touch this package. Exit code + # 183 cancels the build and reports success to the Git provider. Scoped + # to PR builds ("external" versions) because origin/main is only a + # meaningful diff base there. + - | + if [ "$READTHEDOCS_VERSION_TYPE" = "external" ] && git diff --quiet origin/main -- packages/zarr-indexing; + then + exit 183; + fi + install: + - pip install --upgrade pip + - pip install ./packages/zarr-indexing --group packages/zarr-indexing/pyproject.toml:docs + build: + html: + - mkdocs build --strict -f packages/zarr-indexing/mkdocs.yml --site-dir $READTHEDOCS_OUTPUT/html + +mkdocs: + configuration: packages/zarr-indexing/mkdocs.yml diff --git a/packages/zarr-indexing/CHANGELOG.md b/packages/zarr-indexing/CHANGELOG.md new file mode 100644 index 0000000000..7c4bc92cad --- /dev/null +++ b/packages/zarr-indexing/CHANGELOG.md @@ -0,0 +1,3 @@ +# Release notes + + diff --git a/packages/zarr-indexing/LICENSE.txt b/packages/zarr-indexing/LICENSE.txt new file mode 100644 index 0000000000..1e8da4d242 --- /dev/null +++ b/packages/zarr-indexing/LICENSE.txt @@ -0,0 +1,21 @@ +The MIT License (MIT) + +Copyright (c) 2015-2025 Zarr Developers + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/packages/zarr-indexing/README.md b/packages/zarr-indexing/README.md new file mode 100644 index 0000000000..ccdfe595a5 --- /dev/null +++ b/packages/zarr-indexing/README.md @@ -0,0 +1,53 @@ +# zarr-indexing + +Composable, lazy coordinate transforms for Zarr array indexing. + +Documentation: + +This package implements TensorStore-inspired index transforms. The core idea: +every indexing operation (slicing, fancy indexing, etc.) produces a coordinate +mapping from user space to storage space. These mappings compose lazily — no +I/O until you explicitly read or write. + +Key types: + +- `IndexDomain` — a rectangular region of integer coordinates +- `IndexTransform` — maps input coordinates to storage coordinates +- `ConstantMap`, `DimensionMap`, `ArrayMap` — the three ways a single output + dimension can depend on the input +- `compose` — chain two transforms into one + +The package depends only on NumPy and the standard library; it does not import +`zarr`. It is developed in the [zarr-python](https://github.com/zarr-developers/zarr-python) +repository and consumed by `zarr` to resolve array indexing operations. + +## Installation + +```bash +pip install zarr-indexing +``` + +## Developing + +Package-scoped development commands live in the [`justfile`](./justfile) +(requires [just](https://github.com/casey/just)): + +``` +just test # run the test suite (extra args go to pytest) +just lint # ruff, same invocation as CI +just typecheck # pyright, same invocation as CI +just docs-check # strict build of the docs site +just check # all of the above +just docs-serve # serve the docs site locally +``` + +Run them from this directory, or from anywhere in the repository as +`just packages/zarr-indexing/`. + +The test recipe runs against the workspace-root environment, because the +chunk-resolution tests exercise this package against `zarr`'s chunk grids and +`zarr` is deliberately not a dependency of this package. + +## License + +MIT diff --git a/packages/zarr-indexing/changes/3906.feature.md b/packages/zarr-indexing/changes/3906.feature.md new file mode 100644 index 0000000000..fa51b4438e --- /dev/null +++ b/packages/zarr-indexing/changes/3906.feature.md @@ -0,0 +1 @@ +Reworked the JSON layer to conform to the [ndsel](https://github.com/zarr-developers/ndsel) draft wire format, which adapts TensorStore's `IndexTransform`. A new `zarr_indexing.messages` module (`parse_ndsel`, `normalize_ndsel`, `NdselError`) is a pure JSON-to-JSON layer that accepts all five message kinds (`point`/`box`/`slice`/`points`/`transform`) and normalizes them to the canonical transform body, enforcing the full ndsel error taxonomy. The package is checked against the vendored, language-agnostic ndsel conformance corpus. `index_transform_to_json`/`index_transform_from_json` (and the domain variants) now produce and consume the canonical body. On serialization, orthogonal (`oindex`) `index_array` maps no longer emit `input_dimension` alongside `index_array` (a combination both ndsel and TensorStore reject), and degenerate all-singleton index arrays collapse to constant maps; the in-memory `input_dimension` is reconstructed from the array's dependency axes on load. diff --git a/packages/zarr-indexing/changes/README.md b/packages/zarr-indexing/changes/README.md new file mode 100644 index 0000000000..feb3f8674e --- /dev/null +++ b/packages/zarr-indexing/changes/README.md @@ -0,0 +1,25 @@ +Writing a changelog entry for `zarr-indexing` +----------------------------------------------- + +Fragments in **this** directory are release notes for the `zarr-indexing` +package only — kept separate from the parent zarr-python `changes/` +directory so a PR touching only `packages/zarr-indexing/` produces a +release note for this package only. + +Please put a new file in this directory named `xxxx..md`, where + +- `xxxx` is the pull request number associated with this entry +- `` is one of: + - feature + - bugfix + - doc + - removal + - misc + +Inside the file, please write a short description of what you have +changed, and how it impacts users of `zarr-indexing`. + +A `zarr-indexing` release runs `towncrier build` in `packages/zarr-indexing/`, +which consumes the fragments here and updates `CHANGELOG.md`. Fragments +that describe parent zarr-python changes (not the transforms package) +belong in the top-level `changes/` directory, not here. diff --git a/packages/zarr-indexing/docs/_static/favicon-96x96.png b/packages/zarr-indexing/docs/_static/favicon-96x96.png new file mode 100644 index 0000000000..e77977ccf4 Binary files /dev/null and b/packages/zarr-indexing/docs/_static/favicon-96x96.png differ diff --git a/packages/zarr-indexing/docs/_static/logo_bw.png b/packages/zarr-indexing/docs/_static/logo_bw.png new file mode 100644 index 0000000000..df1979d3cc Binary files /dev/null and b/packages/zarr-indexing/docs/_static/logo_bw.png differ diff --git a/packages/zarr-indexing/docs/api/chunk_resolution.md b/packages/zarr-indexing/docs/api/chunk_resolution.md new file mode 100644 index 0000000000..0d81ec3829 --- /dev/null +++ b/packages/zarr-indexing/docs/api/chunk_resolution.md @@ -0,0 +1,5 @@ +--- +title: chunk_resolution +--- + +::: zarr_indexing.chunk_resolution diff --git a/packages/zarr-indexing/docs/api/composition.md b/packages/zarr-indexing/docs/api/composition.md new file mode 100644 index 0000000000..59affe016c --- /dev/null +++ b/packages/zarr-indexing/docs/api/composition.md @@ -0,0 +1,5 @@ +--- +title: composition +--- + +::: zarr_indexing.composition diff --git a/packages/zarr-indexing/docs/api/domain.md b/packages/zarr-indexing/docs/api/domain.md new file mode 100644 index 0000000000..b039d3a7f4 --- /dev/null +++ b/packages/zarr-indexing/docs/api/domain.md @@ -0,0 +1,5 @@ +--- +title: domain +--- + +::: zarr_indexing.domain diff --git a/packages/zarr-indexing/docs/api/errors.md b/packages/zarr-indexing/docs/api/errors.md new file mode 100644 index 0000000000..994a74248a --- /dev/null +++ b/packages/zarr-indexing/docs/api/errors.md @@ -0,0 +1,5 @@ +--- +title: errors +--- + +::: zarr_indexing.errors diff --git a/packages/zarr-indexing/docs/api/grid.md b/packages/zarr-indexing/docs/api/grid.md new file mode 100644 index 0000000000..c4c9cadb4f --- /dev/null +++ b/packages/zarr-indexing/docs/api/grid.md @@ -0,0 +1,5 @@ +--- +title: grid +--- + +::: zarr_indexing.grid diff --git a/packages/zarr-indexing/docs/api/index.md b/packages/zarr-indexing/docs/api/index.md new file mode 100644 index 0000000000..b9a58b70fa --- /dev/null +++ b/packages/zarr-indexing/docs/api/index.md @@ -0,0 +1,47 @@ +--- +title: API reference +--- + +# API reference + +The modules are layered: the transform algebra at the bottom, chunk resolution +and the wire format built on top of it. + +**The transform algebra** + +- [`zarr_indexing.domain`](domain.md) — `IndexDomain`, a rectangular region of + integer coordinates with an explicit (possibly non-zero) origin +- [`zarr_indexing.output_map`](output_map.md) — `ConstantMap`, `DimensionMap`, + and `ArrayMap`: three representations of a set of integer coordinates, one + per storage dimension +- [`zarr_indexing.transform`](transform.md) — `IndexTransform`, which pairs a + domain with output maps, plus the indexing (`[...]`, `.oindex`, `.vindex`), + `intersect`, and `translate` operations, and `selection_to_transform` +- [`zarr_indexing.composition`](composition.md) — `compose`, which chains two + transforms into one + +**Chunk resolution** + +- [`zarr_indexing.chunk_resolution`](chunk_resolution.md) — + `iter_chunk_transforms` (transform + chunk grid → per-chunk transforms) and + `sub_transform_to_selections` (the bridge back to the selection tuples the + current codec pipeline expects) +- [`zarr_indexing.grid`](grid.md) — `DimensionGridLike`, the Protocol + describing the narrow chunk-grid surface chunk resolution consumes, so that + nothing here imports `zarr` + +**The ndsel wire format** (see [the guide](../ndsel.md)) + +- [`zarr_indexing.messages`](messages.md) — `parse_ndsel` / `normalize_ndsel`, + the pure JSON→JSON message layer, and `NdselError` +- [`zarr_indexing.json`](json.md) — lowering between canonical ndsel bodies and + in-memory transforms + +**Errors** + +- [`zarr_indexing.errors`](errors.md) — the canonical index-error types, which + `zarr.errors` re-exports by identity + +Every name listed in `zarr_indexing.__all__` is re-exported at the top level, +so `from zarr_indexing import IndexTransform` and +`from zarr_indexing.transform import IndexTransform` are equivalent. diff --git a/packages/zarr-indexing/docs/api/json.md b/packages/zarr-indexing/docs/api/json.md new file mode 100644 index 0000000000..0183ab30e7 --- /dev/null +++ b/packages/zarr-indexing/docs/api/json.md @@ -0,0 +1,5 @@ +--- +title: json +--- + +::: zarr_indexing.json diff --git a/packages/zarr-indexing/docs/api/messages.md b/packages/zarr-indexing/docs/api/messages.md new file mode 100644 index 0000000000..6c2a434540 --- /dev/null +++ b/packages/zarr-indexing/docs/api/messages.md @@ -0,0 +1,5 @@ +--- +title: messages +--- + +::: zarr_indexing.messages diff --git a/packages/zarr-indexing/docs/api/output_map.md b/packages/zarr-indexing/docs/api/output_map.md new file mode 100644 index 0000000000..55114a6997 --- /dev/null +++ b/packages/zarr-indexing/docs/api/output_map.md @@ -0,0 +1,5 @@ +--- +title: output_map +--- + +::: zarr_indexing.output_map diff --git a/packages/zarr-indexing/docs/api/transform.md b/packages/zarr-indexing/docs/api/transform.md new file mode 100644 index 0000000000..bd754f5ec5 --- /dev/null +++ b/packages/zarr-indexing/docs/api/transform.md @@ -0,0 +1,5 @@ +--- +title: transform +--- + +::: zarr_indexing.transform diff --git a/packages/zarr-indexing/docs/index.md b/packages/zarr-indexing/docs/index.md new file mode 100644 index 0000000000..d123cbd32b --- /dev/null +++ b/packages/zarr-indexing/docs/index.md @@ -0,0 +1,150 @@ +# zarr-indexing + +Composable, lazy coordinate transforms for Zarr array indexing. + +`zarr-indexing` is developed in the +[zarr-python repository](https://github.com/zarr-developers/zarr-python/tree/main/packages/zarr-indexing) +and released independently of `zarr` itself. Install it with: + +``` +pip install zarr-indexing +``` + +## What this is + +An indexing operation — a slice, an integer, a fancy index array — is a +*mapping* from the coordinates a user asks for to the coordinates that live in +storage. This library makes that mapping a first-class value: an +[`IndexTransform`](api/transform.md). Transforms compose, so a view of a view +of an array is still a single transform, and nothing is read until someone +asks for data. + +Three pieces do the work: + +- **The transform algebra** ([`zarr_indexing.transform`](api/transform.md), + [`zarr_indexing.domain`](api/domain.md), + [`zarr_indexing.output_map`](api/output_map.md), + [`zarr_indexing.composition`](api/composition.md)): an `IndexTransform` + pairs an input [`IndexDomain`](api/domain.md) — a rectangular region of + integer coordinates, which unlike NumPy may have a non-zero origin — with + one output map per storage dimension. `ConstantMap`, `DimensionMap`, and + `ArrayMap` are three representations of the same thing, a set of integer + coordinates, traded off against each other for efficiency. +- **Chunk resolution** ([`zarr_indexing.chunk_resolution`](api/chunk_resolution.md)): + given a transform and a chunk grid, which chunks does this selection touch, + which coordinates does it touch *inside* each chunk, and where do the values + land in the output buffer? The resolver is dependency-aware: correlated + (`vindex`) array maps are enumerated jointly rather than as a cartesian + product, and orthogonal (`oindex`) array maps contribute only the chunks + their index arrays actually land in, so resolution scales with the number of + selected coordinates instead of with the size of the grid. +- **A wire format** ([`zarr_indexing.messages`](api/messages.md), + [`zarr_indexing.json`](api/json.md)): selections serialize to and from + [ndsel](https://github.com/zarr-developers/ndsel), a JSON representation of + NumPy-style n-dimensional selections. See [the ndsel wire format](ndsel.md). + +The package depends only on NumPy and the standard library. In particular it +does not import `zarr`: the chunk-grid surface chunk resolution needs is +described by the [`DimensionGridLike`](api/grid.md) Protocol, which zarr's +per-dimension grids satisfy structurally. + +## Relationship to TensorStore + +The model is [TensorStore's](https://google.github.io/tensorstore/index_space.html) +index transform, reimplemented in Python against NumPy: index domains with +explicit origins, output index maps of constant / single-input-dimension / +index-array flavour, and composition as the single operation that stacks +views. Names and semantics follow TensorStore where they overlap — notably, +negative indices are literal coordinates, not Python-style offsets from the +end, and it is the caller's job to normalize them. + +The differences are the ones NumPy compatibility forces. `ArrayMap` records +the input dimension an *orthogonal* (`oindex`) index array varies over, which +TensorStore's format has no field for; the +[serializer collapses or reconstructs that field](api/json.md) so the wire +format stays TensorStore-loadable. Chunk resolution and the `oindex`/`vindex` +helpers exist to serve NumPy-shaped selection semantics, which TensorStore +does not have to model. + +## Quickstart + +Indexing a transform produces a new transform. No I/O happens, and no +coordinates are materialized: + +```python +from zarr_indexing import IndexTransform + +transform = IndexTransform.from_shape((100, 100)) + +view = transform[10:50, 5] +view.domain # IndexDomain(inclusive_min=(10,), exclusive_max=(50,)) +view.selection_repr # '{ [10, 50), 5 }' +``` + +The domain describes what the *user* sees (here a single dimension, 40 long, +with origin 10); the output maps describe what *storage* sees (a stride-1 +`DimensionMap` and the `ConstantMap` for the dropped dimension). + +Fancy indexing works the same way, in both flavours, and still materializes +nothing but the index arrays themselves: + +```python +import numpy as np + +transform.oindex[np.array([3, 1, 90]), 0:4] # '{ {3, 1, 90}, [0, 4) }', shape (3, 4) +transform.vindex[np.array([0, 40, 99]), np.array([1, 2, 3])] # shape (3,) +``` + +Transforms built independently stack with +[`compose`](api/composition.md), which is what `transform[...]` uses +internally when you index an already-indexed view: + +```python +from zarr_indexing import compose + +inner = IndexTransform.from_shape((100,))[::2] # storage 0, 2, 4, ... over domain [0, 50) +outer = IndexTransform.from_shape((50,))[10:20] + +compose(outer, inner).selection_repr # '{ [20, 40) step 2 }' +``` + +Resolution against a chunk grid is where a transform finally meets storage. +Chunk resolution asks the grid only for the per-dimension index-to-chunk +mapping described by [`DimensionGridLike`](api/grid.md), so any object with +those four methods will do: + +```python +from dataclasses import dataclass + +from zarr_indexing import iter_chunk_transforms + + +@dataclass(frozen=True) +class RegularDimensionGrid: + chunk: int + + def index_to_chunk(self, idx): return idx // self.chunk + def chunk_offset(self, chunk_ix): return chunk_ix * self.chunk + def chunk_size(self, chunk_ix): return self.chunk + def indices_to_chunks(self, indices): return indices // self.chunk + + +grids = [RegularDimensionGrid(32), RegularDimensionGrid(32)] + +for chunk_coords, sub_transform, out_indices in iter_chunk_transforms(view, grids): + print(chunk_coords, sub_transform.selection_repr) +# (0, 0) { [10, 32), 5 } +# (1, 0) { [0, 18), 5 } +``` + +Each yielded `sub_transform` is the original transform restricted to one chunk +and translated into chunk-local coordinates — exactly what a codec pipeline +needs to decode that chunk and scatter the result. `out_indices` carries the +output scatter indices for array selections, and is `None` for basic indexing. + +## Reference + +- [The ndsel wire format](ndsel.md) +- [API reference](api/index.md) +- [Changelog](https://github.com/zarr-developers/zarr-python/blob/main/packages/zarr-indexing/CHANGELOG.md) +- [License (MIT)](https://github.com/zarr-developers/zarr-python/blob/main/packages/zarr-indexing/LICENSE.txt) diff --git a/packages/zarr-indexing/docs/ndsel.md b/packages/zarr-indexing/docs/ndsel.md new file mode 100644 index 0000000000..74f394b758 --- /dev/null +++ b/packages/zarr-indexing/docs/ndsel.md @@ -0,0 +1,156 @@ +--- +title: The ndsel wire format +--- + +# The ndsel wire format + +[ndsel](https://github.com/zarr-developers/ndsel) is a draft JSON +representation of NumPy-style n-dimensional selections, adapted from +TensorStore's `IndexTransform` model. `zarr-indexing` implements it in two +layers, and the split between them is the thing worth understanding: + +| Layer | Module | Depends on | Job | +| --- | --- | --- | --- | +| Message | [`zarr_indexing.messages`](api/messages.md) | stdlib only | JSON in, canonical JSON out. Validates and desugars. Never rounds, clamps, or drops information. | +| Engine | [`zarr_indexing.json`](api/json.md) | NumPy | Lowers a *canonical* body into an in-memory [`IndexTransform`](api/transform.md), and back. | + +Constraints that only make sense for a real array — finite bounds, index +arrays as `ndarray`s — live in the engine layer and nowhere else. That is why +`messages` can happily normalize a message with `"-inf"` bounds that +`json.transform_from_canonical` will refuse to lower. + +## Two entry points + +[`parse_ndsel`](api/messages.md#zarr_indexing.messages.parse_ndsel) +structurally validates a message of any kind and returns it **unchanged** — +use it when you want to keep a message in its compact shorthand form but +confirm it is well formed. + +[`normalize_ndsel`](api/messages.md#zarr_indexing.messages.normalize_ndsel) +desugars a message into the single deterministic **canonical transform body** +of the spec (section 4.3): a bare `IndexTransform` body without the `kind` +discriminator. + +```python +from zarr_indexing import normalize_ndsel + +normalize_ndsel({"kind": "box", "inclusive_min": [10, 5], "shape": [40, 1]}) +# {'input_rank': 2, +# 'input_inclusive_min': [10, 5], +# 'input_exclusive_max': [50, 6], +# 'input_labels': ['', ''], +# 'output': [{'offset': 0, 'stride': 1, 'input_dimension': 0}, +# {'offset': 0, 'stride': 1, 'input_dimension': 1}]} +``` + +Normalization is idempotent: re-tag the output with `kind: "transform"` and +normalizing it again returns the same body. Because the canonical body is +field-for-field a TensorStore `IndexTransform` minus `kind`, a normalized +message loads directly into `tensorstore.IndexTransform(json=...)`. + +Both entry points raise +[`NdselError`](api/messages.md#zarr_indexing.messages.NdselError), which +carries the spec `reason` code (`unknown_kind`, `rank_mismatch`, `step_zero`, +`output_map_conflict`, …) alongside a human-readable detail, so callers can +branch on the code rather than on message text. + +## The five message kinds + +Four are shorthands; the fifth is the canonical form itself. + +| `kind` | Fields | Selects | +| --- | --- | --- | +| `point` | `coords` | A single element. Normalizes to rank 0 with one `constant` output map per dimension. | +| `box` | `inclusive_min`, one of `exclusive_max` / `inclusive_max` / `shape`, `labels` | A rectangular region. Exactly one upper-bound spelling may appear. | +| `slice` | `start`, `stop`, `step`, `labels` | A strided region, one Python-style slice per dimension. | +| `points` | `coords` (a list of coordinate rows) | An explicit list of points — the `vindex` case. Normalizes to one `index_array` output map per dimension over a shared rank-1 input domain. | +| `transform` | `input_rank`, `input_inclusive_min`, one of the three `input_*` upper bounds, `input_labels`, `output` | The full canonical form. | + +Value rules the message layer enforces throughout: every integer is a 64-bit +signed value; JSON booleans are **not** integers (Python's +`isinstance(True, int)` is guarded against explicitly); the `"-inf"` / `"+inf"` +sentinels are legal only in bound positions; and an implicit bound is the +one-element `[n]`-bracket form, whose implicit/explicit flag survives +normalization intact. + +## Lowering to a transform + +The engine layer converts between canonical bodies and `IndexTransform`s: + +```python +from zarr_indexing import transform_from_canonical, transform_to_canonical + +t = transform_from_canonical(canonical) +transform_to_canonical(t) == canonical +``` + +`index_transform_to_json` / `index_transform_from_json` (and the +`index_domain_*` variants) are these same converters under their historical +names. + +Two engine constraints apply here and only here. A canonical body carrying a +`"-inf"` or `"+inf"` bound cannot be lowered — an `IndexDomain` addresses a +finite array — so `transform_from_canonical` raises. And implicit bounds lower +*by value*: the `[n]`-bracket flag is a message-layer concern, and the engine +keeps only the integer. + +### The `index_array` round trip + +ndsel and TensorStore both **reject** an output map that carries both +`input_dimension` and `index_array`. The in-memory +[`ArrayMap`](api/output_map.md#zarr_indexing.output_map.ArrayMap), though, +records an `input_dimension` to pin the axis an orthogonal (`oindex`) array +varies over. The serializer bridges that gap in both directions: + +- **On serialize**, a non-degenerate `index_array` map is emitted *without* + `input_dimension`. +- **On load**, the in-memory `input_dimension` is reconstructed from the + full-rank array's dependency axes — its non-singleton axes. An array that + solely owns a single non-singleton axis is orthogonal; arrays that share + non-singleton axes, or vary over several, are correlated (`vindex`), and get + `input_dimension = None`. A single 1-D array over a rank-1 domain is + inherently ambiguous between the two flavours and reconstructs as + orthogonal, which is behaviorally identical in that case. + +There is one deliberate exception, worth calling out because it is the one +place a round trip changes representation rather than preserving it. An +all-singleton `index_array` — size 1 — selects the same coordinate no matter +what the input is, so it is **collapsed to a `constant` map** on serialize: + +```python +from zarr_indexing import IndexTransform, transform_to_canonical + +transform_to_canonical(IndexTransform.from_shape((100, 100)).oindex[[5], 0:2]) +# {'input_rank': 2, +# 'input_inclusive_min': [0, 0], +# 'input_exclusive_max': [1, 2], +# 'input_labels': ['', ''], +# 'output': [{'offset': 5}, +# {'offset': 0, 'stride': 1, 'input_dimension': 1}]} +``` + +The size-1 input dimension stays in the domain, unconsumed by any output map — +still a valid transform, and still the right output shape. A length-1 `oindex` +selection therefore round-trips *behaviorally* (an `ArrayMap` comes back as a +`ConstantMap`) rather than by object identity. + +## Conformance + +The package is checked against the language-agnostic ndsel conformance corpus, +vendored unmodified under +[`tests/conformance/`](https://github.com/zarr-developers/zarr-python/tree/main/packages/zarr-indexing/tests/conformance) +— one JSON file per message kind plus `errors.json`, with the source commit +recorded in `PROVENANCE.md`. Each fixture is either a *success* case +(`input` + expected `normalized` body) or an *error* case (`input` + expected +reason code), and an implementation is conformant iff `normalize` reproduces +every one. `tests/test_conformance.py` runs the whole corpus as one +parametrized test per fixture, so a corpus update reports failures fixture by +fixture rather than as a single opaque assertion. + +Do not edit the vendored files; to pick up spec changes, re-vendor from a newer +ndsel commit and update the recorded SHA. + +A second, optional test (`tests/test_ndsel_tensorstore.py`, skipped unless +`tensorstore` is installed) closes the loop against a real TensorStore by +loading canonical bodies into `tensorstore.IndexTransform` and re-loading +TensorStore's own `to_json()` output back through the engine layer. diff --git a/packages/zarr-indexing/justfile b/packages/zarr-indexing/justfile new file mode 100644 index 0000000000..20e1b2b837 --- /dev/null +++ b/packages/zarr-indexing/justfile @@ -0,0 +1,58 @@ +# Development verbs for the zarr-indexing package. Recipes run with this +# directory as the working directory regardless of where `just` is invoked. + +# List available recipes +default: + @just --list + +# The chunk-resolution tests exercise this package against zarr's ChunkGrid, so +# they need an environment that has both `zarr` and this package installed. +# `zarr` is deliberately not a dependency of this package, and the repo is not +# a uv workspace, so run against the repo-root environment (which provides +# `zarr`) with this package layered in as an editable overlay — the same +# invocation CI uses. +# Run the test suite; extra args are passed to pytest +test *args: + uv run --project ../.. --group test --with-editable . python -m pytest tests {{ args }} + +# Lint with the same invocation CI uses +lint: + uvx ruff check . + +# Type-check the package sources +typecheck: + uv run --group test --with pyright pyright src + +# Run everything CI runs for this package +check: lint typecheck test docs-check + +# Preview the changelog that the next release would generate +changelog-draft: + uvx towncrier build --draft --version Unreleased + +# Build this package's documentation site, warnings as errors +docs-check: + env DISABLE_MKDOCS_2_WARNING=true uv run --group docs mkdocs build --strict + +# With no argument, uses port 8000 if free, otherwise an ephemeral free port; +# an explicitly requested port is used as-is so a conflict fails loudly. +# Serve this package's documentation site +docs-serve port="": + #!/usr/bin/env bash + set -euo pipefail + port="{{ port }}" + if [ -z "$port" ]; then + port=$(uv run --group docs python -c ' + import socket + s = socket.socket() + try: + s.bind(("127.0.0.1", 8000)) + except OSError: + s.close() + s = socket.socket() + s.bind(("127.0.0.1", 0)) + print(s.getsockname()[1]) + s.close() + ') + fi + exec env DISABLE_MKDOCS_2_WARNING=true uv run --group docs mkdocs serve -a "localhost:$port" diff --git a/packages/zarr-indexing/mkdocs.yml b/packages/zarr-indexing/mkdocs.yml new file mode 100644 index 0000000000..d7261f32e1 --- /dev/null +++ b/packages/zarr-indexing/mkdocs.yml @@ -0,0 +1,110 @@ +site_name: zarr-indexing +# The package lives in the zarr-python monorepo; point the header source +# widget at the package directory rather than the repository root. +repo_name: zarr-python/packages/zarr-indexing +repo_url: https://github.com/zarr-developers/zarr-python/tree/main/packages/zarr-indexing +# Absolute because mkdocs would otherwise append this to repo_url's subpath. +edit_uri: https://github.com/zarr-developers/zarr-python/edit/main/packages/zarr-indexing/docs/ +site_description: Composable, lazy coordinate transforms for Zarr array indexing. +site_author: Davis Bennett +site_url: !ENV [READTHEDOCS_CANONICAL_URL, 'https://zarr-indexing.readthedocs.io/'] +docs_dir: docs +use_directory_urls: true + +nav: + - index.md + - ndsel.md + - API Reference: + - api/index.md + - ' zarr_indexing.transform': api/transform.md + - ' zarr_indexing.domain': api/domain.md + - ' zarr_indexing.output_map': api/output_map.md + - ' zarr_indexing.composition': api/composition.md + - ' zarr_indexing.chunk_resolution': api/chunk_resolution.md + - ' zarr_indexing.grid': api/grid.md + - ' zarr_indexing.json': api/json.md + - ' zarr_indexing.messages': api/messages.md + - ' zarr_indexing.errors': api/errors.md + - Changelog: https://github.com/zarr-developers/zarr-python/blob/main/packages/zarr-indexing/CHANGELOG.md + +watch: + - src + +theme: + language: en + name: material + logo: _static/logo_bw.png + favicon: _static/favicon-96x96.png + + palette: + # Light mode + - media: "(prefers-color-scheme: light)" + scheme: default + toggle: + icon: material/brightness-7 + name: Switch to dark mode + + # Dark mode + - media: "(prefers-color-scheme: dark)" + scheme: slate + toggle: + icon: material/brightness-4 + name: Switch to light mode + + font: + text: Roboto + code: Roboto Mono + + features: + - content.code.annotate + - content.code.copy + - navigation.indexes + - navigation.instant + - navigation.tracking + - search.suggest + - search.share + +plugins: + - autorefs + - search + - mkdocstrings: + enable_inventory: true + handlers: + python: + paths: [src] + options: + allow_inspection: true + docstring_section_style: list + docstring_style: numpy + inherited_members: true + line_length: 60 + separate_signature: true + show_root_heading: true + show_signature_annotations: true + show_source: true + show_symbol_type_toc: true + signature_crossrefs: true + show_if_no_docstring: true + extensions: + - griffe_inherited_docstrings + + inventories: + - https://docs.python.org/3/objects.inv + - https://numpy.org/doc/stable/objects.inv + - https://zarr.readthedocs.io/en/stable/objects.inv + +markdown_extensions: + - admonition + - attr_list + - def_list + - footnotes + - md_in_html + - pymdownx.details + - pymdownx.superfences + - toc: + permalink: true + - pymdownx.highlight: + anchor_linenums: true + line_spans: __span + pygments_lang_class: true + - pymdownx.inlinehilite diff --git a/packages/zarr-indexing/pyproject.toml b/packages/zarr-indexing/pyproject.toml new file mode 100644 index 0000000000..21ba4ef7e7 --- /dev/null +++ b/packages/zarr-indexing/pyproject.toml @@ -0,0 +1,124 @@ +[build-system] +requires = ["hatchling>=1.29.0", "hatch-vcs"] +build-backend = "hatchling.build" + +[project] +name = "zarr-indexing" +dynamic = ["version"] +description = "Composable, lazy coordinate transforms for Zarr array indexing." +readme = "README.md" +requires-python = ">=3.12" +license = "MIT" +license-files = ["LICENSE.txt"] +authors = [ + { name = "Davis Bennett", email = "davis.v.bennett@gmail.com" }, +] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "Intended Audience :: Information Technology", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: MIT License", + "Operating System :: OS Independent", + "Programming Language :: Python", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", + "Topic :: Scientific/Engineering", + "Topic :: Software Development :: Libraries :: Python Modules", + "Typing :: Typed", +] +keywords = ["zarr"] +dependencies = [ + "numpy>=2", +] + +[project.urls] +Homepage = "https://github.com/zarr-developers/zarr-python" +Source = "https://github.com/zarr-developers/zarr-python/tree/main/packages/zarr-indexing" +Issues = "https://github.com/zarr-developers/zarr-python/issues" +Changelog = "https://github.com/zarr-developers/zarr-python/blob/main/packages/zarr-indexing/CHANGELOG.md" +Documentation = "https://zarr-indexing.readthedocs.io/" + +[dependency-groups] +# The transform tests exercise chunk resolution against zarr's ChunkGrid +# (tests/test_chunk_resolution.py) and are collected by the parent zarr-python +# test suite, which already has zarr installed. `zarr` is intentionally NOT +# listed here to avoid a workspace dependency cycle; run these tests from the +# repo root (`uv run pytest packages/zarr-indexing/tests`), not in isolation. +test = ["pytest"] +docs = [ + # Pins match the zarr-python docs environment in the repo-root + # pyproject.toml so the two sites render with the same toolchain. + "mkdocs-material==9.7.6", + "mkdocs==1.6.1", + "mkdocstrings==1.0.4", + "mkdocstrings-python==2.0.5", + "griffe-inherited-docstrings==1.1.3", + # mkdocstrings uses ruff to format rendered signatures + "ruff==0.15.20", +] + +[tool.hatch.version] +source = "vcs" +tag-pattern = '^zarr_indexing-v(?P.+)$' +# `git_describe_command` ensures we get the zarr_indexing tags instead of latest. +# `local_scheme` strips the git commit info so the appending info is just a counter from latest tag. +# test-pypi doesn't accept git commit info in tags, and the count should be enough to distinguish unique runs. +raw-options = { root = "../..", git_describe_command = "git describe --dirty --tags --long --match zarr_indexing-v*", local_scheme = "no-local-version" } + +[tool.hatch.build.targets.wheel] +packages = ["src/zarr_indexing"] + +[tool.ruff] +extend = "../../pyproject.toml" +target-version = "py312" + +[tool.pytest.ini_options] +minversion = "7" +testpaths = ["tests"] +xfail_strict = true +addopts = ["-ra", "--strict-config", "--strict-markers"] +filterwarnings = [ + "error", +] + +[tool.pyright] +include = ["src"] +enableExperimentalFeatures = true +typeCheckingMode = "strict" +pythonVersion = "3.12" +# This strict config was written for zarr-metadata's JSON/dataclass-shaped +# code. zarr-indexing is numpy-heavy, and numpy's stubs return partially +# unknown types (e.g. `ndarray[Unknown, Unknown]`, `dtype[Unknown]`) even for +# fully-typed call sites, so the reportUnknown* family below cannot reasonably +# be satisfied here. Downgraded to warnings (not silenced) rather than +# disabled outright, and CI (which only fails the pyright job on errors, not +# warnings) still surfaces them for visibility. +reportUnknownVariableType = "warning" +reportUnknownArgumentType = "warning" +reportUnknownMemberType = "warning" +reportUnknownParameterType = "warning" + +[tool.numpydoc_validation] +checks = [ + "GL10", + "SS04", + "PR02", + "PR03", + "PR05", + "PR06", +] + +[tool.towncrier] +# Fragments for this package live alongside the package source, separate +# from the parent zarr-python `changes/` directory, so a PR touching only +# `packages/zarr-indexing/` produces a release note for this package only. +directory = "changes" +filename = "CHANGELOG.md" +package = "zarr_indexing" +underlines = ["", "", ""] +title_format = "## {version} ({project_date})" +issue_format = "[#{issue}](https://github.com/zarr-developers/zarr-python/issues/{issue})" +start_string = "\n" diff --git a/packages/zarr-indexing/src/zarr_indexing/__init__.py b/packages/zarr-indexing/src/zarr_indexing/__init__.py new file mode 100644 index 0000000000..effe38ca88 --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/__init__.py @@ -0,0 +1,74 @@ +"""Composable, lazy coordinate transforms for zarr array indexing. + +This package implements TensorStore-inspired index transforms. The core idea: +every indexing operation (slicing, fancy indexing, etc.) produces a coordinate +mapping from user space to storage space. These mappings compose lazily — no +I/O until you explicitly read or write. + +Key types: + +- `IndexDomain` — a rectangular region of integer coordinates +- `IndexTransform` — maps input coordinates to storage coordinates +- `ConstantMap`, `DimensionMap`, `ArrayMap` — the three ways a single + output dimension can depend on the input (see `output_map.py`) +- `compose` — chain two transforms into one + +The chunk-resolution helpers (`iter_chunk_transforms`, +`sub_transform_to_selections`) and `selection_to_transform` are also exported +here: they form the surface the zarr integration layer (array indexing) depends +on. The `*Like` grid Protocols describe the chunk-grid surface chunk resolution +consumes without importing zarr. +""" + +from importlib.metadata import version + +from zarr_indexing.chunk_resolution import ( + iter_chunk_transforms, + sub_transform_to_selections, +) +from zarr_indexing.composition import compose +from zarr_indexing.domain import IndexDomain +from zarr_indexing.grid import DimensionGridLike +from zarr_indexing.json import ( + IndexDomainJSON, + IndexTransformJSON, + OutputIndexMapJSON, + index_domain_from_json, + index_domain_to_json, + index_transform_from_json, + index_transform_to_json, + transform_from_canonical, + transform_to_canonical, +) +from zarr_indexing.messages import NdselError, normalize_ndsel, parse_ndsel +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap +from zarr_indexing.transform import IndexTransform, selection_to_transform + +__version__ = version("zarr-indexing") + +__all__ = [ + "ArrayMap", + "ConstantMap", + "DimensionGridLike", + "DimensionMap", + "IndexDomain", + "IndexDomainJSON", + "IndexTransform", + "IndexTransformJSON", + "NdselError", + "OutputIndexMap", + "OutputIndexMapJSON", + "__version__", + "compose", + "index_domain_from_json", + "index_domain_to_json", + "index_transform_from_json", + "index_transform_to_json", + "iter_chunk_transforms", + "normalize_ndsel", + "parse_ndsel", + "selection_to_transform", + "sub_transform_to_selections", + "transform_from_canonical", + "transform_to_canonical", +] diff --git a/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py new file mode 100644 index 0000000000..7aea86ad02 --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py @@ -0,0 +1,380 @@ +"""Chunk resolution — mapping transforms to chunk-level I/O. + +Given an `IndexTransform` (which coordinates a user wants to access) and a +`ChunkGrid` (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 output buffer? + +The algorithm is: + +1. **Enumerate candidate chunks** — determine which chunks could possibly + be touched by the transform's output coordinate ranges. + +2. **Intersect** — for each candidate chunk, call + `transform.intersect(chunk_domain)` to restrict the transform to + coordinates within that chunk. If the intersection is empty, skip it. + +3. **Translate** — shift the restricted transform to chunk-local coordinates + via `transform.translate(-chunk_origin)`. + +4. **Yield** — produce `(chunk_coords, local_transform, surviving_indices)` + triples that the codec pipeline consumes. + +Sorted one-dimensional correlated array maps can be partitioned directly +because every touched chunk owns a contiguous slice of the index array. That +case bypasses candidate enumeration and repeated intersection. + +`sub_transform_to_selections` bridges from the transform representation +back to the raw `(chunk_selection, out_selection, drop_axes)` tuples that +the current codec pipeline expects. This bridge will go away when the codec +pipeline accepts transforms natively. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import numpy as np + +from zarr_indexing.domain import IndexDomain +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr_indexing.transform import IndexTransform + +if TYPE_CHECKING: + from collections.abc import Iterator, Sequence + + from zarr_indexing.grid import DimensionGridLike + +OutIndices = ( + dict[int, np.ndarray[Any, np.dtype[np.intp]]] | np.ndarray[Any, np.dtype[np.intp]] | None +) + +ChunkTransformResult = tuple[ + tuple[int, ...], + IndexTransform, + OutIndices, +] + + +def _one_dimensional_correlated_array_map( + transform: IndexTransform, +) -> tuple[ArrayMap, np.ndarray[Any, np.dtype[np.intp]]] | None: + """Return a nonempty correlated 1-D ArrayMap and its storage coordinates. + + A one-dimensional array selection has no cross-dimensional correlation to + preserve. The computed storage coordinates are also reused by general + resolution when they are unsorted. + """ + if transform.input_rank != 1 or transform.output_rank != 1: + return None + + m = transform.output[0] + if ( + not isinstance(m, ArrayMap) + or m.input_dimension is not None + or m.index_array.ndim != 1 + or m.index_array.size == 0 + ): + return None + + return m, m.offset + m.stride * m.index_array + + +def _iter_sorted_1d_array_map( + m: ArrayMap, + storage: np.ndarray[Any, np.dtype[np.intp]], + dim_grid: DimensionGridLike, +) -> Iterator[ChunkTransformResult]: + """Resolve a sorted 1-D ArrayMap one touched chunk at a time.""" + start = 0 + while start < storage.size: + chunk = dim_grid.index_to_chunk(int(storage[start])) + chunk_start = dim_grid.chunk_offset(chunk) + chunk_stop = chunk_start + dim_grid.chunk_size(chunk) + stop = int(np.searchsorted(storage, chunk_stop, side="left")) + + restricted = IndexTransform( + domain=IndexDomain(inclusive_min=(0,), exclusive_max=(stop - start,)), + output=( + ArrayMap( + index_array=m.index_array[start:stop], + offset=m.offset, + stride=m.stride, + input_dimension=m.input_dimension, + ), + ), + ) + local = restricted.translate((-chunk_start,)) + surviving = np.arange(start, stop, dtype=np.intp) + + yield (chunk,), local, surviving + start = stop + + +def iter_chunk_transforms( + transform: IndexTransform, + dim_grids: Sequence[DimensionGridLike], +) -> Iterator[ChunkTransformResult]: + """Resolve a composed IndexTransform against per-dimension chunk grids. + + `dim_grids` holds one `DimensionGridLike` per output (storage) dimension — + for zarr this is the chunk grid's per-dimension sequence. Yields + `(chunk_coords, sub_transform, out_indices)` triples: + + - `chunk_coords`: which chunk to access. + - `sub_transform`: maps output buffer coords to chunk-local coords. + - `out_indices`: for vectorized/array indexing, the output scatter + indices (integer array). `None` for basic/slice indexing. + """ + + array_map_1d = _one_dimensional_correlated_array_map(transform) + if array_map_1d is not None: + sorted_map, storage = array_map_1d + if storage[0] <= storage[-1] and bool(np.all(storage[1:] >= storage[:-1])): + dim_grid = dim_grids[0] + first_chunk = dim_grid.index_to_chunk(int(storage[0])) + if dim_grid.chunk_size(first_chunk) > 0: + yield from _iter_sorted_1d_array_map(sorted_map, storage, dim_grid) + return + + # Enumerate candidate chunks via the cartesian product of per-slot candidate + # chunk ids, then for each candidate intersect the transform with the chunk + # domain (`transform.intersect` handles orthogonal and vectorized cases + # alike, filtering out combinations it does not actually touch). + # + # A slot covers one or more output dimensions and contributes exactly the + # chunk-coordinate tuples those dimensions can touch: + # + # - `ConstantMap`/`DimensionMap` dims each form their own slot with a + # contiguous range — a single chunk for a constant, and the span between + # the first and last chunk for a slice. These are already tight (or + # nearly so). + # - Orthogonal `ArrayMap` (fancy) dims each form their own slot with only + # the *distinct* chunk ids the index array actually lands in + # (`np.unique`), never the dense `range(min_chunk, max_chunk + 1)` + # between them. A sparse fancy selection (e.g. two far-apart coordinates) + # would otherwise enumerate every chunk in the bounding box, making + # resolution scale with grid size instead of with the number of selected + # coordinates. + # - Correlated (vindex) `ArrayMap` dims share one *joint* slot holding the + # distinct chunk-coordinate tuples the points actually land in. The + # cartesian product of their per-dimension distinct sets would include + # combinations no point touches — quadratic in the number of selected + # points for a diagonal selection — while the joint distinct set is + # bounded by the point count (see zarr-python gh-4174). + correlated_dims: list[int] = [] + correlated_chunk_ids: list[np.ndarray[Any, np.dtype[np.intp]]] = [] + slot_dims: list[tuple[int, ...]] = [] + slot_candidates: list[Sequence[tuple[int, ...]]] = [] + for out_dim, m in enumerate(transform.output): + dg = dim_grids[out_dim] + if isinstance(m, ConstantMap): + # Single chunk + c = dg.index_to_chunk(m.offset) + slot_dims.append((out_dim,)) + slot_candidates.append(((c,),)) + elif isinstance(m, DimensionMap): + d = m.input_dimension + dim_lo = transform.domain.inclusive_min[d] + dim_hi = transform.domain.exclusive_max[d] + if dim_lo >= dim_hi: + return # empty domain + if m.stride > 0: + s_min = m.offset + m.stride * dim_lo + s_max = m.offset + m.stride * (dim_hi - 1) + else: + s_min = m.offset + m.stride * (dim_hi - 1) + s_max = m.offset + m.stride * dim_lo + first = dg.index_to_chunk(s_min) + last = dg.index_to_chunk(s_max) + slot_dims.append((out_dim,)) + slot_candidates.append([(c,) for c in range(first, last + 1)]) + else: + # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap). + # Storage coordinates were already computed for a correlated 1-D map. + storage = ( + array_map_1d[1] if array_map_1d is not None else m.offset + m.stride * m.index_array + ) + if storage.size == 0: + # Empty fancy selection: no coordinates, so no chunks are touched. + return + # Keep the index-array shape: correlated maps broadcast against each + # other below, and raveling first would lose the singleton axes. + chunk_ids = dg.indices_to_chunks(storage.astype(np.intp)) + if m.input_dimension is None: + correlated_dims.append(out_dim) + correlated_chunk_ids.append(chunk_ids) + else: + slot_dims.append((out_dim,)) + slot_candidates.append([(int(c),) for c in np.unique(chunk_ids)]) + + if len(correlated_dims) == 1: + slot_dims.append((correlated_dims[0],)) + slot_candidates.append([(int(c),) for c in np.unique(correlated_chunk_ids[0])]) + elif len(correlated_dims) >= 2: + # Group the points jointly: distinct rows of the per-point chunk + # coordinates, O(points log points) regardless of grid size. + broadcast = np.broadcast_arrays(*correlated_chunk_ids) + stacked = np.stack([b.ravel() for b in broadcast], axis=1) + joint = np.unique(stacked, axis=0) + slot_dims.append(tuple(correlated_dims)) + slot_candidates.append([tuple(int(c) for c in row) for row in joint]) + + import itertools + + output_rank = len(transform.output) + for combo in itertools.product(*slot_candidates): + chunk_coords_list = [0] * output_rank + for dims, part in zip(slot_dims, combo, strict=True): + for d, c in zip(dims, part, strict=True): + chunk_coords_list[d] = c + chunk_coords = tuple(chunk_coords_list) + + # Build the chunk domain in storage space + chunk_min: list[int] = [] + chunk_max: list[int] = [] + chunk_shift: list[int] = [] + for out_dim, c in enumerate(chunk_coords): + dg = dim_grids[out_dim] + c_start = dg.chunk_offset(c) + c_size = dg.chunk_size(c) + chunk_min.append(c_start) + chunk_max.append(c_start + c_size) + chunk_shift.append(-c_start) + + chunk_domain = IndexDomain( + inclusive_min=tuple(chunk_min), + exclusive_max=tuple(chunk_max), + ) + + # Intersect transform with chunk domain + result = transform.intersect(chunk_domain) + if result is None: + continue + + restricted, surviving = result + + # Translate to chunk-local coordinates + local = restricted.translate(tuple(chunk_shift)) + + yield (chunk_coords, local, surviving) + + +def sub_transform_to_selections( + sub_transform: IndexTransform, + out_indices: OutIndices = None, +) -> tuple[ + tuple[int | slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]], ...], + tuple[slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]], ...], + tuple[int, ...], +]: + """Convert a chunk-local sub-transform to raw selections for the codec pipeline. + + Parameters + ---------- + sub_transform + A chunk-local IndexTransform (output maps already translated to + chunk-local coordinates). + out_indices + For vectorized indexing: the output scatter indices for this chunk. + None for orthogonal/basic indexing. + + Returns + ------- + tuple + `(chunk_selection, out_selection, drop_axes)` + """ + inclusive_min = sub_transform.domain.inclusive_min + exclusive_max = sub_transform.domain.exclusive_max + + # Orthogonal outer product: >= 2 ArrayMaps each bound to a distinct input + # dimension. out_indices is a per-output-dim dict of surviving positions. The + # codec applies chunk_array[chunk_sel] / out[out_sel] with NumPy semantics, so + # build np.ix_-style selections (mirroring the legacy OrthogonalIndexer): one + # 1-D selector per dimension, expanded to an open mesh. ConstantMap dims are + # size-1 in chunk space and squeezed out via drop_axes. + if isinstance(out_indices, dict): + chunk_arrays: list[np.ndarray[Any, np.dtype[np.intp]]] = [] + out_arrays: list[np.ndarray[Any, np.dtype[np.intp]]] = [] + drop_axes: list[int] = [] + for out_dim, m in enumerate(sub_transform.output): + if isinstance(m, ConstantMap): + chunk_arrays.append(np.array([m.offset], dtype=np.intp)) + drop_axes.append(out_dim) + elif isinstance(m, DimensionMap): + rng = np.arange(inclusive_min[m.input_dimension], exclusive_max[m.input_dimension]) + chunk_arrays.append((m.offset + m.stride * rng).astype(np.intp)) + out_arrays.append(rng.astype(np.intp)) + else: # ArrayMap + idx = m.index_array.ravel() + chunk_arrays.append((m.offset + m.stride * idx).astype(np.intp)) + out_arrays.append(out_indices[out_dim]) + return np.ix_(*chunk_arrays), np.ix_(*out_arrays), tuple(drop_axes) + + # Correlated (vindex) sub-transforms carry ArrayMaps with `input_dimension` + # None. They scatter through a single flat index (`out_indices`) into the + # row-major-flattened output buffer; the chunk selection reads a + # (points, residual-slice) block via the raveled coordinate arrays and any + # residual DimensionMap slices. + correlated = any( + isinstance(m, ArrayMap) and m.input_dimension is None for m in sub_transform.output + ) + if correlated: + chunk_sel: list[int | slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]]] = [] + for m in sub_transform.output: + if isinstance(m, ConstantMap): + chunk_sel.append(m.offset) + elif isinstance(m, DimensionMap): + d = m.input_dimension + start = m.offset + m.stride * inclusive_min[d] + stop = m.offset + m.stride * exclusive_max[d] + if m.stride < 0: + start, stop = stop + 1, start + 1 + chunk_sel.append(slice(start, stop, m.stride)) + else: # ArrayMap + idx = m.index_array.reshape(-1) + chunk_sel.append((m.offset + m.stride * idx).astype(np.intp)) + # Chunk resolution always supplies the flat scatter index for a + # correlated transform. Absent one (a bare sub-transform), fall back to an + # identity scatter over the whole flattened output buffer. + # `out_indices` is narrowed to a flat scatter array or None here (the + # per-dimension dict is an orthogonal outer product, handled above). + out_scatter: slice | np.ndarray[Any, np.dtype[np.intp]] + if out_indices is None: + n = 1 + for s in sub_transform.domain.shape: + n *= s + out_scatter = slice(0, n) + else: + out_scatter = out_indices + return tuple(chunk_sel), (out_scatter,), () + + chunk_sel = [] # annotated in the correlated branch above (same function scope) + out_sel: list[slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]]] = [] + + # Single-pass build for the basic / single-orthogonal-array cases. + # ConstantMap dims are dropped (no out_sel entry). + for m in sub_transform.output: + if isinstance(m, ConstantMap): + chunk_sel.append(m.offset) + elif isinstance(m, DimensionMap): + d = m.input_dimension + dim_lo = inclusive_min[d] + dim_hi = exclusive_max[d] + start = m.offset + m.stride * dim_lo + stop = m.offset + m.stride * dim_hi + if m.stride < 0: + start, stop = stop + 1, start + 1 + chunk_sel.append(slice(start, stop, m.stride)) + out_sel.append(slice(dim_lo, dim_hi)) + else: # ArrayMap (orthogonal: full-rank, raveled to its 1-D fancy coords) + idx = m.index_array.reshape(-1) + if m.offset == 0 and m.stride == 1: + chunk_sel.append(idx) + else: + chunk_sel.append((m.offset + m.stride * idx).astype(np.intp)) + # Orthogonal ArrayMap: out_indices holds the surviving positions. + out_sel.append(out_indices if out_indices is not None else slice(0, idx.size)) + + return tuple(chunk_sel), tuple(out_sel), () diff --git a/packages/zarr-indexing/src/zarr_indexing/composition.py b/packages/zarr-indexing/src/zarr_indexing/composition.py new file mode 100644 index 0000000000..f5cc82599c --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/composition.py @@ -0,0 +1,133 @@ +"""Composition — chaining two transforms into one. + +`compose(outer, inner)` is the operation that makes views stack. `outer` maps +user coordinates to intermediate coordinates, `inner` maps those intermediate +coordinates to storage, and the result maps user coordinates straight to +storage — so a view of a view of an array is still a single +`IndexTransform`, and indexing never accumulates layers to walk at read time. + +Composition works one output map at a time, and each case reduces to +substituting the outer map into the inner one: + +- A `ConstantMap` inner map ignores its input, so it survives unchanged. +- A `DimensionMap` inner map is affine, so composing it with an outer + `ConstantMap` or `DimensionMap` folds into new `offset`/`stride` values; + composing it with an outer `ArrayMap` leaves the index array alone and + rescales around it. +- An `ArrayMap` inner map must be *evaluated* at the coordinates the outer + transform produces, which is the only case that touches array data. +""" + +from __future__ import annotations + +import numpy as np + +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap +from zarr_indexing.transform import IndexTransform + + +def compose(outer: IndexTransform, inner: IndexTransform) -> IndexTransform: + """Compose two IndexTransforms. + + `outer` maps user coords (rank m) to intermediate coords (rank n). + `inner` maps intermediate coords (rank n) to storage coords (rank p). + The result maps user coords (rank m) to storage coords (rank p). + + Precondition: `outer.output_rank == inner.domain.ndim`. + """ + if outer.output_rank != inner.domain.ndim: + raise ValueError( + f"outer output rank ({outer.output_rank}) must match inner input rank " + f"({inner.domain.ndim})" + ) + + result_output = [_compose_single(outer, inner_map) for inner_map in inner.output] + + return IndexTransform(domain=outer.domain, output=tuple(result_output)) + + +def _compose_single(outer: IndexTransform, inner_map: OutputIndexMap) -> OutputIndexMap: + """Compose a single inner output map with the full outer transform.""" + if isinstance(inner_map, ConstantMap): + return ConstantMap(offset=inner_map.offset) + + if isinstance(inner_map, DimensionMap): + return _compose_dimension(outer, inner_map) + + # inner_map: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap) + return _compose_array(outer, inner_map) + + +def _compose_dimension(outer: IndexTransform, inner_map: DimensionMap) -> OutputIndexMap: + """Compose when inner is a DimensionMap. + + storage = offset_i + stride_i * intermediate[dim_i] + where intermediate[dim_i] = outer.output[dim_i](user_input) + """ + dim_i = inner_map.input_dimension + offset_i = inner_map.offset + stride_i = inner_map.stride + outer_map = outer.output[dim_i] + + if isinstance(outer_map, ConstantMap): + return ConstantMap(offset=offset_i + stride_i * outer_map.offset) + + if isinstance(outer_map, DimensionMap): + return DimensionMap( + input_dimension=outer_map.input_dimension, + offset=offset_i + stride_i * outer_map.offset, + stride=stride_i * outer_map.stride, + ) + + # outer_map: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap) + # Affine post-composition leaves the index array (and hence its full + # input rank and dependency axes) untouched; carry the orthogonal + # binding through unchanged. + return ArrayMap( + index_array=outer_map.index_array, + offset=offset_i + stride_i * outer_map.offset, + stride=stride_i * outer_map.stride, + input_dimension=outer_map.input_dimension, + ) + + +def _compose_array(outer: IndexTransform, inner_map: ArrayMap) -> OutputIndexMap: + """Compose when inner is an ArrayMap. + + storage = offset_i + stride_i * arr_i[intermediate] + We need to evaluate arr_i at the intermediate coordinates produced by outer. + """ + arr_i = inner_map.index_array + offset_i = inner_map.offset + stride_i = inner_map.stride + + # Check if all outer outputs are constant + all_constant = all(isinstance(m, ConstantMap) for m in outer.output) + + if all_constant: + # Evaluate arr_i at the single constant point + idx = tuple(m.offset for m in outer.output if isinstance(m, ConstantMap)) + value = int(arr_i[idx]) + return ConstantMap(offset=offset_i + stride_i * value) + + # For 1D inner array with a single outer output (simple case) + if arr_i.ndim == 1 and len(outer.output) == 1: + outer_map = outer.output[0] + + if isinstance(outer_map, DimensionMap): + dim_size = outer.domain.shape[outer_map.input_dimension] + user_indices = np.arange(dim_size, dtype=np.intp) + intermediate_vals = outer_map.offset + outer_map.stride * user_indices + new_arr = arr_i[intermediate_vals] + return ArrayMap(index_array=new_arr, offset=offset_i, stride=stride_i) + + if isinstance(outer_map, ArrayMap): + intermediate_vals = outer_map.offset + outer_map.stride * outer_map.index_array + new_arr = arr_i[intermediate_vals] + return ArrayMap(index_array=new_arr, offset=offset_i, stride=stride_i) + + # General multi-dim case: not yet implemented + raise NotImplementedError( + "Composing a multi-dimensional inner array map with non-constant outer maps " + "is not yet supported." + ) diff --git a/packages/zarr-indexing/src/zarr_indexing/domain.py b/packages/zarr-indexing/src/zarr_indexing/domain.py new file mode 100644 index 0000000000..f20d5bf7bd --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/domain.py @@ -0,0 +1,189 @@ +"""Index domains — rectangular regions in N-dimensional integer space. + +An `IndexDomain` represents the set of valid coordinates for an array or +array view. It is the cartesian product of per-dimension integer ranges:: + + IndexDomain(inclusive_min=(2, 5), exclusive_max=(10, 20)) + # represents {(i, j) : 2 <= i < 10, 5 <= j < 20} + +Unlike NumPy, domains can have **non-zero origins**. After slicing +`arr[5:10]`, the result has origin 5 and shape 5 — coordinates 5 through +9 are valid. This follows the TensorStore convention. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + + +@dataclass(frozen=True, slots=True) +class IndexDomain: + """A rectangular region in N-dimensional index space. + + The valid coordinates are the integers in + `[inclusive_min[d], exclusive_max[d])` for each dimension `d`. + """ + + inclusive_min: tuple[int, ...] + exclusive_max: tuple[int, ...] + labels: tuple[str, ...] | None = None + # Lazily-memoized shape. Excluded from init/repr/eq/hash: it is derived + # state, not part of the domain's identity. The domain is frozen, so the + # value is computed at most once (see `shape`). `None` is the unset + # sentinel; an empty shape caches as `()`. + _shape: tuple[int, ...] | None = field(default=None, init=False, repr=False, compare=False) + + def __post_init__(self) -> None: + if len(self.inclusive_min) != len(self.exclusive_max): + raise ValueError( + f"inclusive_min and exclusive_max must have the same length. " + f"Got {len(self.inclusive_min)} and {len(self.exclusive_max)}." + ) + for i, (lo, hi) in enumerate(zip(self.inclusive_min, self.exclusive_max, strict=True)): + if lo > hi: + raise ValueError( + f"inclusive_min must be <= exclusive_max for all dimensions. " + f"Dimension {i}: {lo} > {hi}" + ) + if self.labels is not None and len(self.labels) != len(self.inclusive_min): + raise ValueError( + f"labels must have the same length as dimensions. " + f"Got {len(self.labels)} labels for {len(self.inclusive_min)} dimensions." + ) + + @classmethod + def from_shape(cls, shape: tuple[int, ...]) -> IndexDomain: + """Create a domain with origin at zero.""" + return cls( + inclusive_min=(0,) * len(shape), + exclusive_max=shape, + ) + + @property + def ndim(self) -> int: + return len(self.inclusive_min) + + @property + def origin(self) -> tuple[int, ...]: + return self.inclusive_min + + @property + def shape(self) -> tuple[int, ...]: + cached = self._shape + if cached is None: + cached = tuple( + hi - lo for lo, hi in zip(self.inclusive_min, self.exclusive_max, strict=True) + ) + object.__setattr__(self, "_shape", cached) + return cached + + def contains(self, index: tuple[int, ...]) -> bool: + if len(index) != self.ndim: + return False + return all( + lo <= idx < hi + for lo, hi, idx in zip(self.inclusive_min, self.exclusive_max, index, strict=True) + ) + + def contains_domain(self, other: IndexDomain) -> bool: + if other.ndim != self.ndim: + return False + return all( + self_lo <= other_lo and other_hi <= self_hi + for self_lo, self_hi, other_lo, other_hi in zip( + self.inclusive_min, + self.exclusive_max, + other.inclusive_min, + other.exclusive_max, + strict=True, + ) + ) + + def intersect(self, other: IndexDomain) -> IndexDomain | None: + if other.ndim != self.ndim: + raise ValueError( + f"Cannot intersect domains with different ranks: {self.ndim} vs {other.ndim}" + ) + new_min = tuple( + max(a, b) for a, b in zip(self.inclusive_min, other.inclusive_min, strict=True) + ) + new_max = tuple( + min(a, b) for a, b in zip(self.exclusive_max, other.exclusive_max, strict=True) + ) + if any(lo >= hi for lo, hi in zip(new_min, new_max, strict=True)): + return None + return IndexDomain(inclusive_min=new_min, exclusive_max=new_max) + + def translate(self, offset: tuple[int, ...]) -> IndexDomain: + if len(offset) != self.ndim: + raise ValueError( + f"Offset must have same length as domain dimensions. " + f"Domain has {self.ndim} dimensions, offset has {len(offset)}." + ) + new_min = tuple(lo + off for lo, off in zip(self.inclusive_min, offset, strict=True)) + new_max = tuple(hi + off for hi, off in zip(self.exclusive_max, offset, strict=True)) + return IndexDomain(inclusive_min=new_min, exclusive_max=new_max) + + def narrow(self, selection: Any) -> IndexDomain: + """Apply a basic selection and return a narrowed domain. + Indices are absolute coordinates. Integer indices produce length-1 extent. + Strided slices are not supported — use IndexTransform for strides. + """ + normalized = _normalize_selection(selection, self.ndim) + new_inclusive_min: list[int] = [] + new_exclusive_max: list[int] = [] + for dim_idx, (sel, dim_lo, dim_hi) in enumerate( + zip(normalized, self.inclusive_min, self.exclusive_max, strict=True) + ): + if isinstance(sel, int): + if sel < dim_lo or sel >= dim_hi: + raise IndexError( + f"index {sel} is out of bounds for dimension {dim_idx} " + f"with domain [{dim_lo}, {dim_hi})" + ) + new_inclusive_min.append(sel) + new_exclusive_max.append(sel + 1) + else: + start, stop, step = sel.start, sel.stop, sel.step + if step is not None and step != 1: + raise IndexError( + "IndexDomain.narrow only supports step=1 slices. " + f"Got step={step}. Use IndexTransform for strided access." + ) + abs_start = dim_lo if start is None else start + abs_stop = dim_hi if stop is None else stop + abs_start = max(abs_start, dim_lo) + abs_stop = min(abs_stop, dim_hi) + abs_stop = max(abs_stop, abs_start) + new_inclusive_min.append(abs_start) + new_exclusive_max.append(abs_stop) + return IndexDomain( + inclusive_min=tuple(new_inclusive_min), + exclusive_max=tuple(new_exclusive_max), + ) + + +def _normalize_selection(selection: Any, ndim: int) -> tuple[int | slice, ...]: + """Normalize a basic selection to a tuple of ints/slices with length ndim.""" + if not isinstance(selection, tuple): + selection = (selection,) + result: list[int | slice] = [] + ellipsis_seen = False + for sel in selection: + if sel is Ellipsis: + if ellipsis_seen: + raise IndexError("an index can only have a single ellipsis ('...')") + ellipsis_seen = True + num_missing = ndim - (len(selection) - 1) + result.extend([slice(None)] * num_missing) + else: + result.append(sel) + while len(result) < ndim: + result.append(slice(None)) + if len(result) > ndim: + raise IndexError( + f"too many indices for array: array has {ndim} dimensions, " + f"but {len(result)} were indexed" + ) + return tuple(result) diff --git a/packages/zarr-indexing/src/zarr_indexing/errors.py b/packages/zarr-indexing/src/zarr_indexing/errors.py new file mode 100644 index 0000000000..fa2f6fc5d3 --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/errors.py @@ -0,0 +1,21 @@ +"""Canonical index-error types raised by the transform algebra. + +These are the authoritative class definitions. `zarr.errors` re-exports the +same objects (`from zarr_indexing.errors import ...`) so that, e.g., +`zarr.errors.BoundsCheckError is zarr_indexing.errors.BoundsCheckError`. +Both subclass the built-in `IndexError`, so existing `except IndexError` (or +`except zarr.errors.BoundsCheckError`) catch sites keep working unchanged. +""" + +from __future__ import annotations + +__all__ = [ + "BoundsCheckError", + "VindexInvalidSelectionError", +] + + +class VindexInvalidSelectionError(IndexError): ... + + +class BoundsCheckError(IndexError): ... diff --git a/packages/zarr-indexing/src/zarr_indexing/grid.py b/packages/zarr-indexing/src/zarr_indexing/grid.py new file mode 100644 index 0000000000..de1dae2dfc --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/grid.py @@ -0,0 +1,25 @@ +"""Structural typing for the chunk-grid surface used by chunk resolution. + +`chunk_resolution` needs only a narrow slice of a chunk grid: the per-dimension +mapping between storage indices and chunk coordinates, passed as one +`DimensionGridLike` per storage dimension. Rather than import zarr's concrete +grid types, we type against this Protocol; zarr's per-dimension grids satisfy +it structurally, so no zarr import is needed here. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Protocol + +if TYPE_CHECKING: + import numpy as np + import numpy.typing as npt + + +class DimensionGridLike(Protocol): + """The per-dimension chunk-mapping surface consumed by chunk resolution.""" + + def index_to_chunk(self, idx: int) -> int: ... + def chunk_offset(self, chunk_ix: int) -> int: ... + def chunk_size(self, chunk_ix: int) -> int: ... + def indices_to_chunks(self, indices: npt.NDArray[np.intp]) -> npt.NDArray[np.intp]: ... diff --git a/packages/zarr-indexing/src/zarr_indexing/json.py b/packages/zarr-indexing/src/zarr_indexing/json.py new file mode 100644 index 0000000000..c95696f309 --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/json.py @@ -0,0 +1,325 @@ +"""Lowering between canonical ndsel bodies and in-memory `IndexTransform`s. + +This is the **engine layer**. Where `messages.py` is pure JSON→JSON and imposes +no array constraints, this module converts a *canonical* ndsel transform body +(spec section 4.3, as produced by `zarr_indexing.messages.normalize_ndsel`) +into the numpy-backed `IndexTransform` the chunk engine runs on, and back. + +Two engine constraints live **here and only here**: + +- **Finite bounds.** An `IndexDomain` addresses a finite array, so a canonical + body carrying a `"-inf"`/`"+inf"` bound cannot be lowered; `from_json` raises. +- **Implicit bounds lower by value.** The `[n]`-bracket implicit/explicit flag + is a message-layer concern; the engine keeps only the integer value. + +## The `index_array` wire format (and the degenerate-collapse it documents) + +ndsel and TensorStore both **reject** an output map that carries *both* +`input_dimension` and `index_array`. The in-memory `ArrayMap`, however, records +an `input_dimension` to pin the axis an orthogonal (`oindex`) array varies over. +This module bridges the gap: + +- **On serialize** (`transform_to_canonical`): + 1. An all-singleton `index_array` (size 1) selects a single coordinate + regardless of input, so it is **collapsed to a `constant` map** + `{offset: offset + stride*value}`. The size-1 input dimension stays in the + domain, unconsumed — a valid transform. This makes a length-1 `oindex` + selection round-trip *behaviorally* (an `ArrayMap` becomes a `ConstantMap`) + rather than by object identity. + 2. Non-degenerate `index_array` maps are emitted **without** `input_dimension`. + +- **On load** (`transform_from_canonical`): the in-memory `input_dimension` is + reconstructed from the full-rank array's dependency axes (its non-singleton + axes, see `transform._array_map_dependency_axes`). An array that solely owns a + single non-singleton axis is orthogonal (`input_dimension = that axis`); arrays + that share non-singleton axes, or vary over several, are correlated (`vindex`, + `input_dimension = None`). A single 1-D array over a rank-1 domain is + inherently ambiguous between the two flavours; it reconstructs as orthogonal, + which is behaviorally identical for the single-array case. + +`index_transform_to_json` / `index_transform_from_json` (and the `*_domain_*` +variants) are these canonical converters under their historical names. +""" + +from __future__ import annotations + +from collections import Counter +from typing import Any, Required, TypedDict + +import numpy as np + +from zarr_indexing.domain import IndexDomain +from zarr_indexing.messages import normalize_ndsel +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap +from zarr_indexing.transform import ( + IndexTransform, + _array_map_dependency_axes, # pyright: ignore[reportPrivateUsage] +) + +# `_array_map_dependency_axes` is a leading-underscore helper in `transform.py`, +# but it is deliberately shared with this module (the engine-level JSON <-> +# `IndexTransform` lowering below needs the same dependency-axis logic that +# `transform.py`'s own array-reindexing helpers use). It is not part of the +# package's public API; pyright's `reportPrivateUsage` flags the cross-module +# import anyway. See `chunk_resolution.py`'s `_dimensions` suppression for the +# analogous rationale — whether to promote either symbol out of "private" is +# an open pre-publish API decision, not resolved here. + +# --------------------------------------------------------------------------- +# TypedDict definitions (canonical JSON shapes) +# --------------------------------------------------------------------------- + +# An `index_array` serializes via `ndarray.tolist()`, so it is a nested list of +# ints whose nesting depth equals the array rank. +NestedIntList = list[Any] + +# A canonical *lowered* body carries only finite integer bounds, but the JSON +# shape admits the full ndsel `bound` grammar: an explicit int / sentinel, or a +# one-element implicit `[value]` array. +IndexValueJSON = int | str +BoundJSON = int | str | list[IndexValueJSON] + + +class IndexDomainJSON(TypedDict, total=False): + """Canonical JSON representation of an IndexDomain.""" + + input_inclusive_min: Required[list[BoundJSON]] + input_exclusive_max: Required[list[BoundJSON]] + input_labels: Required[list[str]] + + +class OutputIndexMapJSON(TypedDict, total=False): + """Canonical JSON representation of a single output index map. + + Exactly one of three forms (distinguished by which fields are present): + + - `{"offset": 5}` — constant + - `{"offset": 0, "stride": 1, "input_dimension": 0}` — single_input_dimension + - `{"offset": 0, "stride": 1, "index_array": [...], + "index_array_bounds": ["-inf", "+inf"]}` — index_array + """ + + offset: int + stride: int + input_dimension: int + index_array: NestedIntList + index_array_bounds: list[IndexValueJSON] + + +class IndexTransformJSON(TypedDict, total=False): + """Canonical JSON representation of an IndexTransform (spec section 4.3).""" + + input_rank: Required[int] + input_inclusive_min: Required[list[BoundJSON]] + input_exclusive_max: Required[list[BoundJSON]] + input_labels: Required[list[str]] + output: Required[list[OutputIndexMapJSON]] + + +# --------------------------------------------------------------------------- +# Bound / label lowering (engine constraints) +# --------------------------------------------------------------------------- + + +def _lower_bound(bound: BoundJSON, where: str) -> int: + """Lower a canonical bound to a finite integer, rejecting infinities.""" + value = bound[0] if isinstance(bound, list) else bound + if value == "-inf" or value == "+inf": + raise ValueError( + f"{where} is infinite ({value!r}); an IndexDomain addresses a finite " + f"array and cannot lower an infinite bound" + ) + return int(value) + + +def _lower_labels(labels: list[str]) -> tuple[str, ...] | None: + """All-empty labels collapse to `None` so a label-free domain round-trips.""" + return None if all(label == "" for label in labels) else tuple(labels) + + +def _emit_labels(labels: tuple[str, ...] | None, rank: int) -> list[str]: + """Emit canonical labels: `[""]*rank` when the domain is unlabeled.""" + return [""] * rank if labels is None else list(labels) + + +# --------------------------------------------------------------------------- +# IndexDomain serialization +# --------------------------------------------------------------------------- + + +def index_domain_to_json(domain: IndexDomain) -> IndexDomainJSON: + """Convert an IndexDomain to its canonical JSON representation.""" + return { + "input_inclusive_min": list(domain.inclusive_min), + "input_exclusive_max": list(domain.exclusive_max), + "input_labels": _emit_labels(domain.labels, domain.ndim), + } + + +def index_domain_from_json(data: IndexDomainJSON) -> IndexDomain: + """Construct an IndexDomain from its canonical JSON representation.""" + inclusive_min = tuple( + _lower_bound(b, f"input_inclusive_min[{i}]") + for i, b in enumerate(data["input_inclusive_min"]) + ) + exclusive_max = tuple( + _lower_bound(b, f"input_exclusive_max[{i}]") + for i, b in enumerate(data["input_exclusive_max"]) + ) + labels = _lower_labels(list(data["input_labels"])) + return IndexDomain(inclusive_min=inclusive_min, exclusive_max=exclusive_max, labels=labels) + + +# --------------------------------------------------------------------------- +# OutputIndexMap serialization +# --------------------------------------------------------------------------- + + +def output_index_map_to_json(m: OutputIndexMap) -> OutputIndexMapJSON: + """Convert an output index map to its canonical JSON representation. + + A degenerate all-singleton `ArrayMap` collapses to a `constant` map; a + non-degenerate one is emitted without `input_dimension` (see the module + docstring on the wire format). + """ + if isinstance(m, ConstantMap): + return {"offset": m.offset} + + if isinstance(m, DimensionMap): + return {"offset": m.offset, "stride": m.stride, "input_dimension": m.input_dimension} + + # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap) + if m.index_array.size == 1: + value = int(m.index_array.reshape(-1)[0]) + return {"offset": m.offset + m.stride * value} + return { + "offset": m.offset, + "stride": m.stride, + "index_array": m.index_array.tolist(), + "index_array_bounds": ["-inf", "+inf"], + } + + +def output_index_map_from_json(data: OutputIndexMapJSON) -> OutputIndexMap: + """Construct an output index map from its canonical JSON representation. + + An `index_array` map's `input_dimension` is reconstructed from the array's + dependency axes in isolation (single non-singleton axis → orthogonal). The + transform-level loader classifies globally; use it when several maps may + share axes. + """ + if "index_array" in data: + arr = np.asarray(data["index_array"], dtype=np.intp) + return ArrayMap( + index_array=arr, + offset=data.get("offset", 0), + stride=data.get("stride", 1), + input_dimension=_solo_dependency_axis(arr), + ) + + if "input_dimension" in data: + return DimensionMap( + input_dimension=data["input_dimension"], + offset=data.get("offset", 0), + stride=data.get("stride", 1), + ) + + return ConstantMap(offset=data.get("offset", 0)) + + +def _solo_dependency_axis(arr: np.ndarray[Any, Any]) -> int | None: + """The single axis a lone `index_array` varies over, or `None` if not exactly one.""" + dep = _array_map_dependency_axes(arr) + return dep[0] if len(dep) == 1 else None + + +# --------------------------------------------------------------------------- +# IndexTransform serialization +# --------------------------------------------------------------------------- + + +def transform_to_canonical(transform: IndexTransform) -> IndexTransformJSON: + """Convert an IndexTransform to its canonical ndsel transform body. + + The result is fully explicit (spec section 4.3): `input_rank`, fully written + bounds and labels, and an explicit `output` with `offset`/`stride` present + on every affine and array map. + """ + return { + "input_rank": transform.domain.ndim, + "input_inclusive_min": list(transform.domain.inclusive_min), + "input_exclusive_max": list(transform.domain.exclusive_max), + "input_labels": _emit_labels(transform.domain.labels, transform.domain.ndim), + "output": [output_index_map_to_json(m) for m in transform.output], + } + + +def transform_from_canonical(data: IndexTransformJSON) -> IndexTransform: + """Construct an IndexTransform from a canonical (or canonicalizable) body. + + The body is first run through the message layer (`normalize_ndsel`) so that + omitted fields — identity `output`, default bounds/labels — are filled and + validated, then lowered to the engine representation. `index_array` maps' + `input_dimension` values are reconstructed by global dependency-axis + ownership (see the module docstring). + """ + body = normalize_ndsel({"kind": "transform", **data}) + + inclusive_min = tuple( + _lower_bound(b, f"input_inclusive_min[{i}]") + for i, b in enumerate(body["input_inclusive_min"]) + ) + exclusive_max = tuple( + _lower_bound(b, f"input_exclusive_max[{i}]") + for i, b in enumerate(body["input_exclusive_max"]) + ) + domain = IndexDomain( + inclusive_min=inclusive_min, + exclusive_max=exclusive_max, + labels=_lower_labels(body["input_labels"]), + ) + + output_raw: list[dict[str, Any]] = body["output"] + + # Classify index_array maps globally: an axis owned by exactly one array map + # (and the map's sole non-singleton axis) marks that map orthogonal; shared + # or multiple non-singleton axes mark the maps correlated (vindex). + array_axes: dict[int, tuple[int, ...]] = {} + axis_owners: Counter[int] = Counter() + for i, om in enumerate(output_raw): + if "index_array" in om: + arr = np.asarray(om["index_array"], dtype=np.intp) + dep = _array_map_dependency_axes(arr) + array_axes[i] = dep + axis_owners.update(dep) + + output: list[OutputIndexMap] = [] + for i, om in enumerate(output_raw): + if "index_array" in om: + dep = array_axes[i] + input_dim = dep[0] if len(dep) == 1 and axis_owners[dep[0]] == 1 else None + output.append( + ArrayMap( + index_array=np.asarray(om["index_array"], dtype=np.intp), + offset=om.get("offset", 0), + stride=om.get("stride", 1), + input_dimension=input_dim, + ) + ) + elif "input_dimension" in om: + output.append( + DimensionMap( + input_dimension=om["input_dimension"], + offset=om.get("offset", 0), + stride=om.get("stride", 1), + ) + ) + else: + output.append(ConstantMap(offset=om.get("offset", 0))) + + return IndexTransform(domain=domain, output=tuple(output)) + + +# Historical names, now pointing at the canonical converters. +index_transform_to_json = transform_to_canonical +index_transform_from_json = transform_from_canonical diff --git a/packages/zarr-indexing/src/zarr_indexing/messages.py b/packages/zarr-indexing/src/zarr_indexing/messages.py new file mode 100644 index 0000000000..d5761d1a38 --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/messages.py @@ -0,0 +1,657 @@ +"""The ndsel message layer — pure JSON in, canonical JSON out. + +This module implements the [ndsel](https://github.com/zarr-developers/ndsel) draft wire +format: a JSON-serializable representation of NumPy-style n-dimensional +selections that adapts TensorStore's `IndexTransform` model. It is a **pure +JSON→JSON** layer: it depends on nothing but the standard library, imposes no +engine (numpy/array) constraints, and never rounds, clamps, or drops +information. Engine constraints (finite bounds, in-memory `IndexTransform` +construction) live one layer up, in `json.py`. + +Two entry points: + +- `parse_ndsel(obj)` — structurally validate an ndsel message of any of the + five kinds (`point`/`box`/`slice`/`points`/`transform`), returning it + unchanged. Raises `NdselError` (carrying a spec reason code) on any defect. +- `normalize_ndsel(obj)` — desugar and canonicalize a message to the single + deterministic **canonical transform body** of the spec (section 4.3): a bare + `IndexTransform` JSON body, without the `kind` discriminator. `normalize` is + idempotent when its output is re-tagged with `kind: "transform"`. + +The canonical body is, field-for-field, a TensorStore `IndexTransform` (minus +`kind`), so a normalized `transform` loads directly into TensorStore once +`kind` is stripped. + +Value rules enforced here: every integer is a 64-bit signed value; JSON +booleans are **not** integers (Python's `isinstance(True, int)` is guarded +against explicitly); the `"-inf"`/`"+inf"` sentinels are legal only in bound +positions; an implicit bound is the one-element `[n]`-bracket form, and its +implicit/explicit flag is preserved through normalization. +""" + +from __future__ import annotations + +from typing import Any + +__all__ = [ + "NdselError", + "normalize_ndsel", + "parse_ndsel", +] + +# --------------------------------------------------------------------------- +# Error taxonomy +# --------------------------------------------------------------------------- + +#: The complete set of ndsel reason codes (spec section 6). +REASON_CODES = frozenset( + { + "invalid_json", + "unknown_kind", + "unknown_field", + "multiple_upper_bounds", + "bounds_out_of_order", + "output_map_conflict", + "rank_mismatch", + "step_zero", + "negative_step_unsupported", + } +) + + +class NdselError(ValueError): + """An ndsel message failed validation. + + Carries the spec `reason` code (one of `REASON_CODES`) so callers and the + conformance harness can assert on it directly, plus a human-readable + `detail`. + """ + + def __init__(self, reason: str, detail: str = "") -> None: + self.reason = reason + self.detail = detail + super().__init__(f"{reason}: {detail}" if detail else reason) + + +# --------------------------------------------------------------------------- +# 64-bit signed integer range (spec section 3.5) +# --------------------------------------------------------------------------- + +_I64_MIN = -(2**63) +_I64_MAX = 2**63 - 1 + +_KNOWN_KINDS = frozenset({"point", "box", "slice", "points", "transform"}) + +# The two upper-bound spellings, keyed by message prefix. Only one of the three +# per group may appear (spec section 4.1 / 5.2). +_BOX_UPPER = ("exclusive_max", "inclusive_max", "shape") +_TRANSFORM_UPPER = ("input_exclusive_max", "input_inclusive_max", "input_shape") + +_OUTPUT_MAP_FIELDS = frozenset( + {"offset", "stride", "input_dimension", "index_array", "index_array_bounds"} +) + + +# --------------------------------------------------------------------------- +# Leaf value validators +# --------------------------------------------------------------------------- + + +def _is_int(value: Any) -> bool: + """True iff `value` is a JSON integer — an `int` that is not a `bool`. + + JSON has no boolean-as-integer: `True`/`False` are rejected even though + Python makes `bool` a subclass of `int` (spec section 3.6). + """ + return isinstance(value, int) and not isinstance(value, bool) + + +def _check_int(value: Any, where: str) -> int: + """Validate a plain-integer position: an in-range i64, never a sentinel.""" + if not _is_int(value): + raise NdselError("invalid_json", f"{where} must be an integer, got {value!r}") + if value < _I64_MIN or value > _I64_MAX: + raise NdselError("invalid_json", f"{where} is outside the 64-bit signed range: {value}") + return int(value) + + +def _is_sentinel(value: Any) -> bool: + return value in ("-inf", "+inf") + + +def _check_index_value(value: Any, where: str) -> int | str: + """Validate an `index-value`: an in-range i64 or a `"-inf"`/`"+inf"` sentinel.""" + if _is_sentinel(value): + return str(value) + return _check_int(value, where) + + +def _check_bound(value: Any, where: str) -> int | str | list[int | str]: + """Validate a `bound`: an explicit `index-value`, or a one-element implicit `[index-value]`.""" + if isinstance(value, list): + if len(value) != 1: + raise NdselError( + "invalid_json", + f"{where} implicit bound must be a one-element array, got {value!r}", + ) + return [_check_index_value(value[0], where)] + return _check_index_value(value, where) + + +def _check_int_list(value: Any, where: str) -> list[int]: + if not isinstance(value, list): + raise NdselError("invalid_json", f"{where} must be an array, got {value!r}") + return [_check_int(v, f"{where}[{i}]") for i, v in enumerate(value)] + + +def _check_bound_list(value: Any, where: str) -> list[Any]: + if not isinstance(value, list): + raise NdselError("invalid_json", f"{where} must be an array, got {value!r}") + return [_check_bound(v, f"{where}[{i}]") for i, v in enumerate(value)] + + +def _check_label_list(value: Any, where: str) -> list[str]: + if not isinstance(value, list): + raise NdselError("invalid_json", f"{where} must be an array, got {value!r}") + for i, v in enumerate(value): + if not isinstance(v, str): + raise NdselError("invalid_json", f"{where}[{i}] must be a string, got {v!r}") + return list(value) + + +# --------------------------------------------------------------------------- +# Extended-integer order for bounds (spec section 4.1) +# --------------------------------------------------------------------------- + + +def _bound_value(bound: int | str | list[int | str]) -> int | str: + """The underlying `index-value` of a bound, dropping the implicit bracket.""" + return bound[0] if isinstance(bound, list) else bound + + +def _bound_is_implicit(bound: int | str | list[int | str]) -> bool: + return isinstance(bound, list) + + +def _ext_key(value: int | str) -> tuple[int, int]: + """A sort key giving the extended-integer order `-inf < n < +inf` exactly. + + Uses an integer tier plus the value, so no float rounding of near-`2**63` + integers can misorder the `inclusive_min <= exclusive_max` check. + """ + if value == "-inf": + return (0, 0) + if value == "+inf": + return (2, 0) + assert isinstance(value, int) + return (1, value) + + +def _rewrap(value: int | str, *, implicit: bool) -> int | str | list[int | str]: + return [value] if implicit else value + + +# --------------------------------------------------------------------------- +# Message-level helpers +# --------------------------------------------------------------------------- + + +def _require_object(obj: Any) -> dict[str, Any]: + if not isinstance(obj, dict): + raise NdselError("invalid_json", f"message must be a JSON object, got {type(obj).__name__}") + return obj + + +def _message_kind(obj: dict[str, Any]) -> str: + kind = obj.get("kind") + if not isinstance(kind, str): + raise NdselError("invalid_json", "message must have a string 'kind' field") + if kind not in _KNOWN_KINDS: + raise NdselError("unknown_kind", f"unknown kind {kind!r}") + return kind + + +def _check_membership(obj: dict[str, Any], allowed: frozenset[str], what: str) -> None: + """Strict membership (spec section 3.7): reject any undefined member.""" + for key in obj: + if key not in allowed: + raise NdselError("unknown_field", f"{what} has undefined member {key!r}") + + +def _single_upper_bound(obj: dict[str, Any], fields: tuple[str, str, str]) -> str | None: + present = [f for f in fields if f in obj] + if len(present) > 1: + raise NdselError( + "multiple_upper_bounds", + f"at most one of {fields} may be present; got {present}", + ) + return present[0] if present else None + + +def _resolve_upper_bound( + upper_field: str | None, + upper_raw: list[Any] | None, + inclusive_min: list[Any], + rank: int, + *, + kind_of: str, +) -> list[int | str | list[int | str]]: + """Produce `exclusive_max` from whichever upper-bound spelling was supplied. + + - `exclusive_max`/`input_exclusive_max` → used directly. + - `inclusive_max`/`input_inclusive_max` → each element `+1`. + - `shape`/`input_shape` → `inclusive_min + shape` per element. + - none → an **implicit `+inf`** in every dimension. + + The implicit/explicit bracket travels with the extent-bearing field (the + upper bound, or `shape`), matching the spec's `[n]`-bracket convention. + """ + if upper_field is None: + return [["+inf"] for _ in range(rank)] + + assert upper_raw is not None + if kind_of == "exclusive": + return list(upper_raw) + + result: list[int | str | list[int | str]] = [] + for k in range(rank): + raw = upper_raw[k] + implicit = _bound_is_implicit(raw) + value = _bound_value(raw) + if kind_of == "inclusive": + new = _inclusive_to_exclusive(value) + else: # shape + new = _shape_to_exclusive(_bound_value(inclusive_min[k]), value) + result.append(_rewrap(new, implicit=implicit)) + return result + + +def _inclusive_to_exclusive(value: int | str) -> int | str: + if value == "+inf" or value == "-inf": + return value + assert isinstance(value, int) + return value + 1 + + +def _shape_to_exclusive(min_value: int | str, shape_value: int | str) -> int | str: + if shape_value == "+inf" or min_value == "+inf": + return "+inf" + if min_value == "-inf": + return "-inf" + assert isinstance(min_value, int) + assert isinstance(shape_value, int) + return min_value + shape_value + + +def _validate_domain(inclusive_min: list[Any], exclusive_max: list[Any], *, prefix: str) -> None: + """Every dimension must satisfy `inclusive_min <= exclusive_max` (empty is valid).""" + for k, (lo, hi) in enumerate(zip(inclusive_min, exclusive_max, strict=True)): + if _ext_key(_bound_value(lo)) > _ext_key(_bound_value(hi)): + raise NdselError( + "bounds_out_of_order", + f"{prefix}[{k}]: inclusive_min {_bound_value(lo)!r} > " + f"exclusive_max {_bound_value(hi)!r}", + ) + + +def _identity_output(rank: int) -> list[dict[str, Any]]: + return [{"offset": 0, "stride": 1, "input_dimension": k} for k in range(rank)] + + +# --------------------------------------------------------------------------- +# Per-kind desugaring +# --------------------------------------------------------------------------- + + +def _normalize_point(obj: dict[str, Any]) -> dict[str, Any]: + _check_membership(obj, frozenset({"kind", "coords"}), "point") + if "coords" not in obj: + raise NdselError("invalid_json", "point requires 'coords'") + coords = _check_int_list(obj["coords"], "coords") + return { + "input_rank": 0, + "input_inclusive_min": [], + "input_exclusive_max": [], + "input_labels": [], + "output": [{"offset": c} for c in coords], + } + + +def _infer_rank( + obj: dict[str, Any], + named_lengths: list[tuple[str, int]], + *, + declared: int | None, +) -> int: + """Reconcile a declared rank (if any) with every present array's length.""" + rank = declared + for name, length in named_lengths: + if rank is None: + rank = length + elif rank != length: + raise NdselError( + "rank_mismatch", + f"{name} has length {length}, inconsistent with rank {rank}", + ) + return rank if rank is not None else 0 + + +def _normalize_box(obj: dict[str, Any]) -> dict[str, Any]: + allowed = frozenset( + {"kind", "inclusive_min", "exclusive_max", "inclusive_max", "shape", "labels"} + ) + _check_membership(obj, allowed, "box") + + inclusive_min_raw = ( + _check_bound_list(obj["inclusive_min"], "inclusive_min") if "inclusive_min" in obj else None + ) + upper_field = _single_upper_bound(obj, _BOX_UPPER) + upper_raw = _check_bound_list(obj[upper_field], upper_field) if upper_field else None + labels_raw = _check_label_list(obj["labels"], "labels") if "labels" in obj else None + + named_lengths: list[tuple[str, int]] = [] + if inclusive_min_raw is not None: + named_lengths.append(("inclusive_min", len(inclusive_min_raw))) + if upper_raw is not None: + named_lengths.append((upper_field or "", len(upper_raw))) + if labels_raw is not None: + named_lengths.append(("labels", len(labels_raw))) + rank = _infer_rank(obj, named_lengths, declared=None) + + inclusive_min = inclusive_min_raw if inclusive_min_raw is not None else [0] * rank + exclusive_max = _resolve_upper_bound( + upper_field, upper_raw, inclusive_min, rank, kind_of=_upper_kind(upper_field, _BOX_UPPER) + ) + labels = labels_raw if labels_raw is not None else [""] * rank + _validate_domain(inclusive_min, exclusive_max, prefix="box") + + return { + "input_rank": rank, + "input_inclusive_min": inclusive_min, + "input_exclusive_max": exclusive_max, + "input_labels": labels, + "output": _identity_output(rank), + } + + +def _upper_kind(upper_field: str | None, fields: tuple[str, str, str]) -> str: + if upper_field is None or upper_field == fields[0]: + return "exclusive" + if upper_field == fields[1]: + return "inclusive" + return "shape" + + +def _normalize_slice(obj: dict[str, Any]) -> dict[str, Any]: + allowed = frozenset({"kind", "start", "stop", "step", "labels"}) + _check_membership(obj, allowed, "slice") + if "start" not in obj: + raise NdselError("invalid_json", "slice requires 'start'") + if "stop" not in obj: + raise NdselError("invalid_json", "slice requires 'stop'") + start = _check_int_list(obj["start"], "start") + stop = _check_int_list(obj["stop"], "stop") + step = _check_int_list(obj["step"], "step") if "step" in obj else [1] * len(start) + labels_raw = _check_label_list(obj["labels"], "labels") if "labels" in obj else None + + n = len(start) + for name, arr in (("stop", stop), ("step", step)): + if len(arr) != n: + raise NdselError( + "rank_mismatch", f"{name} has length {len(arr)}, expected {n} (from start)" + ) + if labels_raw is not None and len(labels_raw) != n: + raise NdselError( + "rank_mismatch", f"labels has length {len(labels_raw)}, expected {n} (from start)" + ) + + for k, s in enumerate(step): + if s == 0: + raise NdselError("step_zero", f"step[{k}] is zero") + if s < 0: + raise NdselError("negative_step_unsupported", f"step[{k}] is negative ({s})") + + inclusive_min: list[Any] = [] + exclusive_max: list[Any] = [] + output: list[dict[str, Any]] = [] + for k in range(n): + a, b, s = start[k], stop[k], step[k] + m = max(0, -(-(b - a) // s)) # ceil((b - a) / s) + o = _trunc_div(a, s) # trunc(a / s), toward zero + offset = a - s * o # lattice phase, in (-s, s) + inclusive_min.append(o) + exclusive_max.append(o + m) + output.append({"offset": offset, "stride": s, "input_dimension": k}) + + labels = labels_raw if labels_raw is not None else [""] * n + return { + "input_rank": n, + "input_inclusive_min": inclusive_min, + "input_exclusive_max": exclusive_max, + "input_labels": labels, + "output": output, + } + + +def _normalize_points(obj: dict[str, Any]) -> dict[str, Any]: + _check_membership(obj, frozenset({"kind", "coords"}), "points") + if "coords" not in obj: + raise NdselError("invalid_json", "points requires 'coords'") + coords = obj["coords"] + if not isinstance(coords, list): + raise NdselError("invalid_json", f"points coords must be an array, got {coords!r}") + + rows: list[list[int]] = [] + n: int | None = None + for i, row in enumerate(coords): + if not isinstance(row, list): + raise NdselError("invalid_json", f"points coords[{i}] must be an array, got {row!r}") + row_ints = [_check_int(v, f"coords[{i}][{j}]") for j, v in enumerate(row)] + if n is None: + n = len(row_ints) + elif len(row_ints) != n: + raise NdselError( + "rank_mismatch", + f"points coords[{i}] has length {len(row_ints)}, expected {n} (ragged)", + ) + rows.append(row_ints) + + m = len(rows) + n = n if n is not None else 0 + output = [ + { + "offset": 0, + "stride": 1, + "index_array": [rows[i][k] for i in range(m)], + "index_array_bounds": ["-inf", "+inf"], + } + for k in range(n) + ] + return { + "input_rank": 1, + "input_inclusive_min": [0], + "input_exclusive_max": [m], + "input_labels": [""], + "output": output, + } + + +def _normalize_output_map(raw: Any, where: str) -> dict[str, Any]: + if not isinstance(raw, dict): + raise NdselError("invalid_json", f"{where} must be a JSON object, got {raw!r}") + _check_membership(raw, _OUTPUT_MAP_FIELDS, where) + + has_index_array = "index_array" in raw + has_input_dim = "input_dimension" in raw + if has_index_array and has_input_dim: + raise NdselError( + "output_map_conflict", + f"{where} carries both 'input_dimension' and 'index_array'", + ) + + offset = _check_int(raw["offset"], f"{where}.offset") if "offset" in raw else 0 + + if has_index_array: + stride = _check_int(raw["stride"], f"{where}.stride") if "stride" in raw else 1 + bounds = ( + _check_index_array_bounds(raw["index_array_bounds"], where) + if "index_array_bounds" in raw + else ["-inf", "+inf"] + ) + # index_array is carried verbatim (spec section 7 defers shape validation). + return { + "offset": offset, + "stride": stride, + "index_array": raw["index_array"], + "index_array_bounds": bounds, + } + + if has_input_dim: + input_dim = _check_int(raw["input_dimension"], f"{where}.input_dimension") + if input_dim < 0: + raise NdselError( + "invalid_json", f"{where}.input_dimension must be >= 0, got {input_dim}" + ) + stride = _check_int(raw["stride"], f"{where}.stride") if "stride" in raw else 1 + return {"offset": offset, "stride": stride, "input_dimension": input_dim} + + # Constant map: only offset survives. A stray `stride`/`index_array_bounds` + # is schema-valid (the output-map schema permits those members on any map), + # so it is silently dropped rather than rejected — a constant carries only + # `offset` in canonical form (spec section 4.3). + return {"offset": offset} + + +def _check_index_array_bounds(value: Any, where: str) -> list[int | str]: + if not isinstance(value, list) or len(value) != 2: + raise NdselError( + "invalid_json", + f"{where}.index_array_bounds must be a two-element array, got {value!r}", + ) + return [ + _check_index_value(value[0], f"{where}.index_array_bounds[0]"), + _check_index_value(value[1], f"{where}.index_array_bounds[1]"), + ] + + +def _normalize_transform(obj: dict[str, Any]) -> dict[str, Any]: + allowed = frozenset( + { + "kind", + "input_rank", + "input_inclusive_min", + "input_exclusive_max", + "input_inclusive_max", + "input_shape", + "input_labels", + "output", + } + ) + _check_membership(obj, allowed, "transform") + + declared_rank: int | None = None + if "input_rank" in obj: + declared_rank = _check_int(obj["input_rank"], "input_rank") + if declared_rank < 0: + raise NdselError("invalid_json", f"input_rank must be >= 0, got {declared_rank}") + + inclusive_min_raw = ( + _check_bound_list(obj["input_inclusive_min"], "input_inclusive_min") + if "input_inclusive_min" in obj + else None + ) + upper_field = _single_upper_bound(obj, _TRANSFORM_UPPER) + upper_raw = _check_bound_list(obj[upper_field], upper_field) if upper_field else None + labels_raw = ( + _check_label_list(obj["input_labels"], "input_labels") if "input_labels" in obj else None + ) + + named_lengths: list[tuple[str, int]] = [] + if inclusive_min_raw is not None: + named_lengths.append(("input_inclusive_min", len(inclusive_min_raw))) + if upper_raw is not None: + named_lengths.append((upper_field or "", len(upper_raw))) + if labels_raw is not None: + named_lengths.append(("input_labels", len(labels_raw))) + rank = _infer_rank(obj, named_lengths, declared=declared_rank) + + inclusive_min = inclusive_min_raw if inclusive_min_raw is not None else [0] * rank + exclusive_max = _resolve_upper_bound( + upper_field, + upper_raw, + inclusive_min, + rank, + kind_of=_upper_kind(upper_field, _TRANSFORM_UPPER), + ) + labels = labels_raw if labels_raw is not None else [""] * rank + _validate_domain(inclusive_min, exclusive_max, prefix="input") + + if "output" in obj: + if not isinstance(obj["output"], list): + raise NdselError("invalid_json", f"output must be an array, got {obj['output']!r}") + output = [_normalize_output_map(m, f"output[{i}]") for i, m in enumerate(obj["output"])] + else: + output = _identity_output(rank) + + return { + "input_rank": rank, + "input_inclusive_min": inclusive_min, + "input_exclusive_max": exclusive_max, + "input_labels": labels, + "output": output, + } + + +_NORMALIZERS = { + "point": _normalize_point, + "box": _normalize_box, + "slice": _normalize_slice, + "points": _normalize_points, + "transform": _normalize_transform, +} + + +# --------------------------------------------------------------------------- +# trunc division (spec section 5.3 correction, matches _trunc_div in transform.py) +# --------------------------------------------------------------------------- + + +def _trunc_div(a: int, b: int) -> int: + """Integer division rounded toward zero (C semantics).""" + q = a // b + if q < 0 and q * b != a: + q += 1 + return q + + +# --------------------------------------------------------------------------- +# Public entry points +# --------------------------------------------------------------------------- + + +def normalize_ndsel(obj: Any) -> dict[str, Any]: + """Desugar and canonicalize an ndsel message to its canonical transform body. + + Accepts any of the five message kinds and returns the bare canonical + `IndexTransform` body of spec section 4.3 — no `kind` field. Raises + `NdselError` (carrying a reason code) for any invalid input. + """ + message = _require_object(obj) + kind = _message_kind(message) + return _NORMALIZERS[kind](message) + + +def parse_ndsel(obj: Any) -> dict[str, Any]: + """Structurally validate an ndsel message, returning it unchanged. + + A lighter gate than `normalize_ndsel`: it confirms the message is a + well-formed ndsel message of a recognized kind (correct field membership, + JSON types, upper-bound exclusivity, domain ordering, step signs) and + raises `NdselError` otherwise, but does not desugar it. Useful for + validating a message you intend to keep in its compact shorthand form. + """ + message = _require_object(obj) + _message_kind(message) + # Validation and desugaring share one pass; run it and discard the body. + normalize_ndsel(message) + return message diff --git a/packages/zarr-indexing/src/zarr_indexing/output_map.py b/packages/zarr-indexing/src/zarr_indexing/output_map.py new file mode 100644 index 0000000000..581229bd22 --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/output_map.py @@ -0,0 +1,105 @@ +"""Output index maps — three representations of a set of integer coordinates. + +An output index map describes, for one dimension of storage, which coordinates +an array access will touch. Conceptually it is a **set of integers**. Three +representations cover the cases that arise in practice: + +- `ConstantMap(offset=5)` — a singleton set: `{5}` +- `DimensionMap(input_dimension=0, offset=3, stride=2)` over input `[0, 5)` + — an arithmetic progression: `{3, 5, 7, 9, 11}` +- `ArrayMap(index_array=[1, 5, 9])` — an explicit enumeration: `{1, 5, 9}` + +Every output map supports two set-theoretic operations (defined on +`IndexTransform`, which provides the input domain context these maps lack): + +- **intersect** — restrict to coordinates within a range (e.g., a chunk). + `{3, 5, 7, 9, 11} ∩ [4, 8) = {5, 7}` +- **translate** — shift every coordinate by a constant (e.g., make chunk-local). + `{5, 7} - 4 = {1, 3}` + +These two operations are the foundation of chunk resolution: for each chunk, +intersect the map with the chunk's range, then translate to chunk-local +coordinates. + +The three types exist because they trade off generality for efficiency: + +- `ConstantMap`: O(1) storage, O(1) intersection +- `DimensionMap`: O(1) storage, O(1) intersection (analytical) +- `ArrayMap`: O(n) storage, O(n) intersection (must scan the array) + +Collapsing everything to `ArrayMap` would be correct but wasteful — a +billion-element slice would materialize a billion coordinates just to group +them by chunk, when `DimensionMap` does it with three integers. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import numpy as np + import numpy.typing as npt + + +@dataclass(frozen=True, slots=True) +class ConstantMap: + """A singleton set: one storage coordinate. + + Represents `{offset}`. Arises from integer indexing (e.g., `arr[5]` + fixes one dimension to coordinate 5). + """ + + offset: int = 0 + + +@dataclass(frozen=True, slots=True) +class DimensionMap: + """An arithmetic progression of storage coordinates. + + Represents `{offset + stride * i : i in input_range}`, where the input + range comes from the enclosing `IndexTransform`'s domain. Arises from + slice indexing (e.g., `arr[2:10:3]` gives offset=2, stride=3). + """ + + input_dimension: int + offset: int = 0 + stride: int = 1 + + +@dataclass(frozen=True, slots=True) +class ArrayMap: + """An explicit enumeration of storage coordinates. + + Represents `{offset + stride * index_array[i] : i in input_range}`. + Arises from fancy indexing (e.g., `arr[[1, 5, 9]]` or boolean masks). + + Freshly constructed maps are normalized to the **full input rank** of their + enclosing transform: `index_array` has the enclosing domain's rank, sized + fully on the axes it varies over and singleton (size 1) elsewhere. The + dependency axes are therefore derivable from the shape (see + `transform._array_map_dependency_axes`), which distinguishes the two flavours + of multi-array fancy indexing: + + - **orthogonal** (`oindex`): each array varies along a single, *distinct* + axis (all others singleton); the result is their outer product. + - **vectorized** (`vindex`): the arrays are correlated and share the same + non-singleton (broadcast) axes; the result is a pointwise scatter. + + `input_dimension` records the single axis an orthogonal array varies over + (`None` for vectorized), binding it the way `DimensionMap` is bound. It is + usually redundant with the shape-derived classifier, but stays authoritative + for the shapes the classifier cannot distinguish: a length-1 orthogonal + selection normalizes to an all-singleton array (no non-singleton axis), and + length-1 vectorized arrays are equally degenerate. `None` therefore marks a + map as correlated, and an integer pins the dependency axis of a degenerate + orthogonal map (see `transform._array_map_dependent_axis`). + """ + + index_array: npt.NDArray[np.intp] + offset: int = 0 + stride: int = 1 + input_dimension: int | None = None + + +OutputIndexMap = ConstantMap | DimensionMap | ArrayMap diff --git a/packages/zarr-indexing/src/zarr_indexing/py.typed b/packages/zarr-indexing/src/zarr_indexing/py.typed new file mode 100644 index 0000000000..e69de29bb2 diff --git a/packages/zarr-indexing/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py new file mode 100644 index 0000000000..e1a3898b1d --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/transform.py @@ -0,0 +1,1311 @@ +"""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. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any, Literal, cast + +import numpy as np + +from zarr_indexing.domain import IndexDomain +from zarr_indexing.errors import BoundsCheckError, VindexInvalidSelectionError +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap + + +@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) + + +def _intersect( + transform: IndexTransform, output_domain: IndexDomain +) -> ( + tuple[ + IndexTransform, + dict[int, np.ndarray[Any, np.dtype[np.intp]]] | np.ndarray[Any, np.dtype[np.intp]] | None, + ] + | None +): + """Intersect a transform with an output domain (e.g., a chunk's bounds). + + For each output dimension, restrict to storage coordinates within + `[output_domain.inclusive_min[d], output_domain.exclusive_max[d])`. + + Two flavours of fancy indexing require different treatment, distinguished by + the ArrayMaps' dependency axes (see `_array_map_dependency_axes`): + + - **orthogonal** (`oindex`): each ArrayMap varies over a single, distinct + input axis, forming an outer product. Every output dimension is intersected + independently and the input domain narrowed per axis. + - **correlated** (`vindex`): the ArrayMaps share their (broadcast) dependency + axes and scatter through a single flat index. A point survives only if ALL + its storage coordinates fall within the output domain; residual slice + dimensions are intersected independently, as in the orthogonal case. + + A `None` `input_dimension` marks a correlated map, so any such map routes the + whole transform through the correlated intersection. + + Returns `None` if the intersection is empty. + """ + if output_domain.ndim != transform.output_rank: + raise ValueError( + f"output_domain rank ({output_domain.ndim}) != " + f"transform output rank ({transform.output_rank})" + ) + + correlated_dims = [ + i + for i, m in enumerate(transform.output) + if isinstance(m, ArrayMap) and m.input_dimension is None + ] + if len(correlated_dims) > 0: + return _intersect_correlated(transform, output_domain, correlated_dims) + return _intersect_orthogonal(transform, output_domain) + + +def _intersect_dimension_map( + m: DimensionMap, input_lo: int, input_hi: int, lo: int, hi: int +) -> tuple[int, int] | None: + """Narrow a DimensionMap's input range to storage coordinates in `[lo, hi)`. + + `input_lo`/`input_hi` are the current (possibly already narrowed) input + range for the map's axis. Returns the new `(input_lo, input_hi)` or `None` + if no input produces an in-bounds storage coordinate. + """ + if input_lo >= input_hi: + return None + if m.stride > 0: + new_input_lo = max(input_lo, math.ceil((lo - m.offset) / m.stride)) + new_input_hi = min(input_hi, math.ceil((hi - m.offset) / m.stride)) + elif m.stride < 0: + new_input_lo = max(input_lo, math.ceil((hi - 1 - m.offset) / m.stride)) + new_input_hi = min(input_hi, math.ceil((lo - 1 - m.offset) / m.stride)) + else: + if lo <= m.offset < hi: + new_input_lo, new_input_hi = input_lo, input_hi + else: + return None + if new_input_lo >= new_input_hi: + return None + return new_input_lo, new_input_hi + + +def _intersect_orthogonal( + transform: IndexTransform, output_domain: IndexDomain +) -> ( + tuple[ + IndexTransform, + dict[int, np.ndarray[Any, np.dtype[np.intp]]] | np.ndarray[Any, np.dtype[np.intp]] | None, + ] + | None +): + """Intersect a transform with no correlated ArrayMaps. + + Every output dimension is intersected independently. Multiple ArrayMaps bound + to distinct input dimensions form an outer product, so each array's surviving + *output* positions are tracked separately. + """ + new_min = list(transform.domain.inclusive_min) + new_max = list(transform.domain.exclusive_max) + new_output: list[OutputIndexMap] = [] + out_positions: dict[int, np.ndarray[Any, np.dtype[np.intp]]] = {} + + for out_dim, m in enumerate(transform.output): + lo = output_domain.inclusive_min[out_dim] + hi = output_domain.exclusive_max[out_dim] + + if isinstance(m, ConstantMap): + if lo <= m.offset < hi: + new_output.append(m) + else: + return None + + elif isinstance(m, DimensionMap): + d = m.input_dimension + narrowed = _intersect_dimension_map(m, new_min[d], new_max[d], lo, hi) + if narrowed is None: + return None + new_min[d], new_max[d] = narrowed + new_output.append(m) + + else: + # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap) + # Orthogonal: the array varies over a single axis (its dependency + # axis, or `input_dimension` for a degenerate length-1 array). Filter + # along that axis and keep the array at full input rank so the + # singleton axes it broadcasts over are preserved. + d = _array_map_dependent_axis(m) + storage = m.offset + m.stride * m.index_array + mask = (storage >= lo) & (storage < hi) + # The array is singleton on every axis but `d`, so its mask reduces + # to a 1-D vector along `d`. + survivors = np.nonzero(mask.reshape(-1))[0].astype(np.intp) + if survivors.size == 0: + return None + filtered = np.take(m.index_array, survivors, axis=d) + new_output.append( + ArrayMap( + index_array=np.asarray(filtered, dtype=np.intp), + offset=m.offset, + stride=m.stride, + input_dimension=m.input_dimension, + ) + ) + new_max[d] = new_min[d] + int(survivors.size) + out_positions[out_dim] = survivors + + new_domain = IndexDomain( + inclusive_min=tuple(new_min), + exclusive_max=tuple(new_max), + ) + result = IndexTransform(domain=new_domain, output=tuple(new_output)) + + # Hand back the surviving output positions in the shape the bridge expects: + # None (no arrays), a single vector (one array), or a per-output-dim dict + # (>= 2 orthogonal arrays → outer product). + out_indices: ( + dict[int, np.ndarray[Any, np.dtype[np.intp]]] | np.ndarray[Any, np.dtype[np.intp]] | None + ) + if len(out_positions) == 0: + out_indices = None + elif len(out_positions) == 1: + out_indices = next(iter(out_positions.values())) + else: + out_indices = out_positions + return (result, out_indices) + + +def _intersect_correlated( + transform: IndexTransform, + output_domain: IndexDomain, + correlated_dims: list[int], +) -> tuple[IndexTransform, np.ndarray[Any, np.dtype[np.intp]]] | None: + """Intersect a correlated (vindex) transform with an output domain. + + The correlated ArrayMaps share their broadcast (dependency) axes; a broadcast + point survives only if ALL its storage coordinates fall within the output + domain. Residual DimensionMap dimensions are intersected independently (as in + the orthogonal case) and preserved, so a partial vindex — e.g. two coordinate + arrays over a 3-D array, leaving one slice dimension — resolves correctly. + + The surviving broadcast axes collapse to a single axis; the returned + `out_indices` is the flat scatter index into the (row-major flattened) + output buffer, of shape `(surviving_points,) + (residual slice sizes)`. + """ + corr_maps = [cast("ArrayMap", transform.output[i]) for i in correlated_dims] + + # Mixing correlated and orthogonal ArrayMaps in one transform is not produced + # by any single selection and is not supported here. + orthogonal_array_dims = [ + i + for i, m in enumerate(transform.output) + if isinstance(m, ArrayMap) and m.input_dimension is not None + ] + if len(orthogonal_array_dims) > 0: + raise NotImplementedError( + "intersecting a transform with both correlated and orthogonal " + "ArrayMaps is not supported" + ) + + # The broadcast (dependency) axes are shared by every correlated map; they are + # the leading axes of the domain, followed by the residual slice axes. + broadcast_axes = _array_map_dependency_axes(corr_maps[0].index_array) + broadcast_shape = tuple(corr_maps[0].index_array.shape[a] for a in broadcast_axes) + + # Joint bounds mask over the broadcast block. + combined: np.ndarray[Any, np.dtype[np.bool_]] | None = None + for out_dim in correlated_dims: + cm = cast("ArrayMap", transform.output[out_dim]) + storage = cm.offset + cm.stride * cm.index_array + lo = output_domain.inclusive_min[out_dim] + hi = output_domain.exclusive_max[out_dim] + mask = (storage >= lo) & (storage < hi) + combined = mask if combined is None else (combined & mask) + assert combined is not None + # The correlated maps are singleton on every non-broadcast axis, so the mask + # collapses (C-order) to the broadcast block. + combined_bcast = combined.reshape(broadcast_shape) + surviving = np.nonzero(combined_bcast.reshape(-1))[0].astype(np.intp) + if surviving.size == 0: + return None + + # Intersect residual (slice / constant) dimensions independently. Slice dims + # are ordered by input dimension so their flat-buffer strides are row-major. + slice_dims: list[tuple[int, int, int, int, DimensionMap]] = [] # (in_dim, lo, hi, full, m) + for out_dim, m in enumerate(transform.output): + if out_dim in correlated_dims: + continue + lo = output_domain.inclusive_min[out_dim] + hi = output_domain.exclusive_max[out_dim] + if isinstance(m, ConstantMap): + if not (lo <= m.offset < hi): + return None + elif isinstance(m, DimensionMap): + d = m.input_dimension + input_lo = transform.domain.inclusive_min[d] + input_hi = transform.domain.exclusive_max[d] + narrowed = _intersect_dimension_map(m, input_lo, input_hi, lo, hi) + if narrowed is None: + return None + slice_dims.append((d, narrowed[0], narrowed[1], input_hi - input_lo, m)) + slice_dims.sort(key=lambda item: item[0]) + + n_points = int(surviving.size) + n_slice = len(slice_dims) + corr_values = { + out_dim: cast("ArrayMap", transform.output[out_dim]) + .index_array.reshape(broadcast_shape) + .reshape(-1)[surviving] + for out_dim in correlated_dims + } + + # New domain: the collapsed broadcast axis, then one axis per residual slice. + new_min = [0] + new_max = [n_points] + new_input_dim_of = {} + for new_axis, (d, nlo, nhi, _full, _m) in enumerate(slice_dims, start=1): + new_min.append(nlo) + new_max.append(nhi) + new_input_dim_of[d] = new_axis + new_domain = IndexDomain(inclusive_min=tuple(new_min), exclusive_max=tuple(new_max)) + + corr_shape = (n_points,) + (1,) * n_slice + new_output: list[OutputIndexMap] = [] + for out_dim, m in enumerate(transform.output): + if out_dim in correlated_dims: + corr = cast("ArrayMap", m) + new_output.append( + ArrayMap( + index_array=corr_values[out_dim].reshape(corr_shape).astype(np.intp), + offset=corr.offset, + stride=corr.stride, + ) + ) + elif isinstance(m, ConstantMap): + new_output.append(m) + else: + assert isinstance(m, DimensionMap) + new_output.append( + DimensionMap( + input_dimension=new_input_dim_of[m.input_dimension], + offset=m.offset, + stride=m.stride, + ) + ) + result = IndexTransform(domain=new_domain, output=tuple(new_output)) + + # Flat scatter index into the row-major output buffer of shape + # (broadcast points, residual slice sizes...): flat = point * prod(slice) + + # (row-major offset within the slice block). + prod_slice = 1 + for _d, _lo, _hi, full, _m in slice_dims: + prod_slice *= full + out_indices: np.ndarray[Any, np.dtype[np.intp]] = (surviving * prod_slice).reshape( + (n_points,) + (1,) * n_slice + ) + running = 1 + for j in range(n_slice - 1, -1, -1): + _d, nlo, nhi, full, _m = slice_dims[j] + coords = np.arange(nlo, nhi, dtype=np.intp) * running + shape = [1] * (1 + n_slice) + shape[1 + j] = coords.size + out_indices = out_indices + coords.reshape(shape) + running *= full + return (result, out_indices.astype(np.intp)) + + +def _normalize_basic_selection(selection: Any, ndim: int) -> tuple[int | slice | None, ...]: + """Normalize a selection to a tuple of int, slice, or None (newaxis), + expanding ellipsis and padding with slice(None) as needed. + """ + if not isinstance(selection, tuple): + selection = (selection,) + + # Count non-newaxis, non-ellipsis entries to determine how many real dims are addressed + n_newaxis = sum(1 for s in selection if s is None) + has_ellipsis = any(s is Ellipsis for s in selection) + n_real = len(selection) - n_newaxis - (1 if has_ellipsis else 0) + + if n_real > ndim: + raise IndexError( + f"too many indices for array: array has {ndim} dimensions, but {n_real} were indexed" + ) + + result: list[int | slice | None] = [] + ellipsis_seen = False + for sel in selection: + if sel is Ellipsis: + if ellipsis_seen: + raise IndexError("an index can only have a single ellipsis ('...')") + ellipsis_seen = True + num_missing = ndim - n_real + result.extend([slice(None)] * num_missing) + elif isinstance(sel, (int, np.integer)): + result.append(int(sel)) + elif isinstance(sel, slice) or sel is None: + result.append(sel) + else: + raise IndexError(f"unsupported selection type for basic indexing: {type(sel)!r}") + + # Pad remaining dimensions with slice(None) + while sum(1 for s in result if s is not None) < ndim: + result.append(slice(None)) + + return tuple(result) + + +def _reindex_array( + m: ArrayMap, + normalized: tuple[int | slice | None, ...], + domain: IndexDomain, +) -> np.ndarray[Any, np.dtype[np.intp]]: + """Apply basic indexing operations to an ArrayMap's index_array. + + The array's axes correspond to the transform's input dimensions (0-indexed + over the domain shape). Each axis is either a **dependency axis** — the array + genuinely varies with that input dimension — or a **singleton** axis it + broadcasts over. Integer indexing, slicing, or newaxis is applied to the + array only along its dependency axes; a selection on a singleton axis does not + touch the array's values (it just narrows or drops that broadcast axis). + """ + dependent = set(_array_map_dependency_axes(m.index_array)) + if m.input_dimension is not None: + # Degenerate length-1 orthogonal selection: the recorded axis is a + # dependency even though its size (1) makes it look singleton. + dependent.add(m.input_dimension) + arr = m.index_array + + # Build a numpy indexing tuple: one entry per old input dimension + idx: list[Any] = [] + old_dim = 0 + newaxis_positions: list[int] = [] + result_axis = 0 + + for sel in normalized: + if sel is None: + newaxis_positions.append(result_axis) + result_axis += 1 + elif isinstance(sel, int): + if old_dim < arr.ndim: + if old_dim in dependent: + # Convert absolute domain coordinate to 0-based array index + idx.append(sel - domain.inclusive_min[old_dim]) + else: + # Broadcast axis: keep the single element and drop the axis. + idx.append(0) + old_dim += 1 + else: + # sel: slice (normalized: tuple[int | slice | None, ...]) + if old_dim < arr.ndim: + if old_dim in dependent: + lo = domain.inclusive_min[old_dim] + hi = domain.exclusive_max[old_dim] + # Bounds are literal domain coordinates; the stored array is + # indexed positionally, so shift by the domain origin. + start, step, _origin, size = _resolve_slice_ts(sel, old_dim, lo, hi) + pos = start - lo + idx.append(slice(pos, pos + size * step, step)) + else: + # Broadcast axis: preserve the singleton (it still broadcasts + # over the narrowed domain), regardless of the slice bounds. + idx.append(slice(None)) + old_dim += 1 + result_axis += 1 + + result = arr[tuple(idx)] if idx else arr + + for pos in newaxis_positions: + result = np.expand_dims(result, axis=pos) + + return np.asarray(result, dtype=np.intp) + + +_FANCY_AFTER_FANCY_MSG = ( + "applying a fancy (orthogonal/vectorized) selection to a view that already " + "has a fancy-indexed axis is not supported (fancy-after-fancy composition): " + "the new coordinates would index a broadcast axis of the existing selection. " + "Materialize the view first with `.result()` and index the array, or reorder " + "the selections so the fancy step is applied last." +) + + +def _guard_fancy_after_fancy(m: ArrayMap, fancy_dims: set[int] | list[int]) -> None: + """Reject a fancy step that lands on a broadcast axis of an existing ArrayMap. + + A new orthogonal/vectorized selection can only be absorbed into an existing + ArrayMap along the axes that map genuinely varies over (its dependency axes, + plus the recorded `input_dimension` for a degenerate length-1 orthogonal + selection). A fancy index targeting any other axis — a singleton axis the map + merely broadcasts over — cannot be reindexed and used to leak a raw NumPy + `IndexError` at resolve time. Raise a clear `NotImplementedError` instead. + """ + dependent = set(_array_map_dependency_axes(m.index_array)) + if m.input_dimension is not None: + dependent.add(m.input_dimension) + for d in fancy_dims: + if d < m.index_array.ndim and d not in dependent: + raise NotImplementedError(_FANCY_AFTER_FANCY_MSG) + + +def _reindex_array_oindex( + arr: np.ndarray[Any, np.dtype[np.intp]], + normalized: tuple[Any, ...] | list[Any], + domain: IndexDomain, +) -> np.ndarray[Any, np.dtype[np.intp]]: + """Apply oindex/vindex selection to an existing ArrayMap's index_array. + + Each old input dimension gets either an array (fancy index that axis) + or a slice applied to the corresponding array axis. + """ + idx: list[Any] = [] + for old_dim, sel in enumerate(normalized): + if old_dim >= arr.ndim: + break + lo = domain.inclusive_min[old_dim] + if isinstance(sel, np.ndarray): + # Values are literal domain coordinates; the stored array is + # indexed positionally, so shift by the domain origin. + idx.append(sel - lo) + elif isinstance(sel, slice): + hi = domain.exclusive_max[old_dim] + start, step, _origin, size = _resolve_slice_ts(sel, old_dim, lo, hi) + pos = start - lo + idx.append(slice(pos, pos + size * step, step)) + else: + idx.append(slice(None)) + + result = arr[tuple(idx)] if idx else arr + return np.asarray(result, dtype=np.intp) + + +def _apply_basic_indexing(transform: IndexTransform, selection: Any) -> IndexTransform: + """Apply basic indexing (int, slice, ellipsis, newaxis) to an IndexTransform.""" + normalized = _normalize_basic_selection(selection, transform.domain.ndim) + + new_inclusive_min: list[int] = [] + new_exclusive_max: list[int] = [] + old_dim = 0 + new_dim_idx = 0 + old_to_new_dim: dict[int, int] = {} + dropped_dims: set[int] = set() + + # Per old-dim: the slice parameters (for computing new output maps) + dim_slice_params: dict[int, tuple[int, int, int]] = {} # old_dim -> (start, stop, step) + dim_int_val: dict[int, int] = {} # old_dim -> integer index value + + for sel in normalized: + if sel is None: + # newaxis: add a size-1 dimension + new_inclusive_min.append(0) + new_exclusive_max.append(1) + new_dim_idx += 1 + elif isinstance(sel, int): + # Integer index: drop this input dimension. + # Negative indices are literal coordinates (TensorStore convention), + # NOT "from the end" like NumPy. The Array layer handles conversion. + lo = transform.domain.inclusive_min[old_dim] + hi = transform.domain.exclusive_max[old_dim] + idx = sel + if idx < lo or idx >= hi: + hint = _LITERAL_HINT if sel < 0 else "" + raise BoundsCheckError( + f"index {sel} is out of bounds for dimension {old_dim} " + f"(valid indices [{lo}, {hi})){hint}" + ) + dropped_dims.add(old_dim) + dim_int_val[old_dim] = idx + old_dim += 1 + else: + # sel: slice (normalized: tuple[int | slice | None, ...]) + lo = transform.domain.inclusive_min[old_dim] + hi = transform.domain.exclusive_max[old_dim] + + # TensorStore semantics: bounds are literal coordinates; a step-1 + # slice keeps them as the new domain, a strided slice's domain is + # [trunc(start/step), trunc(start/step) + size). + start, step, origin, size = _resolve_slice_ts(sel, old_dim, lo, hi) + new_inclusive_min.append(origin) + new_exclusive_max.append(origin + size) + dim_slice_params[old_dim] = (start, step, origin) + old_to_new_dim[old_dim] = new_dim_idx + new_dim_idx += 1 + old_dim += 1 + + new_domain = IndexDomain( + inclusive_min=tuple(new_inclusive_min), + exclusive_max=tuple(new_exclusive_max), + ) + + # Now update output maps + new_output: list[OutputIndexMap] = [] + for m in transform.output: + if isinstance(m, ConstantMap): + new_output.append(m) + elif isinstance(m, DimensionMap): + d = m.input_dimension + if d in dropped_dims: + # Integer index: this output becomes constant + new_offset = m.offset + m.stride * dim_int_val[d] + new_output.append(ConstantMap(offset=new_offset)) + elif d in old_to_new_dim: + # Slice: new coordinate `origin + k` maps to old coordinate + # `start + k*step`, i.e. old = start - step*origin + step*new. + start, step, origin = dim_slice_params[d] + new_offset = m.offset + m.stride * (start - step * origin) + new_stride = m.stride * step + new_input_dim = old_to_new_dim[d] + new_output.append( + DimensionMap( + input_dimension=new_input_dim, offset=new_offset, stride=new_stride + ) + ) + else: + raise RuntimeError(f"unexpected: dimension {d} not handled") + else: + # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap) + new_arr = _reindex_array(m, normalized, transform.domain) + array_input_dim: int | None = None + if m.input_dimension is not None: + array_input_dim = old_to_new_dim.get(m.input_dimension, m.input_dimension) + new_output.append( + ArrayMap( + index_array=new_arr, + offset=m.offset, + stride=m.stride, + input_dimension=array_input_dim, + ) + ) + + return IndexTransform(domain=new_domain, output=tuple(new_output)) + + +def _array_map_dependency_axes(index_array: np.ndarray[Any, Any]) -> tuple[int, ...]: + """Return the input axes on which a normalized index array varies. + + Normalized `ArrayMap` index arrays carry the full input rank of their + enclosing transform: an axis the array varies over has its full size, while + an axis the array is independent of is a singleton (size 1). The dependency + axes are therefore exactly the non-singleton axes. An orthogonal (`oindex`) + array depends on a single axis; a vectorized (`vindex`) array depends on all + of the (shared) broadcast axes. + """ + return tuple(axis for axis, size in enumerate(index_array.shape) if size != 1) + + +def _array_map_dependent_axis(m: ArrayMap) -> int: + """Return the single input axis an orthogonal `ArrayMap` varies over. + + Normally this is the array's one non-singleton axis. A degenerate length-1 + orthogonal selection normalizes to an all-singleton shape (its dependency + axes are empty and indistinguishable by shape from a scalar), so + `input_dimension` breaks the tie — it records the axis the map binds. + """ + dep = _array_map_dependency_axes(m.index_array) + if len(dep) == 1: + return dep[0] + if m.input_dimension is not None: + return m.input_dimension + raise ValueError( + f"orthogonal ArrayMap must vary over exactly one axis; got dependency " + f"axes {dep} with input_dimension={m.input_dimension}" + ) + + +def _reshape_to_axis( + values: np.ndarray[Any, np.dtype[np.intp]], axis: int, ndim: int +) -> np.ndarray[Any, np.dtype[np.intp]]: + """Reshape a 1-D selection to full rank `ndim` varying only along `axis`. + + The result has `values` laid out along `axis` and singleton (size-1) axes + everywhere else, so its dependency axis is derivable from its shape. + """ + flat = np.asarray(values, dtype=np.intp).ravel() + shape = [1] * ndim + shape[axis] = flat.shape[0] + return flat.reshape(shape) + + +class _OIndexHelper: + """Helper that provides orthogonal (outer) indexing via `transform.oindex[...]`.""" + + def __init__(self, transform: IndexTransform) -> None: + self._transform = transform + + def __getitem__(self, selection: Any) -> IndexTransform: + return _apply_oindex(self._transform, selection) + + +def _normalize_oindex_selection( + selection: Any, ndim: int +) -> tuple[np.ndarray[Any, np.dtype[np.intp]] | slice, ...]: + """Normalize an oindex selection: arrays, slices, booleans, integers.""" + if not isinstance(selection, tuple): + selection = (selection,) + + # Expand ellipsis + has_ellipsis = any(s is Ellipsis for s in selection) + n_ellipsis = 1 if has_ellipsis else 0 + n_real = len(selection) - n_ellipsis + + result: list[np.ndarray[Any, np.dtype[np.intp]] | slice] = [] + for sel in selection: + if sel is Ellipsis: + num_missing = ndim - n_real + result.extend([slice(None)] * num_missing) + elif isinstance(sel, np.ndarray) and sel.dtype == np.bool_: + # Boolean array -> integer indices + (indices,) = np.nonzero(sel) + result.append(indices.astype(np.intp)) + elif isinstance(sel, np.ndarray): + result.append(sel.astype(np.intp)) + elif isinstance(sel, slice): + result.append(sel) + elif isinstance(sel, (int, np.integer)): + # Convert integer scalars to 1-element arrays for orthogonal indexing + result.append(np.array([int(sel)], dtype=np.intp)) + elif isinstance(sel, (list, tuple)): + result.append(np.asarray(sel, dtype=np.intp)) + else: + result.append(sel) + + # Pad with slice(None) + while len(result) < ndim: + result.append(slice(None)) + + return tuple(result) + + +def _apply_oindex(transform: IndexTransform, selection: Any) -> IndexTransform: + """Apply orthogonal indexing to an IndexTransform. + + Each index array is applied independently per dimension (outer product). + """ + normalized = _normalize_oindex_selection(selection, transform.domain.ndim) + + new_inclusive_min: list[int] = [] + new_exclusive_max: list[int] = [] + new_dim_idx = 0 + old_to_new_dim: dict[int, int] = {} + + # Info per old dim + dim_array: dict[int, np.ndarray[Any, np.dtype[np.intp]]] = {} + dim_slice_params: dict[int, tuple[int, int, int]] = {} + + for old_dim, sel in enumerate(normalized): + if isinstance(sel, np.ndarray): + lo = transform.domain.inclusive_min[old_dim] + hi = transform.domain.exclusive_max[old_dim] + # Index-array values are literal domain coordinates; the fancy dim + # they create gets a fresh zero-origin [0, n) domain (TensorStore). + _check_array_in_bounds(sel, lo, hi) + dim_array[old_dim] = sel + new_inclusive_min.append(0) + new_exclusive_max.append(len(sel)) + old_to_new_dim[old_dim] = new_dim_idx + new_dim_idx += 1 + else: + # sel: slice (_normalize_oindex_selection returns + # tuple[np.ndarray | slice, ...]) + lo = transform.domain.inclusive_min[old_dim] + hi = transform.domain.exclusive_max[old_dim] + start, step, origin, size = _resolve_slice_ts(sel, old_dim, lo, hi) + new_inclusive_min.append(origin) + new_exclusive_max.append(origin + size) + dim_slice_params[old_dim] = (start, step, origin) + old_to_new_dim[old_dim] = new_dim_idx + new_dim_idx += 1 + + new_domain = IndexDomain( + inclusive_min=tuple(new_inclusive_min), + exclusive_max=tuple(new_exclusive_max), + ) + + new_output: list[OutputIndexMap] = [] + for m in transform.output: + if isinstance(m, ConstantMap): + new_output.append(m) + elif isinstance(m, DimensionMap): + d = m.input_dimension + if d in dim_array: + new_axis = old_to_new_dim[d] + # Normalize to full input rank: the selection varies along its own + # new axis and is singleton on every other axis. The dependency + # axis is then derivable from the shape (a single non-singleton + # axis marks the selection orthogonal / outer-product rather than + # vectorized). `input_dimension` is kept populated as a + # compatibility shim for consumers not yet migrated to the + # shape-derived classifier. + full_arr = _reshape_to_axis(dim_array[d], new_axis, new_dim_idx) + new_output.append( + ArrayMap( + index_array=full_arr, + offset=m.offset, + stride=m.stride, + input_dimension=new_axis, + ) + ) + elif d in dim_slice_params: + start, step, origin = dim_slice_params[d] + new_offset = m.offset + m.stride * (start - step * origin) + new_stride = m.stride * step + new_input_dim = old_to_new_dim[d] + new_output.append( + DimensionMap( + input_dimension=new_input_dim, offset=new_offset, stride=new_stride + ) + ) + else: + raise RuntimeError(f"unexpected: dimension {d} not handled") + else: + # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap) + _guard_fancy_after_fancy(m, list(dim_array.keys())) + new_arr = _reindex_array_oindex(m.index_array, normalized, transform.domain) + array_input_dim: int | None = None + if m.input_dimension is not None: + array_input_dim = old_to_new_dim.get(m.input_dimension, m.input_dimension) + new_output.append( + ArrayMap( + index_array=new_arr, + offset=m.offset, + stride=m.stride, + input_dimension=array_input_dim, + ) + ) + + return IndexTransform(domain=new_domain, output=tuple(new_output)) + + +class _VIndexHelper: + """Helper that provides vectorized (fancy) indexing via `transform.vindex[...]`.""" + + def __init__(self, transform: IndexTransform) -> None: + self._transform = transform + + def __getitem__(self, selection: Any) -> IndexTransform: + return _apply_vindex(self._transform, selection) + + +def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform: + """Apply vectorized indexing to an IndexTransform. + + All array indices are broadcast together. Broadcast dimensions are prepended, + followed by non-array (slice) dimensions. + """ + if not isinstance(selection, tuple): + selection = (selection,) + + # Expand ellipsis and count consumed dimensions + # Boolean arrays with ndim > 1 consume ndim dims + n_consumed = 0 + for s in selection: + if s is Ellipsis: + continue + if isinstance(s, np.ndarray) and s.dtype == np.bool_ and s.ndim > 1: + n_consumed += s.ndim + else: + n_consumed += 1 + ndim = transform.domain.ndim + + expanded: list[Any] = [] + for sel in selection: + if sel is Ellipsis: + num_missing = ndim - n_consumed + expanded.extend([slice(None)] * num_missing) + else: + expanded.append(sel) + # Count dimensions already consumed by expanded entries + n_expanded_dims = 0 + for sel in expanded: + if isinstance(sel, np.ndarray) and sel.dtype == np.bool_ and sel.ndim > 1: + n_expanded_dims += sel.ndim + else: + n_expanded_dims += 1 + while n_expanded_dims < ndim: + expanded.append(slice(None)) + n_expanded_dims += 1 + + # Convert booleans, lists, ints to integer arrays + processed: list[np.ndarray[Any, np.dtype[np.intp]] | slice] = [] + for sel in expanded: + if isinstance(sel, np.ndarray) and sel.dtype == np.bool_: + indices_tuple = np.nonzero(sel) + processed.extend(indices.astype(np.intp) for indices in indices_tuple) + elif isinstance(sel, np.ndarray): + processed.append(sel.astype(np.intp)) + elif isinstance(sel, (list, tuple)): + processed.append(np.asarray(sel, dtype=np.intp)) + elif isinstance(sel, (int, np.integer)): + processed.append(np.array([int(sel)], dtype=np.intp)) + else: + processed.append(sel) + + # Separate array dims and slice dims + array_dims: list[int] = [] + slice_dims: list[int] = [] + arrays: list[np.ndarray[Any, np.dtype[np.intp]]] = [] + + for i, sel in enumerate(processed): + if isinstance(sel, np.ndarray): + lo = transform.domain.inclusive_min[i] + hi = transform.domain.exclusive_max[i] + _check_array_in_bounds(sel, lo, hi) + array_dims.append(i) + arrays.append(sel) + else: + slice_dims.append(i) + + # Broadcast all arrays together + broadcast_arrays: list[np.ndarray[Any, np.dtype[np.intp]]] + if len(arrays) > 0: + broadcast_arrays = list(np.broadcast_arrays(*arrays)) + broadcast_shape = broadcast_arrays[0].shape + else: + broadcast_arrays = [] + broadcast_shape = () + + # Build new domain: broadcast dims first, then slice dims + new_inclusive_min: list[int] = [] + new_exclusive_max: list[int] = [] + + # Broadcast dimensions + for s in broadcast_shape: + new_inclusive_min.append(0) + new_exclusive_max.append(s) + + # Slice dimensions (preserved-domain literal semantics, like basic indexing) + slice_dim_params: dict[int, tuple[int, int, int]] = {} + for old_dim in slice_dims: + sel = processed[old_dim] + assert isinstance(sel, slice) + lo = transform.domain.inclusive_min[old_dim] + hi = transform.domain.exclusive_max[old_dim] + start, step, origin, size = _resolve_slice_ts(sel, old_dim, lo, hi) + new_inclusive_min.append(origin) + new_exclusive_max.append(origin + size) + slice_dim_params[old_dim] = (start, step, origin) + + new_domain = IndexDomain( + inclusive_min=tuple(new_inclusive_min), + exclusive_max=tuple(new_exclusive_max), + ) + + # Build output maps + array_dim_to_broadcast: dict[int, np.ndarray[Any, np.dtype[np.intp]]] = {} + for i, d in enumerate(array_dims): + array_dim_to_broadcast[d] = broadcast_arrays[i] + + # New dim index for slice dims starts after broadcast dims + n_broadcast_dims = len(broadcast_shape) + + new_output: list[OutputIndexMap] = [] + for m in transform.output: + if isinstance(m, ConstantMap): + new_output.append(m) + elif isinstance(m, DimensionMap): + d = m.input_dimension + if d in array_dim_to_broadcast: + # Normalize to full input rank: the broadcast (correlated) axes + # come first, followed by a singleton axis per slice dimension. + # Every vectorized array shares the same broadcast axes, so the + # dependency axes derived from the shape coincide — the signature + # of a pointwise scatter rather than an outer product. + broadcast_arr = array_dim_to_broadcast[d] + full_arr = broadcast_arr.reshape(broadcast_shape + (1,) * len(slice_dims)) + new_output.append( + ArrayMap( + index_array=full_arr, + offset=m.offset, + stride=m.stride, + ) + ) + else: + # Slice dim: new coord `origin + k` maps to old `start + k*step` + start, step, origin = slice_dim_params[d] + new_offset = m.offset + m.stride * (start - step * origin) + new_stride = m.stride * step + new_input_dim = n_broadcast_dims + slice_dims.index(d) + new_output.append( + DimensionMap( + input_dimension=new_input_dim, offset=new_offset, stride=new_stride + ) + ) + else: + # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap) + _guard_fancy_after_fancy(m, array_dims) + new_arr = _reindex_array_oindex(m.index_array, processed, transform.domain) + new_output.append( + ArrayMap( + index_array=new_arr, + offset=m.offset, + stride=m.stride, + input_dimension=m.input_dimension, + ) + ) + + return IndexTransform(domain=new_domain, output=tuple(new_output)) + + +_LITERAL_HINT = ( + "; within this transform layer, indices are literal domain coordinates (the " + "public Array boundary wraps NumPy-style negatives before they reach here)" +) + + +def _trunc_div(a: int, b: int) -> int: + """Integer division rounded toward zero (C semantics), as TensorStore uses + for strided-slice domain origins — distinct from Python's floor division + for negative operands (`trunc(-9/2) == -4` where `-9 // 2 == -5`).""" + q = a // b + if q < 0 and q * b != a: + q += 1 + return q + + +def _resolve_slice_ts(sel: slice, dim: int, lo: int, hi: int) -> tuple[int, int, int, int]: + """Resolve a slice against domain `[lo, hi)` with TensorStore semantics. + + Slice bounds are **literal domain coordinates** — never from-the-end, never + clamped. Rules (each verified against tensorstore 0.1.84): + + - defaults: `start = lo`, `stop = hi`; + - a non-empty interval must be contained in the domain (no clamping — a + NumPy-style out-of-range or negative bound is an error, not a shorter or + wrapped result); + - an **empty** interval (`start == stop`) is valid anywhere; + - reversed bounds (`start > stop` with positive step) are an error, not + an empty result; + - the result's domain origin is `trunc(start/step)` (rounded toward + zero) and coordinate `origin + k` maps to input `start + k*step`. + + Returns `(start, step, origin, size)` in domain coordinates. + """ + step = 1 if sel.step is None else sel.step + if step <= 0: + # Negative steps are valid in TensorStore but not yet supported here; + # step 0 is invalid everywhere. + raise IndexError("slice step must be positive") + start = lo if sel.start is None else sel.start + stop = hi if sel.stop is None else sel.stop + if stop < start: + raise IndexError( + f"slice interval [{start}, {stop}) with step {step} does not specify " + f"a valid interval for dimension {dim} (start > stop)" + ) + size = -(-(stop - start) // step) # ceil((stop - start) / step) + if size > 0 and (start < lo or stop > hi): + hint = _LITERAL_HINT if (start < 0 or stop < 0) and lo >= 0 else "" + raise BoundsCheckError( + f"slice interval [{start}, {stop}) is not contained within domain " + f"[{lo}, {hi}) for dimension {dim}{hint}" + ) + origin = _trunc_div(start, step) + return start, step, origin, size + + +def _check_array_in_bounds(arr: np.ndarray[Any, np.dtype[np.intp]], lo: int, hi: int) -> None: + """Reject index-array values outside the domain `[lo, hi)`. + + Index-array values are literal domain coordinates (TensorStore semantics): + a value below `inclusive_min` is out of bounds rather than counting from + the end. Out-of-range values raise instead of silently wrapping. + """ + if arr.size == 0: + return + lo_val, hi_val = int(arr.min()), int(arr.max()) + if lo_val < lo: + hint = _LITERAL_HINT if lo_val < 0 and lo >= 0 else "" + raise BoundsCheckError( + f"index {lo_val} is out of bounds (valid indices [{lo}, {hi})){hint}" + ) + if hi_val >= hi: + raise BoundsCheckError(f"index {hi_val} is out of bounds (valid indices [{lo}, {hi}))") + + +def _validate_array_selection(selection: Any, shape: tuple[int, ...], mode: str) -> None: + """Validate array-based selections (orthogonal, vectorized). + + Rejects types that are not valid for coordinate/vectorized indexing. + Does not check bounds — the transform operations handle that. + """ + items = selection if isinstance(selection, tuple) else (selection,) + for sel in items: + if isinstance(sel, slice): + # vindex is coordinate-only (matches eager zarr): every axis needs an + # integer/boolean array, never a slice. Orthogonal (oindex) allows slices. + if mode == "vectorized": + raise VindexInvalidSelectionError( + "unsupported selection type for vectorized indexing; only " + "coordinate selection (tuple of integer arrays) and mask selection " + f"(single Boolean array) are supported; got {selection!r}" + ) + continue + if sel is Ellipsis or isinstance(sel, (int, np.integer)): + continue + if isinstance(sel, (list, np.ndarray)): + continue + raise IndexError(f"unsupported selection type for {mode} indexing: {type(sel)!r}") + + +def _validate_basic_selection(selection: Any) -> None: + """Validate that a selection only contains basic indexing types (int, slice, Ellipsis). + + Rejects None (newaxis), arrays, lists, floats, strings, etc. + """ + items = selection if isinstance(selection, tuple) else (selection,) + for s in items: + if s is Ellipsis or isinstance(s, (int, np.integer, slice)): + continue + raise IndexError(f"unsupported selection type for basic indexing: {type(s)!r}") + + +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}") diff --git a/packages/zarr-indexing/tests/conformance/PROVENANCE.md b/packages/zarr-indexing/tests/conformance/PROVENANCE.md new file mode 100644 index 0000000000..0a6faef60b --- /dev/null +++ b/packages/zarr-indexing/tests/conformance/PROVENANCE.md @@ -0,0 +1,20 @@ +# Provenance of the ndsel conformance corpus + +The JSON fixtures in this directory (`point.json`, `box.json`, `slice.json`, +`points.json`, `transform.json`, `errors.json`) and `README.md` are **vendored, +unmodified**, from the ndsel reference repository. + +- **Source:** +- **Branch:** `main` (merge of d-v-b/ndsel#1, `fix/slice-origin-trunc`) +- **Commit:** `c59bc556c` (fixtures byte-identical to the previously vendored + `c132b4c1caa3205830ce35a42502363171f650a7`) +- **Path in source:** `conformance/` + +**Do not edit these files.** They are vendored as-is so that +`zarr_indexing`' ndsel message layer can be checked against the same +language-agnostic corpus every other ndsel implementation runs. To update the +corpus, re-vendor from a newer ndsel commit and update the commit SHA above. + +ndsel PR #1 (merged) corrected the `slice` desugaring origin from +`floor(a/s)` to `trunc(a/s)` (rounding toward zero), which matches +`zarr_indexing`' existing `_trunc_div` semantics. diff --git a/packages/zarr-indexing/tests/conformance/README.md b/packages/zarr-indexing/tests/conformance/README.md new file mode 100644 index 0000000000..ecb0c57ca2 --- /dev/null +++ b/packages/zarr-indexing/tests/conformance/README.md @@ -0,0 +1,16 @@ +# ndsel conformance corpus + +Language-agnostic fixtures. Each file is a JSON array of cases. + +A **success** case: + { "name": "...", "input": , "normalized": } + +An **error** case: + { "name": "...", "input": , "error": "" } + +An implementation is conformant iff, for every success case, +`normalize(input)` equals `normalized` by structural JSON equality, and for +every error case, `normalize(input)` is rejected with the given reason code. + +The `normalized` value is a canonical `transform` body (the `kind` field is +omitted; implementations compare the transform structure). diff --git a/packages/zarr-indexing/tests/conformance/box.json b/packages/zarr-indexing/tests/conformance/box.json new file mode 100644 index 0000000000..e847872f86 --- /dev/null +++ b/packages/zarr-indexing/tests/conformance/box.json @@ -0,0 +1,50 @@ +[ + { + "name": "box/2d-min-max", + "input": { "kind": "box", "inclusive_min": [0, 0], "exclusive_max": [3, 4] }, + "normalized": { + "input_rank": 2, + "input_inclusive_min": [0, 0], + "input_exclusive_max": [3, 4], + "input_labels": ["", ""], + "output": [ + { "offset": 0, "stride": 1, "input_dimension": 0 }, + { "offset": 0, "stride": 1, "input_dimension": 1 } + ] + } + }, + { + "name": "box/shape-only-origin-zero", + "input": { "kind": "box", "shape": [5] }, + "normalized": { + "input_rank": 1, + "input_inclusive_min": [0], + "input_exclusive_max": [5], + "input_labels": [""], + "output": [ { "offset": 0, "stride": 1, "input_dimension": 0 } ] + } + }, + { + "name": "box/inclusive-max", + "input": { "kind": "box", "inclusive_min": [2], "inclusive_max": [9] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [2], "input_exclusive_max": [10], + "input_labels": [""], + "output": [ { "offset": 0, "stride": 1, "input_dimension": 0 } ] + } + }, + { + "name": "box/implicit-and-infinite-bounds", + "input": { "kind": "box", "inclusive_min": [["-inf"], 0], "exclusive_max": [["+inf"], 4], "labels": ["t", ""] }, + "normalized": { + "input_rank": 2, + "input_inclusive_min": [["-inf"], 0], + "input_exclusive_max": [["+inf"], 4], + "input_labels": ["t", ""], + "output": [ + { "offset": 0, "stride": 1, "input_dimension": 0 }, + { "offset": 0, "stride": 1, "input_dimension": 1 } + ] + } + } +] diff --git a/packages/zarr-indexing/tests/conformance/errors.json b/packages/zarr-indexing/tests/conformance/errors.json new file mode 100644 index 0000000000..e5e08bab3c --- /dev/null +++ b/packages/zarr-indexing/tests/conformance/errors.json @@ -0,0 +1,23 @@ +[ + { "name": "error/step-zero", "input": { "kind": "slice", "start": [0], "stop": [4], "step": [0] }, "error": "step_zero" }, + { "name": "error/negative-step", "input": { "kind": "slice", "start": [9], "stop": [0], "step": [-2] }, "error": "negative_step_unsupported" }, + { "name": "error/multiple-upper-bounds", "input": { "kind": "box", "shape": [3], "exclusive_max": [3] }, "error": "multiple_upper_bounds" }, + { "name": "error/rank-mismatch", "input": { "kind": "slice", "start": [0, 0], "stop": [4] }, "error": "rank_mismatch" }, + { "name": "error/unknown-kind", "input": { "kind": "bogus" }, "error": "unknown_kind" }, + { "name": "error/transform-multiple-upper-bounds", "input": { "kind": "transform", "input_shape": [3], "input_exclusive_max": [3] }, "error": "multiple_upper_bounds" }, + { "name": "error/transform-rank-mismatch", "input": { "kind": "transform", "input_rank": 2, "input_inclusive_min": [0] }, "error": "rank_mismatch" }, + { "name": "error/missing-kind", "input": { "coords": [1, 2] }, "error": "invalid_json" }, + { "name": "error/point-missing-coords", "input": { "kind": "point" }, "error": "invalid_json" }, + { "name": "error/point-bool-coord", "input": { "kind": "point", "coords": [true] }, "error": "invalid_json" }, + { "name": "error/slice-missing-stop", "input": { "kind": "slice", "start": [0] }, "error": "invalid_json" }, + { "name": "error/box-non-list-bound", "input": { "kind": "box", "inclusive_min": 5 }, "error": "invalid_json" }, + { "name": "error/points-bool-coord", "input": { "kind": "points", "coords": [[true]] }, "error": "invalid_json" }, + { "name": "error/integer-out-of-i64-range", "input": { "kind": "point", "coords": [99999999999999999999] }, "error": "invalid_json" }, + { "name": "error/box-inverted-bounds", "input": { "kind": "box", "inclusive_min": [5], "exclusive_max": [3] }, "error": "bounds_out_of_order" }, + { "name": "error/box-negative-shape", "input": { "kind": "box", "shape": [-3] }, "error": "bounds_out_of_order" }, + { "name": "error/transform-inverted-bounds", "input": { "kind": "transform", "input_inclusive_min": [0], "input_exclusive_max": [-1] }, "error": "bounds_out_of_order" }, + { "name": "error/output-map-conflict", "input": { "kind": "transform", "output": [{ "input_dimension": 0, "index_array": [1, 2] }] }, "error": "output_map_conflict" }, + { "name": "error/box-unknown-field", "input": { "kind": "box", "shapee": [3] }, "error": "unknown_field" }, + { "name": "error/point-unknown-field", "input": { "kind": "point", "coords": [1], "extra": true }, "error": "unknown_field" }, + { "name": "error/output-map-unknown-field", "input": { "kind": "transform", "output": [{ "offset": 0, "bogus": 1 }] }, "error": "unknown_field" } +] diff --git a/packages/zarr-indexing/tests/conformance/point.json b/packages/zarr-indexing/tests/conformance/point.json new file mode 100644 index 0000000000..99a5ea16d8 --- /dev/null +++ b/packages/zarr-indexing/tests/conformance/point.json @@ -0,0 +1,30 @@ +[ + { + "name": "point/2d", + "input": { "kind": "point", "coords": [4, 7] }, + "normalized": { + "input_rank": 0, + "input_inclusive_min": [], + "input_exclusive_max": [], + "input_labels": [], + "output": [ { "offset": 4 }, { "offset": 7 } ] + } + }, + { + "name": "point/scalar-0d", + "input": { "kind": "point", "coords": [] }, + "normalized": { + "input_rank": 0, "input_inclusive_min": [], "input_exclusive_max": [], + "input_labels": [], "output": [] + } + }, + { + "name": "point/large-i64", + "input": { "kind": "point", "coords": [1152921504606846976] }, + "normalized": { + "input_rank": 0, "input_inclusive_min": [], "input_exclusive_max": [], + "input_labels": [], + "output": [ { "offset": 1152921504606846976 } ] + } + } +] diff --git a/packages/zarr-indexing/tests/conformance/points.json b/packages/zarr-indexing/tests/conformance/points.json new file mode 100644 index 0000000000..1ad92e12b1 --- /dev/null +++ b/packages/zarr-indexing/tests/conformance/points.json @@ -0,0 +1,34 @@ +[ + { + "name": "points/three-2d", + "input": { "kind": "points", "coords": [[1, 10], [2, 20], [3, 30]] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [0], "input_exclusive_max": [3], + "input_labels": [""], + "output": [ + { "offset": 0, "stride": 1, "index_array": [1, 2, 3], "index_array_bounds": ["-inf", "+inf"] }, + { "offset": 0, "stride": 1, "index_array": [10, 20, 30], "index_array_bounds": ["-inf", "+inf"] } + ] + } + }, + { + "name": "points/1d", + "input": { "kind": "points", "coords": [[5], [9], [2]] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [0], "input_exclusive_max": [3], + "input_labels": [""], + "output": [ + { "offset": 0, "stride": 1, "index_array": [5, 9, 2], "index_array_bounds": ["-inf", "+inf"] } + ] + } + }, + { + "name": "points/empty", + "input": { "kind": "points", "coords": [] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [0], "input_exclusive_max": [0], + "input_labels": [""], + "output": [] + } + } +] diff --git a/packages/zarr-indexing/tests/conformance/slice.json b/packages/zarr-indexing/tests/conformance/slice.json new file mode 100644 index 0000000000..2f1a0694ce --- /dev/null +++ b/packages/zarr-indexing/tests/conformance/slice.json @@ -0,0 +1,61 @@ +[ + { + "name": "slice/unit-step-preserves-frame", + "input": { "kind": "slice", "start": [5], "stop": [10] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [5], "input_exclusive_max": [10], + "input_labels": [""], + "output": [ { "offset": 0, "stride": 1, "input_dimension": 0 } ] + } + }, + { + "name": "slice/divisible-stride", + "input": { "kind": "slice", "start": [4], "stop": [10], "step": [2] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [2], "input_exclusive_max": [5], + "input_labels": [""], + "output": [ { "offset": 0, "stride": 2, "input_dimension": 0 } ] + } + }, + { + "name": "slice/nondivisible-stride-phase-offset", + "input": { "kind": "slice", "start": [5], "stop": [10], "step": [2] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [2], "input_exclusive_max": [5], + "input_labels": [""], + "output": [ { "offset": 1, "stride": 2, "input_dimension": 0 } ] + } + }, + { + "name": "slice/2d-mixed-step", + "input": { "kind": "slice", "start": [0, 5], "stop": [10, 10], "step": [2, 1] }, + "normalized": { + "input_rank": 2, + "input_inclusive_min": [0, 5], + "input_exclusive_max": [5, 10], + "input_labels": ["", ""], + "output": [ + { "offset": 0, "stride": 2, "input_dimension": 0 }, + { "offset": 0, "stride": 1, "input_dimension": 1 } + ] + } + }, + { + "name": "slice/negative-start-trunc-origin", + "input": { "kind": "slice", "start": [-9], "stop": [5], "step": [2] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [-4], "input_exclusive_max": [3], + "input_labels": [""], + "output": [ { "offset": -1, "stride": 2, "input_dimension": 0 } ] + } + }, + { + "name": "slice/negative-start-trunc-origin-step3", + "input": { "kind": "slice", "start": [-8], "stop": [6], "step": [3] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [-2], "input_exclusive_max": [3], + "input_labels": [""], + "output": [ { "offset": -2, "stride": 3, "input_dimension": 0 } ] + } + } +] diff --git a/packages/zarr-indexing/tests/conformance/transform.json b/packages/zarr-indexing/tests/conformance/transform.json new file mode 100644 index 0000000000..f26157aaab --- /dev/null +++ b/packages/zarr-indexing/tests/conformance/transform.json @@ -0,0 +1,57 @@ +[ + { + "name": "transform/omitted-output-identity", + "input": { "kind": "transform", "input_inclusive_min": [0, 0], "input_exclusive_max": [3, 4] }, + "normalized": { + "input_rank": 2, "input_inclusive_min": [0, 0], "input_exclusive_max": [3, 4], + "input_labels": ["", ""], + "output": [ + { "offset": 0, "stride": 1, "input_dimension": 0 }, + { "offset": 0, "stride": 1, "input_dimension": 1 } + ] + } + }, + { + "name": "transform/implicit-bounds-and-labels", + "input": { + "kind": "transform", + "input_inclusive_min": [["-inf"], 7], + "input_exclusive_max": [["+inf"], 11], + "input_labels": ["x", "y"] + }, + "normalized": { + "input_rank": 2, + "input_inclusive_min": [["-inf"], 7], + "input_exclusive_max": [["+inf"], 11], + "input_labels": ["x", "y"], + "output": [ + { "offset": 0, "stride": 1, "input_dimension": 0 }, + { "offset": 0, "stride": 1, "input_dimension": 1 } + ] + } + }, + { + "name": "transform/explicit-output-all-three-map-kinds", + "input": { + "kind": "transform", + "input_inclusive_min": [0], + "input_exclusive_max": [3], + "output": [ + { "offset": 7 }, + { "input_dimension": 0, "stride": 2 }, + { "index_array": [1, 2, 3] } + ] + }, + "normalized": { + "input_rank": 1, + "input_inclusive_min": [0], + "input_exclusive_max": [3], + "input_labels": [""], + "output": [ + { "offset": 7 }, + { "offset": 0, "stride": 2, "input_dimension": 0 }, + { "offset": 0, "stride": 1, "index_array": [1, 2, 3], "index_array_bounds": ["-inf", "+inf"] } + ] + } + } +] diff --git a/packages/zarr-indexing/tests/test_chunk_resolution.py b/packages/zarr-indexing/tests/test_chunk_resolution.py new file mode 100644 index 0000000000..0738384e2b --- /dev/null +++ b/packages/zarr-indexing/tests/test_chunk_resolution.py @@ -0,0 +1,521 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +from zarr.core.chunk_grids import ChunkGrid, FixedDimension, VaryingDimension + +from zarr_indexing import chunk_resolution +from zarr_indexing.chunk_resolution import iter_chunk_transforms, sub_transform_to_selections +from zarr_indexing.domain import IndexDomain +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr_indexing.transform import IndexTransform + +if TYPE_CHECKING: + import pytest + + +class TestChunkResolutionIdentity: + def test_single_chunk(self) -> None: + """Array fits in one chunk.""" + t = IndexTransform.from_shape((10,)) + grid = ChunkGrid(dimensions=(FixedDimension(size=10, extent=10),)) + results = list(iter_chunk_transforms(t, grid._dimensions)) + assert len(results) == 1 + coords, sub_t, _ = results[0] + assert coords == (0,) + assert sub_t.domain.shape == (10,) + + def test_multiple_chunks_1d(self) -> None: + """1D array spanning 3 chunks.""" + t = IndexTransform.from_shape((30,)) + grid = ChunkGrid(dimensions=(FixedDimension(size=10, extent=30),)) + results = list(iter_chunk_transforms(t, grid._dimensions)) + assert len(results) == 3 + coords_list = [r[0] for r in results] + assert (0,) in coords_list + assert (1,) in coords_list + assert (2,) in coords_list + + def test_multiple_chunks_2d(self) -> None: + """2D array spanning 2x3 chunks.""" + t = IndexTransform.from_shape((20, 30)) + grid = ChunkGrid( + dimensions=( + FixedDimension(size=10, extent=20), + FixedDimension(size=10, extent=30), + ) + ) + results = list(iter_chunk_transforms(t, grid._dimensions)) + assert len(results) == 6 + coords_list = [r[0] for r in results] + assert (0, 0) in coords_list + assert (1, 2) in coords_list + + +class TestChunkResolutionSliced: + def test_slice_within_chunk(self) -> None: + """Slice that falls within a single chunk.""" + # Chunk resolution consumes zero-origin transforms: the I/O layer + # normalizes preserved (user-facing) domains via translate_domain_to + # before resolving, so mirror that contract here. + t = IndexTransform.from_shape((100,))[5:8].translate_domain_to((0,)) + grid = ChunkGrid(dimensions=(FixedDimension(size=10, extent=100),)) + results = list(iter_chunk_transforms(t, grid._dimensions)) + assert len(results) == 1 + coords, sub_t, _ = results[0] + assert coords == (0,) + assert isinstance(sub_t.output[0], DimensionMap) + assert sub_t.output[0].offset == 5 + + def test_slice_across_chunks(self) -> None: + """Slice that spans two chunks.""" + t = IndexTransform.from_shape((100,))[8:15] + grid = ChunkGrid(dimensions=(FixedDimension(size=10, extent=100),)) + results = list(iter_chunk_transforms(t, grid._dimensions)) + assert len(results) == 2 + coords_list = [r[0] for r in results] + assert (0,) in coords_list + assert (1,) in coords_list + + +class TestChunkResolutionConstant: + def test_integer_index(self) -> None: + """Integer index produces constant map — single chunk per constant dim.""" + t = IndexTransform.from_shape((100, 100))[25, :] + grid = ChunkGrid( + dimensions=( + FixedDimension(size=10, extent=100), + FixedDimension(size=10, extent=100), + ) + ) + results = list(iter_chunk_transforms(t, grid._dimensions)) + assert len(results) == 10 + for coords, _, _ in results: + assert coords[0] == 2 + + +class TestChunkResolutionArray: + def test_array_index(self) -> None: + """Array index map — chunks determined by array values.""" + idx = np.array([5, 15, 25], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=idx),), + ) + grid = ChunkGrid(dimensions=(FixedDimension(size=10, extent=30),)) + results = list(iter_chunk_transforms(t, grid._dimensions)) + coords_list = [r[0] for r in results] + assert (0,) in coords_list + assert (1,) in coords_list + assert (2,) in coords_list + + +class TestChunkResolutionSorted1D: + def test_matches_general_resolution_for_randomized_sorted_selections( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Direct partitioning matches the original resolver across varied inputs.""" + rng = np.random.default_rng(0) + grids = ( + ChunkGrid(dimensions=(FixedDimension(size=7, extent=30),)), + ChunkGrid(dimensions=(VaryingDimension(edges=(3, 4, 8, 5, 10), extent=30),)), + ) + + for grid in grids: + for _ in range(50): + idx = np.sort(rng.integers(0, 30, size=int(rng.integers(1, 80)))).astype(np.intp) + transform = IndexTransform.from_shape((30,)).vindex[idx] + direct = list(iter_chunk_transforms(transform, grid._dimensions)) + + with monkeypatch.context() as context: + context.setattr( + chunk_resolution, + "_one_dimensional_correlated_array_map", + lambda _transform: None, + ) + general = list(iter_chunk_transforms(transform, grid._dimensions)) + + assert [result[0] for result in direct] == [result[0] for result in general] + for direct_result, general_result in zip(direct, general, strict=True): + _, direct_t, direct_out = direct_result + _, general_t, general_out = general_result + assert direct_t.domain == general_t.domain + + direct_chunk_sel, direct_out_sel, direct_drop = sub_transform_to_selections( + direct_t, direct_out + ) + general_chunk_sel, general_out_sel, general_drop = sub_transform_to_selections( + general_t, general_out + ) + assert direct_drop == general_drop + np.testing.assert_array_equal(direct_chunk_sel[0], general_chunk_sel[0]) + np.testing.assert_array_equal(direct_out_sel[0], general_out_sel[0]) + + def test_sorted_vindex_partitions_chunks_without_intersection( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Sorted vectorized coordinates are sliced directly per touched chunk.""" + idx = np.array([0, 3, 4, 4, 9, 11], dtype=np.intp) + t = IndexTransform.from_shape((12,)).vindex[idx] + grid = ChunkGrid(dimensions=(FixedDimension(size=4, extent=12),)) + + calls = _count_intersect_calls(monkeypatch) + results = list(iter_chunk_transforms(t, grid._dimensions)) + + assert [result[0] for result in results] == [(0,), (1,), (2,)] + assert calls["n"] == 0 + + expected_chunk_indices = ([0, 3], [0, 0], [1, 3]) + expected_out_indices = ([0, 1], [2, 3], [4, 5]) + for result, expected_chunk, expected_out in zip( + results, expected_chunk_indices, expected_out_indices, strict=True + ): + _, sub_t, out_indices = result + chunk_sel, out_sel, drop_axes = sub_transform_to_selections(sub_t, out_indices) + np.testing.assert_array_equal(chunk_sel[0], expected_chunk) + np.testing.assert_array_equal(out_sel[0], expected_out) + assert drop_axes == () + + def test_sorted_array_map_preserves_offset_and_stride(self) -> None: + """Storage partitioning retains the ArrayMap's offset and stride.""" + t = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=( + ArrayMap( + index_array=np.array([0, 1, 2], dtype=np.intp), + offset=1, + stride=3, + ), + ), + ) + grid = ChunkGrid(dimensions=(FixedDimension(size=4, extent=8),)) + + results = list(iter_chunk_transforms(t, grid._dimensions)) + + assert [result[0] for result in results] == [(0,), (1,)] + expected_chunk_indices = ([1], [0, 3]) + expected_out_indices = ([0], [1, 2]) + for result, expected_chunk, expected_out in zip( + results, expected_chunk_indices, expected_out_indices, strict=True + ): + _, sub_t, out_indices = result + chunk_sel, out_sel, _ = sub_transform_to_selections(sub_t, out_indices) + np.testing.assert_array_equal(chunk_sel[0], expected_chunk) + np.testing.assert_array_equal(out_sel[0], expected_out) + + def test_sorted_vindex_with_varying_chunks(self) -> None: + """Touched-boundary searches also support a non-uniform 1-D grid.""" + idx = np.array([0, 1, 2, 3, 5, 9], dtype=np.intp) + t = IndexTransform.from_shape((10,)).vindex[idx] + grid = ChunkGrid(dimensions=(VaryingDimension(edges=(2, 3, 5), extent=10),)) + + results = list(iter_chunk_transforms(t, grid._dimensions)) + + assert [result[0] for result in results] == [(0,), (1,), (2,)] + expected_chunk_indices = ([0, 1], [0, 1], [0, 4]) + for result, expected_chunk in zip(results, expected_chunk_indices, strict=True): + _, sub_t, out_indices = result + chunk_sel, _, _ = sub_transform_to_selections(sub_t, out_indices) + np.testing.assert_array_equal(chunk_sel[0], expected_chunk) + + def test_sorted_vindex_with_zero_sized_dimension_uses_general_resolution( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A zero-sized grid cannot be partitioned by touched boundaries.""" + t = IndexTransform.from_shape((10,)).vindex[np.array([1], dtype=np.intp)] + grid = ChunkGrid(dimensions=(FixedDimension(size=0, extent=10),)) + + calls = _count_intersect_calls(monkeypatch) + results = list(iter_chunk_transforms(t, grid._dimensions)) + + assert results == [] + assert calls["n"] == 1 + + def test_unsorted_vindex_uses_general_resolution(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Unsorted coordinates continue through the general intersection logic.""" + t = IndexTransform.from_shape((12,)).vindex[np.array([9, 0, 4], dtype=np.intp)] + grid = ChunkGrid(dimensions=(FixedDimension(size=4, extent=12),)) + + calls = _count_intersect_calls(monkeypatch) + results = list(iter_chunk_transforms(t, grid._dimensions)) + + assert [result[0] for result in results] == [(0,), (1,), (2,)] + assert calls["n"] == 3 + + def test_sorted_oindex_uses_general_resolution(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Orthogonal ArrayMaps retain their existing domain-aware resolution.""" + t = IndexTransform.from_shape((12,)).oindex[np.array([0, 4, 9], dtype=np.intp)] + grid = ChunkGrid(dimensions=(FixedDimension(size=4, extent=12),)) + + calls = _count_intersect_calls(monkeypatch) + results = list(iter_chunk_transforms(t, grid._dimensions)) + + assert [result[0] for result in results] == [(0,), (1,), (2,)] + assert calls["n"] == 3 + + +def _count_intersect_calls(monkeypatch: pytest.MonkeyPatch) -> dict[str, int]: + """Wrap `IndexTransform.intersect` with a call counter. + + Returns a mutable dict whose `"n"` entry is the number of times + `intersect` is invoked. Used to assert that candidate-chunk enumeration is + proportional to the *touched* chunks, not the dense bounding box between the + min and max touched chunk. + """ + calls = {"n": 0} + original = IndexTransform.intersect + + def counting(self: IndexTransform, output_domain: IndexDomain) -> object: + calls["n"] += 1 + return original(self, output_domain) + + monkeypatch.setattr(IndexTransform, "intersect", counting) + return calls + + +class TestChunkResolutionTouchedOnly: + """`iter_chunk_transforms` must enumerate only the chunks a fancy selection + actually touches — never the dense `range(min_chunk, max_chunk + 1)` bounding + box. These guard against a regression to bounding-box enumeration, whose cost + scales with grid size rather than with the number of selected coordinates. + """ + + def test_1d_sparse_vindex_enumerates_only_touched_chunks( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Two far-apart coordinates on a 1000-chunk grid touch exactly 2 chunks. + + A dense bounding-box enumeration would intersect ~1000 candidate chunks; + touched-only enumeration intersects exactly 2. + """ + # 4000 elements, chunk size 4 -> 1000 chunks. coords 1 and 3997 land in + # chunk 0 and chunk 999 respectively (998 empty chunks between them). + grid = ChunkGrid(dimensions=(FixedDimension(size=4, extent=4000),)) + t = IndexTransform.from_shape((4000,)).vindex[np.array([1, 3997], dtype=np.intp)] + + calls = _count_intersect_calls(monkeypatch) + results = list(iter_chunk_transforms(t, grid._dimensions)) + + coords = sorted(r[0] for r in results) + assert coords == [(0,), (999,)] + # Sorted 1-D coordinates are partitioned directly, without intersecting + # either the touched chunks or the 998 empty chunks between them. + assert calls["n"] == 0 + + def test_2d_orthogonal_enumerates_only_touched_chunks( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Orthogonal outer product of two 2-coordinate arrays touches 2x2 chunks. + + Per-dimension distinct touched chunks: {0, 999} on each axis. The outer + product is 2*2 = 4 candidate chunks (all survive), versus ~1e6 for a + dense 1000x1000 bounding box. + """ + grid = ChunkGrid( + dimensions=( + FixedDimension(size=4, extent=4000), + FixedDimension(size=4, extent=4000), + ) + ) + t = IndexTransform.from_shape((4000, 4000)).oindex[ + np.array([1, 3997], dtype=np.intp), np.array([2, 3998], dtype=np.intp) + ] + + calls = _count_intersect_calls(monkeypatch) + results = list(iter_chunk_transforms(t, grid._dimensions)) + + coords = sorted(r[0] for r in results) + assert coords == [(0, 0), (0, 999), (999, 0), (999, 999)] + assert calls["n"] == 4 + + def test_2d_correlated_vindex_enumerates_joint_touched_chunks( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Two correlated (vindex) coordinate arrays scatter to 2 diagonal chunks. + + The two points (1, 2) and (3997, 3998) touch chunks (0, 0) and + (999, 999). Correlated coordinate arrays are grouped *jointly*, so + enumeration intersects exactly the 2 touched chunks — never the 2x2 + cartesian product of per-dimension distinct chunks, and never the dense + 1e6 grid. + """ + grid = ChunkGrid( + dimensions=( + FixedDimension(size=4, extent=4000), + FixedDimension(size=4, extent=4000), + ) + ) + t = IndexTransform.from_shape((4000, 4000)).vindex[ + np.array([1, 3997], dtype=np.intp), np.array([2, 3998], dtype=np.intp) + ] + + calls = _count_intersect_calls(monkeypatch) + results = list(iter_chunk_transforms(t, grid._dimensions)) + + coords = sorted(r[0] for r in results) + assert coords == [(0, 0), (999, 999)] + assert calls["n"] == 2 + + def test_2d_correlated_vindex_diagonal_is_linear_in_points( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A diagonal of P correlated points touches P chunks with O(P) intersections. + + Enumerating the cartesian product of per-dimension distinct chunk sets + would cost P**2 intersections (2500 here) — quadratic in the number of + selected points for the scattered selections of zarr-python gh-4174. + Joint grouping keeps resolution work proportional to the touched chunks. + """ + p = 50 + grid = ChunkGrid( + dimensions=( + FixedDimension(size=4, extent=4000), + FixedDimension(size=4, extent=4000), + ) + ) + # point i lands in chunk (2i, 2i): all per-dimension chunks distinct + coords_1d = np.arange(p, dtype=np.intp) * 8 + t = IndexTransform.from_shape((4000, 4000)).vindex[coords_1d, coords_1d] + + calls = _count_intersect_calls(monkeypatch) + results = list(iter_chunk_transforms(t, grid._dimensions)) + + assert sorted(r[0] for r in results) == [(2 * i, 2 * i) for i in range(p)] + assert calls["n"] == p + + +class TestSubTransformToSelections: + def test_constant_map(self) -> None: + """ConstantMap produces int selection + drop axis.""" + t = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(ConstantMap(offset=5),), + ) + chunk_sel, out_sel, drop_axes = sub_transform_to_selections(t) + assert chunk_sel == (5,) + assert out_sel == () + assert drop_axes == () + + def test_dimension_map_stride_1(self) -> None: + """DimensionMap with stride=1 produces contiguous slice.""" + t = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(DimensionMap(input_dimension=0, offset=3, stride=1),), + ) + chunk_sel, out_sel, drop_axes = sub_transform_to_selections(t) + assert chunk_sel == (slice(3, 13, 1),) + assert out_sel == (slice(0, 10),) + assert drop_axes == () + + def test_dimension_map_strided(self) -> None: + """DimensionMap with stride>1 produces strided slice.""" + t = IndexTransform( + domain=IndexDomain.from_shape((5,)), + output=(DimensionMap(input_dimension=0, offset=2, stride=3),), + ) + chunk_sel, out_sel, drop_axes = sub_transform_to_selections(t) + assert chunk_sel == (slice(2, 17, 3),) + assert out_sel == (slice(0, 5),) + assert drop_axes == () + + def test_array_map(self) -> None: + """ArrayMap produces integer array selection.""" + arr = np.array([1, 5, 9], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=arr, offset=0, stride=1),), + ) + chunk_sel, out_sel, drop_axes = sub_transform_to_selections(t) + assert isinstance(chunk_sel[0], np.ndarray) + np.testing.assert_array_equal(chunk_sel[0], arr) + # Without chunk_mask, out_sel falls back to domain-based slices + assert out_sel == (slice(0, 3),) + assert drop_axes == () + + def test_array_map_with_offset_stride(self) -> None: + """ArrayMap with offset and stride computes storage coords.""" + arr = np.array([0, 1, 2], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=arr, offset=10, stride=5),), + ) + chunk_sel, _out_sel, drop_axes = sub_transform_to_selections(t) + assert isinstance(chunk_sel[0], np.ndarray) + np.testing.assert_array_equal(chunk_sel[0], np.array([10, 15, 20])) + assert drop_axes == () + + def test_mixed_maps_2d(self) -> None: + """Mix of ConstantMap and DimensionMap.""" + t = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=( + ConstantMap(offset=5), + DimensionMap(input_dimension=0, offset=0, stride=1), + ), + ) + chunk_sel, _out_sel, drop_axes = sub_transform_to_selections(t) + assert chunk_sel[0] == 5 + assert chunk_sel[1] == slice(0, 10, 1) + # drop_axes is empty — integer in chunk_sel naturally drops the dim via numpy + assert drop_axes == () + + +class TestChunkResolutionArrayMapFlavours: + """Chunk resolution must yield outer-product (np.ix_) selectors for + orthogonal ArrayMaps and shared flat-scatter selectors for correlated ones, + and must return early for empty fancy selections.""" + + def test_empty_array_selection_yields_nothing(self) -> None: + """An empty ArrayMap selection produces no chunk transforms (no crash).""" + t = IndexTransform( + domain=IndexDomain.from_shape((0,)), + output=(ArrayMap(index_array=np.array([], dtype=np.intp)),), + ) + grid = ChunkGrid(dimensions=(FixedDimension(size=3, extent=10),)) + assert list(iter_chunk_transforms(t, grid._dimensions)) == [] + + def test_orthogonal_outer_product_selectors(self) -> None: + """Two independent arrays produce np.ix_-style (mesh) chunk/out selectors.""" + t = IndexTransform.from_shape((10, 10)).oindex[np.array([1, 3]), np.array([2, 4, 6])] + grid = ChunkGrid( + dimensions=(FixedDimension(size=10, extent=10), FixedDimension(size=10, extent=10)) + ) + results = list(iter_chunk_transforms(t, grid._dimensions)) + assert len(results) == 1 + _coords, sub_t, out_indices = results[0] + chunk_sel, out_sel, drop_axes = sub_transform_to_selections(sub_t, out_indices) + # np.ix_ produces one 2-D open-mesh selector per axis, for both sides. + assert len(chunk_sel) == 2 + assert len(out_sel) == 2 + assert isinstance(chunk_sel[0], np.ndarray) + assert isinstance(chunk_sel[1], np.ndarray) + assert chunk_sel[0].shape == (2, 1) + assert chunk_sel[1].shape == (1, 3) + assert drop_axes == () + + def test_correlated_scatter_with_residual_slice(self) -> None: + """Correlated arrays + a residual slice dim scatter through a single flat + index whose shape matches the (points, slice) block read from the chunk.""" + t = IndexTransform.from_shape((4, 3, 5)).vindex[np.array([1, 3]), np.array([2, 0])] + grid = ChunkGrid( + dimensions=( + FixedDimension(size=4, extent=4), + FixedDimension(size=3, extent=3), + FixedDimension(size=5, extent=5), + ) + ) + # One chunk holds everything: both points survive, slice dim spans [0,5). + results = list(iter_chunk_transforms(t, grid._dimensions)) + assert len(results) == 1 + _coords, sub_t, out_indices = results[0] + chunk_sel, out_sel, _drop = sub_transform_to_selections(sub_t, out_indices) + # Chunk side: flat coordinate arrays for the two correlated dims plus a + # slice for the residual dim. + assert len(chunk_sel) == 3 + np.testing.assert_array_equal(np.asarray(chunk_sel[0]), [1, 3]) + np.testing.assert_array_equal(np.asarray(chunk_sel[1]), [2, 0]) + assert chunk_sel[2] == slice(0, 5, 1) + # Output side: a single flat scatter index of shape (points, slice) = (2, 5). + assert len(out_sel) == 1 + assert np.asarray(out_sel[0]).shape == (2, 5) diff --git a/packages/zarr-indexing/tests/test_composition.py b/packages/zarr-indexing/tests/test_composition.py new file mode 100644 index 0000000000..dd92f59b80 --- /dev/null +++ b/packages/zarr-indexing/tests/test_composition.py @@ -0,0 +1,166 @@ +from __future__ import annotations + +import numpy as np +import pytest + +from zarr_indexing.composition import compose +from zarr_indexing.domain import IndexDomain +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr_indexing.transform import IndexTransform + + +class TestComposeConstantInner: + """Inner = constant. Result is always constant.""" + + def test_constant_inner_any_outer(self) -> None: + outer = IndexTransform.from_shape((5,)) + inner = IndexTransform( + domain=IndexDomain.from_shape((5,)), + output=(ConstantMap(offset=42),), + ) + result = compose(outer, inner) + assert isinstance(result.output[0], ConstantMap) + assert result.output[0].offset == 42 + + +class TestComposeDimensionInner: + """Inner = DimensionMap.""" + + def test_dimension_inner_constant_outer(self) -> None: + outer = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(ConstantMap(offset=5),), + ) + inner = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(DimensionMap(input_dimension=0, offset=10, stride=3),), + ) + result = compose(outer, inner) + assert isinstance(result.output[0], ConstantMap) + assert result.output[0].offset == 25 + + def test_dimension_inner_dimension_outer(self) -> None: + outer = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(DimensionMap(input_dimension=0, offset=5, stride=2),), + ) + inner = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(DimensionMap(input_dimension=0, offset=10, stride=3),), + ) + result = compose(outer, inner) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == 25 + assert result.output[0].stride == 6 + assert result.output[0].input_dimension == 0 + + def test_dimension_inner_array_outer(self) -> None: + arr = np.array([0, 2, 4], dtype=np.intp) + outer = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=arr, offset=5, stride=2),), + ) + inner = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(DimensionMap(input_dimension=0, offset=10, stride=3),), + ) + result = compose(outer, inner) + assert isinstance(result.output[0], ArrayMap) + assert result.output[0].offset == 25 + assert result.output[0].stride == 6 + np.testing.assert_array_equal(result.output[0].index_array, arr) + + +class TestComposeArrayInner: + """Inner = ArrayMap.""" + + def test_array_inner_constant_outer(self) -> None: + inner_arr = np.array([10, 20, 30], dtype=np.intp) + outer = IndexTransform( + domain=IndexDomain.from_shape((5,)), + output=(ConstantMap(offset=1),), + ) + inner = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=inner_arr, offset=0, stride=1),), + ) + result = compose(outer, inner) + assert isinstance(result.output[0], ConstantMap) + assert result.output[0].offset == 20 + + def test_array_inner_array_outer(self) -> None: + outer_arr = np.array([0, 2, 1], dtype=np.intp) + inner_arr = np.array([10, 20, 30], dtype=np.intp) + outer = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=outer_arr, offset=0, stride=1),), + ) + inner = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=inner_arr, offset=0, stride=1),), + ) + result = compose(outer, inner) + assert isinstance(result.output[0], ArrayMap) + expected = np.array([10, 30, 20], dtype=np.intp) + np.testing.assert_array_equal(result.output[0].index_array, expected) + + +class TestComposeMultiDim: + def test_2d_identity_compose(self) -> None: + a = IndexTransform.from_shape((10, 20)) + b = IndexTransform.from_shape((10, 20)) + result = compose(a, b) + assert result.domain.shape == (10, 20) + for i in range(2): + m = result.output[i] + assert isinstance(m, DimensionMap) + assert m.input_dimension == i + assert m.offset == 0 + assert m.stride == 1 + + def test_mixed_map_types(self) -> None: + outer = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=( + ConstantMap(offset=5), + DimensionMap(input_dimension=0, offset=0, stride=1), + ), + ) + inner = IndexTransform( + domain=IndexDomain.from_shape((10, 10)), + output=( + DimensionMap(input_dimension=0, offset=2, stride=3), + DimensionMap(input_dimension=1, offset=0, stride=1), + ), + ) + result = compose(outer, inner) + assert isinstance(result.output[0], ConstantMap) + assert result.output[0].offset == 17 + assert isinstance(result.output[1], DimensionMap) + assert result.output[1].input_dimension == 0 + assert result.output[1].offset == 0 + assert result.output[1].stride == 1 + + def test_rank_mismatch_raises(self) -> None: + outer = IndexTransform.from_shape((10,)) + inner = IndexTransform.from_shape((10, 20)) + with pytest.raises(ValueError, match="rank"): + compose(outer, inner) + + +class TestComposeChain: + def test_three_transforms(self) -> None: + a = IndexTransform.from_shape((100,)) + b = IndexTransform( + domain=IndexDomain.from_shape((100,)), + output=(DimensionMap(input_dimension=0, offset=10, stride=1),), + ) + c = IndexTransform( + domain=IndexDomain.from_shape((100,)), + output=(DimensionMap(input_dimension=0, offset=5, stride=2),), + ) + bc = compose(b, c) + abc = compose(a, bc) + assert isinstance(abc.output[0], DimensionMap) + assert abc.output[0].offset == 25 + assert abc.output[0].stride == 2 diff --git a/packages/zarr-indexing/tests/test_conformance.py b/packages/zarr-indexing/tests/test_conformance.py new file mode 100644 index 0000000000..207a9d8236 --- /dev/null +++ b/packages/zarr-indexing/tests/test_conformance.py @@ -0,0 +1,55 @@ +"""ndsel conformance corpus harness. + +Runs the vendored, language-agnostic ndsel fixtures (see +`tests/conformance/PROVENANCE.md`) against this package's message layer +(`zarr_indexing.messages`). An implementation is conformant iff: + +- for every *success* fixture, `normalize_ndsel(input)` equals the fixture's + `normalized` value by structural JSON equality; +- for every *error* fixture, `normalize_ndsel(input)` is rejected with an + `NdselError` carrying the fixture's `error` reason code. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import pytest + +from zarr_indexing.messages import NdselError, normalize_ndsel + +_CONFORMANCE_DIR = Path(__file__).parent / "conformance" + + +def _load_cases() -> list[tuple[str, dict[str, Any]]]: + cases: list[tuple[str, dict[str, Any]]] = [] + for path in sorted(_CONFORMANCE_DIR.glob("*.json")): + data = json.loads(path.read_text()) + cases.extend((f"{path.stem}::{case['name']}", case) for case in data) + return cases + + +_CASES = _load_cases() +_SUCCESS = [(name, c) for name, c in _CASES if "normalized" in c] +_ERROR = [(name, c) for name, c in _CASES if "error" in c] + + +def test_corpus_is_present() -> None: + # Guard against an empty/missing vendored corpus silently passing. + assert len(_SUCCESS) > 0 + assert len(_ERROR) > 0 + + +@pytest.mark.parametrize(("name", "case"), _SUCCESS, ids=[name for name, _ in _SUCCESS]) +def test_success_fixture(name: str, case: dict[str, Any]) -> None: + result = normalize_ndsel(case["input"]) + assert result == case["normalized"] + + +@pytest.mark.parametrize(("name", "case"), _ERROR, ids=[name for name, _ in _ERROR]) +def test_error_fixture(name: str, case: dict[str, Any]) -> None: + with pytest.raises(NdselError) as excinfo: + normalize_ndsel(case["input"]) + assert excinfo.value.reason == case["error"] diff --git a/packages/zarr-indexing/tests/test_domain.py b/packages/zarr-indexing/tests/test_domain.py new file mode 100644 index 0000000000..9664a0b08a --- /dev/null +++ b/packages/zarr-indexing/tests/test_domain.py @@ -0,0 +1,202 @@ +from __future__ import annotations + +import pytest + +from zarr_indexing.domain import IndexDomain + + +class TestIndexDomainConstruction: + def test_from_shape(self) -> None: + d = IndexDomain.from_shape((10, 20)) + assert d.inclusive_min == (0, 0) + assert d.exclusive_max == (10, 20) + assert d.ndim == 2 + assert d.origin == (0, 0) + assert d.shape == (10, 20) + + def test_from_shape_0d(self) -> None: + d = IndexDomain.from_shape(()) + assert d.ndim == 0 + assert d.shape == () + + def test_non_zero_origin(self) -> None: + d = IndexDomain(inclusive_min=(5, 10), exclusive_max=(15, 30)) + assert d.origin == (5, 10) + assert d.shape == (10, 20) + assert d.ndim == 2 + + def test_validation_mismatched_lengths(self) -> None: + with pytest.raises(ValueError, match="same length"): + IndexDomain(inclusive_min=(0,), exclusive_max=(10, 20)) + + def test_validation_min_greater_than_max(self) -> None: + with pytest.raises(ValueError, match="inclusive_min must be <="): + IndexDomain(inclusive_min=(10,), exclusive_max=(5,)) + + def test_empty_domain(self) -> None: + d = IndexDomain(inclusive_min=(5,), exclusive_max=(5,)) + assert d.shape == (0,) + + def test_labels(self) -> None: + d = IndexDomain(inclusive_min=(0, 0), exclusive_max=(10, 20), labels=("x", "y")) + assert d.labels == ("x", "y") + + def test_labels_none(self) -> None: + d = IndexDomain.from_shape((10,)) + assert d.labels is None + + +class TestIndexDomainContains: + def test_contains_inside(self) -> None: + d = IndexDomain.from_shape((10, 20)) + assert d.contains((0, 0)) is True + assert d.contains((9, 19)) is True + assert d.contains((5, 10)) is True + + def test_contains_outside(self) -> None: + d = IndexDomain.from_shape((10, 20)) + assert d.contains((10, 0)) is False + assert d.contains((-1, 0)) is False + assert d.contains((0, 20)) is False + + def test_contains_non_zero_origin(self) -> None: + d = IndexDomain(inclusive_min=(5,), exclusive_max=(10,)) + assert d.contains((5,)) is True + assert d.contains((9,)) is True + assert d.contains((4,)) is False + assert d.contains((10,)) is False + + def test_contains_wrong_ndim(self) -> None: + d = IndexDomain.from_shape((10, 20)) + assert d.contains((5,)) is False + + def test_contains_domain_inside(self) -> None: + outer = IndexDomain.from_shape((10, 20)) + inner = IndexDomain(inclusive_min=(2, 3), exclusive_max=(8, 15)) + assert outer.contains_domain(inner) is True + + def test_contains_domain_outside(self) -> None: + outer = IndexDomain.from_shape((10, 20)) + inner = IndexDomain(inclusive_min=(2, 3), exclusive_max=(11, 15)) + assert outer.contains_domain(inner) is False + + def test_contains_domain_wrong_ndim(self) -> None: + outer = IndexDomain.from_shape((10, 20)) + inner = IndexDomain.from_shape((5,)) + assert outer.contains_domain(inner) is False + + +class TestIndexDomainIntersect: + def test_overlapping(self) -> None: + a = IndexDomain(inclusive_min=(0, 0), exclusive_max=(10, 10)) + b = IndexDomain(inclusive_min=(5, 5), exclusive_max=(15, 15)) + result = a.intersect(b) + assert result is not None + assert result.inclusive_min == (5, 5) + assert result.exclusive_max == (10, 10) + + def test_disjoint(self) -> None: + a = IndexDomain(inclusive_min=(0,), exclusive_max=(5,)) + b = IndexDomain(inclusive_min=(10,), exclusive_max=(15,)) + assert a.intersect(b) is None + + def test_touching_boundary(self) -> None: + a = IndexDomain(inclusive_min=(0,), exclusive_max=(5,)) + b = IndexDomain(inclusive_min=(5,), exclusive_max=(10,)) + assert a.intersect(b) is None + + def test_contained(self) -> None: + a = IndexDomain.from_shape((20,)) + b = IndexDomain(inclusive_min=(5,), exclusive_max=(10,)) + result = a.intersect(b) + assert result is not None + assert result.inclusive_min == (5,) + assert result.exclusive_max == (10,) + + def test_wrong_ndim(self) -> None: + a = IndexDomain.from_shape((10,)) + b = IndexDomain.from_shape((10, 20)) + with pytest.raises(ValueError, match="different ranks"): + a.intersect(b) + + +class TestIndexDomainTranslate: + def test_translate_positive(self) -> None: + d = IndexDomain.from_shape((10, 20)) + result = d.translate((5, 10)) + assert result.inclusive_min == (5, 10) + assert result.exclusive_max == (15, 30) + + def test_translate_negative(self) -> None: + d = IndexDomain(inclusive_min=(10, 20), exclusive_max=(30, 40)) + result = d.translate((-10, -20)) + assert result.inclusive_min == (0, 0) + assert result.exclusive_max == (20, 20) + + def test_translate_wrong_length(self) -> None: + d = IndexDomain.from_shape((10,)) + with pytest.raises(ValueError, match="same length"): + d.translate((1, 2)) + + +class TestIndexDomainNarrow: + def test_narrow_slice(self) -> None: + d = IndexDomain.from_shape((10, 20)) + result = d.narrow((slice(2, 8), slice(5, 15))) + assert result.inclusive_min == (2, 5) + assert result.exclusive_max == (8, 15) + + def test_narrow_int(self) -> None: + d = IndexDomain.from_shape((10, 20)) + result = d.narrow((3, slice(None))) + assert result.inclusive_min == (3, 0) + assert result.exclusive_max == (4, 20) + + def test_narrow_ellipsis(self) -> None: + d = IndexDomain.from_shape((10, 20, 30)) + result = d.narrow((slice(1, 5), ...)) + assert result.inclusive_min == (1, 0, 0) + assert result.exclusive_max == (5, 20, 30) + + def test_narrow_slice_none(self) -> None: + d = IndexDomain.from_shape((10,)) + result = d.narrow((slice(None),)) + assert result == d + + def test_narrow_non_zero_origin(self) -> None: + d = IndexDomain(inclusive_min=(10,), exclusive_max=(20,)) + result = d.narrow((slice(12, 18),)) + assert result.inclusive_min == (12,) + assert result.exclusive_max == (18,) + + def test_narrow_int_out_of_bounds(self) -> None: + d = IndexDomain.from_shape((10,)) + with pytest.raises(IndexError, match="out of bounds"): + d.narrow((10,)) + + def test_narrow_int_below_origin(self) -> None: + d = IndexDomain(inclusive_min=(5,), exclusive_max=(10,)) + with pytest.raises(IndexError, match="out of bounds"): + d.narrow((4,)) + + def test_narrow_clamps_to_domain(self) -> None: + d = IndexDomain.from_shape((10,)) + result = d.narrow((slice(-5, 100),)) + assert result.inclusive_min == (0,) + assert result.exclusive_max == (10,) + + def test_narrow_bare_slice(self) -> None: + d = IndexDomain.from_shape((10,)) + result = d.narrow(slice(2, 8)) + assert result.inclusive_min == (2,) + assert result.exclusive_max == (8,) + + def test_narrow_too_many_indices(self) -> None: + d = IndexDomain.from_shape((10,)) + with pytest.raises(IndexError, match="too many indices"): + d.narrow((1, 2)) + + def test_narrow_step_not_one(self) -> None: + d = IndexDomain.from_shape((10,)) + with pytest.raises(IndexError, match="step=1"): + d.narrow((slice(0, 10, 2),)) diff --git a/packages/zarr-indexing/tests/test_json.py b/packages/zarr-indexing/tests/test_json.py new file mode 100644 index 0000000000..42b59b2c30 --- /dev/null +++ b/packages/zarr-indexing/tests/test_json.py @@ -0,0 +1,336 @@ +from __future__ import annotations + +import numpy as np +import pytest + +from zarr_indexing.domain import IndexDomain +from zarr_indexing.json import ( + IndexTransformJSON, + index_domain_from_json, + index_domain_to_json, + index_transform_from_json, + index_transform_to_json, + output_index_map_from_json, + output_index_map_to_json, +) +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr_indexing.transform import IndexTransform + + +def _maps_equal(a: object, b: object) -> bool: + if type(a) is not type(b): + return False + if isinstance(a, ConstantMap): + assert isinstance(b, ConstantMap) + return a.offset == b.offset + if isinstance(a, DimensionMap): + assert isinstance(b, DimensionMap) + return (a.input_dimension, a.offset, a.stride) == (b.input_dimension, b.offset, b.stride) + assert isinstance(a, ArrayMap) + assert isinstance(b, ArrayMap) + return ( + a.offset == b.offset + and a.stride == b.stride + and a.input_dimension == b.input_dimension + and np.array_equal(a.index_array, b.index_array) + ) + + +def _transforms_equal(a: IndexTransform, b: IndexTransform) -> bool: + """Structural equality that compares `ArrayMap` index arrays element-wise + (`IndexTransform`'s dataclass `__eq__` cannot, as numpy `==` is ambiguous).""" + return ( + a.domain == b.domain + and len(a.output) == len(b.output) + and all(_maps_equal(x, y) for x, y in zip(a.output, b.output, strict=True)) + ) + + +class TestIndexDomainJSON: + def test_roundtrip(self) -> None: + domain = IndexDomain(inclusive_min=(2, 5), exclusive_max=(10, 20)) + json = index_domain_to_json(domain) + assert json == { + "input_inclusive_min": [2, 5], + "input_exclusive_max": [10, 20], + "input_labels": ["", ""], + } + restored = index_domain_from_json(json) + assert restored == domain + + def test_with_labels(self) -> None: + domain = IndexDomain(inclusive_min=(0, 0), exclusive_max=(10, 20), labels=("x", "y")) + json = index_domain_to_json(domain) + assert json["input_labels"] == ["x", "y"] + restored = index_domain_from_json(json) + assert restored.labels == ("x", "y") + + def test_without_labels_emits_empty_and_round_trips_to_none(self) -> None: + domain = IndexDomain.from_shape((5,)) + json = index_domain_to_json(domain) + # Canonical form always writes labels; an unlabeled domain gets [""]*rank. + assert json["input_labels"] == [""] + restored = index_domain_from_json(json) + assert restored.labels is None + + def test_zero_origin(self) -> None: + domain = IndexDomain.from_shape((10, 20, 30)) + json = index_domain_to_json(domain) + assert json == { + "input_inclusive_min": [0, 0, 0], + "input_exclusive_max": [10, 20, 30], + "input_labels": ["", "", ""], + } + assert index_domain_from_json(json) == domain + + +class TestOutputIndexMapJSON: + def test_constant(self) -> None: + m = ConstantMap(offset=42) + json = output_index_map_to_json(m) + assert json == {"offset": 42} + restored = output_index_map_from_json(json) + assert isinstance(restored, ConstantMap) + assert restored.offset == 42 + + def test_constant_zero(self) -> None: + m = ConstantMap(offset=0) + json = output_index_map_to_json(m) + assert json == {"offset": 0} + restored = output_index_map_from_json(json) + assert isinstance(restored, ConstantMap) + assert restored.offset == 0 + + def test_dimension(self) -> None: + m = DimensionMap(input_dimension=1, offset=10, stride=3) + json = output_index_map_to_json(m) + assert json == {"offset": 10, "stride": 3, "input_dimension": 1} + restored = output_index_map_from_json(json) + assert isinstance(restored, DimensionMap) + assert restored.input_dimension == 1 + assert restored.offset == 10 + assert restored.stride == 3 + + def test_dimension_stride_1_written(self) -> None: + """Canonical form writes stride even at its default of 1.""" + m = DimensionMap(input_dimension=0) + json = output_index_map_to_json(m) + assert json == {"offset": 0, "stride": 1, "input_dimension": 0} + restored = output_index_map_from_json(json) + assert isinstance(restored, DimensionMap) + assert restored.stride == 1 + + def test_array(self) -> None: + arr = np.array([1, 5, 9], dtype=np.intp) + m = ArrayMap(index_array=arr, offset=2, stride=3) + json = output_index_map_to_json(m) + # Canonical: stride/offset present, index_array_bounds present, and + # no input_dimension (ndsel/TensorStore reject it beside index_array). + assert json == { + "offset": 2, + "stride": 3, + "index_array": [1, 5, 9], + "index_array_bounds": ["-inf", "+inf"], + } + restored = output_index_map_from_json(json) + assert isinstance(restored, ArrayMap) + np.testing.assert_array_equal(restored.index_array, arr) + assert restored.offset == 2 + assert restored.stride == 3 + + def test_array_stride_1_written(self) -> None: + arr = np.array([0, 1, 2], dtype=np.intp) + m = ArrayMap(index_array=arr) + json = output_index_map_to_json(m) + assert json["stride"] == 1 + restored = output_index_map_from_json(json) + assert isinstance(restored, ArrayMap) + assert restored.stride == 1 + + def test_array_2d(self) -> None: + arr = np.array([[1, 2], [3, 4]], dtype=np.intp) + m = ArrayMap(index_array=arr) + json = output_index_map_to_json(m) + assert json["index_array"] == [[1, 2], [3, 4]] + restored = output_index_map_from_json(json) + assert isinstance(restored, ArrayMap) + np.testing.assert_array_equal(restored.index_array, arr) + + def test_degenerate_singleton_array_collapses_to_constant(self) -> None: + """An all-singleton index_array selects one coordinate -> constant map.""" + m = ArrayMap(index_array=np.array([[4]], dtype=np.intp), offset=1, stride=2) + json = output_index_map_to_json(m) + assert json == {"offset": 1 + 2 * 4} + restored = output_index_map_from_json(json) + assert isinstance(restored, ConstantMap) + assert restored.offset == 9 + + +class TestIndexTransformJSON: + def test_identity(self) -> None: + t = IndexTransform.from_shape((10, 20)) + json = index_transform_to_json(t) + assert json == { + "input_rank": 2, + "input_inclusive_min": [0, 0], + "input_exclusive_max": [10, 20], + "input_labels": ["", ""], + "output": [ + {"offset": 0, "stride": 1, "input_dimension": 0}, + {"offset": 0, "stride": 1, "input_dimension": 1}, + ], + } + restored = index_transform_from_json(json) + assert restored.domain == t.domain + assert len(restored.output) == 2 + for orig, rest in zip(t.output, restored.output, strict=True): + assert type(orig) is type(rest) + + def test_sliced(self) -> None: + t = IndexTransform.from_shape((100,))[10:50:2] + json = index_transform_to_json(t) + restored = index_transform_from_json(json) + assert restored.domain.shape == t.domain.shape + assert isinstance(restored.output[0], DimensionMap) + orig = t.output[0] + assert isinstance(orig, DimensionMap) + assert restored.output[0].offset == orig.offset + assert restored.output[0].stride == orig.stride + + def test_with_constant(self) -> None: + t = IndexTransform.from_shape((10, 20))[3] + json = index_transform_to_json(t) + restored = index_transform_from_json(json) + assert isinstance(restored.output[0], ConstantMap) + assert restored.output[0].offset == 3 + assert isinstance(restored.output[1], DimensionMap) + + def test_with_array(self) -> None: + idx = np.array([1, 5, 9], dtype=np.intp) + t = IndexTransform.from_shape((10, 20)).oindex[idx, :] + json = index_transform_to_json(t) + # The oindex array must not carry input_dimension on the wire. + assert "input_dimension" not in json["output"][0] + restored = index_transform_from_json(json) + assert isinstance(restored.output[0], ArrayMap) + # Orthogonal arrays are normalized to full input rank with a singleton + # axis on the dimension they do not vary over. + assert restored.output[0].index_array.shape == (3, 1) + np.testing.assert_array_equal(restored.output[0].index_array, idx.reshape(3, 1)) + # input_dimension is reconstructed from the sole non-singleton axis. + assert restored.output[0].input_dimension == 0 + assert isinstance(restored.output[1], DimensionMap) + + def test_roundtrip_preserves_singleton_axes(self) -> None: + """Full-rank orthogonal arrays keep their singleton axes across JSON.""" + t = IndexTransform.from_shape((10, 20)).oindex[np.array([1, 3]), np.array([2, 4, 6])] + restored = index_transform_from_json(index_transform_to_json(t)) + orig0, orig1 = t.output[0], t.output[1] + rest0, rest1 = restored.output[0], restored.output[1] + assert isinstance(orig0, ArrayMap) + assert isinstance(orig1, ArrayMap) + assert isinstance(rest0, ArrayMap) + assert isinstance(rest1, ArrayMap) + assert rest0.index_array.shape == (2, 1) + assert rest1.index_array.shape == (1, 3) + np.testing.assert_array_equal(rest0.index_array, orig0.index_array) + np.testing.assert_array_equal(rest1.index_array, orig1.index_array) + # Distinct, exclusively-owned axes -> reconstructed as orthogonal. + assert rest0.input_dimension == 0 + assert rest1.input_dimension == 1 + + def test_with_labels(self) -> None: + domain = IndexDomain(inclusive_min=(0, 0), exclusive_max=(10, 20), labels=("x", "y")) + t = IndexTransform.identity(domain) + json = index_transform_to_json(t) + assert json["input_labels"] == ["x", "y"] + restored = index_transform_from_json(json) + assert restored.domain.labels == ("x", "y") + + def test_tensorstore_compatible_format(self) -> None: + """A canonical body loads and round-trips through the engine layer.""" + json: IndexTransformJSON = { + "input_rank": 3, + "input_inclusive_min": [0, 0, 0], + "input_exclusive_max": [100, 200, 3], + "input_labels": ["x", "y", "channel"], + "output": [ + {"offset": 5}, + {"offset": 10, "stride": 2, "input_dimension": 1}, + {"offset": 0, "stride": 1, "index_array": [1, 2, 0]}, + ], + } + t = index_transform_from_json(json) + assert t.domain.shape == (100, 200, 3) + assert t.domain.labels == ("x", "y", "channel") + assert isinstance(t.output[0], ConstantMap) + assert t.output[0].offset == 5 + assert isinstance(t.output[1], DimensionMap) + assert t.output[1].offset == 10 + assert t.output[1].stride == 2 + assert t.output[1].input_dimension == 1 + assert isinstance(t.output[2], ArrayMap) + np.testing.assert_array_equal(t.output[2].index_array, [1, 2, 0]) + + # Roundtrip + json_rt = index_transform_to_json(t) + t_rt = index_transform_from_json(json_rt) + assert t_rt.domain == t.domain + + +class TestCanonicalRoundTrips: + """Round-trip `transform == from(to(transform))`, up to the documented + degenerate-collapse (all-singleton ArrayMap -> ConstantMap).""" + + def test_oindex_multi_axis(self) -> None: + t = IndexTransform.from_shape((10, 20, 30)).oindex[np.array([1, 3]), :, np.array([2, 4, 6])] + rt = index_transform_from_json(index_transform_to_json(t)) + assert _transforms_equal(rt, t) + + def test_oindex_with_slice(self) -> None: + t = IndexTransform.from_shape((10, 20))[2:8].oindex[np.array([3, 5, 7]), :] + rt = index_transform_from_json(index_transform_to_json(t)) + assert _transforms_equal(rt, t) + + def test_vindex(self) -> None: + t = IndexTransform.from_shape((10, 20)).vindex[np.array([1, 3, 5]), np.array([2, 4, 6])] + rt = index_transform_from_json(index_transform_to_json(t)) + assert _transforms_equal(rt, t) + + def test_vindex_with_residual_slice(self) -> None: + t = IndexTransform.from_shape((10, 20, 30)).vindex[np.array([1, 3]), np.array([2, 4]), :] + rt = index_transform_from_json(index_transform_to_json(t)) + assert _transforms_equal(rt, t) + + def test_length1_degenerate_oindex_collapses(self) -> None: + """A length-1 oindex array becomes an all-singleton ArrayMap; the JSON + round-trip collapses it to a ConstantMap (behaviorally identical).""" + t = IndexTransform.from_shape((10, 20)).oindex[np.array([7]), :] + m = t.output[0] + assert isinstance(m, ArrayMap) + assert m.index_array.size == 1 + + rt = index_transform_from_json(index_transform_to_json(t)) + # The degenerate array collapsed to a constant selecting the same cell. + rm = rt.output[0] + assert isinstance(rm, ConstantMap) + assert rm.offset == 7 + # The size-1 input dimension survives, unconsumed, in the domain. + assert rt.domain == t.domain + + def test_slices_and_constants(self) -> None: + t = IndexTransform.from_shape((10, 20, 30))[2:8:2, 5, :] + rt = index_transform_from_json(index_transform_to_json(t)) + assert _transforms_equal(rt, t) + + +def test_infinite_bound_rejected_on_lowering() -> None: + body: IndexTransformJSON = { + "input_rank": 1, + "input_inclusive_min": [0], + "input_exclusive_max": [["+inf"]], + "input_labels": [""], + "output": [{"offset": 0, "stride": 1, "input_dimension": 0}], + } + with pytest.raises(ValueError, match="infinite"): + index_transform_from_json(body) diff --git a/packages/zarr-indexing/tests/test_messages.py b/packages/zarr-indexing/tests/test_messages.py new file mode 100644 index 0000000000..14bed66448 --- /dev/null +++ b/packages/zarr-indexing/tests/test_messages.py @@ -0,0 +1,91 @@ +"""Message-layer tests beyond the vendored conformance corpus. + +The corpus (see `test_conformance.py`) covers the desugaring matrix and error +codes. These tests pin behaviors the corpus does not: `normalize` idempotence, +`parse_ndsel`, 64-bit boundary handling, and schema-valid-but-redundant maps. +""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from zarr_indexing.messages import NdselError, normalize_ndsel, parse_ndsel + +_MESSAGES = [ + {"kind": "point", "coords": [4, 7]}, + {"kind": "box", "inclusive_min": [0, 0], "exclusive_max": [3, 4]}, + {"kind": "box", "inclusive_min": [["-inf"], 0], "exclusive_max": [["+inf"], 4]}, + {"kind": "slice", "start": [5], "stop": [10], "step": [2]}, + {"kind": "points", "coords": [[1, 10], [2, 20]]}, + { + "kind": "transform", + "input_inclusive_min": [0], + "input_exclusive_max": [3], + "output": [{"offset": 7}, {"input_dimension": 0, "stride": 2}, {"index_array": [1, 2, 3]}], + }, +] + + +@pytest.mark.parametrize("message", _MESSAGES) +def test_normalize_is_idempotent(message: dict[str, Any]) -> None: + once = normalize_ndsel(message) + twice = normalize_ndsel({"kind": "transform", **once}) + assert twice == once + + +@pytest.mark.parametrize("message", _MESSAGES) +def test_parse_returns_message_unchanged(message: dict[str, Any]) -> None: + assert parse_ndsel(message) == message + + +def test_parse_rejects_invalid() -> None: + with pytest.raises(NdselError) as excinfo: + parse_ndsel({"kind": "slice", "start": [0]}) + assert excinfo.value.reason == "invalid_json" + + +def test_constant_map_drops_redundant_stride() -> None: + # A constant map (no input_dimension, no index_array) is schema-valid even + # with a stray stride; it canonicalizes to offset-only. + result = normalize_ndsel( + {"kind": "transform", "input_rank": 0, "output": [{"offset": 5, "stride": 9}]} + ) + assert result["output"] == [{"offset": 5}] + + +def test_i64_min_and_max_round_trip() -> None: + i64_min, i64_max = -(2**63), 2**63 - 1 + result = normalize_ndsel({"kind": "point", "coords": [i64_min, i64_max]}) + assert result["output"] == [{"offset": i64_min}, {"offset": i64_max}] + + +def test_i64_overflow_rejected() -> None: + with pytest.raises(NdselError) as excinfo: + normalize_ndsel({"kind": "point", "coords": [2**63]}) + assert excinfo.value.reason == "invalid_json" + + +def test_bool_in_output_offset_rejected() -> None: + with pytest.raises(NdselError) as excinfo: + normalize_ndsel({"kind": "transform", "input_rank": 0, "output": [{"offset": True}]}) + assert excinfo.value.reason == "invalid_json" + + +def test_sentinel_not_allowed_in_plain_integer_position() -> None: + with pytest.raises(NdselError) as excinfo: + normalize_ndsel({"kind": "point", "coords": ["+inf"]}) + assert excinfo.value.reason == "invalid_json" + + +def test_not_an_object_rejected() -> None: + with pytest.raises(NdselError) as excinfo: + normalize_ndsel([1, 2, 3]) + assert excinfo.value.reason == "invalid_json" + + +def test_empty_string_kind_is_unknown_kind() -> None: + with pytest.raises(NdselError) as excinfo: + normalize_ndsel({"kind": ""}) + assert excinfo.value.reason == "unknown_kind" diff --git a/packages/zarr-indexing/tests/test_ndsel_tensorstore.py b/packages/zarr-indexing/tests/test_ndsel_tensorstore.py new file mode 100644 index 0000000000..794ac5f3d5 --- /dev/null +++ b/packages/zarr-indexing/tests/test_ndsel_tensorstore.py @@ -0,0 +1,52 @@ +"""Cross-check canonical ndsel bodies against a real TensorStore. + +A normalized ndsel `transform` body is, field-for-field, a TensorStore +`IndexTransform` (minus the `kind` discriminator, which the canonical body never +carries). This test loads a handful of finite-bound canonical bodies into +`tensorstore.IndexTransform(json=...)` and confirms that TensorStore's own +`to_json()` re-loads, through our engine layer, into an equivalent transform. + +Skipped when tensorstore is not installed. Run it explicitly with: + + uv run --with tensorstore pytest \ + packages/zarr-indexing/tests/test_ndsel_tensorstore.py -q +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from zarr_indexing.json import transform_from_canonical, transform_to_canonical +from zarr_indexing.transform import IndexTransform + +ts = pytest.importorskip("tensorstore") + + +def _canonical_transforms() -> list[IndexTransform]: + base = IndexTransform.from_shape((10, 20)) + return [ + base, # identity + base[2:8:2, :], # strided DimensionMap + identity + base[3, :], # integer index -> ConstantMap + DimensionMap + base.oindex[np.array([1, 5, 9]), :], # orthogonal index_array + IndexTransform.from_shape((10, 20, 30)).vindex[ + np.array([1, 3]), np.array([2, 4]), : + ], # correlated index_arrays + residual slice + ] + + +@pytest.mark.parametrize("transform", _canonical_transforms()) +def test_body_loads_in_tensorstore_and_round_trips(transform: IndexTransform) -> None: + body = transform_to_canonical(transform) + + # (1) The canonical body loads directly as a TensorStore IndexTransform. + ts_transform = ts.IndexTransform(json=body) + + # (2) TensorStore's own JSON re-loads, through our engine, to an equivalent + # transform. Comparing via our canonical form normalizes away + # representational choices (index_array_bounds, default omissions) that + # both sides make differently but that denote the same selection. + ts_json = ts_transform.to_json() + reloaded = transform_from_canonical(ts_json) + assert transform_to_canonical(reloaded) == transform_to_canonical(transform) diff --git a/packages/zarr-indexing/tests/test_output_map.py b/packages/zarr-indexing/tests/test_output_map.py new file mode 100644 index 0000000000..498101444e --- /dev/null +++ b/packages/zarr-indexing/tests/test_output_map.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import numpy as np + +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap + + +class TestConstantMap: + def test_construction(self) -> None: + m = ConstantMap(offset=42) + assert m.offset == 42 + + def test_default_offset(self) -> None: + m = ConstantMap() + assert m.offset == 0 + + def test_frozen(self) -> None: + m = ConstantMap(offset=5) + assert isinstance(m, ConstantMap) + + +class TestDimensionMap: + def test_construction(self) -> None: + m = DimensionMap(input_dimension=3, offset=5, stride=2) + assert m.input_dimension == 3 + assert m.offset == 5 + assert m.stride == 2 + + def test_defaults(self) -> None: + m = DimensionMap(input_dimension=0) + assert m.offset == 0 + assert m.stride == 1 + + def test_frozen(self) -> None: + m = DimensionMap(input_dimension=0) + assert isinstance(m, DimensionMap) + + +class TestArrayMap: + def test_construction(self) -> None: + arr = np.array([1, 3, 5], dtype=np.intp) + m = ArrayMap(index_array=arr, offset=10, stride=2) + assert m.offset == 10 + assert m.stride == 2 + np.testing.assert_array_equal(m.index_array, arr) + + def test_defaults(self) -> None: + arr = np.array([0, 1], dtype=np.intp) + m = ArrayMap(index_array=arr) + assert m.offset == 0 + assert m.stride == 1 + + def test_frozen(self) -> None: + arr = np.array([0], dtype=np.intp) + m = ArrayMap(index_array=arr) + assert isinstance(m, ArrayMap) diff --git a/packages/zarr-indexing/tests/test_tensorstore_parity.py b/packages/zarr-indexing/tests/test_tensorstore_parity.py new file mode 100644 index 0000000000..1ed99046e9 --- /dev/null +++ b/packages/zarr-indexing/tests/test_tensorstore_parity.py @@ -0,0 +1,263 @@ +"""TensorStore-parity oracle tests for IndexTransform semantics. + +Every case in this module was executed against tensorstore 0.1.84 (see the +lazy-indexing design notes): the expected domains, values, and error conditions +are TensorStore's observed behavior, which zarr's lazy indexing matches by +design. Core rules pinned here: + +- **Domain preservation**: a step-1 slice keeps the literal coordinates of the + selected interval (`a[2:10]` has domain `[2, 10)`); nothing re-zeros + implicitly. Re-zeroing is explicit via `translate_to`. +- **Strided-domain rule**: for step ``k``, ``origin = trunc(start/k)`` (rounded + toward zero), ``shape = ceil((stop - start)/k)``, and coordinate + ``origin + i`` maps to base cell ``start + i*k``. +- **Strict containment**: non-empty slice intervals must lie within the domain + — no clamping, no negative-wrapping; empty intervals are valid anywhere; + reversed non-empty bounds are an error, not an empty result. +- **Fancy-dim rule**: index-array dims get fresh explicit ``[0, n)`` domains; + index-array values are absolute domain coordinates. +- **Translate rules**: ``translate_by``/``translate_to`` shift the input domain + while preserving which cells are addressed. +""" + +from __future__ import annotations + +from typing import ClassVar + +import numpy as np +import pytest + +from zarr_indexing.domain import IndexDomain +from zarr_indexing.errors import BoundsCheckError +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr_indexing.transform import IndexTransform + + +def _identity(lo: int, hi: int) -> IndexTransform: + """Identity transform over the 1-D domain [lo, hi).""" + return IndexTransform.identity(IndexDomain(inclusive_min=(lo,), exclusive_max=(hi,))) + + +def _a() -> IndexTransform: + """The oracle's base fixture: identity over [0, 12).""" + return _identity(0, 12) + + +def _w() -> IndexTransform: + """The oracle's translated fixture: identity over [-10, 2), cell c -> base c + 10.""" + return _a().translate_domain_by((-10,)) + + +def _dim(t: IndexTransform) -> DimensionMap: + m = t.output[0] + assert isinstance(m, DimensionMap) + return m + + +def _base_cells(t: IndexTransform) -> list[int]: + """The base cells a 1-D single-DimensionMap transform addresses, in order.""" + m = _dim(t) + lo, hi = t.domain.inclusive_min[0], t.domain.exclusive_max[0] + return [m.offset + m.stride * c for c in range(lo, hi)] + + +class TestDomainPreservation: + """Oracle section 1-2: step-1 slices keep literal coordinates.""" + + def test_slice_preserves_domain(self) -> None: + t = _a()[2:10] + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((2,), (10,)) + assert _base_cells(t) == list(range(2, 10)) + + def test_integer_on_preserved_domain_is_a_coordinate(self) -> None: + v = _a()[2:10] + assert isinstance(v[3].output[0], ConstantMap) + assert v[3].output[0].offset == 3 # coordinate 3 = base cell 3 + assert v[2].output[0].offset == 2 + assert v[9].output[0].offset == 9 + + @pytest.mark.parametrize("bad", [0, -1, 10]) + def test_out_of_domain_integer_raises(self, bad: int) -> None: + with pytest.raises(BoundsCheckError, match=r"valid indices \[2, 10\)"): + _a()[2:10][bad] + + def test_slice_of_slice_is_literal(self) -> None: + v = _a()[2:10] + t = v[3:7] + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((3,), (7,)) + assert _base_cells(t) == [3, 4, 5, 6] + + def test_ellipsis_preserves_domain(self) -> None: + v = _a()[2:10] + t = v[...] + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((2,), (10,)) + + +class TestNegativeOriginDomain: + """Oracle section 3: on domain [-10, 2), -1 is just another index.""" + + def test_translated_domain(self) -> None: + w = _w() + assert (w.domain.inclusive_min, w.domain.exclusive_max) == ((-10,), (2,)) + assert _base_cells(w) == list(range(12)) + + @pytest.mark.parametrize(("coord", "base"), [(-5, 5), (-10, 0), (-1, 9), (1, 11)]) + def test_negative_coordinates_address_cells(self, coord: int, base: int) -> None: + t = _w()[coord] + assert isinstance(t.output[0], ConstantMap) + assert t.output[0].offset == base + + @pytest.mark.parametrize("bad", [-11, 2]) + def test_out_of_domain_raises(self, bad: int) -> None: + with pytest.raises(BoundsCheckError, match=r"valid indices \[-10, 2\)"): + _w()[bad] + + def test_negative_slice_bounds_are_coordinates(self) -> None: + t = _w()[-5:] + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((-5,), (2,)) + assert _base_cells(t) == [5, 6, 7, 8, 9, 10, 11] + t2 = _w()[-5:-2] + assert (t2.domain.inclusive_min, t2.domain.exclusive_max) == ((-5,), (-2,)) + assert _base_cells(t2) == [5, 6, 7] + + +class TestStridedDomains: + """Oracle section 5: origin = trunc(start/step), coord origin+i -> start + i*step.""" + + # (slice, expected (lo, hi), expected base cells) — verbatim oracle rows. + CASES: ClassVar[list[tuple[slice, tuple[int, int], list[int]]]] = [ + (slice(1, 10, 3), (0, 3), [1, 4, 7]), + (slice(None, None, 2), (0, 6), [0, 2, 4, 6, 8, 10]), + (slice(2, 11, 3), (0, 3), [2, 5, 8]), + (slice(0, 12, 4), (0, 3), [0, 4, 8]), + (slice(5, 12, 2), (2, 6), [5, 7, 9, 11]), + (slice(6, 12, 2), (3, 6), [6, 8, 10]), + (slice(7, 12, 3), (2, 4), [7, 10]), + ] + + @pytest.mark.parametrize(("sel", "dom", "cells"), CASES) + def test_strided_domain_and_cells( + self, sel: slice, dom: tuple[int, int], cells: list[int] + ) -> None: + t = _a()[sel] + assert (t.domain.inclusive_min[0], t.domain.exclusive_max[0]) == dom + assert _base_cells(t) == cells + + def test_strided_on_negative_origin(self) -> None: + # w[-9:2:2] -> domain [-4, 2), base cells 1,3,5,7,9,11 + t = _w()[-9:2:2] + assert (t.domain.inclusive_min[0], t.domain.exclusive_max[0]) == (-4, 2) + assert _base_cells(t) == [1, 3, 5, 7, 9, 11] + # w[::2] -> domain [-5, 1), base cells 0,2,4,6,8,10 + t2 = _w()[::2] + assert (t2.domain.inclusive_min[0], t2.domain.exclusive_max[0]) == (-5, 1) + assert _base_cells(t2) == [0, 2, 4, 6, 8, 10] + + def test_strided_composition(self) -> None: + s = _a()[1:10:3] # domain [0, 3), cells 1,4,7 + assert [s[k].output[0].offset for k in range(3)] == [1, 4, 7] + t = s[1:3] + assert (t.domain.inclusive_min[0], t.domain.exclusive_max[0]) == (1, 3) + assert _base_cells(t) == [4, 7] + t2 = _a()[::2][1:4] + assert (t2.domain.inclusive_min[0], t2.domain.exclusive_max[0]) == (1, 4) + assert _base_cells(t2) == [2, 4, 6] + t3 = _a()[::2][::2] + assert (t3.domain.inclusive_min[0], t3.domain.exclusive_max[0]) == (0, 3) + assert _base_cells(t3) == [0, 4, 8] + + @pytest.mark.parametrize("bad", [-2, -1, 3, 4]) + def test_strided_bounds(self, bad: int) -> None: + with pytest.raises(BoundsCheckError, match=r"valid indices \[0, 3\)"): + _a()[1:10:3][bad] + + +class TestStrictContainment: + """Oracle section 11: no clamping, no wrapping; empty intervals valid anywhere.""" + + @pytest.mark.parametrize( + "sel", + [ + slice(5, 100), + slice(-3, None), + slice(-3, -1), + slice(0, 13), + slice(12, 14), + slice(100, 200), + ], + ) + def test_uncontained_interval_raises(self, sel: slice) -> None: + with pytest.raises(BoundsCheckError, match="not contained"): + _a()[sel] + + def test_uncontained_on_negative_origin(self) -> None: + with pytest.raises(BoundsCheckError, match="not contained"): + _w()[-20:] + + @pytest.mark.parametrize( + ("sel", "pos"), [(slice(5, 5), 5), (slice(0, 0), 0), (slice(13, 13), 13)] + ) + def test_empty_interval_valid_anywhere(self, sel: slice, pos: int) -> None: + t = _a()[sel] + assert t.domain.shape == (0,) + assert t.domain.inclusive_min[0] == pos + + @pytest.mark.parametrize("sel", [slice(5, 2), slice(100, 50)]) + def test_reversed_bounds_raise(self, sel: slice) -> None: + with pytest.raises(IndexError, match="valid.*interval|interval"): + _a()[sel] + + +class TestTranslate: + """Oracle sections 4 and 12: translate_by / translate_to preserve the cell mapping.""" + + def test_translate_to_zero(self) -> None: + t = _a()[2:10].translate_domain_to((0,)) + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((0,), (8,)) + assert _base_cells(t) == list(range(2, 10)) + + def test_translate_to_offset(self) -> None: + t = _a().translate_domain_to((5,)) + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((5,), (17,)) + assert _base_cells(t) == list(range(12)) + + def test_translate_by_composes_with_stride(self) -> None: + # a[::2].translate_by[5] -> domain [5, 11), base = 2*(coord-5) + t = _a()[::2].translate_domain_by((5,)) + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((5,), (11,)) + assert _base_cells(t) == [0, 2, 4, 6, 8, 10] + assert t[5].output[0].offset == 0 + assert t[10].output[0].offset == 10 + with pytest.raises(BoundsCheckError, match=r"valid indices \[5, 11\)"): + t[0] + + def test_translate_strided_to(self) -> None: + t = _a()[1:10:3].translate_domain_to((100,)) + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((100,), (103,)) + assert _base_cells(t) == [1, 4, 7] + + +class TestFancyDims: + """Oracle section 7: fancy dims get fresh [0, n); values are absolute coordinates.""" + + def test_index_array_values_are_coordinates(self) -> None: + v = _a()[2:10] + t = v.oindex[(np.array([3, 5], dtype=np.intp),)] + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((0,), (2,)) + m = t.output[0] + assert isinstance(m, ArrayMap) + storage = m.offset + m.stride * m.index_array + np.testing.assert_array_equal(np.asarray(storage).ravel(), [3, 5]) + + def test_index_array_on_negative_origin(self) -> None: + t = _w().oindex[(np.array([-10, -1], dtype=np.intp),)] + m = t.output[0] + assert isinstance(m, ArrayMap) + storage = m.offset + m.stride * m.index_array + np.testing.assert_array_equal(np.asarray(storage).ravel(), [0, 9]) + + def test_index_array_out_of_domain_raises(self) -> None: + v = _a()[2:10] + for bad in ([0, 3], [-1, 3], [3, 10]): + with pytest.raises(BoundsCheckError): + v.oindex[(np.array(bad, dtype=np.intp),)] diff --git a/packages/zarr-indexing/tests/test_transform.py b/packages/zarr-indexing/tests/test_transform.py new file mode 100644 index 0000000000..baecd9ada2 --- /dev/null +++ b/packages/zarr-indexing/tests/test_transform.py @@ -0,0 +1,628 @@ +from __future__ import annotations + +import numpy as np +import pytest + +from zarr_indexing.domain import IndexDomain +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr_indexing.transform import IndexTransform, selection_to_transform + + +class TestIndexTransformConstruction: + def test_from_shape(self) -> None: + t = IndexTransform.from_shape((10, 20)) + assert t.input_rank == 2 + assert t.output_rank == 2 + assert t.domain.shape == (10, 20) + assert t.domain.origin == (0, 0) + for i, m in enumerate(t.output): + assert isinstance(m, DimensionMap) + assert m.input_dimension == i + assert m.offset == 0 + assert m.stride == 1 + + def test_identity(self) -> None: + domain = IndexDomain(inclusive_min=(5,), exclusive_max=(15,)) + t = IndexTransform.identity(domain) + assert t.input_rank == 1 + assert t.output_rank == 1 + assert t.domain == domain + assert isinstance(t.output[0], DimensionMap) + assert t.output[0].input_dimension == 0 + + def test_from_shape_0d(self) -> None: + t = IndexTransform.from_shape(()) + assert t.input_rank == 0 + assert t.output_rank == 0 + assert t.domain.shape == () + + def test_custom_output_maps(self) -> None: + domain = IndexDomain.from_shape((10,)) + maps = (ConstantMap(offset=42), DimensionMap(input_dimension=0, offset=5, stride=2)) + t = IndexTransform(domain=domain, output=maps) + assert t.input_rank == 1 + assert t.output_rank == 2 + + def test_validation_input_dimension_out_of_range(self) -> None: + domain = IndexDomain.from_shape((10,)) + maps = (DimensionMap(input_dimension=5),) + with pytest.raises(ValueError, match="input_dimension"): + IndexTransform(domain=domain, output=maps) + + +class TestIndexTransformBasicIndexing: + def test_slice_identity(self) -> None: + """slice(None) on identity transform is a no-op.""" + t = IndexTransform.from_shape((10, 20)) + result = t[slice(None), slice(None)] + assert result.domain.shape == (10, 20) + assert result.input_rank == 2 + assert result.output_rank == 2 + + def test_slice_narrows(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = t[2:8, 5:15] + # Domains are preserved (TensorStore): the slice keeps its literal + # coordinates, so the map stays the identity (out = in). + assert result.domain.shape == (6, 10) + assert result.domain.origin == (2, 5) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == 0 + assert result.output[0].stride == 1 + assert result.output[0].input_dimension == 0 + assert isinstance(result.output[1], DimensionMap) + assert result.output[1].offset == 0 + assert result.output[1].input_dimension == 1 + + def test_strided_slice(self) -> None: + t = IndexTransform.from_shape((10,)) + result = t[::2] + assert result.domain.shape == (5,) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == 0 + assert result.output[0].stride == 2 + + def test_strided_slice_with_start(self) -> None: + t = IndexTransform.from_shape((10,)) + result = t[1:9:3] + # indices: 1, 4, 7 -> 3 elements + assert result.domain.shape == (3,) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == 1 + assert result.output[0].stride == 3 + + def test_int_drops_dimension(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = t[3] + assert result.input_rank == 1 + assert result.output_rank == 2 + assert isinstance(result.output[0], ConstantMap) + assert result.output[0].offset == 3 + assert isinstance(result.output[1], DimensionMap) + assert result.output[1].input_dimension == 0 + + def test_int_middle_dimension(self) -> None: + t = IndexTransform.from_shape((10, 20, 30)) + result = t[:, 5, :] + assert result.input_rank == 2 + assert result.output_rank == 3 + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].input_dimension == 0 + assert isinstance(result.output[1], ConstantMap) + assert result.output[1].offset == 5 + assert isinstance(result.output[2], DimensionMap) + assert result.output[2].input_dimension == 1 + + def test_ellipsis(self) -> None: + t = IndexTransform.from_shape((10, 20, 30)) + result = t[2:8, ...] + assert result.input_rank == 3 + assert result.domain.shape == (6, 20, 30) + + def test_newaxis(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = t[np.newaxis, :, :] + assert result.input_rank == 3 + assert result.domain.shape == (1, 10, 20) + assert result.output_rank == 2 + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].input_dimension == 1 + assert isinstance(result.output[1], DimensionMap) + assert result.output[1].input_dimension == 2 + + def test_int_out_of_bounds(self) -> None: + t = IndexTransform.from_shape((10,)) + with pytest.raises(IndexError): + t[10] + + def test_negative_int_is_literal(self) -> None: + """Negative indices are literal coordinates (TensorStore convention), + not 'from the end' like NumPy.""" + t = IndexTransform.from_shape((10,)) + with pytest.raises(IndexError): + t[-1] # -1 is out of bounds for domain [0, 10) + + def test_negative_int_valid_with_negative_origin(self) -> None: + """Negative index is valid if the domain includes negative coordinates.""" + domain = IndexDomain(inclusive_min=(-5,), exclusive_max=(5,)) + t = IndexTransform.identity(domain) + result = t[-3] + assert isinstance(result.output[0], ConstantMap) + assert result.output[0].offset == -3 + + def test_composition_of_slices(self) -> None: + """Slicing a sliced transform re-selects in literal domain coordinates.""" + t = IndexTransform.from_shape((100,)) + result = t[10:50][15:30] + assert result.domain.shape == (15,) + assert result.domain.origin == (15,) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == 0 + assert result.output[0].stride == 1 + + def test_composition_of_strides(self) -> None: + t = IndexTransform.from_shape((100,)) + result = t[::2][::3] + # t[::2] -> shape (50,), offset=0, stride=2 + # [::3] -> shape ceil(50/3)=17, offset=0, stride=2*3=6 + assert result.domain.shape == (17,) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].stride == 6 + + def test_bare_int(self) -> None: + """Non-tuple selection.""" + t = IndexTransform.from_shape((10, 20)) + result = t[3] + assert result.input_rank == 1 + + def test_bare_slice(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = t[2:8] + assert result.domain.shape == (6, 20) + + +class TestBasicIndexingOnArrayMaps: + """When a transform already has ArrayMap outputs, basic indexing must + apply the corresponding operation to the index_array's axes.""" + + def test_int_on_array_map_drops_axis(self) -> None: + """Integer index on a dimension referenced by an ArrayMap should + index into the array on that axis.""" + arr = np.array([[10, 20], [30, 40], [50, 60]], dtype=np.intp) + # 2D input domain (3, 2), one ArrayMap output + t = IndexTransform( + domain=IndexDomain.from_shape((3, 2)), + output=(ArrayMap(index_array=arr),), + ) + # Index with int on dim 0 -> pick row 1 -> arr[1, :] = [30, 40] + result = t[1] + assert result.input_rank == 1 + assert result.domain.shape == (2,) + assert isinstance(result.output[0], ArrayMap) + np.testing.assert_array_equal(result.output[0].index_array, np.array([30, 40])) + + def test_slice_on_array_map(self) -> None: + """Slice on a dimension referenced by an ArrayMap should slice the array.""" + arr = np.array([10, 20, 30, 40, 50], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((5,)), + output=(ArrayMap(index_array=arr),), + ) + result = t[1:4] + assert result.domain.shape == (3,) + assert isinstance(result.output[0], ArrayMap) + np.testing.assert_array_equal(result.output[0].index_array, np.array([20, 30, 40])) + + def test_strided_slice_on_array_map(self) -> None: + """Strided slice on ArrayMap should stride the array.""" + arr = np.array([10, 20, 30, 40, 50], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((5,)), + output=(ArrayMap(index_array=arr),), + ) + result = t[::2] + assert result.domain.shape == (3,) + assert isinstance(result.output[0], ArrayMap) + np.testing.assert_array_equal(result.output[0].index_array, np.array([10, 30, 50])) + + def test_newaxis_on_array_map(self) -> None: + """Newaxis should insert an axis in the index_array.""" + arr = np.array([10, 20, 30], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=arr),), + ) + result = t[np.newaxis, :] + assert result.input_rank == 2 + assert result.domain.shape == (1, 3) + assert isinstance(result.output[0], ArrayMap) + assert result.output[0].index_array.shape == (1, 3) + np.testing.assert_array_equal(result.output[0].index_array, np.array([[10, 20, 30]])) + + def test_int_drops_one_of_two_array_dims(self) -> None: + """2D array map, int on dim 0, slice on dim 1.""" + arr = np.array([[10, 20, 30], [40, 50, 60]], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((2, 3)), + output=(ArrayMap(index_array=arr),), + ) + result = t[0, 1:3] + assert result.input_rank == 1 + assert result.domain.shape == (2,) + assert isinstance(result.output[0], ArrayMap) + # arr[0, 1:3] = [20, 30] + np.testing.assert_array_equal(result.output[0].index_array, np.array([20, 30])) + + +class TestIndexTransformOindex: + def test_oindex_int_array(self) -> None: + t = IndexTransform.from_shape((10, 20)) + idx = np.array([1, 3, 5], dtype=np.intp) + result = t.oindex[idx, :] + assert result.input_rank == 2 + assert result.domain.shape == (3, 20) + assert isinstance(result.output[0], ArrayMap) + # Full input rank: the array varies along its own axis (0), singleton on 1. + assert result.output[0].index_array.shape == (3, 1) + np.testing.assert_array_equal(result.output[0].index_array, idx.reshape(3, 1)) + assert result.output[0].offset == 0 + assert result.output[0].stride == 1 + assert isinstance(result.output[1], DimensionMap) + assert result.output[1].input_dimension == 1 + + def test_oindex_bool_array(self) -> None: + t = IndexTransform.from_shape((5,)) + mask = np.array([True, False, True, False, True]) + result = t.oindex[mask] + assert result.domain.shape == (3,) + assert isinstance(result.output[0], ArrayMap) + np.testing.assert_array_equal( + result.output[0].index_array, np.array([0, 2, 4], dtype=np.intp) + ) + + def test_oindex_mixed(self) -> None: + t = IndexTransform.from_shape((10, 20)) + idx = np.array([2, 4], dtype=np.intp) + result = t.oindex[idx, 5:15] + assert result.input_rank == 2 + assert result.domain.shape == (2, 10) + # fancy dim: fresh zero-origin; slice dim: preserved literal coords + assert result.domain.origin == (0, 5) + assert isinstance(result.output[0], ArrayMap) + assert isinstance(result.output[1], DimensionMap) + assert result.output[1].offset == 0 + + def test_oindex_multiple_arrays(self) -> None: + t = IndexTransform.from_shape((10, 20, 30)) + idx0 = np.array([1, 3], dtype=np.intp) + idx1 = np.array([5, 10, 15], dtype=np.intp) + result = t.oindex[idx0, :, idx1] + assert result.input_rank == 3 + assert result.domain.shape == (2, 20, 3) + assert isinstance(result.output[0], ArrayMap) + assert isinstance(result.output[1], DimensionMap) + assert isinstance(result.output[2], ArrayMap) + + def test_oindex_multiple_arrays_preserves_independent_axes(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = t.oindex[np.array([1, 3]), np.array([2, 4, 6])] + assert result.domain.shape == (2, 3) + assert isinstance(result.output[0], ArrayMap) + assert isinstance(result.output[1], ArrayMap) + assert result.output[0].index_array.shape == (2, 1) + assert result.output[1].index_array.shape == (1, 3) + + +class TestIndexTransformVindex: + def test_vindex_single_array(self) -> None: + t = IndexTransform.from_shape((10,)) + idx = np.array([1, 3, 5], dtype=np.intp) + result = t.vindex[idx] + assert result.input_rank == 1 + assert result.domain.shape == (3,) + assert isinstance(result.output[0], ArrayMap) + np.testing.assert_array_equal(result.output[0].index_array, idx) + + def test_vindex_broadcast(self) -> None: + t = IndexTransform.from_shape((10, 20)) + idx0 = np.array([[1, 2], [3, 4]], dtype=np.intp) + idx1 = np.array([[10, 11], [12, 13]], dtype=np.intp) + result = t.vindex[idx0, idx1] + assert result.input_rank == 2 + assert result.domain.shape == (2, 2) + assert isinstance(result.output[0], ArrayMap) + assert isinstance(result.output[1], ArrayMap) + np.testing.assert_array_equal(result.output[0].index_array, idx0) + np.testing.assert_array_equal(result.output[1].index_array, idx1) + + def test_vindex_with_slice(self) -> None: + t = IndexTransform.from_shape((10, 20, 30)) + idx = np.array([1, 3, 5], dtype=np.intp) + result = t.vindex[idx, :, :] + assert result.input_rank == 3 + assert result.domain.shape == (3, 20, 30) + assert isinstance(result.output[0], ArrayMap) + + def test_vindex_bool_mask(self) -> None: + t = IndexTransform.from_shape((5,)) + mask = np.array([True, False, True, False, True]) + result = t.vindex[mask] + assert result.domain.shape == (3,) + assert isinstance(result.output[0], ArrayMap) + + def test_vindex_broadcast_different_shapes(self) -> None: + t = IndexTransform.from_shape((10, 20)) + idx0 = np.array([1, 2, 3], dtype=np.intp) + idx1 = np.array([[10], [11]], dtype=np.intp) + result = t.vindex[idx0, idx1] + assert result.input_rank == 2 + assert result.domain.shape == (2, 3) + + def test_vindex_multiple_arrays_preserves_shared_axes(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = t.vindex[np.array([1, 3]), np.array([2, 4])] + assert result.domain.shape == (2,) + assert isinstance(result.output[0], ArrayMap) + assert isinstance(result.output[1], ArrayMap) + assert result.output[0].index_array.shape == (2,) + assert result.output[1].index_array.shape == (2,) + + +class TestSelectionToTransform: + def test_basic_slice(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = selection_to_transform((slice(2, 8), slice(5, 15)), t, "basic") + assert result.domain.shape == (6, 10) + assert result.domain.origin == (2, 5) # preserved literal coordinates + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == 0 + + def test_basic_int(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = selection_to_transform((3, slice(None)), t, "basic") + assert result.input_rank == 1 + assert isinstance(result.output[0], ConstantMap) + assert result.output[0].offset == 3 + + def test_basic_ellipsis(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = selection_to_transform(Ellipsis, t, "basic") + assert result.domain.shape == (10, 20) + + def test_orthogonal(self) -> None: + t = IndexTransform.from_shape((10, 20)) + idx = np.array([1, 3, 5], dtype=np.intp) + result = selection_to_transform((idx, slice(None)), t, "orthogonal") + assert result.domain.shape == (3, 20) + assert isinstance(result.output[0], ArrayMap) + + def test_vectorized(self) -> None: + t = IndexTransform.from_shape((10, 20)) + idx0 = np.array([1, 3], dtype=np.intp) + idx1 = np.array([5, 7], dtype=np.intp) + result = selection_to_transform((idx0, idx1), t, "vectorized") + assert result.domain.shape == (2,) + assert isinstance(result.output[0], ArrayMap) + assert isinstance(result.output[1], ArrayMap) + + def test_composition_with_non_identity(self) -> None: + """Indexing a sliced transform uses literal domain coordinates. + + The slice [10:50] preserves its domain, so a follow-up [15:30] + re-selects coordinates 15..29 of the base (TensorStore semantics), and + the composed map stays the identity (out = in). + """ + t = IndexTransform.from_shape((100,))[10:50] + result = selection_to_transform(slice(15, 30), t, "basic") + assert (result.domain.inclusive_min, result.domain.exclusive_max) == ((15,), (30,)) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == 0 + assert result.output[0].stride == 1 + + +class TestIndexTransformIntersect: + def test_constant_inside(self) -> None: + t = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(ConstantMap(offset=5),), + ) + result = t.intersect(IndexDomain(inclusive_min=(0,), exclusive_max=(10,))) + assert result is not None + restricted, surviving = result + assert isinstance(restricted.output[0], ConstantMap) + assert restricted.output[0].offset == 5 + assert surviving is None + + def test_constant_outside(self) -> None: + t = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(ConstantMap(offset=5),), + ) + result = t.intersect(IndexDomain(inclusive_min=(10,), exclusive_max=(20,))) + assert result is None + + def test_dimension_partial(self) -> None: + """DimensionMap over [0,10) intersected with [5,15) narrows input to [5,10).""" + t = IndexTransform.from_shape((10,)) + result = t.intersect(IndexDomain(inclusive_min=(5,), exclusive_max=(15,))) + assert result is not None + restricted, surviving = result + assert restricted.domain.inclusive_min == (5,) + assert restricted.domain.exclusive_max == (10,) + assert surviving is None + + def test_dimension_no_overlap(self) -> None: + t = IndexTransform.from_shape((10,)) + result = t.intersect(IndexDomain(inclusive_min=(20,), exclusive_max=(30,))) + assert result is None + + def test_dimension_strided(self) -> None: + """stride=2, offset=1 over [0,5): storage 1,3,5,7,9. Chunk [4,8).""" + t = IndexTransform( + domain=IndexDomain.from_shape((5,)), + output=(DimensionMap(input_dimension=0, offset=1, stride=2),), + ) + result = t.intersect(IndexDomain(inclusive_min=(4,), exclusive_max=(8,))) + assert result is not None + restricted, _surviving = result + # input 2->5, input 3->7. Both in [4,8). + assert restricted.domain.inclusive_min == (2,) + assert restricted.domain.exclusive_max == (4,) + + def test_array_partial(self) -> None: + arr = np.array([3, 8, 15, 22], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((4,)), + output=(ArrayMap(index_array=arr),), + ) + result = t.intersect(IndexDomain(inclusive_min=(5,), exclusive_max=(20,))) + assert result is not None + restricted, surviving = result + assert isinstance(restricted.output[0], ArrayMap) + np.testing.assert_array_equal(restricted.output[0].index_array, np.array([8, 15])) + assert surviving is not None + np.testing.assert_array_equal(surviving, np.array([1, 2])) + + def test_array_none_inside(self) -> None: + arr = np.array([1, 2, 3], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=arr),), + ) + assert t.intersect(IndexDomain(inclusive_min=(10,), exclusive_max=(20,))) is None + + def test_2d_mixed(self) -> None: + """2D: ConstantMap on dim 0, DimensionMap on dim 1.""" + t = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=( + ConstantMap(offset=5), + DimensionMap(input_dimension=0, offset=0, stride=1), + ), + ) + chunk = IndexDomain(inclusive_min=(0, 5), exclusive_max=(10, 15)) + result = t.intersect(chunk) + assert result is not None + restricted, _ = result + assert isinstance(restricted.output[0], ConstantMap) + assert restricted.output[0].offset == 5 + assert isinstance(restricted.output[1], DimensionMap) + assert restricted.domain.inclusive_min == (5,) + assert restricted.domain.exclusive_max == (10,) + + +class TestIndexTransformTranslate: + def test_translate_constant(self) -> None: + t = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(ConstantMap(offset=5),), + ) + result = t.translate((-5,)) + assert isinstance(result.output[0], ConstantMap) + assert result.output[0].offset == 0 + + def test_translate_dimension(self) -> None: + t = IndexTransform.from_shape((10,)) + result = t.translate((-3,)) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == -3 + assert result.output[0].stride == 1 + + def test_translate_array(self) -> None: + arr = np.array([5, 10], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((2,)), + output=(ArrayMap(index_array=arr, offset=3),), + ) + result = t.translate((-3,)) + assert isinstance(result.output[0], ArrayMap) + assert result.output[0].offset == 0 + np.testing.assert_array_equal(result.output[0].index_array, arr) + + def test_translate_2d(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = t.translate((-5, -10)) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == -5 + assert isinstance(result.output[1], DimensionMap) + assert result.output[1].offset == -10 + + +class TestArrayMapDependencyAxes: + """`_array_map_dependency_axes` derives the input axes an array varies on + from its (full-rank) shape: non-singleton axes vary, singleton axes do not.""" + + def test_orthogonal_single_axis(self) -> None: + from zarr_indexing.transform import _array_map_dependency_axes + + t = IndexTransform.from_shape((10, 20)).oindex[np.array([1, 3]), np.array([2, 4, 6])] + m0, m1 = t.output[0], t.output[1] + assert isinstance(m0, ArrayMap) + assert isinstance(m1, ArrayMap) + assert _array_map_dependency_axes(m0.index_array) == (0,) + assert _array_map_dependency_axes(m1.index_array) == (1,) + + def test_vectorized_shares_axes(self) -> None: + from zarr_indexing.transform import _array_map_dependency_axes + + t = IndexTransform.from_shape((10, 20)).vindex[np.array([1, 3]), np.array([2, 4])] + m0, m1 = t.output[0], t.output[1] + assert isinstance(m0, ArrayMap) + assert isinstance(m1, ArrayMap) + assert _array_map_dependency_axes(m0.index_array) == (0,) + assert _array_map_dependency_axes(m1.index_array) == (0,) + + def test_scalar_array_has_no_dependency(self) -> None: + from zarr_indexing.transform import _array_map_dependency_axes + + assert _array_map_dependency_axes(np.ones((1, 1), dtype=np.intp)) == () + + +class TestIntersectArrayMapClassification: + """`_intersect` must distinguish orthogonal (outer-product) ArrayMaps from + correlated (vectorized) ones by their dependency axes, keep surviving arrays + at full input rank, and preserve residual (slice) dimensions.""" + + def test_orthogonal_outer_product_keeps_full_rank(self) -> None: + """Two arrays on distinct axes narrow independently and stay full rank; + out_indices is a per-output-dim dict of surviving positions.""" + t = IndexTransform.from_shape((10, 10)).oindex[np.array([1, 3, 8]), np.array([2, 6, 9])] + # Chunk covering storage [0,5) x [0,5): rows 1,3 survive (out pos 0,1), + # cols 2 survives (out pos 0). + chunk = IndexDomain(inclusive_min=(0, 0), exclusive_max=(5, 5)) + result = t.intersect(chunk) + assert result is not None + restricted, out_indices = result + assert isinstance(restricted.output[0], ArrayMap) + assert isinstance(restricted.output[1], ArrayMap) + # Full input rank preserved (not raveled to 1-D). + assert restricted.output[0].index_array.ndim == 2 + assert restricted.output[1].index_array.ndim == 2 + assert restricted.domain.ndim == 2 + assert isinstance(out_indices, dict) + np.testing.assert_array_equal(out_indices[0], np.array([0, 1])) + np.testing.assert_array_equal(out_indices[1], np.array([0])) + + def test_correlated_with_residual_slice_preserves_slice_dim(self) -> None: + """A vindex transform with two correlated arrays plus a residual slice + dim intersects without a rank error and keeps the DimensionMap.""" + t = IndexTransform.from_shape((4, 3, 5)).vindex[np.array([1, 3]), np.array([2, 0])] + # Chunk covering storage [0,2) x [2,3) x [0,5): only point (1,2,*) is in + # bounds on both array dims -> one surviving broadcast point. + chunk = IndexDomain(inclusive_min=(0, 2, 0), exclusive_max=(2, 3, 5)) + result = t.intersect(chunk) + assert result is not None + restricted, out_indices = result + # A DimensionMap for the residual slice dim survives (no post-init error). + assert any(isinstance(m, DimensionMap) for m in restricted.output) + assert out_indices is not None + + def test_length1_orthogonal_not_treated_as_correlated(self) -> None: + """A length-1 orthogonal array (all-singleton shape) is still an outer + product with the length-3 axis: out_indices is a dict, not a flat array.""" + t = IndexTransform.from_shape((6, 6)).oindex[np.array([2]), np.array([1, 3, 5])] + chunk = IndexDomain(inclusive_min=(0, 0), exclusive_max=(6, 6)) + result = t.intersect(chunk) + assert result is not None + _restricted, out_indices = result + assert isinstance(out_indices, dict)