From 38cd4cf7af7f2a6d3986e0b8ba15d9ce1cae6a68 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 26 Mar 2026 08:28:44 +0000 Subject: [PATCH] Add signal_ids filtering to SequenceInterpolator with lazy memmap strategy Co-authored-by: reneburghardt <7131281+reneburghardt@users.noreply.github.com> Agent-Logs-Url: https://github.com/reneburghardt/experanto/sessions/51fac970-0fc7-4299-892e-120fec1c57f7 --- experanto/interpolators.py | 126 ++++++++++++++++++++--- tests/test_sequence_interpolator.py | 154 +++++++++++++++++++++++++++- 2 files changed, 265 insertions(+), 15 deletions(-) diff --git a/experanto/interpolators.py b/experanto/interpolators.py index 853ea4d..92db722 100644 --- a/experanto/interpolators.py +++ b/experanto/interpolators.py @@ -20,6 +20,13 @@ logger = logging.getLogger(__name__) +def _is_contiguous_range(ids: np.ndarray) -> bool: + """Return True if *ids* (sorted 1-D integer array) form a contiguous range.""" + if len(ids) == 0: + return True + return bool(ids[-1] - ids[0] == len(ids) - 1) + + class Interpolator: """Abstract base class for time series interpolation. @@ -177,6 +184,20 @@ class SequenceInterpolator(Interpolator): If True, subtracts mean during normalization. normalize_std_threshold : float, optional Minimum std threshold to prevent division by near-zero values. + signal_ids : array-like of int, optional + Indices of the signals to select. When omitted all signals are used. + + For memory-mapped data (``cache_data=False``) the selection is applied + lazily inside :meth:`interpolate` to avoid materialising the full array: + + * **Contiguous** indices (e.g. ``[3, 4, 5, 6]``) → a strided *view* + of the underlying file is used, so no data is copied. + * **Non-contiguous** indices → the full memmap is kept on disk; only + the rows required by each query are read and then the column subset + is taken in RAM. + + For already-in-memory data (``cache_data=True`` or ``.npy`` files) the + subset is applied once during initialisation. **kwargs Additional keyword arguments (ignored). @@ -187,7 +208,8 @@ class SequenceInterpolator(Interpolator): time_delta : float Time between samples (1 / sampling_rate). n_signals : int - Number of signals (e.g., neurons, behavior channels). + Number of selected signals (equals ``len(signal_ids)`` when provided, + otherwise the total number of signals in the file). Notes ----- @@ -209,6 +231,7 @@ def __init__( normalize: bool = False, normalize_subtract_mean: bool = False, normalize_std_threshold: float | None = None, # or 0.01 + signal_ids: np.ndarray | None = None, **kwargs, ) -> None: super().__init__(root_folder) @@ -226,7 +249,22 @@ def __init__( # Valid interval can be different to start time and end time. self.valid_interval = TimeInterval(self.start_time, self.end_time) - self.n_signals = meta["n_signals"] + # Normalise and sort the requested signal indices (sorted order gives + # better sequential disk-access patterns). + if signal_ids is not None: + signal_ids = np.sort(np.asarray(signal_ids, dtype=np.intp)) + + # _signal_ids is non-None only for the lazy-filter path: a memmap that + # is not fully cached AND whose signal selection is non-contiguous. + # In that case self._data remains the full memmap and the column subset + # is applied per-query inside interpolate(). + self._signal_ids: np.ndarray | None = None + + # Expose the final sorted signal selection so subclasses (e.g. + # PhaseShiftedSequenceInterpolator) can subset their own per-signal + # arrays without re-sorting the input. + self._sorted_signal_ids: np.ndarray | None = signal_ids + # read .mem (memmap) or .npy file if self.is_mem_mapped: self._data = np.memmap( @@ -237,18 +275,50 @@ def __init__( ) if cache_data: - self._data = np.array(self._data).astype( - np.float32 - ) # Convert memmap to ndarray + # Materialise into RAM; subset first to avoid a large + # intermediate allocation when signal_ids is given. + if signal_ids is not None: + self._data = np.asarray( + self._data[:, signal_ids], dtype=np.float32 + ) + else: + self._data = np.array(self._data).astype( + np.float32 + ) # Convert memmap to ndarray + elif signal_ids is not None: + if _is_contiguous_range(signal_ids): + # Strided view — zero-copy, still lazy. + self._data = self._data[ + :, signal_ids[0] : signal_ids[-1] + 1 + ] + else: + # Keep the full memmap on disk; filter lazily per query. + self._signal_ids = signal_ids else: self._data = np.load(self.root_folder / "data.npy") + if signal_ids is not None: + self._data = self._data[:, signal_ids] + + # n_signals reflects the number of signals returned by interpolate(). + self.n_signals = ( + len(self._signal_ids) + if self._signal_ids is not None + else self._data.shape[1] + ) if self.normalize: self.normalize_init() def normalize_init(self): - self.mean = np.load(self.root_folder / "meta/means.npy") - self.std = np.load(self.root_folder / "meta/stds.npy") + mean = np.load(self.root_folder / "meta/means.npy") + std = np.load(self.root_folder / "meta/stds.npy") + # Subset statistics to the selected signals when lazy filtering is + # active (self._data still holds the full array in that case). + if self._signal_ids is not None: + mean = mean[self._signal_ids] + std = std[self._signal_ids] + self.mean = mean + self.std = std assert ( self.mean.shape[0] == self.n_signals ), f"mean shape does not match: {self.mean.shape} vs {self._data.shape}" @@ -284,9 +354,9 @@ def interpolate( stacklevel=2, ) return ( - (np.empty((0, self._data.shape[1])), valid) + (np.empty((0, self.n_signals)), valid) if return_valid - else np.empty((0, self._data.shape[1])) + else np.empty((0, self.n_signals)) ) idx_lower = np.floor((valid_times - self.start_time) / self.time_delta).astype( @@ -295,6 +365,8 @@ def interpolate( if self.interpolation_mode == "nearest_neighbor": data = self._data[idx_lower] + if self._signal_ids is not None: + data = data[:, self._signal_ids] return (data, valid) if return_valid else data @@ -326,6 +398,9 @@ def interpolate( data_lower = self._data[idx_lower] data_upper = self._data[idx_upper] + if self._signal_ids is not None: + data_lower = data_lower[:, self._signal_ids] + data_upper = data_upper[:, self._signal_ids] interpolated = ( lower_signal_ratio * data_lower + upper_signal_ratio * data_upper @@ -359,13 +434,17 @@ class PhaseShiftedSequenceInterpolator(SequenceInterpolator): ---------- root_folder : str Path to the modality directory. Must contain ``meta/phase_shifts.npy``. + signal_ids : array-like of int, optional + See :class:`SequenceInterpolator`. Phase shifts are subsetted to the + same selection. **kwargs All parameters from :class:`SequenceInterpolator`. Attributes ---------- _phase_shifts : numpy.ndarray - Per-signal phase shift values in seconds. + Per-signal phase shift values in seconds (subset when *signal_ids* is + given). """ def __init__( @@ -377,6 +456,7 @@ def __init__( normalize: bool = False, normalize_subtract_mean: bool = False, normalize_std_threshold: float | None = None, # or 0.01 + signal_ids: np.ndarray | None = None, **kwargs, ) -> None: super().__init__( @@ -387,10 +467,14 @@ def __init__( normalize=normalize, normalize_subtract_mean=normalize_subtract_mean, normalize_std_threshold=normalize_std_threshold, + signal_ids=signal_ids, **kwargs, ) - self._phase_shifts = np.load(self.root_folder / "meta/phase_shifts.npy") + phase_shifts_all = np.load(self.root_folder / "meta/phase_shifts.npy") + if self._sorted_signal_ids is not None: + phase_shifts_all = phase_shifts_all[self._sorted_signal_ids] + self._phase_shifts = phase_shifts_all self.valid_interval = TimeInterval( self.start_time + (np.max(self._phase_shifts) if len(self._phase_shifts) > 0 else 0), @@ -411,9 +495,9 @@ def interpolate( stacklevel=2, ) return ( - (np.empty((0, self._data.shape[1])), valid) + (np.empty((0, self.n_signals)), valid) if return_valid - else np.empty((0, self._data.shape[1])) + else np.empty((0, self.n_signals)) ) idx_lower = np.floor( @@ -426,7 +510,17 @@ def interpolate( ).astype(int) if self.interpolation_mode == "nearest_neighbor": - data = np.take_along_axis(self._data, idx_lower, axis=0) + if self._signal_ids is not None: + # self._data is the full memmap with N_all columns, while + # idx_lower has only N_subset columns. np.take_along_axis + # requires matching shapes, so we use explicit fancy indexing + # to map each subset position to its original column index. + cols = np.broadcast_to( + self._signal_ids[np.newaxis, :], idx_lower.shape + ) + data = self._data[idx_lower, np.asarray(cols)] + else: + data = np.take_along_axis(self._data, idx_lower, axis=0) return (data, valid) if return_valid else data elif self.interpolation_mode == "linear": @@ -456,6 +550,10 @@ def interpolate( upper_signal_ratio = upper_numerator / denom _, cols = np.indices(idx_lower.shape) + if self._signal_ids is not None: + # Remap 0-based column positions to original signal indices in + # the full memmap. + cols = self._signal_ids[cols] data_lower = self._data[idx_lower, cols] data_upper = self._data[idx_upper, cols] diff --git a/tests/test_sequence_interpolator.py b/tests/test_sequence_interpolator.py index 8ec6993..6695f85 100644 --- a/tests/test_sequence_interpolator.py +++ b/tests/test_sequence_interpolator.py @@ -1,12 +1,14 @@ import numpy as np import pytest +from contextlib import closing from experanto.interpolators import ( + Interpolator, PhaseShiftedSequenceInterpolator, SequenceInterpolator, ) -from .create_sequence_data import sequence_data_and_interpolator +from .create_sequence_data import create_sequence_data, sequence_data_and_interpolator DEFAULT_SEQUENCE_LENGTH = 10 @@ -547,5 +549,155 @@ def test_interpolation_mode_not_implemented(): seq_interp.interpolate(np.array([0.0, 1.0, 2.0]), return_valid=True) +# --------------------------------------------------------------------------- +# Tests for signal_ids filtering +# --------------------------------------------------------------------------- + +# Non-contiguous ids exercise the lazy-filter path on memmaps and the +# direct-subset path on npy/cached data. +NON_CONTIGUOUS_IDS = np.array([0, 2, 7, 9]) # 4 out of 10 signals +CONTIGUOUS_IDS = np.array([3, 4, 5, 6]) # 4 contiguous signals + + +@pytest.mark.parametrize("interpolation_mode", ["nearest_neighbor", "linear"]) +@pytest.mark.parametrize( + "signal_ids", [NON_CONTIGUOUS_IDS, CONTIGUOUS_IDS, np.array([5])] +) +@pytest.mark.parametrize( + "use_mem_mapped,cache_data", + [(False, False), (True, False), (True, True)], +) +def test_signal_ids_output_matches_full_interpolation( + interpolation_mode, signal_ids, use_mem_mapped, cache_data +): + """Filtered output must equal the corresponding columns of the full output.""" + n_signals = 10 + data_kwargs = { + "n_signals": n_signals, + "use_mem_mapped": use_mem_mapped, + "t_end": 5.0, + "sampling_rate": 10.0, + } + with create_sequence_data(**data_kwargs) as (timestamps, _, _): + times = timestamps[1 : DEFAULT_SEQUENCE_LENGTH + 1] + 1e-9 + + with closing( + Interpolator.create("tests/sequence_data", cache_data=cache_data) + ) as full_interp: + full_interp.interpolation_mode = interpolation_mode + full_out, valid_full = full_interp.interpolate( + times=times, return_valid=True + ) + + with closing( + Interpolator.create( + "tests/sequence_data", + cache_data=cache_data, + signal_ids=signal_ids, + ) + ) as sub_interp: + sub_interp.interpolation_mode = interpolation_mode + + assert sub_interp.n_signals == len( + signal_ids + ), "n_signals should reflect the number of selected signals" + + sub_out, valid_sub = sub_interp.interpolate( + times=times, return_valid=True + ) + + sorted_ids = np.sort(signal_ids) + assert sub_out.shape == (len(valid_sub), len(sorted_ids)) + np.testing.assert_array_equal(valid_full, valid_sub) + np.testing.assert_allclose( + sub_out, + full_out[:, sorted_ids], + rtol=1e-5, + err_msg="Filtered output does not match the corresponding columns of full output", + ) + + +@pytest.mark.parametrize("interpolation_mode", ["nearest_neighbor", "linear"]) +@pytest.mark.parametrize( + "signal_ids", [NON_CONTIGUOUS_IDS, CONTIGUOUS_IDS] +) +def test_signal_ids_memmap_not_materialised_when_lazy(interpolation_mode, signal_ids): + """For a non-contiguous selection on an uncached memmap the underlying + ``_data`` must remain a ``np.memmap`` (lazy access, no full copy).""" + n_signals = 10 + with sequence_data_and_interpolator( + data_kwargs={ + "n_signals": n_signals, + "use_mem_mapped": True, + "t_end": 5.0, + "sampling_rate": 10.0, + }, + interp_kwargs={"cache_data": False, "signal_ids": signal_ids}, + ) as (timestamps, _, _, sub_interp): + assert isinstance( + sub_interp._data, np.memmap + ), "_data should remain a memmap when cache_data=False" + + sub_interp.interpolation_mode = interpolation_mode + times = timestamps[1 : DEFAULT_SEQUENCE_LENGTH + 1] + 1e-9 + out, _ = sub_interp.interpolate(times=times, return_valid=True) + assert out.shape[1] == len(signal_ids) + + +@pytest.mark.parametrize("interpolation_mode", ["nearest_neighbor", "linear"]) +@pytest.mark.parametrize( + "signal_ids", [NON_CONTIGUOUS_IDS, CONTIGUOUS_IDS] +) +def test_phase_shifted_signal_ids_output_matches_full_interpolation( + interpolation_mode, signal_ids +): + """PhaseShiftedSequenceInterpolator with signal_ids must return the same + columns as the unfiltered interpolator for those signals.""" + n_signals = 10 + with create_sequence_data( + n_signals=n_signals, + use_mem_mapped=True, + t_end=5.0, + sampling_rate=10.0, + shifts_per_signal=True, + ) as (timestamps, _, _): + times = timestamps[2 : DEFAULT_SEQUENCE_LENGTH + 2] + 1e-9 + + with closing( + Interpolator.create("tests/sequence_data") + ) as full_interp: + full_interp.interpolation_mode = interpolation_mode + full_out, valid_full = full_interp.interpolate( + times=times, return_valid=True + ) + + with closing( + Interpolator.create("tests/sequence_data", signal_ids=signal_ids) + ) as sub_interp: + assert isinstance( + sub_interp, PhaseShiftedSequenceInterpolator + ), "Expected PhaseShiftedSequenceInterpolator" + sub_interp.interpolation_mode = interpolation_mode + + assert sub_interp.n_signals == len(signal_ids) + assert len(sub_interp._phase_shifts) == len(signal_ids), ( + "_phase_shifts should be subsetted to the selected signals" + ) + + sub_out, valid_sub = sub_interp.interpolate( + times=times, return_valid=True + ) + + sorted_ids = np.sort(signal_ids) + assert sub_out.shape == (len(valid_sub), len(sorted_ids)) + np.testing.assert_array_equal(valid_full, valid_sub) + np.testing.assert_allclose( + sub_out, + full_out[:, sorted_ids], + rtol=1e-5, + err_msg="Filtered PhaseShifted output does not match corresponding columns", + ) + + if __name__ == "__main__": pytest.main([__file__])