funlib.persistence 0.6.1__tar.gz → 0.8.0__tar.gz

{funlib_persistence-0.6.1 → funlib_persistence-0.8.0}/PKG-INFO

@@ -1,16 +1,15 @@
 Metadata-Version: 2.4
 Name: funlib.persistence
-Version: 0.6.1
+Version: 0.8.0
 Summary: Interfaces for data (arrays and graphs) and storage formats (databases and file formats)
 Author-email: William Patton <pattonw@hhmi.org>, Jan Funke <funkej@hhmi.org>
 License: MIT
 Classifier: Programming Language :: Python :: 3
-Requires-Python: >=3.10
+Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: zarr<3,>=2
-Requires-Dist: numcodecs<0.16.0,>0.13
-Requires-Dist: iohub>=0.2.0b0
+Requires-Dist: zarr<4,>=3
+Requires-Dist: iohub>=0.3.0a5
 Requires-Dist: funlib.geometry>=0.3.0
 Requires-Dist: networkx>=3.0.0
 Requires-Dist: pymongo>=4.0.0
@@ -18,13 +17,13 @@ Requires-Dist: numpy>=2.0.0
 Requires-Dist: pydantic>=2.0.0
 Requires-Dist: dask>=2024.0.0
 Requires-Dist: toml>=0.10.0
-Requires-Dist: psycopg2-binary>=2.9.5
+Requires-Dist: psycopg2-binary>=2.9.11
 Provides-Extra: dev
 Requires-Dist: coverage>=7.7.1; extra == "dev"
-Requires-Dist: mypy>=1.15.0; extra == "dev"
 Requires-Dist: pytest>=8.3.5; extra == "dev"
 Requires-Dist: pytest-mock>=3.14.0; extra == "dev"
 Requires-Dist: ruff>=0.11.2; extra == "dev"
+Requires-Dist: ty>=0.0.16; extra == "dev"
 Requires-Dist: types-networkx; extra == "dev"
 Requires-Dist: types-psycopg2; extra == "dev"
 Requires-Dist: types-toml; extra == "dev"
