@@ -330,14 +330,14 @@ def fit(self, freqs=None, power_spectrum=None, freq_range=None, ap_range=None):
330
330
Frequency values for the power spectrum, in linear space.
331
331
power_spectrum : 1d array, optional
332
332
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 :
335
336
336
337
Notes
337
338
-----
338
339
Data is optional if data has been already been added to FOOOF object.
339
340
"""
340
-
341
341
# If freqs & power_spectrum provided together, add data to object.
342
342
if freqs is not None and power_spectrum is not None :
343
343
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):
358
358
# In rare cases, the model fails to fit. Therefore it's in a try/except
359
359
# Cause of failure: RuntimeError, failure to find parameters in curve_fit
360
360
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' )
361
369
362
- if ap_range :
363
- ap_inds = (self .freqs >= ap_range [0 ]) & (self .freqs <= ap_range [1 ])
364
370
ap_freqs = self .freqs [ap_inds ]
365
371
ap_spectrum = self .power_spectrum [ap_inds ]
366
372
else :
@@ -374,22 +380,27 @@ def fit(self, freqs=None, power_spectrum=None, freq_range=None, ap_range=None):
374
380
# Flatten the power_spectrum using fit aperiodic fit
375
381
self ._spectrum_flat = self .power_spectrum - self ._ap_fit
376
382
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
379
389
freqs_0 = self .freqs
380
390
self .freqs = self .freqs [per_inds ]
381
- per_spectrum_flat = self ._spectrum_flat [per_inds ]
382
391
if freq_range :
392
+ freq_range_0 = self .freq_range
383
393
self .freq_range = freq_range
384
- else :
385
- per_spectrum_flat = np .copy (self ._spectrum_flat )
394
+
386
395
387
396
388
397
# 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 ) )
390
399
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
393
404
394
405
# Calculate the peak fit
395
406
# 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):
398
409
# Create peak-removed (but not flattened) power spectrum.
399
410
self ._spectrum_peak_rm = self .power_spectrum - self ._peak_fit
400
411
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
+
401
417
# Run final aperiodic fit on peak-removed power spectrum
402
418
# 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 )
404
420
self ._ap_fit = gen_aperiodic (self .freqs , self .aperiodic_params_ )
405
421
406
422
# Create full power_spectrum model fit
0 commit comments