create module for chunk grids, and registry

Author: d-v-b

src/zarr/core/chunk_grids.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+import itertools
+import operator
+from dataclasses import dataclass
+from functools import reduce
+from typing import TYPE_CHECKING, Any, Self
+
+import numpy as np
+import numpy.typing as npt
+
+from zarr.core.chunk_grids.common import ChunkGrid
+from zarr.core.common import (
+    JSON,
+    NamedConfig,
+    ShapeLike,
+    ceildiv,
+    parse_named_configuration,
+    parse_shapelike,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+
+@dataclass(frozen=True)
+class RegularChunkGrid(ChunkGrid):
+    chunk_shape: tuple[int, ...]
+
+    def __init__(self, *, chunk_shape: ShapeLike) -> None:
+        chunk_shape_parsed = parse_shapelike(chunk_shape)
+
+        object.__setattr__(self, "chunk_shape", chunk_shape_parsed)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, JSON] | NamedConfig[str, Any]) -> Self:
+        _, configuration_parsed = parse_named_configuration(data, "regular")
+
+        return cls(**configuration_parsed)  # type: ignore[arg-type]
+
+    def to_dict(self) -> dict[str, JSON]:
+        return {"name": "regular", "configuration": {"chunk_shape": tuple(self.chunk_shape)}}
+
+    def update_shape(self, new_shape: tuple[int, ...]) -> Self:
+        return self
+
+    def all_chunk_coords(self, array_shape: tuple[int, ...]) -> Iterator[tuple[int, ...]]:
+        return itertools.product(
+            *(range(ceildiv(s, c)) for s, c in zip(array_shape, self.chunk_shape, strict=False))
+        )
+
+    def get_nchunks(self, array_shape: tuple[int, ...]) -> int:
+        return reduce(
+            operator.mul,
+            itertools.starmap(ceildiv, zip(array_shape, self.chunk_shape, strict=True)),
+            1,
+        )
+
+    def get_chunk_shape(
+        self, array_shape: tuple[int, ...], chunk_coord: tuple[int, ...]
+    ) -> tuple[int, ...]:
+        """
+        Get the shape of a specific chunk.
+
+        For RegularChunkGrid, all chunks have the same shape except possibly
+        the last chunk in each dimension.
+        """
+        return tuple(
+            int(min(self.chunk_shape[i], array_shape[i] - chunk_coord[i] * self.chunk_shape[i]))
+            for i in range(len(array_shape))
+        )
+
+    def get_chunk_start(
+        self, array_shape: tuple[int, ...], chunk_coord: tuple[int, ...]
+    ) -> tuple[int, ...]:
+        """
+        Get the starting position of a chunk in the array.
+
+        For RegularChunkGrid, this is simply chunk_coord * chunk_shape.
+        """
+        return tuple(
+            coord * size for coord, size in zip(chunk_coord, self.chunk_shape, strict=False)
+        )
+
+    def array_index_to_chunk_coord(
+        self, array_shape: tuple[int, ...], array_index: tuple[int, ...]
+    ) -> tuple[int, ...]:
+        """
+        Map an array index to chunk coordinates.
+
+        For RegularChunkGrid, this is simply array_index // chunk_shape.
+        """
+        return tuple(
+            0 if size == 0 else idx // size
+            for idx, size in zip(array_index, self.chunk_shape, strict=False)
+        )
+
+    def array_indices_to_chunk_dim(
+        self, array_shape: tuple[int, ...], dim: int, indices: npt.NDArray[np.intp]
+    ) -> npt.NDArray[np.intp]:
+        """
+        Vectorized mapping of array indices to chunk coordinates along one dimension.
+
+        For RegularChunkGrid, this is simply indices // chunk_size.
+        """
+        chunk_size = self.chunk_shape[dim]
+        if chunk_size == 0:
+            return np.zeros_like(indices)
+        return indices // chunk_size
+
+    def chunks_per_dim(self, array_shape: tuple[int, ...], dim: int) -> int:
+        """
+        Get the number of chunks along a specific dimension.
+
+        For RegularChunkGrid, this is ceildiv(array_shape[dim], chunk_shape[dim]).
+        """
+        return ceildiv(array_shape[dim], self.chunk_shape[dim])
+
+    def get_chunk_grid_shape(self, array_shape: tuple[int, ...]) -> tuple[int, ...]:
+        """
+        Get the shape of the chunk grid (number of chunks along each dimension).
+
+        For RegularChunkGrid, this is computed using ceildiv for each dimension.
+        """
+        return tuple(
+            ceildiv(array_len, chunk_len)
+            for array_len, chunk_len in zip(array_shape, self.chunk_shape, strict=False)
+        )

src/zarr/core/chunk_grids/__init__.py

@@ -28,6 +28,7 @@
     ChunkGrid,
     RectilinearChunkGrid,
     RegularChunkGrid,
