Skip to content

Commit 5109345

Browse files
author
Laurent Mackay
committed
Adds more flexibility in the aperiodic/oscillatory fitting
DESCRIPTION: I have added a keyword `ap_range` to the FOOOF.fit() function. When `ap_range` is specified, the already existing `freq_range` keyword is used to determine the frequencies for the oscillatory fits while the ap_range keyword is used for the aperiodic fit. When it is not specified, everything proceeds as before. This allows the two parts FO and OOF to be more-or-less independent of one another if so-desired. Furthermore, the `ap_range` keyword may also be a set of indices in order to allow for a very simple implementation of exclusion zones. TESTING: Tested on 5 different EEG files, no errors unless inputs are the wrong size.
1 parent 9a7fabb commit 5109345

File tree

1 file changed

+30
-14
lines changed

1 file changed

+30
-14
lines changed

fooof/fit.py

Lines changed: 30 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -330,14 +330,14 @@ def fit(self, freqs=None, power_spectrum=None, freq_range=None, ap_range=None):
330330
Frequency values for the power spectrum, in linear space.
331331
power_spectrum : 1d array, optional
332332
Power values, which must be input in linear space.
333-
freq_range : list of [float, float], optional
334-
Frequency range to restrict power spectrum to. If not provided, keeps the entire range.
333+
freq_range : list of [float, float], or np.ndarray of booleans of the same length as freqs, optional
334+
Frequency range/indices to restrict power spectrum to. If not provided, keeps the entire range.
335+
ap_range :
335336
336337
Notes
337338
-----
338339
Data is optional if data has been already been added to FOOOF object.
339340
"""
340-
341341
# If freqs & power_spectrum provided together, add data to object.
342342
if freqs is not None and power_spectrum is not None:
343343
self.add_data(freqs, power_spectrum, freq_range if ap_range is None else None)
@@ -358,9 +358,15 @@ def fit(self, freqs=None, power_spectrum=None, freq_range=None, ap_range=None):
358358
# In rare cases, the model fails to fit. Therefore it's in a try/except
359359
# Cause of failure: RuntimeError, failure to find parameters in curve_fit
360360
try:
361+
362+
if ap_range is not None:#isolate aperiodic frequencies/spectrum
363+
if not isinstance(ap_range,np.ndarray) or ap_range.shape[-1]==2:
364+
ap_inds = (self.freqs >= ap_range[0]) & (self.freqs <= ap_range[1])
365+
elif ap_range.shape[-1]==self.freqs.shape[-1]:
366+
ap_inds = ap_range
367+
else:
368+
raise ValueError('ap_range must have the same length as freqs - can not proceed')
361369

362-
if ap_range:
363-
ap_inds = (self.freqs >= ap_range[0]) & (self.freqs <= ap_range[1])
364370
ap_freqs = self.freqs[ap_inds]
365371
ap_spectrum = self.power_spectrum[ap_inds]
366372
else:
@@ -374,22 +380,27 @@ def fit(self, freqs=None, power_spectrum=None, freq_range=None, ap_range=None):
374380
# Flatten the power_spectrum using fit aperiodic fit
375381
self._spectrum_flat = self.power_spectrum - self._ap_fit
376382

377-
if ap_range:
378-
per_inds = (self.freqs >= ap_range[0]) & (self.freqs <= ap_range[1])
383+
384+
if ap_range is not None:#isolate periodic frequencies/spectrum
385+
per_inds = (self.freqs >= freq_range[0]) & (self.freqs <= freq_range[1])
386+
per_spectrum_flat = np.copy(self._spectrum_flat[per_inds])
387+
self._spectrum_flat = per_spectrum_flat
388+
#save/set some attributes so peak fitting works properly
379389
freqs_0 = self.freqs
380390
self.freqs = self.freqs[per_inds]
381-
per_spectrum_flat = self._spectrum_flat[per_inds]
382391
if freq_range:
392+
freq_range_0 = self.freq_range
383393
self.freq_range = freq_range
384-
else:
385-
per_spectrum_flat = np.copy(self._spectrum_flat)
394+
386395

387396

388397
# Find peaks, and fit them with gaussians
389-
self.gaussian_params_ = self._fit_peaks(per_spectrum_flat)
398+
self.gaussian_params_ = self._fit_peaks(np.copy(self._spectrum_flat))
390399

391-
if ap_range:
392-
self.freqs=freqs_0
400+
if ap_range is not None:
401+
#restore attributes to initial values
402+
self.freqs = freqs_0
403+
self.freq_range = freq_range_0
393404

394405
# Calculate the peak fit
395406
# Note: if no peaks are found, this creates a flat (all zero) peak fit.
@@ -398,9 +409,14 @@ def fit(self, freqs=None, power_spectrum=None, freq_range=None, ap_range=None):
398409
# Create peak-removed (but not flattened) power spectrum.
399410
self._spectrum_peak_rm = self.power_spectrum - self._peak_fit
400411

412+
if ap_range is not None:
413+
ap_spectrum_peak_rm = self._spectrum_peak_rm[ap_inds]
414+
else:
415+
ap_spectrum_peak_rm = self._spectrum_peak_rm
416+
401417
# Run final aperiodic fit on peak-removed power spectrum
402418
# Note: This overwrites previous aperiodic fit
403-
self.aperiodic_params_ = self._simple_ap_fit(self.freqs, self._spectrum_peak_rm)
419+
self.aperiodic_params_ = self._simple_ap_fit(ap_freqs, ap_spectrum_peak_rm)
404420
self._ap_fit = gen_aperiodic(self.freqs, self.aperiodic_params_)
405421

406422
# Create full power_spectrum model fit

0 commit comments

Comments
 (0)