hoa-tools 1.0.9__py3-none-any.whl → 2.0.0b1__py3-none-any.whl

hoa_tools/_n5.py

@@ -1,10 +1,436 @@
+import codecs
+import json
+import numbers
+import re
+import struct
+import sys
+from collections.abc import AsyncIterator, Iterable, Sequence
 from typing import Any
 
-import zarr.n5
-from zarr.storage import FSStore
+import numpy as np
+import numpy.typing as npt
+from numcodecs.abc import Codec
+from numcodecs.registry import get_codec, register_codec
+from zarr.abc.store import ByteRequest
+from zarr.core.buffer import Buffer, BufferPrototype, default_buffer_prototype
+from zarr.storage import FsspecStore
 
+zarr_group_meta_key = ".zgroup"
+zarr_array_meta_key = ".zarray"
+zarr_attrs_key = ".zattrs"
+n5_attrs_key = "attributes.json"
+N5_FORMAT = "2.0.0"
+ZARR_FORMAT = 2
 
-# Override __init__ to prevent the deprecation warning from showing up
-class N5FSStore(zarr.n5.N5FSStore):  # type: ignore[misc]
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
-        FSStore.__init__(self, *args, **kwargs)
+zarr_to_n5_keys = [
+    ("chunks", "blockSize"),
+    ("dtype", "dataType"),
+    ("compressor", "compression"),
+    ("shape", "dimensions"),
+]
+n5_keywords = ["n5", "dataType", "dimensions", "blockSize", "compression"]
+
+
+class N5FSStore(FsspecStore):
+    async def get(
+        self,
+        key: str,
+        prototype: BufferPrototype,
+        byte_range: ByteRequest | None = None,
+    ) -> Buffer | None:
+        if key.endswith(zarr_group_meta_key):
+            key_new = key.replace(zarr_group_meta_key, n5_attrs_key)
+            value = group_metadata_to_zarr(await self._load_n5_attrs(key_new))
+
+            return prototype.buffer.from_bytes(json_dumps(value))
+
+        if key.endswith(zarr_array_meta_key):
+            key_new = key.replace(zarr_array_meta_key, n5_attrs_key)
+            top_level = key == zarr_array_meta_key
+            value = array_metadata_to_zarr(
+                await self._load_n5_attrs(key_new), top_level=top_level
+            )
+            return prototype.buffer.from_bytes(json_dumps(value))
+
+        if key.endswith(zarr_attrs_key):
+            key_new = key.replace(zarr_attrs_key, n5_attrs_key)
+            value = attrs_to_zarr(await self._load_n5_attrs(key_new))
+
+            if len(value) == 0:
+                raise KeyError(key_new)
+            return prototype.buffer.from_bytes(json_dumps(value))
+
+        key_new = invert_chunk_coords(key) if is_chunk_key(key) else key
+
+        return await super().get(key_new, prototype=prototype, byte_range=byte_range)
+
+    async def get_partial_values(
+        self,
+        prototype: BufferPrototype,
+        key_ranges: Iterable[tuple[str, ByteRequest | None]],
+    ) -> list[Buffer | None]:
+        raise NotImplementedError
+
+    async def exists(self, key: str) -> bool:
+        if key.endswith(zarr_group_meta_key):
+            key_new = key.replace(zarr_group_meta_key, n5_attrs_key)
+            if not await super().exists(key_new):
+                return False
+            # group if not a dataset (attributes do not contain 'dimensions')
+            return "dimensions" not in await self._load_n5_attrs(key_new)
+
+        if key.endswith(zarr_array_meta_key):
+            key_new = key.replace(zarr_array_meta_key, n5_attrs_key)
+            # array if attributes contain 'dimensions'
+            return "dimensions" in await self._load_n5_attrs(key_new)
+
+        if key.endswith(zarr_attrs_key):
+            key_new = key.replace(zarr_attrs_key, n5_attrs_key)
+            return await self._contains_attrs(key_new)
+
+        key_new = invert_chunk_coords(key) if is_chunk_key(key) else key
+
+        return await super().exists(key_new)
+
+    @property
+    def supports_writes(self) -> bool:  # type: ignore[override]
+        return False
+
+    async def set(
+        self,
+        key: str,
+        value: Buffer,
+        byte_range: tuple[int, int] | None = None,
+    ) -> None:
+        raise NotImplementedError
+
+    @property
+    def supports_deletes(self) -> bool:  # type: ignore[override]
+        return False
+
+    async def delete(self, key: str) -> None:
+        raise NotImplementedError
+
+    @property
+    def supports_listing(self) -> bool:  # type: ignore[override]
+        return False
+
+    def list(self) -> AsyncIterator[str]:
+        raise NotImplementedError
+
+    def list_prefix(self, prefix: str) -> AsyncIterator[str]:
+        raise NotImplementedError
+
+    def list_dir(self, prefix: str) -> AsyncIterator[str]:
+        # This method should be async, like overridden methods in child classes.
+        # However, that's not straightforward:
+        # https://stackoverflow.com/questions/68905848
+        raise NotImplementedError
+
+    async def _load_n5_attrs(self, path: str) -> dict[str, Any]:
+        try:
+            s = await super().get(path, prototype=default_buffer_prototype())
+            if s is None:
+                raise RuntimeError(f"No N5 attributes at path {path}")
+            return json_loads(s.to_bytes())
+        except KeyError:
+            return {}
+
+    async def _contains_attrs(self, path: str | None) -> bool:
+        if path is None:
+            attrs_key = n5_attrs_key
+        elif not path.endswith(n5_attrs_key):
+            attrs_key = f"{path}/{n5_attrs_key}"
+        else:
+            attrs_key = path
+
+        attrs = attrs_to_zarr(await self._load_n5_attrs(attrs_key))
+        return len(attrs) > 0
+
+
+# match strings of numbers with "." between
+# (e.g., "1.2.4", "1.2", "5")
+_prog_ckey = re.compile(r"^(\d+)(\.\d+)+$")
+
+
+def is_chunk_key(key: str) -> bool:
+    rv = False
+    segments = list(key.split("/"))
+    if segments:
+        last_segment = segments[-1]
+        rv = bool(_prog_ckey.match(last_segment))
+    return rv
+
+
+def invert_chunk_coords(key: str) -> str:
+    segments = list(key.split("/"))
+    if segments:
+        last_segment = segments[-1]
+        if _prog_ckey.match(last_segment):
+            coords = list(last_segment.split("."))
+            last_segment = "/".join(coords[::-1])
+            segments = [*segments[:-1], last_segment]
+            key = "/".join(segments)
+    return key
+
+
+def group_metadata_to_zarr(group_metadata: dict[str, Any]) -> dict[str, Any]:
+    """Convert group metadata from N5 to zarr format."""
+    return {"zarr_format": ZARR_FORMAT}
+
+
+def array_metadata_to_zarr(
+    array_metadata: dict[str, Any], *, top_level: bool = False
+) -> dict[str, Any]:
+    """
+    Convert array metadata from N5 to zarr format.
+
+    If the `top_level` keyword argument is True,
+    then the `N5` key will be removed from metadata
+    """
+    for t, f in zarr_to_n5_keys:
+        array_metadata[t] = array_metadata.pop(f)
+    if top_level:
+        array_metadata.pop("n5")
+    array_metadata["zarr_format"] = ZARR_FORMAT
+
+    array_metadata["shape"] = array_metadata["shape"][::-1]
+    array_metadata["chunks"] = array_metadata["chunks"][::-1]
+    array_metadata["fill_value"] = 0  # also if None was requested
+    array_metadata["order"] = "C"
+    array_metadata["filters"] = None
+    array_metadata["dimension_separator"] = "."
+    array_metadata["dtype"] = np.dtype(array_metadata["dtype"]).str
+
+    compressor_config = array_metadata["compressor"]
+    compressor_config = compressor_config_to_zarr(compressor_config)
+    array_metadata["compressor"] = {
+        "id": N5ChunkWrapper.codec_id,
+        "compressor_config": compressor_config,
+        "dtype": array_metadata["dtype"],
+        "chunk_shape": array_metadata["chunks"],
+    }
+
+    return array_metadata
+
+
+def attrs_to_zarr(attrs: dict[str, Any]) -> dict[str, Any]:
+    """
+    Get all zarr attributes from an N5 attributes dictionary.
+
+       (i.e.,
+    all non-keyword attributes).
+
+    """
+    # remove all N5 keywords
+    for n5_key in n5_keywords:
+        attrs.pop(n5_key, None)
+
+    return attrs
+
+
+def json_loads(s: bytes | str) -> dict[str, Any]:
+    """Read JSON in a consistent way."""
+    return json.loads(ensure_text(s, "utf-8"))  # type: ignore[no-any-return]
+
+
+def json_dumps(o: Any) -> bytes:
+    """Write JSON in a consistent, human-readable way."""
+    return json.dumps(
+        o,
+        indent=4,
+        sort_keys=True,
+        ensure_ascii=True,
+        separators=(",", ": "),
+        cls=NumberEncoder,
+    ).encode("ascii")
+
+
+class NumberEncoder(json.JSONEncoder):
+    def default(self, o: Any) -> float:
+        # See json.JSONEncoder.default docstring for explanation
+        # This is necessary to encode numpy dtype
+        if isinstance(o, numbers.Integral):
+            return int(o)
+        if isinstance(o, numbers.Real):
+            return float(o)
+        return json.JSONEncoder.default(self, o)  # type: ignore[no-any-return]
+
+
+def ensure_text(s: bytes | str, encoding: str = "utf-8") -> str:
+    if not isinstance(s, str):
+        return codecs.decode(s, encoding)
+    return s
+
+
+def compressor_config_to_zarr(
+    compressor_config: dict[str, Any],
+) -> dict[str, Any] | None:
+    codec_id = compressor_config["type"]
+    zarr_config = {"id": codec_id}
+
+    if codec_id == "bzip2":
+        zarr_config["id"] = "bz2"
+        zarr_config["level"] = compressor_config["blockSize"]
+
+    elif codec_id == "blosc":
+        zarr_config["cname"] = compressor_config["cname"]
+        zarr_config["clevel"] = compressor_config["clevel"]
+        zarr_config["shuffle"] = compressor_config["shuffle"]
+        zarr_config["blocksize"] = compressor_config["blocksize"]
+
+    elif codec_id == "lzma":
+        zarr_config["format"] = compressor_config["format"]
+        zarr_config["check"] = compressor_config["check"]
+        zarr_config["preset"] = compressor_config["preset"]
+        zarr_config["filters"] = compressor_config["filters"]
+
+    elif codec_id == "xz":
+        zarr_config["id"] = "lzma"
+        zarr_config["format"] = 1  # lzma.FORMAT_XZ
+        zarr_config["check"] = -1
+        zarr_config["preset"] = compressor_config["preset"]
+        zarr_config["filters"] = None
+
+    elif codec_id == "gzip":
+        if compressor_config.get("useZlib"):
+            zarr_config["id"] = "zlib"
+            zarr_config["level"] = compressor_config["level"]
+        else:
+            zarr_config["id"] = "gzip"
+            zarr_config["level"] = compressor_config["level"]
+
+    elif codec_id == "raw":
+        return None
+
+    else:
+        zarr_config.update({k: v for k, v in compressor_config.items() if k != "type"})
+
+    return zarr_config
+
+
+class N5ChunkWrapper(Codec):  # type: ignore[misc]
+    codec_id = "n5_wrapper"
+    chunk_shape: tuple[int, ...]
+    dtype: np.dtype
+
+    def __init__(
+        self,
+        dtype: npt.DTypeLike,
+        chunk_shape: Sequence[int],
+        compressor_config: dict[str, Any] | None = None,
+        compressor: Codec | None = None,
+    ):
+        self.dtype = np.dtype(dtype)
+        self.chunk_shape = tuple(chunk_shape)
+        # is the dtype a little endian format?
+        self._little_endian = self.dtype.byteorder == "<" or (
+            self.dtype.byteorder == "=" and sys.byteorder == "little"
+        )
+
+        if compressor is not None:
+            if compressor_config is not None:
+                raise ValueError(
+                    "Only one of compressor_config or compressor should be given."
+                )
+            compressor_config = compressor.get_config()
+
+        if compressor_config is None or compressor_config["id"] == "raw":
+            self.compressor_config = None
+            self._compressor = None
+        else:
+            self._compressor = get_codec(compressor_config)
+            self.compressor_config = self._compressor.get_config()
+
+    def get_config(self) -> dict[str, Any]:
+        return {"id": self.codec_id, "compressor_config": self.compressor_config}
+
+    def encode(self, chunk: npt.NDArray[Any]) -> bytes:
+        assert chunk.flags.c_contiguous, "Chunk is not C contiguous"  # noqa: S101
+
+        header = self._create_header(chunk)
+        chunk = self._to_big_endian(chunk)
+
+        if self._compressor:
+            return header + self._compressor.encode(chunk)  # type: ignore[no-any-return]
+        return header + chunk.tobytes(order="A")
+
+    def decode(
+        self, chunk: bytes, out: npt.NDArray[Any] | None = None
+    ) -> npt.NDArray[Any]:
+        len_header, chunk_shape = self._read_header(chunk)
+        chunk = chunk[len_header:]
+
+        if out is not None:
+            # out should only be used if we read a complete chunk
+            assert chunk_shape == self.chunk_shape, (  # noqa: S101
+                f"Expected chunk of shape {self.chunk_shape}, found {chunk_shape}"
+            )
+
+            if self._compressor:
+                self._compressor.decode(chunk, out)
+            else:
+                raise RuntimeError("Can't handle case with no compressor")
+                # ndarray_copy(chunk, out)
+
+            # we can byteswap in-place
+            if self._little_endian:
+                out.byteswap(inplace=True)
+
+            return out
+
+        if self._compressor:
+            chunk = self._compressor.decode(chunk)
+
+        # more expensive byteswap
+        chunk = self._from_big_endian(chunk)  # type: ignore[assignment]
+
+        # read partial chunk
+        if chunk_shape != self.chunk_shape:
+            chunk = np.frombuffer(chunk, dtype=self.dtype)  # type: ignore[assignment]
+            chunk = chunk.reshape(chunk_shape)  # type: ignore[attr-defined]
+            complete_chunk = np.zeros(self.chunk_shape, dtype=self.dtype)
+            target_slices = tuple(slice(0, s) for s in chunk_shape)
+            complete_chunk[target_slices] = chunk
+            chunk = complete_chunk  # type: ignore[assignment]
+
+        return chunk  # type: ignore[return-value]
+
+    @staticmethod
+    def _create_header(chunk: npt.NDArray[Any]) -> bytes:
+        mode = struct.pack(">H", 0)
+        num_dims = struct.pack(">H", len(chunk.shape))
+        shape = b"".join(struct.pack(">I", d) for d in chunk.shape[::-1])
+
+        return mode + num_dims + shape
+
+    @staticmethod
+    def _read_header(chunk: bytes) -> tuple[int, tuple[int, ...]]:
+        num_dims = struct.unpack(">H", chunk[2:4])[0]
+        shape = tuple(
+            struct.unpack(">I", chunk[i : i + 4])[0]
+            for i in range(4, num_dims * 4 + 4, 4)
+        )[::-1]
+
+        len_header = 4 + num_dims * 4
+
+        return len_header, shape
+
+    def _to_big_endian(self, data: npt.NDArray[Any]) -> npt.NDArray[Any]:
+        # assumes data is ndarray
+
+        if self._little_endian:
+            return data.byteswap()
+        return data
+
+    def _from_big_endian(self, data: bytes) -> npt.NDArray[Any]:
+        # assumes data is byte array in big endian
+
+        if not self._little_endian:
+            return data  # type:ignore[return-value]
+
+        a = np.frombuffer(data, self.dtype.newbyteorder(">"))
+        return a.astype(self.dtype)
+
+
+register_codec(N5ChunkWrapper, N5ChunkWrapper.codec_id)

