Repository: the-siesta-group/edfio

Issue:
Read only specific time segments?
I was wondering if it would be useful to be able to read only specific time segments (as opposed to the entire data), WDYT? This is probably a very niche feature, but I'm bringing it up because if we ever wanted to use edfio as a reader in MNE, then we'd need this functionality (and this would be a separate discussion on its own).

Hints from issue comments:
@hofaflo: I think this would be a good addition! There might need to be different options to handle segment limits which do not coincide with sample times for all signals (in `Edf.slice_between_seconds` we currently take the easiest way by failing if that is the case). Would you be willing to create a PR for this?@cbrnr: I don't have the time to implement this feature currently. It is also not urgent at all, I just wanted to know if you were open to this addition in principle – so thanks, great 👍. If MNE devs decide that it would be worth having edfio replace or supplement their own reader, I can tackle this of course (but feel free to beat me to it if you want/need this feature earlier).@marcoross: This would be a nice feature as well if we ever wanted to implement a viewer or read partial data from a network location for machine learning purposes.@jtmiles: I also think this would be very helpful! I am sometimes given cumbersome (~100GB) clinical iEEG data files that I only need a small time window and select channels from. Additional slicing functionality would be awesome!@hofaflo: As I'm not sure if #58 applies to your use cases @cbrnr and @jtmiles, I'll reopen this for now@cbrnr: I think this looks good for now 👍! If there is anything still missing for an MNE reader using edfio, I'll let you know 😄 (I don't have time to implement it in the near future).@jtmiles: Hi, hope that I haven't held up closing this. Yes, this seems to work for me - thanks very much!

