Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[ENH] - Add interpolate_spectra function #288

Merged
merged 5 commits into from
Jul 21, 2023
Merged
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions doc/api.rst
Original file line number Diff line number Diff line change
Expand Up @@ -341,6 +341,7 @@ Utilities for working with data.

trim_spectrum
interpolate_spectrum
interpolate_spectra
subsample_spectra

Parameter Utilities
Expand Down
10 changes: 8 additions & 2 deletions fooof/core/modutils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -177,6 +177,7 @@ 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

Returns
-------
Expand All @@ -193,7 +194,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
Expand Down
5 changes: 5 additions & 0 deletions fooof/tests/core/test_modutils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 + \
Expand Down
15 changes: 15 additions & 0 deletions fooof/tests/utils/test_data.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
52 changes: 52 additions & 0 deletions fooof/utils/data.py
Original file line number Diff line number Diff line change
@@ -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

###################################################################################################
###################################################################################################

Expand Down Expand Up @@ -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.

Expand Down