Skip to content

Commit

Permalink
add copy param to to_dask to avoid segfaults from closed file (#9)
Browse files Browse the repository at this point in the history
  • Loading branch information
tlambert03 authored Oct 9, 2021
1 parent 5f0a8b8 commit 35fa068
Showing 1 changed file with 5 additions and 4 deletions.
9 changes: 5 additions & 4 deletions nd2/nd2file.py
Original file line number Diff line number Diff line change
Expand Up @@ -181,12 +181,12 @@ def asarray(self) -> np.ndarray:
def __array__(self) -> np.ndarray:
return self.asarray()

def to_dask(self) -> da.Array:
def to_dask(self, copy=True) -> da.Array:
from dask.array import map_blocks

chunks = [(1,) * x for x in self._coord_shape]
chunks += [(x,) for x in self._frame_shape]
return map_blocks(self._dask_block, chunks=chunks, dtype=self.dtype)
return map_blocks(self._dask_block, copy, chunks=chunks, dtype=self.dtype)

_NO_IDX = -1

Expand All @@ -195,7 +195,7 @@ def _seq_index_from_coords(self, coords: Sequence) -> int:
return self._NO_IDX
return np.ravel_multi_index(coords, self._coord_shape)

def _dask_block(self, block_id: Tuple[int]) -> np.ndarray:
def _dask_block(self, copy, block_id: Tuple[int]) -> np.ndarray:
if isinstance(block_id, np.ndarray):
return

Expand All @@ -206,7 +206,8 @@ def _dask_block(self, block_id: Tuple[int]) -> np.ndarray:
if any(block_id):
raise ValueError(f"Cannot get chunk {block_id} for single frame image.")
idx = 0
return self._get_frame(idx)[(np.newaxis,) * ncoords]
data = self._get_frame(idx)[(np.newaxis,) * ncoords]
return data.copy() if copy else data

def to_xarray(self, delayed: bool = True, squeeze=True) -> xr.DataArray:
import xarray as xr
Expand Down

0 comments on commit 35fa068

Please sign in to comment.