@@ -32,7 +31,7 @@ Dynamic: license-file
 
 [![tests](https://github.com/funkelab/funlib.persistence/actions/workflows/tests.yaml/badge.svg)](https://github.com/funkelab/funlib.persistence/actions/workflows/tests.yaml)
 [![ruff](https://github.com/funkelab/funlib.persistence/actions/workflows/ruff.yaml/badge.svg)](https://github.com/funkelab/funlib.persistence/actions/workflows/ruff.yaml)
-[![mypy](https://github.com/funkelab/funlib.persistence/actions/workflows/mypy.yaml/badge.svg)](https://github.com/funkelab/funlib.persistence/actions/workflows/mypy.yaml)
+[![ty](https://github.com/funkelab/funlib.persistence/actions/workflows/ty.yaml/badge.svg)](https://github.com/funkelab/funlib.persistence/actions/workflows/ty.yaml)
 [![pypi](https://github.com/funkelab/funlib.persistence/actions/workflows/publish.yaml/badge.svg)](https://pypi.org/project/funlib.persistence/)
 
 # funlib.persistence

{funlib_persistence-0.6.1 → funlib_persistence-0.8.0}/README.md

@@ -1,6 +1,6 @@
 [![tests](https://github.com/funkelab/funlib.persistence/actions/workflows/tests.yaml/badge.svg)](https://github.com/funkelab/funlib.persistence/actions/workflows/tests.yaml)
 [![ruff](https://github.com/funkelab/funlib.persistence/actions/workflows/ruff.yaml/badge.svg)](https://github.com/funkelab/funlib.persistence/actions/workflows/ruff.yaml)
-[![mypy](https://github.com/funkelab/funlib.persistence/actions/workflows/mypy.yaml/badge.svg)](https://github.com/funkelab/funlib.persistence/actions/workflows/mypy.yaml)
+[![ty](https://github.com/funkelab/funlib.persistence/actions/workflows/ty.yaml/badge.svg)](https://github.com/funkelab/funlib.persistence/actions/workflows/ty.yaml)
 [![pypi](https://github.com/funkelab/funlib.persistence/actions/workflows/publish.yaml/badge.svg)](https://pypi.org/project/funlib.persistence/)
 
 # funlib.persistence

{funlib_persistence-0.6.1 → funlib_persistence-0.8.0}/pyproject.toml

@@ -13,17 +13,13 @@ authors = [
 ]
 dynamic = ['version']
 
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 classifiers = ["Programming Language :: Python :: 3"]
 keywords = []
 
 dependencies = [
-    "zarr>=2,<3",
-    # ImportError: cannot import name 'cbuffer_sizes' from 'numcodecs.blosc'
-    # We can pin zarr to >2.18.7 but then we have to drop python 3.10
-    # pin numcodecs to avoid breaking change
-    "numcodecs>0.13,<0.16.0",
-    "iohub>=0.2.0b0",
+    "zarr>=3,<4",
+    "iohub>=0.3.0a5",
     "funlib.geometry>=0.3.0",
     "networkx>=3.0.0",
     "pymongo>=4.0.0",
@@ -31,7 +27,7 @@ dependencies = [
     "pydantic>=2.0.0",
     "dask>=2024.0.0",
     "toml>=0.10.0",
-    "psycopg2-binary>=2.9.5",
+    "psycopg2-binary>=2.9.11",
 ]
 
 [tool.setuptools.dynamic]
@@ -40,10 +36,10 @@ version = { attr = "funlib.persistence.__version__" }
 [project.optional-dependencies]
 dev = [
     "coverage>=7.7.1",
-    "mypy>=1.15.0",
     "pytest>=8.3.5",
     "pytest-mock>=3.14.0",
     "ruff>=0.11.2",
+    "ty>=0.0.16",
     "types-networkx",
     "types-psycopg2",
     "types-toml",
@@ -56,10 +52,8 @@ lint.select = ["F", "W", "I001"]
 [tool.setuptools.package-data]
 "funlib.persistence" = ["py.typed"]
 
-[tool.mypy]
-explicit_package_bases = true
-
 # # module specific overrides
 [[tool.mypy.overrides]]
 module = ["zarr.*", "iohub.*"]
 ignore_missing_imports = true
+

{funlib_persistence-0.6.1 → funlib_persistence-0.8.0/src}/funlib/persistence/__init__.py

@@ -1,4 +1,4 @@
 from .arrays import Array, open_ds, prepare_ds, open_ome_ds, prepare_ome_ds  # noqa
 
-__version__ = "0.6.1"
+__version__ = "0.8.0"
 __version_info__ = tuple(int(i) for i in __version__.split("."))

{funlib_persistence-0.6.1 → funlib_persistence-0.8.0/src}/funlib/persistence/arrays/array.py

@@ -6,9 +6,8 @@ from typing import Any, Optional, Union
 import dask.array as da
 import numpy as np
 from dask.array.optimization import fuse_slice
-from zarr import Array as ZarrArray
-
 from funlib.geometry import Coordinate, Roi
+from zarr import Array as ZarrArray
 
 from .freezable import Freezable
 from .lazy_ops import LazyOp
@@ -75,6 +74,18 @@ class Array(Freezable):
             offset, voxel_size, axis_names, units, and types. Metadata
             can either be passed in or read from array attributes.
 
+        scheduler (``Optional[str]``):
+
+            The dask scheduler used to compute reads (`__getitem__` and
+            `to_ndarray`); writes do not go through dask. `None` (the default)
+            leaves the choice to dask, which uses its threaded scheduler for
+            arrays. Pass `"single-threaded"` when the caller already
+            parallelises array access (one process or thread per block, as
+            daisy does), where dask's own thread pool only adds contention
+            and can deadlock. See
+            https://docs.dask.org/en/stable/scheduler-overview.html for the
+            available schedulers.
+
     """
 
     data: da.Array
@@ -90,7 +101,9 @@ class Array(Freezable):
         chunks: Optional[Union[int, Sequence[int], str]] = "auto",
         lazy_op: Optional[LazyOp] = None,
         strict_metadata: bool = False,
+        scheduler: Optional[str] = None,
     ):
+        self.scheduler = scheduler
         if not isinstance(data, da.Array):
             self.data = da.from_array(data, chunks=chunks)
         else:
@@ -111,11 +124,13 @@ class Array(Freezable):
         # only used if not reading from zarr and there is no built in `.attrs`
         self._attrs: dict[str, Any] = {}
 
+        # whether the view's chunks still sit on the stored chunk grid (see
+        # `_keeps_chunk_grid`); `chunk_shape` is reported only while they do
+        self._on_chunk_grid = True
+        self.lazy_ops: list[LazyOp] = []
         if lazy_op is not None:
             self.apply_lazy_ops(lazy_op)
-
-        lazy_ops = [] if lazy_op is None else [lazy_op]
-        self.lazy_ops = lazy_ops
+            self.lazy_ops = [lazy_op]
 
         self.freeze()
 
@@ -134,8 +149,66 @@ class Array(Freezable):
             return self._attrs
 
     @property
-    def chunk_shape(self) -> Coordinate:
-        return Coordinate(self.data.chunksize)
+    def chunk_shape(self) -> Optional[Coordinate]:
+        """The write-safe chunk shape: the shard shape if sharded, else the chunk
+        shape (an ``open_ds(chunks=...)`` override is reported as given). ``None``
+        for a view that is not writeable or that lazy operations moved off the
+        stored chunk grid."""
+        if not self.is_writeable or not self._on_chunk_grid:
+            return None
+        shards = self.shard_shape
+        return shards if shards is not None else self._dask_chunk_shape()
+
+    @property
+    def inner_chunk_shape(self) -> Optional[Coordinate]:
+        """Chunk shape for efficient reads. Can be smaller than the write-safe
+        :attr:`chunk_shape` in the case of sharding. ``None`` for a view that
+        lazy operations moved off the stored chunk grid."""
+        return self._dask_chunk_shape() if self._on_chunk_grid else None
+
+    @property
+    def shard_shape(self) -> Optional[Coordinate]:
+        """The shard shape in the case of sharded datasets. ``None`` for non
+        sharded or unwritable datasets, or off the stored chunk grid. See
+        https://zarr.readthedocs.io/en/stable/user-guide/performance.html#sharding
+        """
+        shards = getattr(self._source_data, "shards", None)
+        if shards is None or not self.is_writeable or not self._on_chunk_grid:
+            return None
+        return Coordinate(
+            s for s, uncollapsed in zip(shards, self._uncollapsed_dims) if uncollapsed
+        )
+
+    def _dask_chunk_shape(self) -> Coordinate:
+        return Coordinate(self.data.chunksize)  # ty: ignore[unresolved-attribute]
+
+    def _keeps_chunk_grid(self, lazy_op) -> bool:
+        """Whether a lazy op leaves this view's chunks on the stored chunk grid:
+        dropping or truncating axes does, a start, a step, a selection list, a
+        mask or an arbitrary callable does not. A ``Roi`` does if it begins on
+        a chunk boundary."""
+        if isinstance(lazy_op, Roi):
+            if (
+                not all(self.uncollapsed_dims(physical=True))
+                or lazy_op.dims != self.spatial_dims
+            ):
+                return True  # apply_lazy_ops refuses the op, with its own message
+            unit = self.chunk_shape
+            if unit is None:
+                return False
+            physical = [c for c, t in zip(unit, self.types) if t in ["space", "time"]]
+            begin = (lazy_op.offset - self.offset) / self.voxel_size
+            return all(b % c == 0 for b, c in zip(begin, physical))
+        if not isinstance(lazy_op, tuple):
+            lazy_op = (lazy_op,)
+        for a in lazy_op:
+            if isinstance(a, int):
+                continue
+            if isinstance(a, slice):
+                if a.start in (None, 0) and a.step in (None, 1):
+                    continue
+            return False
+        return True
 
     def uncollapsed_dims(self, physical: bool = False) -> list[bool]:
         """
@@ -254,17 +327,18 @@ class Array(Freezable):
         )
 
     def apply_lazy_ops(self, lazy_op):
+        # decided before the op moves the offset, recorded only once it applied
+        on_chunk_grid = self._on_chunk_grid and self._keeps_chunk_grid(lazy_op)
         if self._is_slice(lazy_op):
             if not isinstance(lazy_op, tuple):
                 lazy_op = (lazy_op,)
-            for ii, a in enumerate(lazy_op):
+            # Pair each entry with the dimension it indexes BEFORE collapsing any: the
+            # previous in-place walk re-counted uncollapsed dims after each collapse,
+            # so a tuple with two or more ints collapsed the wrong dimensions.
+            dims = [i for i, uc in enumerate(self._uncollapsed_dims) if uc]
+            for a, i in zip(lazy_op, dims):
                 if isinstance(a, int):
-                    for i, uc in enumerate(self._uncollapsed_dims):
-                        if uc:
-                            if ii == 0:
-                                self._uncollapsed_dims[i] = False
-                                break
-                            ii -= 1
+                    self._uncollapsed_dims[i] = False
             self.data = self.data[lazy_op]
         elif isinstance(lazy_op, Roi):
             assert all(self.uncollapsed_dims(physical=True)), (
@@ -283,6 +357,7 @@ class Array(Freezable):
                 f"LazyOp {lazy_op} is not a supported lazy_op. "
                 f"Supported lazy_ops are: {LazyOp}"
             )
+        self._on_chunk_grid = on_chunk_grid
 
     def lazy_op(self, lazy_op: LazyOp):
         """Apply an lazy operation to this array.
@@ -330,7 +405,9 @@ class Array(Freezable):
                     % (roi, self.roi)
                 )
 
-            return self.data[self.__slices(roi, use_lazy_slices=False)].compute()
+            return self.data[self.__slices(roi, use_lazy_slices=False)].compute(
+                scheduler=self.scheduler
+            )
 
         elif isinstance(key, Coordinate):
             coordinate = key
@@ -339,23 +416,26 @@ class Array(Freezable):
                 raise IndexError("Requested coordinate is not contained in this array.")
 
             index = self.__index(coordinate)
-            return self.data[index].compute()
+            return self.data[index].compute(scheduler=self.scheduler)
 
         else:
-            return self.data[key].compute()
+            return self.data[key].compute(scheduler=self.scheduler)
 
-    def __setitem__(self, key, value: np.ndarray):
-        """Set the data of this array within the given ROI.
+    def __setitem__(self, key: Roi | slice | tuple, value: np.ndarray | float | int):
+        """Set the data of this array.
 
         Args:
 
-            key (`class:Roi`):
+            key (`class:Roi` or any numpy compatible key):
 
-                The ROI to write to.
+                The region to write to. Can be a `Roi` for world-unit indexing,
+                or any numpy-compatible key (e.g. ``np.s_[:]``, a slice, a tuple
+                of slices).
 
-            value (``ndarray``):
+            value (``ndarray`` or scalar):
 
-                The value to write.
+                The value to write. Can be a numpy array or a scalar that will
+                be broadcast.
         """
 
         if self.is_writeable:
@@ -419,7 +499,7 @@ class Array(Freezable):
         if fill_value != 0:
             data[:] = fill_value
 
-        array = Array(data, roi.offset, self.voxel_size)
+        array = Array(data, roi.offset, self.voxel_size, scheduler=self.scheduler)
 
         shared_roi = self.roi.intersect(roi)
 
@@ -436,18 +516,20 @@ class Array(Freezable):
         voxel_roi = (roi - self.offset) / self.voxel_size
 
         if check_chunk_align:
+            chunk_shape = self.chunk_shape
+            assert chunk_shape is not None, "this view has no chunk grid to align to"
             for d in range(roi.dims):
                 end_of_array = roi.get_end()[d] == self.roi.get_end()[d]
 
-                begin_align_with_chunks = voxel_roi.begin[d] % self.chunk_shape[d] == 0
-                shape_align_with_chunks = voxel_roi.shape[d] % self.chunk_shape[d] == 0
+                begin_align_with_chunks = voxel_roi.begin[d] % chunk_shape[d] == 0
+                shape_align_with_chunks = voxel_roi.shape[d] % chunk_shape[d] == 0
 
                 assert begin_align_with_chunks and (
                     shape_align_with_chunks or end_of_array
                 ), (
                     "ROI %s (in voxels: %s) does not align with chunks of "
                     "size %s (mismatch in dimension %d)"
-                    % (roi, voxel_roi, self.chunk_shape, d)
+                    % (roi, voxel_roi, chunk_shape, d)
                 )
 
         roi_slices = tuple(
@@ -474,7 +556,7 @@ class Array(Freezable):
         elif isinstance(lazy_op, list) and all([isinstance(a, int) for a in lazy_op]):
             return True
         elif isinstance(lazy_op, tuple) and all(
-            [self._is_slice(a, writeable) for a in lazy_op]
+            [self._is_slice(a, writeable) for a in lazy_op]  # type: ignore[arg-type]  # ty can't narrow parameterized tuple iteration
         ):
             return True
         elif (
@@ -506,3 +588,56 @@ class Array(Freezable):
                 f"Chunk shape ({self.chunk_shape}) must have the same "
                 f"number of dimensions as the data ({self.shape})"
             )
+
+    def to_pixel_space(
+        self, world_loc: Roi | Coordinate | Sequence[int | float] | np.ndarray
+    ) -> Roi | Coordinate | np.ndarray:
+        """Convert a point or roi in world space into the pixel space of this array.
+        Works on sequences of floats by returning a numpy array that is not guaranteed
+        to be aligned with pixels.
+
+        Args:
+            world_loc (Roi | Coordinate | Sequence[int | float]):  The world space
+                point or roi to convert to pixel space
+
+        Raises:
+            ValueError: If the world location falls outside the array's roi
+
+        Returns:
+            Roi | Coordinate | np.ndarray: The point or roi converted into the pixel
+                space of the array. If the input was a Coordinate or Roi, returns
+                the same type of object. If the input was a sequence of python numbers,
+                returns a numpy array with the result, which may or may not be integers.
+        """
+        if not self.roi.contains(world_loc):
+            raise ValueError(
+                f"Given world location {world_loc} is not included in this array with world ROI {self.roi}"
+            )
+        if isinstance(world_loc, (Roi, Coordinate)):
+            return (world_loc - self.roi.offset) / self.voxel_size
+        else:
+            return (np.array(world_loc) - np.array(self.roi.offset)) / np.array(
+                self.voxel_size
+            )
+
+    def to_world_space(
+        self, pixel_loc: Roi | Coordinate | Sequence[int | float] | np.ndarray
+    ) -> Roi | Coordinate:
+        """Convert a point or roi from pixel space in this array to the world
+        coordinate system defined by this array's roi and voxel size.
+        Works on Coordiantes and Rois, as well as sequences of ints or floats.
+
+        Args:
+            pixel_loc (Roi | Coordinate | Sequence[int | float]): A location in pixel space
+                of this array's data.
+
+        Returns:
+            Roi | Coordinate: The world location of the given pixel location. If the input
+                was a Roi or Coordinate, return the same type of object. If it was a
+                sequence of python numbers, return a numpy array that may be int or float.
+        """
+        if isinstance(pixel_loc, (Roi, Coordinate)):
+            return pixel_loc * self.voxel_size + self.roi.offset
+        return np.array(pixel_loc) * np.array(self.voxel_size) + np.array(
+            self.roi.offset
+        )