diff --git a/bdpy/dataform/__init__.py b/bdpy/dataform/__init__.py index 30d66a60..fe9ea43a 100644 --- a/bdpy/dataform/__init__.py +++ b/bdpy/dataform/__init__.py @@ -5,6 +5,7 @@ """ from .datastore import * +from .feature_hdf5 import FeatureWriter, convert_features_to_hdf5, save_features from .features import * from .kvs import SQLite3KeyValueStore from .pd import * diff --git a/bdpy/dataform/_feature_chunking.py b/bdpy/dataform/_feature_chunking.py new file mode 100644 index 00000000..24eaf742 --- /dev/null +++ b/bdpy/dataform/_feature_chunking.py @@ -0,0 +1,146 @@ +"""Chunk-shape policy for chunked feature storage. + +HDF5 reads whole chunks even when only a few elements of a chunk are selected, +so the chunk shape - not the file format alone - is what decides the cost of a +partial read. This module turns an array shape and dtype into a chunk shape +under a byte budget, rather than hard-coding an extent along any one axis. + +This file is a part of BdPy. +""" + +import math +from typing import Sequence, Tuple + +import numpy as np + +__all__ = ["DEFAULT_TARGET_CHUNK_BYTES", "choose_chunk_shape"] + +#: Default per-chunk byte budget. h5py's own guidance puts a useful chunk +#: somewhere between 10 KiB and 1 MiB; features are large and read in slabs, so +#: we sit at the top of that range. +DEFAULT_TARGET_CHUNK_BYTES = 1 << 20 + + +def _prod(values: Sequence[int]) -> int: + out = 1 + for v in values: + out *= int(v) + return out + + +def choose_chunk_shape( + shape: Sequence[int], + dtype: np.dtype, + target_bytes: int = DEFAULT_TARGET_CHUNK_BYTES, + n_samples_known: bool = True, +) -> Tuple[int, ...]: + """Choose an HDF5 chunk shape for a feature array under a byte budget. + + The two axes callers actually slice are axis 0 (samples) and axis 1 (the + outermost feature axis, typically channels or units), so both are chunked. + Trailing axes - the spatial axes of a convolutional feature map, say + ``(13, 13)`` - are kept whole: they are small and are never sliced in + practice, and splitting them would only shrink the chunks. + + Within the budget the two chunked axes are split roughly evenly in element + count, so that neither a sample-wise read nor a channel-wise read degenerates + into reading the entire dataset. Each extent is then shrunk to the smallest + one needing the same number of chunks, which keeps HDF5 from padding the + edge chunks and inflating the file. + + Parameters + ---------- + shape : sequence of int + Full array shape, ``(n_samples, *feature_shape)``. + dtype : numpy.dtype + Array dtype; only its itemsize matters. + target_bytes : int, optional + Size budget for one chunk, in bytes + (default: ``DEFAULT_TARGET_CHUNK_BYTES``). It is an upper bound only + when the budget can be met at all: the trailing axes are kept whole, so + a shape whose trailing axes alone exceed `target_bytes` necessarily + produces a larger chunk. + n_samples_known : bool, optional + Whether ``shape[0]`` is the final number of samples. Pass ``False`` when + writing to a resizable dataset whose sample count is not yet known, so + that the sample-axis extent is not capped by the current size. + + Returns + ------- + tuple of int + Chunk shape, same length as ``shape``. Every element is at least 1, and + no larger than the corresponding entry of ``shape`` for axes of + non-zero size (a layer with zero samples still needs a chunk extent of + at least 1 on the sample axis, since HDF5 forbids a zero chunk dim). + + Raises + ------ + ValueError + If `shape` has fewer than two axes, any feature axis has size 0, the + sample axis is negative, or `target_bytes` is not positive. + + Examples + -------- + >>> choose_chunk_shape((1200, 1000), np.dtype(np.float32)) + (400, 500) + >>> choose_chunk_shape((1200, 256, 13, 13), np.dtype(np.float32)) + (39, 37, 13, 13) + >>> choose_chunk_shape((50, 1000), np.dtype(np.float32)) + (50, 1000) + """ + shape = tuple(int(s) for s in shape) + if len(shape) < 2: + raise ValueError( + "Feature arrays need at least a sample axis and a feature axis; " + "got shape {}".format(shape) + ) + if target_bytes < 1: + raise ValueError("target_bytes must be positive, got {}".format(target_bytes)) + if shape[0] < 0: + raise ValueError("sample axis cannot be negative, got {}".format(shape)) + if any(s < 1 for s in shape[1:]): + # A zero-width feature axis carries nothing and cannot be chunked. + # A zero-sample layer, by contrast, is meaningful and is allowed below. + raise ValueError( + "feature axes must all be non-empty, got shape {}".format(shape) + ) + + itemsize = np.dtype(dtype).itemsize + + # The whole array fits in one chunk: nothing to gain from splitting it. + # Clamp to at least 1 per axis -- a zero-sample layer is legal but HDF5 + # rejects a chunk with a zero dimension. + if n_samples_known and _prod(shape) * itemsize <= target_bytes: + return tuple(max(1, s) for s in shape) + + # Bytes taken by one (sample, feature-0) cell, i.e. one element of the two + # chunked axes with the trailing axes kept whole. + inner = max(1, _prod(shape[2:])) * itemsize + budget_cells = max(1, target_bytes // inner) + + # Split the budget evenly between the two chunked axes, then let each axis + # absorb the headroom the other did not need: a short sample axis should not + # force a narrow feature chunk, and vice versa. + c1 = min(shape[1], max(1, math.isqrt(budget_cells))) + c0 = max(1, budget_cells // c1) + if n_samples_known and shape[0] > 0: + c0 = min(shape[0], c0) + c1 = min(shape[1], max(c1, budget_cells // c0)) + + # HDF5 allocates whole chunks, so a chunk extent that divides its axis + # unevenly pads the edge chunks and inflates the file. Shrinking the extent + # to the smallest one that needs the same number of chunks removes most of + # that padding, and can only lower the chunk size, so the budget still holds. + c1 = _snap(shape[1], c1) + if n_samples_known and shape[0] > 0: + c0 = _snap(shape[0], c0) + + return (c0, c1, *shape[2:]) + + +def _snap(length: int, chunk: int) -> int: + """Smallest extent covering `length` in the same number of chunks as `chunk`.""" + if length < 1 or chunk >= length: + return chunk + n_chunks = -(-length // chunk) # ceil + return -(-length // n_chunks) # ceil diff --git a/bdpy/dataform/_feature_store.py b/bdpy/dataform/_feature_store.py new file mode 100644 index 00000000..d57c7ff5 --- /dev/null +++ b/bdpy/dataform/_feature_store.py @@ -0,0 +1,909 @@ +"""Storage backends for DNN feature collections. + +``Features`` used to be welded to one on-disk layout: a directory per layer +holding one ``.mat`` file per stimulus. That layout has a file boundary only on +the sample axis, so reading a few channels still costs a full read of every +stimulus. This module puts a backend interface between ``Features`` and the +bytes, so the legacy layout and the chunked HDF5 layout are interchangeable and +a future backend needs no change to ``Features`` itself. + +Every backend answers the same question -- "give me these labels, sliced this +way along the feature axes" -- and differs only in how much it has to read to do +it. The legacy backend loads whole stimulus files and slices in memory; the HDF5 +backend pushes the slice down to h5py and reads only the chunks it covers. + +This file is a part of BdPy. +""" + +import glob +import os +from abc import ABC, abstractmethod +from functools import partial +from multiprocessing import Pool +from typing import TYPE_CHECKING, Dict, Iterator, List, Optional, Sequence, Tuple, Union + +import h5py +import numpy as np + +from . import _mat_v73 +from ._feature_chunking import DEFAULT_TARGET_CHUNK_BYTES + +__all__ = [ + "FeatureStore", + "HDF5FeatureStore", + "MatFeatureStore", + "detect_format", +] + +#: Root attribute marking a file as bdpy chunked feature storage. +FORMAT_ATTR = "bdpy_format" +#: Root attribute holding the schema version. +FORMAT_VERSION_ATTR = "bdpy_format_version" +#: Value of ``FORMAT_ATTR`` for feature files. +FORMAT_NAME = "features" +#: Highest schema version this module can read. +SUPPORTED_FORMAT_VERSION = 1 +#: Dataset holding the feature array, shape ``(n_samples, *feature_shape)``. +FEATURES_DATASET = "features" +#: Dataset holding the stimulus labels, shape ``(n_samples,)``. +LABELS_DATASET = "labels" +#: Extension of a chunked feature file. +HDF5_EXT = "h5" + +# A slice spec for the feature axes. Deliberately narrow: basic indexing only, +# i.e. what ``numpy.s_[128:256]`` or ``numpy.s_[8:16, 1:4]`` produces. See +# _validate_feature_slice for what is rejected and why. +# `types.EllipsisType` is 3.10+, and a bare "ellipsis" forward reference does not +# resolve at runtime (it is not a Python name), which makes +# typing.get_type_hints() raise. typeshed does define builtins.ellipsis, so the +# checker gets the real type and the interpreter gets an equivalent object. +if TYPE_CHECKING: + from builtins import ellipsis +else: + ellipsis = type(Ellipsis) +BasicIndexer = Union[slice, int, np.integer, ellipsis] +FeatureSlice = Union[BasicIndexer, Tuple[BasicIndexer, ...], None] + + +def _load_array_with_key(key: str, path: str) -> np.ndarray: + """Load one array from a ``.mat`` file. Shared with :mod:`bdpy.dataform.features`. + + v5 ``.mat`` via scipy, v7.3 (HDF5) via h5py; avoids hdf5storage on the load + path, which breaks under NumPy 2.0 (see :mod:`bdpy.dataform._mat_v73`). + `key` comes first so ``partial(..., 'feat')`` is a ``Pool.map`` callable. + """ + return _mat_v73.loadmat_key(path, key) + + +def _determine_num_parallel(num_files: int) -> int: + # NOTE: optimal number of parallel processes is not clear. It could depend + # on several factors such as the number of files, the size of files, the + # number of cores, etc. For now, we use a simple heuristic based on the + # number of files. + num_parallel: int + if num_files < 16: + num_parallel = 1 + elif num_files < 64: + num_parallel = 16 + else: + num_parallel = 64 + return num_parallel + + +def _normalize_feature_slice(feature_slice: FeatureSlice) -> Tuple: + """Normalize and validate a feature-axis slice spec. + + Returns a tuple of per-axis indexers; ``None`` means "no slice" and gives + the empty tuple. Note this makes ``numpy.newaxis`` (which *is* ``None``) + inexpressible, so a bare newaxis reads as "no slice"; inside a tuple it is + rejected explicitly. + """ + if feature_slice is None: + return () + indexers = feature_slice if isinstance(feature_slice, tuple) else (feature_slice,) + _validate_feature_slice(indexers) + return indexers + + +def _validate_feature_slice(indexers: Tuple) -> None: + """Reject anything outside basic, forward indexing. + + `feature_slice` exists to read part of a layer, not to reimplement NumPy + indexing. Keeping it to basic forward indexing is what lets every backend + mean the same thing by it: h5py handles slices with a positive step + (negative start/stop included), integers and one ``Ellipsis`` exactly as + NumPy does, but rejects a negative step outright. Allowing more would make + the legacy and HDF5 backends diverge, or require reimplementing NumPy's + semantics to keep them together. + + Parameters + ---------- + indexers : tuple + Per-axis indexers, as returned by :func:`_normalize_feature_slice`. + + Raises + ------ + ValueError + If any indexer is outside the supported set. + """ + advice = ( + " feature_slice supports basic forward indexing only (slices with a " + "positive step, integers and a single Ellipsis). Read without " + "feature_slice and index the resulting array with NumPy instead." + ) + if sum(ix is Ellipsis for ix in indexers) > 1: + raise ValueError("feature_slice may contain at most one Ellipsis." + advice) + for ix in indexers: + # bool is a subclass of int in Python, so it must be rejected *before* + # the integer check or True would silently be read as index 1. + if isinstance(ix, (bool, np.bool_)): + raise ValueError( + "feature_slice does not support boolean indexing." + advice + ) + if ix is Ellipsis or isinstance(ix, (int, np.integer)): + continue + if isinstance(ix, slice): + for part, value in ( + ("start", ix.start), ("stop", ix.stop), ("step", ix.step), + ): + if value is None: + continue + if isinstance(value, (bool, np.bool_)) or not isinstance( + value, (int, np.integer) + ): + raise ValueError( + "feature_slice slice {} must be an integer or None, " + "got {!r}.".format(part, value) + advice + ) + if ix.step is not None and ix.step < 1: + raise ValueError( + "feature_slice does not support a step below 1 " + "(got {}).".format(ix.step) + advice + ) + continue + if ix is None: + raise ValueError( + "feature_slice does not support numpy.newaxis." + advice + ) + raise ValueError( + "feature_slice does not support {} indexing.".format( + type(ix).__name__ + ) + + advice + ) + + +def _normalize_axis(axis: int, shape: Tuple[int, ...]) -> int: + """Return `axis` as a non-negative axis of an array of `shape`. + + Raises + ------ + ValueError + If `axis` is out of range for `shape`. + """ + normalized = axis + len(shape) if axis < 0 else axis + if not 0 <= normalized < len(shape): + raise ValueError( + "axis {} is out of range for selected shape {}".format(axis, shape) + ) + return normalized + + +def _check_size(size: int) -> None: + """Reject a slab size below 1.""" + if size < 1: + raise ValueError("size must be positive, got {}".format(size)) + + +def _partition(length: int, size: int) -> Iterator[slice]: + """Split ``range(length)`` into consecutive slices of at most `size`. + + The slices cover the range exactly once, in order; only the last may be + shorter. ``length == 0`` yields nothing. `size` must already be positive. + + A generator rather than a list: a materialized one would cost an object per + slab, which for a large axis at size 1 is millions of them. + """ + for start in range(0, length, size): + yield slice(start, min(start + size, length)) + + +def _axis_index(axis: int, sl: slice) -> Tuple[slice, ...]: + """Index tuple that applies `sl` to `axis` and leaves earlier axes whole.""" + return (slice(None),) * axis + (sl,) + + +def _default_iter_size( + shape: Tuple[int, ...], axis: int, itemsize: int, extent: Optional[int] +) -> int: + """Default slab size along `axis` of an array of `shape`. + + The on-disk chunk `extent` along the axis when there is one, so iteration is + chunk-aligned; otherwise as many elements as fit in the byte budget. Always + at least 1 and at most the axis length (when that is positive). + """ + if extent is not None: + return max(1, min(extent, shape[axis])) + other = 1 + for i, n in enumerate(shape): + if i != axis: + other *= int(n) + per_element = max(1, other * itemsize) + # Read the module constant at call time so it can be patched in tests. + return max(1, min(shape[axis], DEFAULT_TARGET_CHUNK_BYTES // per_element)) + + +def _iter_array_slabs( + array: np.ndarray, axis: int, slices: Iterator[slice] +) -> Iterator[Tuple[slice, np.ndarray]]: + """Yield ``(sl, view)`` for each of `slices` applied to `axis` of `array`. + + The slabs are views, so holding one keeps `array` alive. + """ + for sl in slices: + yield sl, array[_axis_index(axis, sl)] + + +class FeatureStore(ABC): + """Interface to a collection of DNN features on disk. + + A store exposes layers and stimulus labels, and reads features for a set of + labels with an optional slice along the feature axes. Implementations differ + only in how much I/O a given request costs. + """ + + @property + @abstractmethod + def layers(self) -> List[str]: + """List of DNN layers held by this store.""" + + @property + @abstractmethod + def labels(self) -> List[str]: + """List of stimulus labels, in the store's own order.""" + + @abstractmethod + def shape(self, layer: str) -> Tuple[int, ...]: + """Full shape ``(n_samples, *feature_shape)`` of `layer`.""" + + @abstractmethod + def dtype(self, layer: str) -> np.dtype: + """Dtype of the features in `layer`.""" + + def read( + self, + layer: str, + labels: Optional[Sequence[str]] = None, + feature_slice: FeatureSlice = None, + ) -> np.ndarray: + """Read features from `layer`. + + `feature_slice` is validated here, before any I/O, so every backend + accepts and rejects exactly the same specs. Backends implement + :meth:`_read` and must not override this method. + + Parameters + ---------- + layer : str + DNN layer. + labels : sequence of str, optional + Stimulus labels to read. ``None`` reads every label in store order. + Otherwise rows come back **in the order given**, and repeated labels + yield repeated rows. + feature_slice : slice, int, Ellipsis or tuple of those, optional + Index applied to the feature axes (axes 1 and up), as produced by + ``numpy.s_[...]``. Basic forward indexing only: slices with a + positive step and integer bounds, integers, and at most one + ``Ellipsis``. Fancy indexing, a negative step, a non-integer slice + bound, booleans, and ``numpy.newaxis`` *inside a tuple* raise + ``ValueError``. A bare ``numpy.newaxis`` is ``None``, the "no + slice" default, so it reads the whole feature tensor. + + Returns + ------- + numpy.ndarray + Array of shape ``(n_labels, *sliced_feature_shape)``. + """ + return self._read(layer, labels, _normalize_feature_slice(feature_slice)) + + @abstractmethod + def _read( + self, + layer: str, + labels: Optional[Sequence[str]], + indexers: Tuple, + ) -> np.ndarray: + """Backend read: `indexers` is already validated, ``()`` for no slice.""" + + def chunk_extent(self, layer: str, axis: int) -> Optional[int]: + """On-disk chunk extent of `layer` along `axis`, if the store has one. + + Returns ``None`` when the backend has no chunking to align with, in + which case :meth:`iter_chunks` falls back to a byte budget. + """ + return None + + def iter_chunks( + self, + layer: str, + labels: Optional[Sequence[str]] = None, + axis: int = 1, + size: Optional[int] = None, + ) -> Iterator[Tuple[slice, np.ndarray]]: + """Iterate over `layer` in slabs along `axis`. + + Yields ``(sl, block)`` pairs, where `sl` is the slice applied to `axis` + of the selected array and `block` is that slab. The slice is yielded so + that callers can place results back without tracking offsets themselves. + + There is deliberately no `feature_slice` here. Composing an arbitrary + slice with the per-slab slice means reimplementing NumPy's index + arithmetic, which is both the source of subtle wrong results and more + generality than reading a layer in slabs needs. Slice the blocks as they + come out instead. + + Parameters + ---------- + layer : str + DNN layer. + labels : sequence of str, optional + Stimulus labels, as in :meth:`read`. + axis : int, optional + Axis to iterate over. Axis 0 is the sample axis; the default, axis + 1, is the outermost feature axis, which is the one partial feature + reads are about. + size : int, optional + Number of elements per slab. Defaults to the on-disk chunk extent + along `axis` when the backend has one, so that iteration is + chunk-aligned and every element is read exactly once; otherwise a + size derived from the same byte budget is used. + + Yields + ------ + tuple of (slice, numpy.ndarray) + The slice applied to `axis`, and the corresponding slab. + + Notes + ----- + How much is resident at once depends on the backend. Along the sample + axis every backend reads one slab at a time. Along a feature axis only + a backend with real partial reads can do so: the legacy ``.mat`` + layout has no file boundary there, so it reads the selection once and + yields views into it, and the blocks keep that array alive for as long + as they are held. + """ + selected_shape = self._selected_shape(layer, labels) + axis = _normalize_axis(axis, selected_shape) + + length = selected_shape[axis] + if size is None: + size = _default_iter_size( + selected_shape, + axis, + np.dtype(self.dtype(layer)).itemsize, + self.chunk_extent(layer, axis), + ) + _check_size(size) + + if length == 0: + # Nothing to iterate. Returning here rather than falling into the + # loop keeps a backend that reads the selection up front from + # reading anything the per-slab path would not have read either. + return + + slices = _partition(length, size) + if axis == 0: + all_labels = list(self.labels) if labels is None else list(labels) + for sl in slices: + yield sl, self.read(layer, all_labels[sl]) + else: + yield from self._iter_feature_slabs(layer, labels, axis, slices) + + def _iter_feature_slabs( + self, + layer: str, + labels: Optional[Sequence[str]], + axis: int, + slices: Iterator[slice], + ) -> Iterator[Tuple[slice, np.ndarray]]: + """Yield the feature-axis slabs of `layer` named by `slices`. + + This is the seam for a backend that cannot read a feature axis + partially. `axis` is already normalized and validated to a feature axis + (``>= 1``), and `slices` partitions that axis, so an implementation + decides only *how* a slab is obtained -- never which slabs there are. + + The default reads each slab on its own, which is what a backend with + real partial reads wants. + """ + for sl in slices: + # axis >= 1 indexes the feature axes, where axis 1 of the array is + # entry 0 of the feature-slice tuple. Go through read(), the single + # entry point that validates. + yield sl, self.read(layer, labels, _axis_index(axis - 1, sl)) + + def _selected_shape( + self, + layer: str, + labels: Optional[Sequence[str]] = None, + ) -> Tuple[int, ...]: + """Shape the selection would have, without reading the data.""" + full = self.shape(layer) + n_samples = full[0] if labels is None else len(labels) + return (n_samples, *full[1:]) + + + +class MatFeatureStore(FeatureStore): + """Legacy per-stimulus ``.mat`` feature directory. + + Layout is ``//