Merge pull request #288 from fooof-tools/interpg

[ENH] - Add `interpolate_spectra` function

Author: TomDonoghue

doc/api.rst

@@ -341,6 +341,7 @@ Utilities for working with data.
 
     trim_spectrum
     interpolate_spectrum
+    interpolate_spectra
     subsample_spectra
 
 Parameter Utilities

fooof/core/modutils.py

@@ -164,7 +164,7 @@ def docs_append_to_section(docstring, section, add):
                         for split in docstring.split('\n\n')])
 
 
-def docs_get_section(docstring, section, output='extract'):
+def docs_get_section(docstring, section, output='extract', end=None):
     """Extract and/or remove a specified section from a docstring.
 
     Parameters
@@ -177,6 +177,9 @@ def docs_get_section(docstring, section, output='extract'):
         Run mode, options:
             'extract' - returns the extracted section from the docstring.
             'remove' - returns the docstring after removing the specified section.
+    end : str, optional
+        Indicates the contents of a line that signals the end of the section to select.
+        If not provided, the section is selected until a blank line.
 
     Returns
     -------
@@ -193,7 +196,12 @@ def docs_get_section(docstring, section, output='extract'):
         # Track whether in the desired section
         if section in line and '--' in docstring_split[ind + 1]:
             in_section = True
-        if in_section and line == '':
+        if end:
+            if in_section and '    ' + end == line:
+                in_section = False
+                # In this approach, an extra newline is caught - so pop it off
+                outs.pop()
+        elif in_section and line == '':
             in_section = False
 
         # Collect desired outputs based on whether extracting or removing section

fooof/tests/core/test_modutils.py

@@ -73,6 +73,11 @@ def test_docs_get_section(tdocstring):
     assert 'Parameters' not in out2
     assert 'Returns' in out2
 
+    # Test with end_selection
+    out3 = docs_get_section(tdocstring, 'Parameters', output='extract', end='Returns')
+    assert 'Parameters' in out3
+    assert 'Returns' not in out3
+
 def test_docs_add_section(tdocstring):
 
     tdocstring = tdocstring + \

fooof/tests/utils/test_data.py

@@ -50,6 +50,21 @@ def test_interpolate_spectrum():
         mask = np.logical_and(freqs >= f_range[0], freqs <= f_range[1])
         assert powers[mask].sum() > powers_out[mask].sum()
 
+def test_interpolate_spectra():
+
+    freqs, powers = gen_group_power_spectra(\
+        5, [1, 150], [1, 100, 1], [[10, 0.5, 1.0], [60, 1, 0.1], [120, 0.5, 0.1]])
+
+    exclude = [[58, 62], [118, 122]]
+    freqs_out, powers_out = interpolate_spectra(freqs, powers, exclude)
+    assert np.array_equal(freqs, freqs_out)
+    assert np.all(powers)
+    assert powers.shape == powers_out.shape
+
+    for f_range in exclude:
+        mask = np.logical_and(freqs >= f_range[0], freqs <= f_range[1])
+        assert powers[:, mask].sum() > powers_out[:, mask].sum()
+
 def test_subsample_spectra():
 
     # Simulate spectra, each with unique osc peak (for checking)

fooof/utils/data.py

@@ -1,9 +1,12 @@
 """Utilities for working with data and models."""
 
 from itertools import repeat
+from functools import partial
 
 import numpy as np
 
+from fooof.core.modutils import docs_get_section, replace_docstring_sections
+
 ###################################################################################################
 ###################################################################################################
 
@@ -138,6 +141,55 @@ def interpolate_spectrum(freqs, powers, interp_range, buffer=3):
     return freqs, powers
 
 
+def wrap_interpolate_spectrum(powers, freqs, interp_range, buffer):
+    """Wraps interpolate function, organizing inputs & outputs to use `partial`."""
+    return interpolate_spectrum(freqs, powers, interp_range, buffer)[1]
+
+
+@replace_docstring_sections(docs_get_section(interpolate_spectrum.__doc__, 'Notes', end='Examples'))
+def interpolate_spectra(freqs, powers, interp_range, buffer=3):
+    """Interpolate a frequency region across a group of power spectra.
+
+    Parameters
+    ----------
+    freqs : 1d array
+        Frequency values for the power spectrum.
+    powers : 2d array
+        Power values for the power spectra.
+    interp_range : list of float or list of list of float
+        Frequency range to interpolate, as [lowest_freq, highest_freq].
+        If a list of lists, applies each as it's own interpolation range.
+    buffer : int or list of int
+        The number of samples to use on either side of the interpolation
+        range, that are then averaged and used to calculate the interpolation.
+
+    Returns
+    -------
+    freqs : 1d array
+        Frequency values for the power spectrum.
+    powers : 2d array
+        Power values, with interpolation, for the power spectra.
+
+    Notes
+    -----
+    % copied in from interpolate_spectrum
+
+    Examples
+    --------
+    Using simulated spectra, interpolate away line noise peaks:
+
+    >>> from fooof.sim import gen_group_power_spectra
+    >>> freqs, powers = gen_group_power_spectra(5, [1, 75], [1, 1], [[10, 0.5, 1.0], [60, 2, 0.1]])
+    >>> freqs, powers = interpolate_spectra(freqs, powers, [58, 62])
+    """
+
+    tfunc = partial(wrap_interpolate_spectrum, freqs=freqs,
+                    interp_range=interp_range, buffer=buffer)
+    powers = np.apply_along_axis(tfunc, 1, powers)
+
+    return freqs,powers
+
+
 def subsample_spectra(spectra, selection, return_inds=False):
     """Subsample a group of power spectra.