[ENH] - Add signal yielder functionality

doc/api.rst

@@ -525,3 +525,13 @@ Outliers
     remove_nans
     restore_nans
     discard_outliers
+
+Yielders
+~~~~~~~~
+
+.. currentmodule:: neurodsp.utils.yielders
+.. autosummary::
+    :toctree: generated/
+
+    step_over_time
+    step_over_signals

neurodsp/tests/utils/test_yielders.py

@@ -0,0 +1,44 @@
+"""Tests for neurodsp.utils.yielders."""
+
+from collections.abc import Iterable
+
+import numpy as np
+
+from neurodsp.utils.yielders import *
+
+###################################################################################################
+###################################################################################################
+
+def test_step_over_time(tsig):
+
+    start = 0
+    size = 100
+    step = 50
+
+    yielder = step_over_time(tsig, start, size, step)
+    assert isinstance(yielder, Iterable)
+
+    it0 = next(yielder)
+    np.array_equal(it0, tsig[0:size])
+
+    it1 = next(yielder)
+    np.array_equal(it0, tsig[0+step:step+size])
+
+    # Test non-zero start
+    nzero_start = 25
+    yielder2 = step_over_time(tsig, nzero_start, size, step)
+    assert np.array_equal(next(yielder2), tsig[nzero_start:nzero_start+size])
+
+    for csig in yielder:
+        pass
+
+def test_step_over_signals(tsig2d):
+
+    yielder = step_over_signals(tsig2d)
+    assert isinstance(yielder, Iterable)
+
+    it0 = next(yielder)
+    assert np.array_equal(tsig2d[0, :], it0)
+
+    it1 = next(yielder)
+    assert np.array_equal(tsig2d[1, :], it1)

neurodsp/utils/yielders.py

@@ -0,0 +1,45 @@
+"""Utilities that create yielders to iterate across signals."""
+
+###################################################################################################
+###################################################################################################
+
+def step_over_time(sig, start=0, size=100, step=1):
+    """Step across time segments of a signal.
+
+    Parameters
+    ----------
+    sig : 1d array
+        Time series to iterate across.
+    start : int, optional, default: 0
+        Staring index.
+    size : int, optional, default: 100
+        Size of each segment to yield.
+    step : int, optional, default: 1
+        Step size of each iteration.
+
+    Yields
+    ------
+    segment : 1d array
+        Extracted segment of the time series.
+    """
+
+    for st in range(start, len(sig)-size, step):
+        yield sig[st:st+size]
+
+
+def step_over_signals(signals):
+    """"Step across signals within an array.
+
+    Parameters
+    ----------
+    signals : 2d array
+        Array of signals to iterate across.
+
+    Yields
+    ------
+    sig : 1d array
+        Extracted signal.
+    """
+
+    for sig in signals:
+        yield sig