hoa_tools/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '1.0.9'
-__version_tuple__ = version_tuple = (1, 0, 9)
+__version__ = version = '2.0.0b1'
+__version_tuple__ = version_tuple = (2, 0, 0, 'b1')
 
 __commit_id__ = commit_id = None

hoa_tools/dataset.py

@@ -20,8 +20,7 @@ import gcsfs
 import networkx as nx
 import numpy as np
 import xarray as xr
-import zarr.core
-import zarr.n5
+import zarr.abc.store
 import zarr.storage
 
 from hoa_tools._n5 import N5FSStore
@@ -63,7 +62,7 @@ class Dataset(HOAMetadata):
         """
         Whether this dataset contains the whole organ or not.
         """
-        return self.voi.startswith("complete")
+        return self.voi.startswith("complete") or self.voi.startswith("overview")
 
     @property
     def is_zoom(self) -> bool:
@@ -148,15 +147,21 @@ class Dataset(HOAMetadata):
 
         # n5://gs://ucl-hip-ct-35a68e99feaae8932b1d44da0358940b/S-20-29/heart/2.5um_VOI-01_bm05/
         bucket, path = gcs_path.split("/", maxsplit=1)
-        fs = gcsfs.GCSFileSystem(project=bucket, token="anon", access="read_only")  # noqa: S106
+        fs = gcsfs.GCSFileSystem(
+            project="ucl-hip-ct",
+            token="anon",  # noqa: S106
+            access="read_only",
+            asynchronous=True,
+        )
+        store: zarr.abc.store.Store
         if self._remote_fmt == "n5":
