feat: add read_frame and loop_indices to public api (#181)

* feat: add stuff to public api

* coverage

src/nd2/_util.py

@@ -14,6 +14,7 @@
     from typing_extensions import Final
 
     from nd2.readers import ND2Reader
+    from nd2.structures import ExpLoop
 
     StrOrPath = Union[str, PathLike]
     FileOrBinaryIO = Union[StrOrPath, BinaryIO]
@@ -234,3 +235,22 @@ def convert_dict_of_lists_to_records(
         }
         for row_data in zip(*columns.values())
     ]
+
+
+def loop_indices(experiment: list[ExpLoop]) -> list[dict[str, int]]:
+    """Return a list of dicts of loop indices for each frame.
+
+    Examples
+    --------
+    >>> with nd2.ND2File("path/to/file.nd2") as f:
+    ...     f.loop_indices()
+    [
+        {'Z': 0, 'T': 0, 'C': 0},
+        {'Z': 0, 'T': 0, 'C': 1},
+        {'Z': 0, 'T': 0, 'C': 2},
+        ...
+    ]
+    """
+    axes = [AXIS._MAP[x.type] for x in experiment]
+    indices = product(*(range(x.count) for x in experiment))
+    return [dict(zip(axes, x)) for x in indices]

src/nd2/nd2file.py

@@ -874,7 +874,7 @@ def asarray(self, position: int | None = None) -> np.ndarray:
                 seqs = self._seq_index_from_coords(coords)  # type: ignore
                 final_shape[pidx] = 1
 
-        arr: np.ndarray = np.stack([self._get_frame(i) for i in seqs])
+        arr: np.ndarray = np.stack([self.read_frame(i) for i in seqs])
         return arr.reshape(final_shape)
 
     def __array__(self) -> np.ndarray:
@@ -954,7 +954,7 @@ def _dask_block(self, copy: bool, block_id: tuple[int]) -> np.ndarray:
                             f"Cannot get chunk {block_id} for single frame image."
                         )
                     idx = 0
-                data = self._get_frame(int(idx))  # type: ignore
+                data = self.read_frame(int(idx))  # type: ignore
                 data = data.copy() if copy else data
                 return data[(np.newaxis,) * ncoords]
             finally:
@@ -1052,11 +1052,36 @@ def _coord_shape(self) -> tuple[int, ...]:
     def _frame_count(self) -> int:
         return int(np.prod(self._coord_shape))
 
-    def _get_frame(self, index: SupportsInt) -> np.ndarray:
-        frame = self._rdr.read_frame(int(index))
+    def _get_frame(self, index: SupportsInt) -> np.ndarray:  # pragma: no cover
+        warnings.warn(
+            'Use of "_get_frame" is deprecated, use the public "read_frame" instead.',
+            stacklevel=2,
+        )
+        return self.read_frame(index)
+
+    def read_frame(self, frame_index: SupportsInt) -> np.ndarray:
+        """Read a single frame from the file, indexed by frame number."""
+        frame = self._rdr.read_frame(int(frame_index))
         frame.shape = self._raw_frame_shape
         return frame.transpose((2, 0, 1, 3)).squeeze()
 
+    @cached_property
+    def loop_indices(self) -> list[dict[str, int]]:
+        """Return a list of dicts of loop indices for each frame.
+
+        Examples
+        --------
+        >>> with nd2.ND2File("path/to/file.nd2") as f:
+        ...     f.loop_indices
+        [
+            {'Z': 0, 'T': 0, 'C': 0},
+            {'Z': 0, 'T': 0, 'C': 1},
+            {'Z': 0, 'T': 0, 'C': 2},
+            ...
+        ]
+        """
+        return _util.loop_indices(self.experiment)
+
     def _expand_coords(self, squeeze: bool = True) -> dict:
         """Return a dict that can be used as the coords argument to xr.DataArray.
 

src/nd2/readers/_modern/modern_reader.py

@@ -3,7 +3,6 @@
 import os
 import warnings
 import zlib
-from itertools import product
 from typing import TYPE_CHECKING, Any, Iterable, Mapping, Sequence, cast
 
 import numpy as np
@@ -79,6 +78,8 @@ def __init__(self, path: FileOrBinaryIO, error_radius: int | None = None) -> Non
         self._raw_text_info: RawTextInfoDict | None = None
         self._raw_image_metadata: RawMetaDict | None = None
 
+        self._loop_indices: list[dict[str, int]] | None = None
+
     @property
     def chunkmap(self) -> ChunkMap:
         """Load and return the chunkmap.
@@ -234,9 +235,9 @@ def loop_indices(self) -> list[dict[str, int]]:
             ...
         ]
         """
-        axes = [_util.AXIS._MAP[x.type] for x in self.experiment()]
-        indices = product(*(range(x.count) for x in self.experiment()))
-        return [dict(zip(axes, x)) for x in indices]
+        if self._loop_indices is None:
+            self._loop_indices = _util.loop_indices(self.experiment())
+        return self._loop_indices
 
     def _img_exp_events(self) -> list[structures.ExperimentEvent]:
         """Parse and return all Image and Experiment events."""

tests/test_metadata.py

@@ -74,6 +74,8 @@ def test_metadata_extraction(new_nd2: Path) -> None:
         for i in range(nd._rdr._seq_count()):
             assert isinstance(nd.frame_metadata(i), structures.FrameMetadata)
         assert isinstance(nd.experiment, list)
+        assert isinstance(nd.loop_indices, list)
+        assert all(isinstance(x, dict) for x in nd.loop_indices)
         assert isinstance(nd.text_info, dict)
         assert isinstance(nd.sizes, dict)
         assert isinstance(nd.custom_data, dict)