From 82e7005808385a5441672e1c57c9f3dc34ab1c71 Mon Sep 17 00:00:00 2001 From: Talley Lambert Date: Fri, 6 Oct 2023 15:39:00 +0200 Subject: [PATCH] feat: add `read_frame` and `loop_indices` to public api (#181) * feat: add stuff to public api * coverage --- src/nd2/_util.py | 20 ++++++++++++++ src/nd2/nd2file.py | 33 +++++++++++++++++++++--- src/nd2/readers/_modern/modern_reader.py | 9 ++++--- tests/test_metadata.py | 2 ++ 4 files changed, 56 insertions(+), 8 deletions(-) diff --git a/src/nd2/_util.py b/src/nd2/_util.py index 4037a67..5c26564 100644 --- a/src/nd2/_util.py +++ b/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] diff --git a/src/nd2/nd2file.py b/src/nd2/nd2file.py index 1c4afc0..6de81c0 100644 --- a/src/nd2/nd2file.py +++ b/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. diff --git a/src/nd2/readers/_modern/modern_reader.py b/src/nd2/readers/_modern/modern_reader.py index 438b0ba..938f10b 100644 --- a/src/nd2/readers/_modern/modern_reader.py +++ b/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.""" diff --git a/tests/test_metadata.py b/tests/test_metadata.py index 6bace47..70ff7ae 100644 --- a/tests/test_metadata.py +++ b/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)