Solution patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 79a8604..d047d58 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+### Added
+- Load a segment of EDF raw data from a signal without reading the entire file using `EdfSignal.get_data_slice` method ([#58](https://github.com/the-siesta-group/edfio/pull/58)).
+- Load annotations from a specific time window of the EDF file without reading the entire file using `Edf.get_annotations` method ([#58](https://github.com/the-siesta-group/edfio/pull/58)).
+
 ## [0.4.4] - 2024-09-25
 
 ### Added
diff --git a/edfio/_lazy_loading.py b/edfio/_lazy_loading.py
new file mode 100644
index 0000000..cc0d943
--- /dev/null
+++ b/edfio/_lazy_loading.py
@@ -0,0 +1,61 @@
+from typing import Any, Optional, Union
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+class LazyLoader:
+    """
+    Class to load data for a single signal from a buffer of EDF data records (array or memmap).
+
+    Parameters
+    ----------
+    buffer : numpy.ndarray or numpy.memmap
+        Buffer of EDF data records.
+    start_sample : int
+        Offset of the signal in the data records (in samples).
+    end_sample : int
+        End of the signal in the data records (in samples).
+    """
+
+    def __init__(
+        self,
+        buffer: Union[NDArray[np.int16], np.memmap[Any, np.dtype[np.int16]]],
+        start_sample: int,
+        end_sample: int,
+    ) -> None:
+        self.buffer = buffer
+        self.start_sample = start_sample
+        self.end_sample = end_sample
+
+    def load(
+        self, start_record: Optional[int] = None, end_record: Optional[int] = None
+    ) -> NDArray[np.int16]:
+        """
+        Load signal data from the buffer.
+
+        Parameters
+        ----------
+        start_record : int, optional
+            The first EDF data record to load samples from. If None, load from the beginning of the buffer.
+        end_record : int, optional
+            The last EDF data record to load samples from. If None, load until the end of the buffer.
+
+        Returns
+        -------
+        numpy.ndarray
+            Signal data (digital).
+        """
+        if start_record is None:
+            start_record = 0
+        if end_record is None:
+            end_record = self.buffer.shape[0]
+        if (
+            end_record < start_record
+            or start_record < 0
+            or end_record > self.buffer.shape[0]
+        ):
+            raise ValueError("Invalid slice: Slice exceeds EDF duration")
+        return self.buffer[
+            start_record:end_record, self.start_sample : self.end_sample
+        ].flatten()
diff --git a/edfio/edf.py b/edfio/edf.py
index 69a0e2c..7a5f521 100644
--- a/edfio/edf.py
+++ b/edfio/edf.py
@@ -16,7 +16,6 @@
 from typing import Any, Literal
 
 import numpy as np
-import numpy.typing as npt
 
 from edfio._header_field import (
     decode_date,
@@ -29,6 +28,7 @@
     encode_str,
     encode_time,
 )
+from edfio._lazy_loading import LazyLoader
 from edfio.edf_annotations import (
     EdfAnnotation,
     _create_annotations_signal,
@@ -248,11 +248,7 @@ def _load_data(
 
         for signal, start, end in zip(self._signals, starts, ends):
             if lazy_load_data:
-
-                def lazy_load(s: int = start, e: int = end) -> npt.NDArray[np.int16]:
-                    return datarecords[:, s:e].flatten()
-
-                signal._lazy_loader = lazy_load
+                signal._lazy_loader = LazyLoader(datarecords, start, end)
             else:
                 signal._digital = datarecords[:, start:end].flatten()
 
@@ -772,16 +768,40 @@ def duration(self) -> float:
         """Recording duration in seconds."""
         return self.num_data_records * self.data_record_duration
 
-    @property
-    def annotations(self) -> tuple[EdfAnnotation, ...]:
+    def get_annotations(
+        self, start_second: float | None = None, stop_second: float | None = None
+    ) -> tuple[EdfAnnotation, ...]:
         """
-        All annotations contained in the Edf, sorted chronologically.
+        All annotations defined (starting) in a specified time region of the Edf, sorted chronologically.
 
-        Does not include timekeeping annotations.
+        Does not include timekeeping annotations. Will not include annotations that start within the specified time region but are
+        defined in data records outside of that window.
+
+        Parameters
+        ----------
+        start_second : float, optional
+            The start of the time region in seconds from recording start. If not provided, the start of the recording is used.
+        stop_second : float, optional
+            The end of the time region in seconds. If not provided, the end of the recording is used.
         """
         annotations: list[EdfAnnotation] = []
+        start_second_defined = start_second is not None
+        stop_second_defined = stop_second is not None
+        start_second = start_second or 0
+        stop_second = stop_second or self.duration
         for i, signal in enumerate(self._annotation_signals):
-            for data_record in signal.digital.reshape(
+            if self.duration == 0 or signal.sampling_frequency == 0:
+                digital_slice = signal.digital
+            else:
+                # Make sure to read full data records
+                start_record = start_second // self.data_record_duration
+                adjusted_start = start_record * self.data_record_duration
+                last_record = ceil(
+                    stop_second / self.data_record_duration - 1e-12
+                )  # avoid floating point errors
+                adjusted_stop = last_record * self.data_record_duration
+                digital_slice = signal.get_digital_slice(adjusted_start, adjusted_stop)
+            for data_record in digital_slice.reshape(
                 (-1, signal.samples_per_data_record)
             ):
                 annot_dr = _EdfAnnotationsDataRecord.from_bytes(data_record.tobytes())
@@ -800,7 +820,23 @@ def annotations(self) -> tuple[EdfAnnotation, ...]:
             )
             for ann in annotations
         ]
-        return tuple(sorted(annotations))
+        annotations = sorted(annotations)
+        if start_second_defined:
+            while annotations and annotations[0].onset < start_second:
+                annotations.pop(0)
+        if stop_second_defined:
+            while annotations and annotations[-1].onset >= stop_second:
+                annotations.pop()
+        return tuple(annotations)
+
+    @property
+    def annotations(self) -> tuple[EdfAnnotation, ...]:
+        """
+        All annotations contained in the Edf, sorted chronologically.
+
+        Does not include timekeeping annotations.
+        """
+        return self.get_annotations()
 
     def drop_annotations(self, text: str) -> None:
         """
diff --git a/edfio/edf_signal.py b/edfio/edf_signal.py
index ee3251c..42526db 100644
--- a/edfio/edf_signal.py
+++ b/edfio/edf_signal.py
@@ -14,6 +14,7 @@
     encode_int,
     encode_str,
 )
+from edfio._lazy_loading import LazyLoader
 
 
 class _IntRange(NamedTuple):
@@ -99,7 +100,7 @@ class EdfSignal:
     )
 
     _digital: npt.NDArray[np.int16] | None = None
-    _lazy_loader: Callable[[], npt.NDArray[np.int16]] | None = None
+    _lazy_loader: LazyLoader | None = None
 
     def __init__(
         self,
@@ -309,19 +310,11 @@ def digital(self) -> npt.NDArray[np.int16]:
         if self._digital is None:
             if self._lazy_loader is None:
                 raise ValueError("Signal data not set")
-            self._digital = self._lazy_loader()
+            self._digital = self._lazy_loader.load()
             self._lazy_loader = None
         return self._digital
 
-    @property
-    def data(self) -> npt.NDArray[np.float64]:
-        """
-        Numpy array containing the physical signal values as floats.
-
-        To simplify avoiding inconsistencies between signal data and header fields,
-        individual values in the returned array can not be modified. Use
-        :meth:`EdfSignal.update_data` to overwrite with new physical data.
-        """
+    def _calibrate(self, digital: npt.NDArray[np.int16]) -> npt.NDArray[np.float64]:
         try:
             gain, offset = _calculate_gain_and_offset(
                 self.digital_min,
@@ -330,7 +323,7 @@ def data(self) -> npt.NDArray[np.float64]:
                 self.physical_max,
             )
         except ZeroDivisionError:
-            data = self.digital.astype(np.float64)
+            data = digital.astype(np.float64)
             if self.digital_max == self.digital_min:
                 warnings.warn(
                     f"Digital minimum equals digital maximum ({self.digital_min}) for {self.label}, returning uncalibrated signal."
@@ -340,12 +333,77 @@ def data(self) -> npt.NDArray[np.float64]:
                     f"Physical minimum equals physical maximum ({self.physical_min}) for {self.label}, returning uncalibrated signal."
                 )
         except ValueError:
-            data = self.digital.astype(np.float64)
+            data = digital.astype(np.float64)
         else:
-            data = (self.digital + offset) * gain
+            data = (digital + offset) * gain
         data.setflags(write=False)
         return data
 
+    @property
+    def data(self) -> npt.NDArray[np.float64]:
+        """
+        Numpy array containing the physical signal values as floats.
+
+        To simplify avoiding inconsistencies between signal data and header fields,
+        individual values in the returned array can not be modified. Use
+        :meth:`EdfSignal.update_data` to overwrite with new physical data.
+        """
+        return self._calibrate(self.digital)
+
+    def get_digital_slice(
+        self, start_second: float, stop_second: float
+    ) -> npt.NDArray[np.int16]:
+        """
+        Get a slice of the digital signal values.
+
+        If the signal has not been loaded into memory so far, only the requested slice will be read.
+
+        Parameters
+        ----------
+        start_second : float
+            The start of the slice in seconds.
+        stop_second : float
+            The end of the slice in seconds.
+        """
+        duration = stop_second - start_second
+        if duration < 0:
+            raise ValueError("Invalid slice: Duration must be non-negative")
+        if start_second < 0:
+            raise ValueError("Invalid slice: Start second must be non-negative")
+        start_index = round(start_second * self.sampling_frequency)
+        end_index = round(stop_second * self.sampling_frequency)
+        if self._digital is not None:
+            if end_index > len(self._digital):
+                raise ValueError("Invalid slice: Slice exceeds EDF duration")
+            return self._digital[start_index:end_index]
+        if self._lazy_loader is None:
+            raise ValueError("Signal data not set")
+        first_data_record = start_index // self.samples_per_data_record
+        last_data_record = (end_index - 1) // self.samples_per_data_record + 1
+        digital_portion = self._lazy_loader.load(first_data_record, last_data_record)
+        offset_within_first_record = start_index % self.samples_per_data_record
+        num_samples = end_index - start_index
+        return digital_portion[
+            offset_within_first_record : offset_within_first_record + num_samples
+        ]
+
+    def get_data_slice(
+        self, start_second: float, stop_second: float
+    ) -> npt.NDArray[np.float64]:
+        """
+        Get a slice of the signal data.
+
+        If the signal has not been loaded into memory so far, only the requested slice will be read.
+
+        Parameters
+        ----------
+        start_second : float
+            The start of the slice in seconds.
+        stop_second : float
+            The end of the slice in seconds.
+        """
+        return self._calibrate(self.get_digital_slice(start_second, stop_second))
+
     def update_data(
         self,
         data: npt.NDArray[np.float64],


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 5e7b6ed..70eeb37 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,8 +1,27 @@
 from pathlib import Path
 
+import numpy as np
 import pytest
 
+from edfio._lazy_loading import LazyLoader
+
 
 @pytest.fixture()
 def tmp_file(tmp_path: Path) -> Path:
     return tmp_path / "target.edf"
+
+
+@pytest.fixture()
+def buffered_lazy_loader() -> LazyLoader:
+    # Buffer of 4 raw EDF records containing 12 samples for the tested signal.
+    # Data for other signals is set to zero.
+    data_records = np.array(
+        [
+            [0, 0, 0, 1, 2, 3, 0, 0],
+            [0, 0, 0, 4, 5, 6, 0, 0],
+            [0, 0, 0, 7, 8, 9, 0, 0],
+            [0, 0, 0, 10, 11, 12, 0, 0],
+        ],
+        dtype=np.int16,
+    )
+    return LazyLoader(data_records, 3, 6)
diff --git a/tests/test_edf.py b/tests/test_edf.py
index 862a827..b4bf7fd 100644
--- a/tests/test_edf.py
+++ b/tests/test_edf.py
@@ -1174,7 +1174,6 @@ def test_sampling_frequencies_leading_to_floating_point_issues_in_signal_duratio
         (0,              11,                  "EDF header indicates 11 data records, but file contains 10 records. Updating header."),
     ],
 )
-# fmt: on
 def test_read_edf_with_invalid_number_of_records(
     tmp_path: Path,
     extra_bytes: int,
@@ -1197,11 +1196,14 @@ def test_read_edf_with_invalid_number_of_records(
         comparison_edf = read_edf(EDF_FILE)
         for signal, comparison_signal in zip(edf.signals, comparison_edf.signals):
             np.testing.assert_array_equal(signal.data, comparison_signal.data)
+# fmt: on
 
 
 def test_fail_lazy_load():
     edf_data = EDF_FILE.read_bytes()
-    with pytest.raises(ValueError, match="Lazy loading is only supported for local file paths"):
+    with pytest.raises(
+        ValueError, match="Lazy loading is only supported for local file paths"
+    ):
         read_edf(edf_data, lazy_load_data=True)
 
 
diff --git a/tests/test_edf_annotations.py b/tests/test_edf_annotations.py
index 1d0fb4d..229e471 100644
--- a/tests/test_edf_annotations.py
+++ b/tests/test_edf_annotations.py
@@ -465,3 +465,14 @@ def test_add_annotations(tmp_file: Path):
     edf.write(tmp_file)
     edf = read_edf(tmp_file)
     assert edf.annotations == tuple(sorted(old_annotations + new_annotations))
+
+
+def test_get_annotations_defined_in_timewindow():
+    all_annotations = [
+        EdfAnnotation(0.15, 3, "A"),
+        EdfAnnotation(3.5, 1, "B"),
+        EdfAnnotation(5, 1, "C"),
+        EdfAnnotation(7.8, 1, "D"),
+    ]
+    edf = Edf(signals=[EdfSignal(np.arange(10), 1)], annotations=all_annotations)
+    assert edf.get_annotations(0.9, 7.6) == tuple(all_annotations[1:3])
diff --git a/tests/test_edf_signal.py b/tests/test_edf_signal.py
index b769c12..87bd821 100644
--- a/tests/test_edf_signal.py
+++ b/tests/test_edf_signal.py
@@ -8,6 +8,7 @@
 import pytest
 
 from edfio import EdfSignal
+from edfio._lazy_loading import LazyLoader
 from edfio.edf_signal import (
     _FloatRange,
     _IntRange,
@@ -45,7 +46,6 @@
         (-1111111.5,    -1111112,         -1111112,         -1111111,),
     ],
 )
-# fmt: on
 def test_round_float_to_8_characters(
     value: float,
     expected_round: float,
@@ -55,6 +55,7 @@ def test_round_float_to_8_characters(
     assert _round_float_to_8_characters(value, round) == expected_round
     assert _round_float_to_8_characters(value, math.floor) == expected_floor
     assert _round_float_to_8_characters(value, math.ceil) == expected_ceil
+# fmt: on
 
 
 def sine(duration, f, fs):
@@ -367,3 +368,96 @@ def test_edf_signal_init_does_not_accept_edf_annotations_as_label():
 def test_edf_signal_label_cannot_be_set_to_edf_annotations():
     with pytest.raises(ValueError, match="must not be 'EDF Annotations'"):
         EdfSignal(np.arange(2), 1).label = "EDF Annotations"
+
+
+def test_get_data_slice_already_loaded():
+    signal = EdfSignal(
+        np.arange(10),
+        sampling_frequency=2,
+        digital_range=(0, 9),
+        physical_range=(0, 9),
+    )
+    slice = signal.get_data_slice(1.5, 4.5)
+    np.testing.assert_array_equal(slice, np.arange(3, 9))
+
+
+@pytest.mark.parametrize(
+    ("start", "stop"),
+    [
+        (-1.0, 3.0),
+        (1.5, 5.5),
+    ],
+)
+def test_get_data_slice_outside_of_bounds_already_loaded(start: float, stop: float):
+    signal = EdfSignal(
+        np.arange(10), sampling_frequency=2, digital_range=(0, 9), physical_range=(0, 9)
+    )
+    with pytest.raises(ValueError, match="Invalid slice"):
+        signal.get_data_slice(start, stop)
+
+
+@pytest.fixture()
+def lazy_loaded_signal(buffered_lazy_loader: LazyLoader) -> EdfSignal:
+    # Simulate initialization of the signal from the buffer.
+    signal = EdfSignal._from_raw_header(
+        sampling_frequency=3,
+        label=b"SpO2",
+        transducer_type=b"Pulse Oximeter",
+        physical_dimension=b"%",
+        physical_min=b"0",
+        physical_max=b"100",
+        digital_min=b"0",
+        digital_max=b"1000",
+        prefiltering=b"",
+        samples_per_data_record=b"3",
+        reserved=b"",
+    )
+    signal._lazy_loader = buffered_lazy_loader
+    return signal
+
+
+@pytest.mark.parametrize(
+    ("start", "stop"),
+    [
+        (0.0, 4.0),
+        (1.0, 3.0),
+        (0.67, 1.0),
+        (2.0, 3.33),
+        (1.33, 2.67),
+        (1.33, 1.33),
+    ],
+)
+def test_lazy_load_get_data_slice(
+    start: float, stop: float, lazy_loaded_signal: EdfSignal
+):
+    expected_digital_values = np.arange(1, 13, dtype=np.int16)
+    expected_digital_slice = expected_digital_values[round(start * 3) : round(stop * 3)]
+    actual_digital_slice = lazy_loaded_signal.get_digital_slice(start, stop)
+    np.testing.assert_array_equal(actual_digital_slice, expected_digital_slice)
+
+    # Expected signal values for the slice.
+    expected_data_slice = expected_digital_slice.astype(np.float64) / 10
+    actual_data_slice = lazy_loaded_signal.get_data_slice(start, stop)
+    np.testing.assert_allclose(actual_data_slice, expected_data_slice, atol=1e-14)
+
+
+@pytest.mark.parametrize(
+    ("start", "stop"),
+    [
+        (-1.0, 3.0),
+        (1.0, 5.0),
+        (4.33, 4.67),
+        (2.0, 1.0),
+    ],
+)
+def test_lazy_load_get_data_slice_outside_of_bounds(
+    start: float, stop: float, lazy_loaded_signal: EdfSignal
+):
+    with pytest.raises(ValueError, match="Invalid slice"):
+        lazy_loaded_signal.get_data_slice(start, stop)
+
+
+def test_get_data_slice_with_no_data_available(lazy_loaded_signal: EdfSignal):
+    lazy_loaded_signal._lazy_loader = None
+    with pytest.raises(ValueError, match="Signal data not set"):
+        lazy_loaded_signal.get_data_slice(0, 1)
diff --git a/tests/test_lazy_loading.py b/tests/test_lazy_loading.py
new file mode 100644
index 0000000..8120d75
--- /dev/null
+++ b/tests/test_lazy_loading.py
@@ -0,0 +1,29 @@
+import numpy as np
+import pytest
+
+from edfio._lazy_loading import LazyLoader
+
+
+def test_load_entire_signal_data(buffered_lazy_loader: LazyLoader):
+    expected_data = np.arange(1, 13, dtype=np.int16)
+    np.testing.assert_array_equal(buffered_lazy_loader.load(), expected_data)
+
+
+def test_load_some_records(buffered_lazy_loader: LazyLoader):
+    expected_data = np.arange(4, 10, dtype=np.int16)
+    np.testing.assert_array_equal(buffered_lazy_loader.load(1, 3), expected_data)
+
+
+@pytest.mark.parametrize(
+    ("start_record", "end_record"),
+    [
+        (-1, 3),
+        (1, 5),
+        (3, 2),
+    ],
+)
+def test_load_invalid_records(
+    start_record: int, end_record: int, buffered_lazy_loader: LazyLoader
+):
+    with pytest.raises(ValueError, match="Invalid slice: Slice exceeds EDF duration"):
+        buffered_lazy_loader.load(start_record, end_record)