+    parse_chunk_grid_from_dict,
 )
 from zarr.core.chunk_key_encodings import (
     ChunkKeyEncoding,
@@ -233,7 +234,7 @@ def __init__(
         """
 
         shape_parsed = parse_shapelike(shape)
-        chunk_grid_parsed = ChunkGrid.from_dict(chunk_grid)
+        chunk_grid_parsed = parse_chunk_grid_from_dict(chunk_grid)
         chunk_key_encoding_parsed = parse_chunk_key_encoding(chunk_key_encoding)
         dimension_names_parsed = parse_dimension_names(dimension_names)
         # Note: relying on a type method is numpy-specific

src/zarr/core/chunk_grids/common.py

@@ -22,17 +22,20 @@
     )
     from zarr.abc.numcodec import Numcodec
     from zarr.core.buffer import Buffer, NDBuffer
+    from zarr.core.chunk_grids import ChunkGrid
     from zarr.core.chunk_key_encodings import ChunkKeyEncoding
     from zarr.core.common import JSON
 
 __all__ = [
     "Registry",
     "get_buffer_class",
+    "get_chunk_grid_class",
     "get_chunk_key_encoding_class",
     "get_codec_class",
     "get_ndbuffer_class",
     "get_pipeline_class",
     "register_buffer",
+    "register_chunk_grid",
     "register_chunk_key_encoding",
     "register_codec",
     "register_ndbuffer",
@@ -63,6 +66,7 @@ def register(self, cls: type[T], qualname: str | None = None) -> None:
 __pipeline_registry: Registry[CodecPipeline] = Registry()
 __buffer_registry: Registry[Buffer] = Registry()
 __ndbuffer_registry: Registry[NDBuffer] = Registry()
+__chunk_grid_registry: Registry[ChunkGrid] = Registry()
 __chunk_key_encoding_registry: Registry[ChunkKeyEncoding] = Registry()
 
 """
@@ -103,6 +107,11 @@ def _collect_entrypoints() -> list[Registry[Any]]:
     data_type_registry._lazy_load_list.extend(entry_points.select(group="zarr.data_type"))
     data_type_registry._lazy_load_list.extend(entry_points.select(group="zarr", name="data_type"))
 
+    __chunk_grid_registry.lazy_load_list.extend(entry_points.select(group="zarr.chunk_grid"))
+    __chunk_grid_registry.lazy_load_list.extend(
+        entry_points.select(group="zarr", name="chunk_grid")
+    )
+
     __chunk_key_encoding_registry.lazy_load_list.extend(
         entry_points.select(group="zarr.chunk_key_encoding")
     )
@@ -125,6 +134,7 @@ def _collect_entrypoints() -> list[Registry[Any]]:
         __pipeline_registry,
         __buffer_registry,
         __ndbuffer_registry,
+        __chunk_grid_registry,
         __chunk_key_encoding_registry,
     ]
 
@@ -156,6 +166,10 @@ def register_buffer(cls: type[Buffer], qualname: str | None = None) -> None:
     __buffer_registry.register(cls, qualname)
 
 
+def register_chunk_grid(key: str, cls: type) -> None:
+    __chunk_grid_registry.register(cls, key)
+
+
 def register_chunk_key_encoding(key: str, cls: type) -> None:
     __chunk_key_encoding_registry.register(cls, key)
 
@@ -296,6 +310,16 @@ def get_ndbuffer_class(reload_config: bool = False) -> type[NDBuffer]:
     )
 
 
+def get_chunk_grid_class(key: str) -> type[ChunkGrid]:
+    __chunk_grid_registry.lazy_load(use_entrypoint_name=True)
+    if key not in __chunk_grid_registry:
+        raise KeyError(
+            f"Chunk grid '{key}' not found in registered chunk grids: "
+            f"{list(__chunk_grid_registry)}."
+        )
+    return __chunk_grid_registry[key]
+
+
 def get_chunk_key_encoding_class(key: str) -> type[ChunkKeyEncoding]:
     __chunk_key_encoding_registry.lazy_load(use_entrypoint_name=True)
     if key not in __chunk_key_encoding_registry:

src/zarr/core/chunk_grids/rectilinear.py

@@ -130,7 +130,7 @@ def test_rectilinear_from_dict_spec_example() -> None:
         },
     }
 
-    grid = RectilinearChunkGrid._from_dict(metadata)  # type: ignore[arg-type]
+    grid = RectilinearChunkGrid.from_dict(metadata)  # type: ignore[arg-type]
 
     assert grid.chunk_shapes == (
         (2, 2, 2),
@@ -151,7 +151,7 @@ def test_rectilinear_from_dict_invalid_kind() -> None:
         },
     }
     with pytest.raises(ValueError, match="Only 'inline' kind is supported"):
-        RectilinearChunkGrid._from_dict(metadata)  # type: ignore[arg-type]
+        RectilinearChunkGrid.from_dict(metadata)  # type: ignore[arg-type]
 
 
 def test_rectilinear_from_dict_missing_chunk_shapes() -> None:
@@ -163,7 +163,7 @@ def test_rectilinear_from_dict_missing_chunk_shapes() -> None:
         },
     }
     with pytest.raises(ValueError, match="must contain 'chunk_shapes'"):
-        RectilinearChunkGrid._from_dict(metadata)  # type: ignore[arg-type]
+        RectilinearChunkGrid.from_dict(metadata)  # type: ignore[arg-type]
 
 
 def test_rectilinear_to_dict() -> None:
@@ -234,7 +234,7 @@ def test_rectilinear_roundtrip() -> None:
     """Test that to_dict and from_dict are inverses"""
     original = RectilinearChunkGrid(chunk_shapes=[[1, 2, 3], [4, 5]])
     metadata = original.to_dict()
-    reconstructed = RectilinearChunkGrid._from_dict(metadata)
+    reconstructed = RectilinearChunkGrid.from_dict(metadata)
 
     assert reconstructed.chunk_shapes == original.chunk_shapes
 
@@ -331,7 +331,7 @@ def test_roundtrip_with_compression() -> None:
     assert metadata["configuration"]["chunk_shapes"] == [[[10, 6]], [[5, 5]]]  # type: ignore[call-overload,index]
 
     # Deserialize from dict
-    grid2 = RectilinearChunkGrid._from_dict(metadata)
+    grid2 = RectilinearChunkGrid.from_dict(metadata)
 
     # Verify the expanded chunk_shapes match
     assert grid2.chunk_shapes == grid1.chunk_shapes