Merge pull request #47 from schmidtfa/minor_fixes

examples/basic_func_fun.py

@@ -34,7 +34,7 @@
 
 #from pyrasa.irasa import irasa
 #%%
-freq_irasa, psd_ap, psd_p = irasa(sig, 
+irasa_result = irasa(sig, 
                                   fs=fs, 
                                   band=(1, 100), 
                                   psd_kwargs={'nperseg': duration*fs, 
@@ -44,26 +44,21 @@
 
 f, axes = plt.subplots(ncols=2, figsize=(8, 4))
 axes[0].set_title('Periodic')
-axes[0].plot(freq_irasa, psd_p[0,:])
+axes[0].plot(irasa_result.freqs, irasa_result.periodic[0,:])
 axes[0].set_ylabel('Power (a.u.)')
 axes[0].set_xlabel('Frequency (Hz)')
 axes[1].set_title('Aperiodic')
-axes[1].loglog(freq_irasa, psd_ap[0,:])
+axes[1].loglog(irasa_result.freqs, irasa_result.aperiodic[0,:])
 axes[1].set_ylabel('Power (a.u.)')
 axes[1].set_xlabel('Frequency (Hz)')
 
 f.tight_layout()
 
 # %% get periodic stuff
-from pyrasa.utils.peak_utils import get_peak_params
-get_peak_params(psd_p, freqs=freq_irasa)
-# %%
-from pyrasa.utils.aperiodic_utils import compute_slope
 
-ap_params, gof_params = compute_slope(aperiodic_spectrum=psd_ap,
-                                       freqs=freq_irasa-1, 
-                                       fit_func='fixed',
-                                       #fit_bounds=[0, 40]
-                                       )
-ap_params
+pe_params = irasa_result.get_peaks()
+pe_params
+#%% get aperiodics
+ap_params = irasa_result.get_slopes(fit_func='knee')
+ap_params.gof
 # %%

examples/custom_fit_functions.py

@@ -0,0 +1,50 @@
+#%%
+import scipy.signal as dsp
+from pyrasa.utils.aperiodic_utils import compute_slope
+from pyrasa.utils.fit_funcs import AbstractFitFun
+import numpy as np
+from neurodsp.sim import sim_powerlaw
+from typing import Any
+
+n_secs = 60
+fs=500
+f_range = [1.5, 300]
+exponent = -1.5
+
+sig = sim_powerlaw(n_seconds=n_secs, fs=fs, exponent=exponent)
+
+# test whether recombining periodic and aperiodic spectrum is equivalent to the original spectrum
+freqs, psd = dsp.welch(sig, fs, nperseg=int(4 * fs))
+freq_logical = np.logical_and(freqs >= f_range[0], freqs <= f_range[1])
+psd, freqs = psd[freq_logical], freqs[freq_logical]
+
+
+class CustomFitFun(AbstractFitFun):
+    def func(self, x: np.ndarray, a: float, b: float) -> np.ndarray:
+        """
+        Specparams fixed fitting function.
+        Use this to model aperiodic activity without a spectral knee
+        """
+        y_hat = a + b * x
+
+        return y_hat
+
+    @property
+    def curve_kwargs(self) -> dict[str, Any]:
+        aperiodic_nolog = 10**self.aperiodic_spectrum
+        off_guess = [aperiodic_nolog[0]]
+        exp_guess = [
+            np.abs(np.log10(aperiodic_nolog[0] / aperiodic_nolog[-1]) / np.log10(self.freq[-1] / self.freq[0]))
+        ]
+        return {
+            'maxfev': 10_000,
+            'ftol': 1e-5,
+            'xtol': 1e-5,
+            'gtol': 1e-5,
+            'p0': np.array(off_guess + exp_guess),
+            'bounds': ((-np.inf, -np.inf), (np.inf, np.inf)),
+        }
+
+#%%
+slope_fit = compute_slope(np.log10(psd), np.log10(freqs), fit_func=CustomFitFun)
+# %%

pyrasa/irasa_mne/mne_objs.py

@@ -1,4 +1,5 @@
 # %% inherit from spectrum array
+
 import matplotlib
 import mne
 import numpy as np
@@ -10,6 +11,8 @@
 from pyrasa.utils.peak_utils import get_peak_params
 from pyrasa.utils.types import SlopeFit
 
+# FutureWarning:
+
 
 class PeriodicSpectrumArray(SpectrumArray):
     """Subclass of SpectrumArray"""
@@ -170,7 +173,10 @@ def __init__(
         )
 
     def get_slopes(
-        self: SpectrumArray, fit_func: str = 'fixed', scale: bool = False, fit_bounds: tuple[float, float] | None = None
+        self: SpectrumArray,
+        fit_func: str = 'fixed',
+        scale: bool = False,
+        fit_bounds: tuple[float, float] | None = None,
     ) -> SlopeFit:
         """
         This method can be used to extract aperiodic parameters from the aperiodic spectrum extracted from IRASA.
@@ -393,7 +399,10 @@ def __init__(
         )
 
     def get_slopes(
-        self: SpectrumArray, fit_func: str = 'fixed', scale: bool = False, fit_bounds: tuple[float, float] | None = None
+        self: SpectrumArray,
+        fit_func: str = 'fixed',
+        scale: bool = False,
+        fit_bounds: tuple[float, float] | None = None,
     ) -> SlopeFit:
         """
         This method can be used to extract aperiodic parameters from the aperiodic spectrum extracted from IRASA.

pyrasa/utils/aperiodic_utils.py

@@ -5,137 +5,37 @@
 
 import numpy as np
 import pandas as pd
-from scipy.optimize import curve_fit
 
+from pyrasa.utils.fit_funcs import AbstractFitFun, FixedFitFun, KneeFitFun
 from pyrasa.utils.types import SlopeFit
 
 
-def fixed_model(x: np.ndarray, b0: float, b: float) -> np.ndarray:
-    """
-    Specparams fixed fitting function.
-    Use this to model aperiodic activity without a spectral knee
-    """
-
-    y_hat = b0 - np.log10(x**b)
-
-    return y_hat
-
-
-def knee_model(x: np.ndarray, b0: float, k: float, b1: float, b2: float) -> np.ndarray:
-    """
-    Model aperiodic activity with a spectral knee and a pre-knee slope.
-    Use this to model aperiodic activity with a spectral knee
-    """
-
-    y_hat = b0 - np.log10(x**b1 * (k + x**b2))
-
-    return y_hat
-
-
-def _get_gof(psd: np.ndarray, psd_pred: np.ndarray, fit_func: str) -> pd.DataFrame:
-    """
-    get goodness of fit (i.e. mean squared error and R2)
-    BIC and AIC currently assume OLS
-    https://machinelearningmastery.com/probabilistic-model-selection-measures/
-    """
-
-    residuals = np.log10(psd) - psd_pred
-    ss_res = np.sum(residuals**2)
-    ss_tot = np.sum((np.log10(psd) - np.mean(np.log10(psd))) ** 2)
-    mse = np.mean(residuals**2)
-
-    if fit_func == 'knee':
-        k = 3  # k -> number of params
-    elif fit_func == 'fixed':
-        k = 1  # k -> number of params
-
-    n = len(psd)
-    bic = n * np.log(mse) + k * np.log(n)
-    aic = n * np.log(mse) + 2 * k
-
-    gof = pd.DataFrame({'mse': mse, 'r_squared': 1 - (ss_res / ss_tot), 'BIC': bic, 'AIC': aic}, index=[0])
-    return gof
-
-
 def _compute_slope(
     aperiodic_spectrum: np.ndarray,
     freq: np.ndarray,
-    fit_func: str,
-    fit_bounds: tuple | None = None,
+    fit_func: str | type[AbstractFitFun],
     scale_factor: float | int = 1,
 ) -> tuple[pd.DataFrame, pd.DataFrame]:
     """get the slope of the aperiodic spectrum"""
 
-    curv_kwargs = {
-        'maxfev': 10_000,
-        'ftol': 1e-5,
-        'xtol': 1e-5,
-        'gtol': 1e-5,
-    }
-
-    off_guess = [aperiodic_spectrum[0]] if fit_bounds is None else fit_bounds[0]
-    exp_guess = (
-        [np.abs(np.log10(aperiodic_spectrum[0] / aperiodic_spectrum[-1]) / np.log10(freq[-1] / freq[0]))]
-        if fit_bounds is None
-        else fit_bounds[1]
-    )
-
-    valid_slope_functions = ['fixed', 'knee']
-    assert fit_func in valid_slope_functions, f'The slope fitting function has to be in {valid_slope_functions}'
-
-    if fit_func == 'fixed':
-        fit_f = fixed_model
-
-        p, _ = curve_fit(fit_f, freq, np.log10(aperiodic_spectrum))
-
-        params = pd.DataFrame(
-            {
-                'Offset': p[0],
-                'Exponent': p[1],
-                'fit_type': 'fixed',
-            },
-            index=[0],
-        )
-        psd_pred = fit_f(freq, *p)
-
-    elif fit_func == 'knee':
-        fit_f = knee_model  # type: ignore
-        # curve_fit_specs
-        cumsum_psd = np.cumsum(aperiodic_spectrum)
-        half_pw_freq = freq[np.abs(cumsum_psd - (0.5 * cumsum_psd[-1])).argmin()]
-        # make the knee guess the point where we have half the power in the spectrum seems plausible to me
-        knee_guess = [half_pw_freq ** (exp_guess[0] + exp_guess[0])]
-        # convert knee freq to knee val which should be 2*exp_1 but this seems good enough
-        curv_kwargs['p0'] = np.array(off_guess + knee_guess + exp_guess + exp_guess)  # type: ignore
-        # print(curv_kwargs['p0'])
-        # make this optional
-        curv_kwargs['bounds'] = ((0, 0, 0, 0), (np.inf, np.inf, np.inf, np.inf))  # type: ignore
-        # knee value should always be positive at least intuitively
-        p, _ = curve_fit(fit_f, freq, np.log10(aperiodic_spectrum), **curv_kwargs)
-
-        params = pd.DataFrame(
-            {
-                'Offset': p[0] / scale_factor,
-                'Knee': p[1],
-                'Exponent_1': p[2],
-                'Exponent_2': p[3],
-                'Knee Frequency (Hz)': p[1] ** (1.0 / (2 * p[2] + p[3])),
-                'fit_type': 'knee',
-            },
-            index=[0],
-        )
-        psd_pred = fit_f(freq, *p)
+    if isinstance(fit_func, str):
+        if fit_func == 'fixed':
+            fit_func = FixedFitFun
+        elif fit_func == 'knee':
+            fit_func = KneeFitFun
+        else:
+            raise ValueError('fit_func should be either "fixed" or "knee"')
 
-    gof = _get_gof(aperiodic_spectrum, psd_pred, fit_func)
-    gof['fit_type'] = fit_func
+    fit_f = fit_func(freq, aperiodic_spectrum, scale_factor=scale_factor)
+    params, gof = fit_f.fit_func()
 
     return params, gof
 
 
 def compute_slope(
     aperiodic_spectrum: np.ndarray,
     freqs: np.ndarray,
-    fit_func: str,
+    fit_func: str | type[AbstractFitFun] = 'fixed',
     ch_names: Iterable | None = None,
     scale: bool = False,
     fit_bounds: tuple[float, float] | None = None,
@@ -186,6 +86,8 @@ def compute_slope(
         fmin, fmax = freqs.min(), freqs.max()
         assert fit_bounds[0] > fmin, f'The selected lower bound is lower than the lowest frequency of {fmin}Hz'
         assert fit_bounds[1] < fmax, f'The selected upper bound is higher than the highest frequency of {fmax}Hz'
+        freq_logical = np.logical_and(freqs >= fit_bounds[0], freqs <= fit_bounds[1])
+        aperiodic_spectrum, freqs = aperiodic_spectrum[:, freq_logical], freqs[freq_logical]
 
     if freqs[0] == 0:
         warnings.warn(
@@ -216,7 +118,6 @@ def num_zeros(decimal: int) -> float:
             freq=freqs,
             fit_func=fit_func,
             scale_factor=scale_factor,
-            fit_bounds=fit_bounds,
         )
 
         params['ch_name'] = ch_name