"""Provide lazy, array-like access to datasets without loading full arrays into memory.
This module provides lightweight wrappers that expose array-like, lazy access to
datasets without loading full arrays into memory.
The interface is defined by :class:`Dataset` and mimics that of h5py
datasets, allowing users to slice and access data on demand. Specific
implementations load the necessary data and apply transformations, such as
stacking, scaling, summing, etc., only for the requested slices.
.. code-block:: python
xyz = LazyStackedDataset([ds1, ds2, ds3])
# Loaded and computed only for this slice
xyz_slice = xyz[100:200]
# Writing is supported for LazyScaledDataset and LazyStackedDataset
xyz[100:200] = new_data
scaled = LazyScaledDataset(ds, scale_factor=2.0)
scaled[50:100] = values # Applies inverse scaling before writing
combined_signal = LazySumDataset([ds1, ds2])
combined_flags = LazyBooleanOrDataset([flags1, flags2])
# Loaded and computed only for this slice (read-only)
signal_slice = combined_signal[0:1024]
flags_slice = combined_flags[0:1024]
Some operations also supports writing. For example, :class:`LazyScaledDataset`
applies the inverse scaling operation before writing to the underlying dataset.
.. code-block:: python
xyz = LazyScaledDataset(ds, scale_factor=2.0)
xyz[100:200] = new_values # Writes new_values / 2.0 to ds[100:200]
Dataset interface
-----------------
.. autoclass:: Dataset
:members:
:special-members: __getitem__, __setitem__, __len__, __repr__
Operations on datasets
----------------------
.. autoclass:: LazyMappedDataset
:members:
:special-members: __getitem__, __setitem__
.. autoclass:: LazyScaledDataset
:members:
:special-members: __getitem__, __setitem__
.. autoclass:: LazyReducedDataset
:members:
:special-members: __getitem__
.. autoclass:: LazySumDataset
:members:
:special-members: __getitem__
.. autoclass:: LazyBooleanOrDataset
:members:
:special-members: __getitem__
Stacking datasets
-----------------
.. autoclass:: LazyStackedDataset
:members:
:special-members: __getitem__, __setitem__
"""
from functools import partial
from typing import Any, Protocol, Sequence, override, runtime_checkable
import numpy as np
from numpy.typing import DTypeLike, NDArray
[docs]
@runtime_checkable
class Dataset(Protocol):
"""Protocol for (potentially lazy) datasets with NumPy-style access.
Implementations must expose shape, dimensionality, dtype and support slicing
through ``__getitem__`` similarly to ``h5py.Dataset``. Writing via
``__setitem__`` is optional and may not be supported by all implementations.
"""
@property
def shape(self) -> tuple[int, ...]:
"""Shape of the dataset."""
...
@property
def ndim(self) -> int:
"""Number of dimensions."""
...
@property
def dtype(self) -> np.dtype[Any]:
"""Data type of the dataset."""
...
[docs]
def __getitem__(self, *args: Any, **kwargs: Any) -> Any:
"""Return a sliced array."""
...
[docs]
def __setitem__(self, *args: Any, **kwargs: Any) -> None:
"""Write to the dataset.
Raises
------
NotImplementedError
If writing is not supported by the implementation.
"""
raise NotImplementedError("Writing is not supported for this dataset")
[docs]
def __len__(self) -> int:
"""Return the length of the first axis.
Raises
------
TypeError
If the dataset is 0-dimensional.
"""
if self.ndim == 0:
raise TypeError("len() of unsized dataset")
return self.shape[0]
def __array__(
self, dtype: np.dtype[Any] | None = None, copy: bool | None = None
) -> NDArray:
"""Return the full dataset as a NumPy array.
Parameters
----------
dtype
Desired data type of the returned array. If None, the dtype of the dataset
is used.
copy
Whether to return a copy of the data. If None, the default behavior of
``np.asarray`` is used.
Returns
-------
The full dataset as a NumPy array.
"""
return np.asarray(self[...], dtype=dtype, copy=copy)
[docs]
def __repr__(self) -> str:
"""Return a concise representation showing class name, shape and dtype."""
return f"{type(self).__name__}(shape={self.shape}, dtype={self.dtype})"
[docs]
class LazyMappedDataset(Dataset):
"""Lazy-loaded dataset that applies a function to values on slicing.
This is a generic helper for dataset wrappers that transform a single dataset by
applying a callable to the sliced values. Writing is supported by providing an
``inverse_func`` that maps values back before writing to the underlying dataset; if
omitted, writing raises ``NotImplementedError``.
Since ``func`` can change the data type of the values (e.g. casting or thresholding
into booleans), ``dtype`` is inferred by applying ``func`` to a zero-sized probe
array unless explicitly provided via ``dtype``.
Supports both scalar-valued and multivariate map functions, with two
constraints: ``func`` must not mix elements across the first (time) dimension, since
the dataset is lazily sliced along it before ``func`` is applied; and ``func`` must
preserve the exact shape of its input, as functions that change the output shape are
not supported. In other words, ``func`` and its inverse (if defined) must accept an
array-like broadcastable to ``dataset.shape[1:]`` and return an array of the exact
same shape as its input.
Parameters
----------
dataset
Dataset to map.
func
Callable applied to sliced values when reading.
inverse_func
Callable applied to values before writing, effectively the inverse of ``func``.
If ``None``, writing is not supported.
dtype
Data type of the mapped values. If ``None``, it is inferred by applying ``func``
to a zero-sized probe array.
is_multivariate
Whether ``func`` is multivariate. This is used for writing to the dataset. If
``None``, this property is inferred by subjecting ``func`` to a series of simple
tests. These tests might in very rare cases fail at identifying a multivariate
function, so it is recommended to always explicitly provide it.
"""
def __init__(
self,
dataset: Dataset,
func: Any,
inverse_func: Any = None,
*,
dtype: DTypeLike | None = None,
is_multivariate: bool | None = None,
) -> None:
self.dataset = dataset
self.func = func
self.inverse_func = inverse_func
self._dtype = np.dtype(dtype) if dtype is not None else self._infer_dtype()
self._is_multivariate = (
is_multivariate
if is_multivariate is not None
else self._infer_is_multivariate()
)
self._validate_func()
def _infer_dtype(self) -> np.dtype:
"""Infer the mapped dtype by probing ``func`` with an empty array.
The probe is of the same shape as self.dataset, except for its first dimension
which is zero-sized.
Returns
-------
Inferred dtype of the mapped values.
Raises
------
ValueError
If the mapped dtype cannot be inferred with the zero-sized probe.
"""
probe = np.empty((0, *self.dataset.shape[1:]), dtype=self.dataset.dtype)
try:
return np.asarray(self.func(probe)).dtype
except ValueError as e:
raise ValueError(
"Cannot infer dtype from func; please provide dtype explicitly"
) from e
def _infer_is_multivariate(self) -> bool:
"""Infer whether ``func`` is multivariate.
In some very rare cases, this method can mistake a multivariate function for a
non-multivariate one. Users are encouraged to explicitly confirm whether
``func`` is multivariate or not.
Returns
-------
False if ``func`` is purely element-wise with no reduction operations
involved, True otherwise.
"""
is_multivariate = False
dtype = np.dtype(self.dataset.dtype)
probe = np.empty((0, *self.dataset.shape[1:]), dtype=dtype)
try:
self.func(probe)
except ValueError:
# this means func requires at least one element, so it contains a reduction
is_multivariate = True
if not is_multivariate:
# test against a scalar probe -> catching an exception there means the
# function expects a multidimensional input
scalar_probe = dtype.type(100.0)
try:
output = self.func(scalar_probe)
# some numpy functions that are multivariate and work on scalars will
# implicitly transform them into length-1 arrays, so this additional
# check avoids mistaking them for non-multivariate functions
is_multivariate = not np.isscalar(output)
except Exception:
is_multivariate = True
if not is_multivariate:
# The one numpy function that is multivariate and does not get flagged
# by the scalar check is np.flip. We check it by perturbing the first
# element in a length-2 probe and evaluating if the second element has
# changed. Hopefully this can also catch any other multivariate
# functions that would have passed all the previous tests
reference_probe = np.array([0.0, 2.0], dtype=dtype)
perturbed_probe = reference_probe.copy()
perturbed_probe[0] = dtype.type(100.0)
try:
reference_output = np.asarray(self.func(reference_probe))
perturbed_output = np.asarray(self.func(perturbed_probe))
is_multivariate = not np.isclose(
reference_output[1], perturbed_output[1]
)
except Exception:
is_multivariate = True
# some functions that broadcast silently can pass all checks in this method,
# but they are not shape-preserving and will be caught by _validate_func
return is_multivariate
def _validate_func(self) -> None:
"""Verify that ``self.func`` preserves the exact shape of the input.
Raises
------
ValueError
If ``func`` is not shape-preserving.
"""
probe_shape = (0, *self.dataset.shape[1:])
probe = np.empty(probe_shape, dtype=self.dataset.dtype)
try:
output = np.asarray(self.func(probe))
except ValueError:
# A ValueError would be raised here because func contains an operation that
# requires at least one element. We create a probe with one line instead to
# ensure it is not empty.
probe_shape = (1, *self.dataset.shape[1:])
probe = np.empty(probe_shape, dtype=self.dataset.dtype)
output = np.asarray(self.func(probe))
if output.shape != probe.shape:
raise ValueError(
f"func must preserve the exact shape of its input: "
f"input shape {probe.shape}, output shape {output.shape}"
)
@property
@override
def shape(self) -> tuple[int, ...]:
"""Shape of the dataset."""
return self.dataset.shape
@property
@override
def ndim(self) -> int:
"""Number of dimensions."""
return self.dataset.ndim
@property
@override
def dtype(self) -> np.dtype:
"""Data type of the dataset."""
return self._dtype
[docs]
@override
def __getitem__(self, key) -> Any:
"""Slice and map the dataset lazily.
In case the mapping function is multivariate, it needs elements from different
dimensions to compute the output. Therefore, the dataset is only sliced along
the first dimension (time) before applying the mapping function. The output is
then sliced with the original key to return the required slice.
Parameters
----------
key
Slicing key (integers, slices, lists, ellipsis).
Returns
-------
NDArray
Result of ``func`` applied to the sliced array from the dataset.
"""
if not self._is_multivariate:
return self.func(self.dataset[key])
# if key is a tuple and func is multivariate, separate key into its component
# along the first dimension and its components along all other dimensions
if isinstance(key, tuple):
first, rest = key[0], key[1:]
else:
first, rest = key, ()
# If the first dimension is collapsed, replace the slicing key along axis 0
# with a size-1 slice to preserve the first dimension
collapse_first_dimension = isinstance(first, (int, np.integer))
if collapse_first_dimension:
idx = first if first >= 0 else first + self.dataset.shape[0]
first = slice(idx, idx + 1)
key = (first, *rest)
# if key is not a tuple, then indexing only occurs along axis 0, so key can be
# used as is
if not isinstance(key, tuple):
result = self.func(self.dataset[key])
if collapse_first_dimension:
# collapse the first dimension back to match the shape of the requested
# slice
result = result[0]
return result
# if first is an Ellipsis, then no slicing occurs along the first dimension. The
# function is applied to the whole dataset before slicing.
if first is Ellipsis:
return self.func(self.dataset[:])[key]
# if first is not an Ellipsis, then the dataset is sliced only along axis 0
# using first before applying func. The output is then sliced with the rest of
# the key to match the shape of the requested slice
result = self.func(self.dataset[first, ...])
if collapse_first_dimension:
# if first was replaced by a size-1 slice, it needs to be collapsed to match
# the shape of the requested slice
return result[(0, *rest)]
return result[(slice(None), *rest)]
[docs]
@override
def __setitem__(self, key: Any, value: Any) -> None:
"""Write mapped values to the dataset lazily.
``inverse_func`` is applied to ``value`` before writing to the
underlying dataset.
For multivariate map functions, partial writes along dimensions other than axis
0 are not supported: either ``value`` must be scalar or ``value.shape[1:]``
must be broadcastable to ``self.dataset.shape[1:]``. This restriction does not
apply to scalar-valued map functions.
Parameters
----------
key
Slicing key (integers, slices, lists, ellipsis).
value
Data to write. Will be passed through ``inverse_func`` before writing.
Raises
------
NotImplementedError
If no ``inverse_func`` was provided.
"""
if self.inverse_func is None:
raise NotImplementedError("Writing is not supported for this dataset")
if self._is_multivariate and isinstance(key, tuple) and len(key[1:]):
raise ValueError(
"Partial writes along dimensions other than axis 0 are not supported "
"for multivariate map functions"
)
self.dataset[key] = self.inverse_func(value)
[docs]
class LazyScaledDataset(LazyMappedDataset):
"""Lazy-loaded dataset that applies scaling on slicing.
This is a thin wrapper around :class:`LazyMappedDataset` with a scaling
function and its inverse.
Parameters
----------
dataset
HDF5 dataset to scale.
scale_factor
Scaling factor to apply to all values.
dtype
Data type of the scaled values. If ``None``, it is inferred automatically.
"""
def __init__(
self,
dataset: Dataset,
scale_factor: float = 1.0,
*,
dtype: DTypeLike | None = None,
) -> None:
self.scale_factor = scale_factor
super().__init__(
dataset, func=self._scale, inverse_func=self._unscale, dtype=dtype
)
def _scale(self, value: Any) -> Any:
return value * self.scale_factor
def _unscale(self, value: Any) -> Any:
"""Invert the scaling before writing.
Raises
------
ValueError
If scale_factor is zero (inverse scaling undefined).
"""
if self.scale_factor == 0:
raise ValueError("Cannot write to dataset with scale_factor=0")
return value / self.scale_factor
[docs]
class LazyReducedDataset(Dataset):
"""Lazy-loaded dataset that applies a NumPy reduction across datasets.
This is a generic helper for dataset wrappers that combine multiple datasets by
loading the corresponding slice from each dataset and reducing the resulting arrays
with a callable such as ``np.sum`` or ``np.logical_or.reduce``.
Since ``reducer`` can change the shape (e.g. collapsing an axis) and data type (e.g.
summing booleans into integers) of the values, ``shape`` and ``dtype`` are inferred
by applying ``reducer`` to zero-sized probe arrays unless explicitly provided via
``shape`` and ``dtype``.
Parameters
----------
datasets
Sequence of datasets to reduce.
reducer
Callable that accepts a list or array-like of sliced values and returns a
reduced result. The function can optionally accept an ``axis`` argument to
specify the axis along which to reduce.
dtype
Data type of the reduced values. If ``None``, it is inferred by applying
``reducer`` to zero-sized probe arrays.
shape
Shape of the reduced dataset. If ``None``, it is inferred the same way.
Raises
------
ValueError
If datasets list is empty.
"""
def __init__(
self,
datasets: Sequence[Dataset],
reducer: Any,
*,
dtype: DTypeLike | None = None,
shape: tuple[int, ...] | None = None,
) -> None:
self.datasets = list(datasets)
if not self.datasets:
raise ValueError("At least one dataset is required")
self.reducer = reducer
resolved_dtype = np.dtype(dtype) if dtype is not None else None
resolved_shape = tuple(shape) if shape is not None else None
if resolved_dtype is None or resolved_shape is None:
inferred_shape, inferred_dtype = self._infer_metadata()
if resolved_dtype is None:
resolved_dtype = inferred_dtype
if resolved_shape is None:
resolved_shape = inferred_shape
self._dtype: np.dtype = resolved_dtype
self._shape: tuple[int, ...] = resolved_shape
def _reduce(self, slices: list[Any]) -> Any:
"""Apply the reducer, passing ``axis=0`` when the reducer supports it.
Parameters
----------
slices
List of slices to be reduced.
Returns
-------
NDArray
Reduced array from the sliced datasets.
"""
try:
return self.reducer(slices, axis=0)
except TypeError:
return self.reducer(slices)
def _infer_metadata(self) -> tuple[tuple[int, ...], np.dtype]:
"""Infer the reduced shape and dtype by probing with zero-sized arrays."""
probes = [np.empty((0,) + ds.shape[1:], dtype=ds.dtype) for ds in self.datasets]
result = np.asarray(self._reduce(probes))
shape = (self.datasets[0].shape[0],) + result.shape[1:]
return shape, result.dtype
@property
@override
def shape(self) -> tuple[int, ...]:
"""Shape of the dataset."""
return self._shape
@property
@override
def ndim(self) -> int:
"""Number of dimensions."""
return len(self._shape)
@property
@override
def dtype(self) -> np.dtype:
"""Data type of the dataset."""
return self._dtype
[docs]
@override
def __getitem__(self, key) -> Any:
"""Slice and reduce datasets lazily.
Parameters
----------
key
Slicing key (integers, slices, lists, ellipsis).
Returns
-------
NDArray
Reduced array from the sliced datasets.
Raises
------
ValueError
If the reduced result does not match the expected shape or dtype.
"""
slices = [ds[key] for ds in self.datasets]
reduced = self._reduce(slices)
probe = np.broadcast_to(np.empty((), dtype=self.dtype), self.shape)
expected = probe[key]
reduced_arr = np.asarray(reduced)
if reduced_arr.shape != expected.shape:
raise ValueError(
f"Reduced result has shape {reduced_arr.shape}, "
f"expected {expected.shape}. If shape has not been "
"specified, it may need to be provided explicitly."
)
if reduced_arr.dtype != expected.dtype:
raise ValueError(
f"Reduced result has dtype {reduced_arr.dtype}, "
f"expected {expected.dtype}. If dtype has not been "
"specified, it may need to be provided explicitly."
)
return reduced
@override
def __setitem__(self, *args: Any, **kwargs: Any) -> None:
raise NotImplementedError
[docs]
class LazySumDataset(LazyReducedDataset):
"""Lazy-loaded dataset that sums values across multiple datasets on slicing.
This is a thin wrapper around :class:`LazyReducedDataset` with a sum reducer.
Parameters
----------
datasets
Sequence of HDF5 datasets to sum.
dtype
Data type of the summed values. If ``None``, it is inferred automatically.
shape
Shape of the summed dataset. If ``None``, it is inferred automatically.
Raises
------
ValueError
If datasets list is empty.
"""
def __init__(
self,
datasets: Sequence[Dataset],
*,
dtype: DTypeLike | None = None,
shape: tuple[int, ...] | None = None,
) -> None:
super().__init__(datasets, reducer=np.sum, dtype=dtype, shape=shape)
def _logical_or_reduce_and_cast(
x: Sequence[np.ndarray], dtype: DTypeLike | None = None
) -> np.ndarray:
"""Apply a logical OR reduction across a sequence of numpy arrays.
Parameters
----------
x
Sequence of numpy arrays to reduce.
dtype
If not ``None``, the resulting array will be cast to this data type.
Returns
-------
The result of applying logical OR across the input arrays.
"""
reduced = np.logical_or.reduce(x)
if dtype is not None:
reduced = reduced.astype(dtype)
return reduced
[docs]
class LazyBooleanOrDataset(LazyReducedDataset):
"""Lazy-loaded dataset that applies Boolean OR across multiple datasets.
This is a thin wrapper around :class:`LazyReducedDataset` with a boolean OR reducer.
Parameters
----------
datasets
Sequence of HDF5 datasets to OR.
dtype
Data type of the combined values. If ``None``, it is inferred automatically.
shape
Shape of the combined dataset. If ``None``, it is inferred automatically.
Raises
------
ValueError
If datasets list is empty.
"""
def __init__(
self,
datasets: Sequence[Dataset],
*,
dtype: DTypeLike | None = None,
shape: tuple[int, ...] | None = None,
) -> None:
reducer = partial(_logical_or_reduce_and_cast, dtype=dtype)
super().__init__(datasets, reducer=reducer, dtype=dtype, shape=shape)
[docs]
class LazyStackedDataset(Dataset):
"""Lazy-loaded stacked dataset that does not load all data into memory.
This is useful to provide access to stacked and normalized datasets without
loading all data into memory at once, similar to h5py datasets.
>>> stacked = LazyStackedDataset([ds1, ds2, ds3])
>>> print(stacked.shape)
(1000, 3)
>>> data_slice = stacked[100:200] # Loads only the requested slice
>>> data_slice2 = stacked[100:200, ..., 0] # Slicing along stacked axis
Parameters
----------
datasets
Sequence of HDF5 datasets to stack.
axis
Axis along which to stack the datasets.
Raises
------
ValueError
If ``axis`` is out of bounds for stacking.
"""
def __init__(
self,
datasets: Sequence[Dataset],
axis: int = -1,
) -> None:
self.datasets = list(datasets)
self._validate_datasets()
self.axis = axis
# Compute the shape of the stacked array
shape = list(self.datasets[0].shape)
if self.axis < -len(shape) - 1 or self.axis > len(shape):
raise ValueError("Axis out of bounds for stacking")
if self.axis < 0:
shape.insert(len(shape) + self.axis + 1, len(datasets))
else:
shape.insert(self.axis, len(datasets))
self._shape = tuple(shape)
def _validate_datasets(self) -> None:
"""Make sure all datasets have the same shape and dtype.
Raises
------
ValueError
If ``self.datasets`` is empty.
ValueError
If all elements in ``self.datasets`` do not have the same shape.
ValueError
If all elements in ``self.datasets`` do not have the same dtype.
"""
# Make sure there is at least one dataset
if not self.datasets:
raise ValueError("At least one dataset is required")
# Make sure all datasets have the same shape and dtype
first_shape = self.datasets[0].shape
first_dtype = self.datasets[0].dtype
for ds in self.datasets:
if ds.shape != first_shape:
raise ValueError("All datasets must have the same shape")
if ds.dtype != first_dtype:
raise ValueError("All datasets must have the same dtype")
@property
def normalized_axis(self) -> int:
"""Normalize axis to a positive value."""
axis = self.axis
if axis < 0:
axis = self.ndim + axis
return axis
@property
@override
def shape(self) -> tuple[int, ...]:
"""Shape of the stacked array."""
return self._shape
@property
@override
def ndim(self) -> int:
"""Number of dimensions of the stacked array."""
return len(self._shape)
@property
@override
def dtype(self) -> np.dtype:
"""Data type of the stacked array."""
return self.datasets[0].dtype
def _normalize_key(self, key: Any) -> tuple[Any, ...]:
"""Normalize an indexing key into an explicit per-axis tuple.
Wraps non-tuple keys, expands ellipsis, and pads missing trailing axes with full
slices, so the result always has exactly ``self.ndim`` entries, one per axis,
with NumPy scalar integers converted to Python ints.
Parameters
----------
key
Slicing key, can include integers (Python int or Numpy integer scalars),
slices, lists of indices, and ellipsis.
Returns
-------
tuple
Canonical per-axis key with exactly ``self.ndim`` entries.
Raises
------
IndexError
If None or np.newaxis is used in the key.
IndexError
If too many indices are provided for the array.
"""
if not isinstance(key, tuple):
key = (key,)
if any(k is None for k in key):
raise IndexError("None or np.newaxis indexing is not supported")
key = tuple(int(k) if isinstance(k, np.integer) else k for k in key)
if Ellipsis in key:
ellipsis_index = key.index(Ellipsis)
num_missing = self.ndim - (len(key) - 1)
key = (
key[:ellipsis_index]
+ (slice(None),) * num_missing
+ key[ellipsis_index + 1 :]
)
if len(key) < self.ndim:
key = key + (slice(None),) * (self.ndim - len(key))
elif len(key) > self.ndim:
raise IndexError("Too many indices for array")
return key
def _dataset_indices(self, dataset_slice: Any) -> list[int]:
"""Resolve the stacked-axis entry of a normalized key into dataset indices."""
if isinstance(dataset_slice, int):
return [dataset_slice]
if isinstance(dataset_slice, slice):
start, stop, step = dataset_slice.indices(len(self.datasets))
return list(range(start, stop, step))
return list(dataset_slice)
@staticmethod
def _output_axis(key: tuple[Any, ...], axis: int) -> int:
"""Find where ``axis`` ends up once integer-sliced dimensions are removed.
Since integer indices remove a dimension from the result, the position of
``axis`` in the sliced/stacked output shifts left by the number of integer
entries that precede it in ``key``.
"""
return axis - sum(1 for idx in key[:axis] if isinstance(idx, int))
[docs]
@override
def __getitem__(self, key) -> Any:
"""Slice the stacked datasets on demand without loading all data.
Supports slicing along all dimensions including the stacked axis. The
stacked axis can be sliced using integers, slices, or lists of indices.
Note that fancy indexing (i.e., using arrays of indices for multiple
axes) for non-stacked axes is not supported due to a limitation in h5py.
However, it is supported for a single non-stacked axis at a time.
Insertion of None or np.newaxis is not supported.
Parameters
----------
key
Slicing key, can include integers (Python int or Numpy integer
scalars), slices, lists of indices, and ellipsis.
Returns
-------
Any
Sliced and stacked array.
Raises
------
IndexError
If None or np.newaxis is used in the key.
IndexError
If too many indices are provided for the array.
"""
key = self._normalize_key(key)
# Select datasets along the stacked axis
stacked_axis = self.normalized_axis
dataset_slice = key[stacked_axis]
dataset_indices = self._dataset_indices(dataset_slice)
# Slice each selected dataset with the remaining (non-stacked) axes
remaining_slices = key[:stacked_axis] + key[stacked_axis + 1 :]
slices = [self.datasets[i][remaining_slices] for i in dataset_indices]
# If only one dataset is selected, collapse the stacked axis
if isinstance(dataset_slice, int):
return slices[0]
# Stack the selected slices at the (possibly shifted) stacked axis
output_axis = self._output_axis(key, stacked_axis)
return np.stack(slices, axis=output_axis)
[docs]
@override
def __setitem__(self, key: Any, value: Any) -> None:
"""Write to stacked datasets on demand.
The value is unstacked along the stacked axis and distributed to the
appropriate underlying datasets.
Parameters
----------
key
Slicing key, can include integers, slices, or lists of indices.
value
Data to write. Must be compatible with the sliced shape.
Raises
------
IndexError
If None or np.newaxis is used in the key, or if too many indices
are provided.
"""
key = self._normalize_key(key)
value = np.asarray(value)
# Select datasets along the stacked axis
stacked_axis = self.normalized_axis
dataset_slice = key[stacked_axis]
dataset_indices = self._dataset_indices(dataset_slice)
# Write to each selected dataset with the remaining (non-stacked) axes
remaining_slices = key[:stacked_axis] + key[stacked_axis + 1 :]
# If only one dataset is selected, write directly
if isinstance(dataset_slice, int):
self.datasets[dataset_indices[0]][remaining_slices] = value
return
# Unstack the value at the (possibly shifted) stacked axis and write
# the corresponding slice to each selected dataset
value_axis = self._output_axis(key, stacked_axis)
for i, ds_idx in enumerate(dataset_indices):
value_slice = np.take(value, i, axis=value_axis)
self.datasets[ds_idx][remaining_slices] = value_slice