@@ -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