diff --git a/doc/api.rst b/doc/api.rst index e76e040b..98f1b981 100644 --- a/doc/api.rst +++ b/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 diff --git a/neurodsp/tests/utils/test_yielders.py b/neurodsp/tests/utils/test_yielders.py new file mode 100644 index 00000000..cf3f5f79 --- /dev/null +++ b/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) diff --git a/neurodsp/utils/yielders.py b/neurodsp/utils/yielders.py new file mode 100644 index 00000000..132a434e --- /dev/null +++ b/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