-            store = N5FSStore(url=bucket, fs=fs, mode="r")
+            store = N5FSStore(fs=fs, path=f"/{bucket}", read_only=True)
         elif self._remote_fmt == "zarr":
-            store = zarr.storage.FSStore(url=bucket, fs=fs, mode="r")
+            store = zarr.storage.FsspecStore(fs=fs, path=f"/{bucket}", read_only=True)
 
-        return zarr.open_group(store, mode="r", path=path)
+        return zarr.open_group(store, mode="r", path=path, zarr_format=2)
 
-    def _remote_array(self, *, downsample_level: int) -> zarr.core.Array:
+    def _remote_array(self, *, downsample_level: int) -> zarr.Array:
         """
         Get an object representing the data array in the remote Google Cloud Store.
         """
@@ -168,7 +173,7 @@ class Dataset(HOAMetadata):
             key = f"s{downsample_level}"
         else:
             key = f"{downsample_level}"
-        return self._remote_store[key]
+        return self._remote_store[key]  # type: ignore[return-value]
 
     def data_array(self, *, downsample_level: int) -> xr.DataArray:
         """
@@ -216,7 +221,7 @@ def _load_datasets_from_files(data_dir: Path) -> dict[str, Dataset]:
     }
     if len(datasets) == 0:
         raise FileNotFoundError(
-            f"Did not find any dataset metadata files at {data_dir}"  # noqa: EM102
+            f"Did not find any dataset metadata files at {data_dir}"
         )
     return datasets
 

{hoa_tools-1.0.9.dist-info → hoa_tools-2.0.0b1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hoa-tools
-Version: 1.0.9
+Version: 2.0.0b1
 Summary: Tools for working with the Human Organ Atlas
 Author-email: David Stansby <d.stansby@ucl.ac.uk>
 License: Copyright (c) 2024, David Stansby
@@ -49,7 +49,7 @@ Requires-Dist: pandas>=2
 Requires-Dist: pydantic>=2
 Requires-Dist: simpleitk
 Requires-Dist: xarray
-Requires-Dist: zarr<3,>2.16
+Requires-Dist: zarr>=3
 Provides-Extra: dev
 Requires-Dist: build; extra == "dev"
 Requires-Dist: mypy; extra == "dev"

{hoa_tools-1.0.9.dist-info → hoa_tools-2.0.0b1.dist-info}/RECORD

@@ -1,7 +1,7 @@
 hoa_tools/__init__.py,sha256=wc7Yn5CXR3hxr2JyV3eEyCx3p8UZdknXDFJPSthKE_g,87
-hoa_tools/_n5.py,sha256=NGDJBbuyp1scGZ9fjuH7QqAQBpax9nxg2MhwhnV698M,310
-hoa_tools/_version.py,sha256=JXTThZsIEQNG8lSfLsQqv8iVrLso3IkPevWFvCathJU,704
-hoa_tools/dataset.py,sha256=VMkp6_YR1DcEASgLeH5bkutlBML5ApeRGX2AlFPNs-A,10028
+hoa_tools/_n5.py,sha256=7hl3Tf5n1hKVTM-l593w7eeYzQRenk_znFdyWnif-kw,14400
+hoa_tools/_version.py,sha256=fz65no5MBu9R7ellXZDEC0QTY0L4mFHp5ebemVVbIqg,712
+hoa_tools/dataset.py,sha256=DiQl-1lBZhG1M-Pbtalji5aqWe6Srb__VpeKo_EHxp0,10228
 hoa_tools/inventory.py,sha256=44IbE1eLFTetAeCadpBFYXXIwcX5fagq76zSMLWBOmE,613
 hoa_tools/metadata.py,sha256=TwcTsXjwnhBdeSsb8TC_9XznQrfwXLOPi-teaZ9-UtI,23011
 hoa_tools/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -306,8 +306,8 @@ hoa_tools/data/metadata/metadata/S-20-29_lung_left_complete-organ_25.31um_bm05.j
 hoa_tools/data/metadata/metadata/S-21-33_brain_complete-organ_19.57um_bm18.json,sha256=rrodLFZXIYWktEifFiRCy1EGZOSZw57Tm6_yliWagW4,7256
 hoa_tools/data/metadata/metadata/S-21-46_brain_complete-organ_19.57um_bm18.json,sha256=p4n2FRkBdoddNnI_fw1veRbMcrhOGjWcMcgPYJ73S6M,7415
 hoa_tools/data/metadata/metadata/S-22-16_brain_complete-organ_19.59um_bm18.json,sha256=aPKtOTz3qTPa0gvVuOi9kDDsBk0qmHNg8HjKb43gNbM,7426
-hoa_tools-1.0.9.dist-info/licenses/LICENSE.md,sha256=CU-KGODbkoeDUXYPeLgpicDrFTffLQECao-wTKJrXFU,1480
-hoa_tools-1.0.9.dist-info/METADATA,sha256=2dlcUrhrNxLbt0uZ_DgGHGq_4n1KQyFbg1p-fC9BHpg,5136
-hoa_tools-1.0.9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-hoa_tools-1.0.9.dist-info/top_level.txt,sha256=NkWbNX-2byZN6hcVyowOVv0aiRNdH0uXpOkA0SFHDgY,10
-hoa_tools-1.0.9.dist-info/RECORD,,
+hoa_tools-2.0.0b1.dist-info/licenses/LICENSE.md,sha256=CU-KGODbkoeDUXYPeLgpicDrFTffLQECao-wTKJrXFU,1480
+hoa_tools-2.0.0b1.dist-info/METADATA,sha256=cSxXHiPS_c1Wn8DK7q6YnbR7yDju-ce7V7TwpDjMCHQ,5133
+hoa_tools-2.0.0b1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+hoa_tools-2.0.0b1.dist-info/top_level.txt,sha256=NkWbNX-2byZN6hcVyowOVv0aiRNdH0uXpOkA0SFHDgY,10
+hoa_tools-2.0.0b1.dist-info/RECORD,,