Merge pull request #299 from fooof-tools/optimize

[MNT] - Optimizations for curve fit

Author: TomDonoghue

fooof/core/funcs.py

@@ -32,9 +32,7 @@ def gaussian_function(xs, *params):
 
     ys = np.zeros_like(xs)
 
-    for ii in range(0, len(params), 3):
-
-        ctr, hgt, wid = params[ii:ii+3]
+    for ctr, hgt, wid in zip(*[iter(params)] * 3):
 
         ys = ys + hgt * np.exp(-(xs-ctr)**2 / (2*wid**2))
 
@@ -60,11 +58,8 @@ def expo_function(xs, *params):
         Output values for exponential function.
     """
 
-    ys = np.zeros_like(xs)
-
     offset, knee, exp = params
-
-    ys = ys + offset - np.log10(knee + xs**exp)
+    ys = offset - np.log10(knee + xs**exp)
 
     return ys
 
@@ -88,11 +83,8 @@ def expo_nk_function(xs, *params):
         Output values for exponential function, without a knee.
     """
 
-    ys = np.zeros_like(xs)
-
     offset, exp = params
-
-    ys = ys + offset - np.log10(xs**exp)
+    ys = offset - np.log10(xs**exp)
 
     return ys
 
@@ -113,11 +105,8 @@ def linear_function(xs, *params):
         Output values for linear function.
     """
 
-    ys = np.zeros_like(xs)
-
     offset, slope = params
-
-    ys = ys + offset + (xs*slope)
+    ys = offset + (xs*slope)
 
     return ys
 
@@ -138,11 +127,8 @@ def quadratic_function(xs, *params):
         Output values for quadratic function.
     """
 
-    ys = np.zeros_like(xs)
-
     offset, slope, curve = params
-
-    ys = ys + offset + (xs*slope) + ((xs**2)*curve)
+    ys = offset + (xs*slope) + ((xs**2)*curve)
 
     return ys
 

fooof/core/jacobians.py

@@ -0,0 +1,103 @@
+""""Functions for computing Jacobian matrices to be used during fitting.
+
+Notes
+-----
+These functions line up with those in `funcs`.
+The parameters in these functions are labeled {a, b, c, ...}, but follow the order in `funcs`.
+These functions are designed to be passed into `curve_fit` to provide a computed Jacobian.
+"""
+
+import numpy as np
+
+###################################################################################################
+###################################################################################################
+
+## Periodic Jacobian functions
+
+def jacobian_gauss(xs, *params):
+    """Create the Jacobian matrix for the Gaussian function.
+
+    Parameters
+    ----------
+    xs : 1d array
+        Input x-axis values.
+    *params : float
+        Parameters for the function.
+
+    Returns
+    -------
+    jacobian : 2d array
+        Jacobian matrix, with shape [len(xs), n_params].
+    """
+
+    jacobian = np.zeros((len(xs), len(params)))
+
+    for i, (a, b, c) in enumerate(zip(*[iter(params)] * 3)):
+
+        ax = -a + xs
+        ax2 = ax**2
+
+        c2 = c**2
+        c3 = c**3
+
+        exp = np.exp(-ax2 / (2 * c2))
+        exp_b = exp * b
+
+        ii = i * 3
+        jacobian[:, ii] = (exp_b * ax) / c2
+        jacobian[:, ii+1] = exp
+        jacobian[:, ii+2] = (exp_b * ax2) / c3
+
+    return jacobian
+
+
+## Aperiodic Jacobian functions
+
+def jacobian_expo(xs, *params):
+    """Create the Jacobian matrix for the exponential function.
+
+    Parameters
+    ----------
+    xs : 1d array
+        Input x-axis values.
+    *params : float
+        Parameters for the function.
+
+    Returns
+    -------
+    jacobian : 2d array
+        Jacobian matrix, with shape [len(xs), n_params].
+    """
+
+    a, b, c = params
+
+    xs_c = xs**c
+    b_xs_c = xs_c + b
+
+    jacobian = np.ones((len(xs), len(params)))
+    jacobian[:, 1] = -1 / b_xs_c
+    jacobian[:, 2] = -(xs_c * np.log10(xs)) / b_xs_c
+
+    return jacobian
+
+
+def jacobian_expo_nk(xs, *params):
+    """Create the Jacobian matrix for the exponential no-knee function.
+
+    Parameters
+    ----------
+    xs : 1d array
+        Input x-axis values.
+    *params : float
+        Parameters for the function.
+
+    Returns
+    -------
+    jacobian : 2d array
+        Jacobian matrix, with shape [len(xs), n_params].
+    """
+
+    jacobian = np.ones((len(xs), len(params)))
+    jacobian[:, 1] = -np.log10(xs)
+
+    return jacobian

fooof/objs/fit.py

@@ -71,6 +71,7 @@
 from fooof.core.modutils import copy_doc_func_to_method
 from fooof.core.utils import group_three, check_array_dim
 from fooof.core.funcs import gaussian_function, get_ap_func, infer_ap_func
+from fooof.core.jacobians import jacobian_gauss
 from fooof.core.errors import (FitError, NoModelError, DataError,
                                NoDataError, InconsistentDataError)
 from fooof.core.strings import (gen_settings_str, gen_results_fm_str,
@@ -192,12 +193,17 @@ def __init__(self, peak_width_limits=(0.5, 12.0), max_n_peaks=np.inf, min_peak_h
         self._gauss_overlap_thresh = 0.75
         # Parameter bounds for center frequency when fitting gaussians, in terms of +/- std dev
         self._cf_bound = 1.5
-        # The maximum number of calls to the curve fitting function
-        self._maxfev = 5000
         # The error metric to calculate, post model fitting. See `_calc_error` for options
         #   Note: this is for checking error post fitting, not an objective function for fitting
         self._error_metric = 'MAE'
 
+        ## PRIVATE CURVE_FIT SETTINGS
+        # The maximum number of calls to the curve fitting function
+        self._maxfev = 5000
+        # The tolerance setting for curve fitting (see scipy.curve_fit - ftol / xtol / gtol)
+        #   Here reduce tolerance to speed fitting. Set value to 1e-8 to match curve_fit default
+        self._tol = 0.00001
+
         ## RUN MODES
         # Set default debug mode - controls if an error is raised if model fitting is unsuccessful
         self._debug = False
@@ -946,7 +952,9 @@ def _simple_ap_fit(self, freqs, power_spectrum):
                 warnings.simplefilter("ignore")
                 aperiodic_params, _ = curve_fit(get_ap_func(self.aperiodic_mode),
                                                 freqs, power_spectrum, p0=guess,
-                                                maxfev=self._maxfev, bounds=ap_bounds)
+                                                maxfev=self._maxfev, bounds=ap_bounds,
+                                                ftol=self._tol, xtol=self._tol, gtol=self._tol,
+                                                check_finite=False)
         except RuntimeError as excp:
             error_msg = ("Model fitting failed due to not finding parameters in "
                          "the simple aperiodic component fit.")
@@ -1003,7 +1011,9 @@ def _robust_ap_fit(self, freqs, power_spectrum):
                 warnings.simplefilter("ignore")
                 aperiodic_params, _ = curve_fit(get_ap_func(self.aperiodic_mode),
                                                 freqs_ignore, spectrum_ignore, p0=popt,
-                                                maxfev=self._maxfev, bounds=ap_bounds)
+                                                maxfev=self._maxfev, bounds=ap_bounds,
+                                                ftol=self._tol, xtol=self._tol, gtol=self._tol,
+                                                check_finite=False)
         except RuntimeError as excp:
             error_msg = ("Model fitting failed due to not finding "
                          "parameters in the robust aperiodic fit.")
@@ -1149,7 +1159,9 @@ def _fit_peak_guess(self, guess):
         # Fit the peaks
         try:
             gaussian_params, _ = curve_fit(gaussian_function, self.freqs, self._spectrum_flat,
-                                           p0=guess, maxfev=self._maxfev, bounds=gaus_param_bounds)
+                                           p0=guess, maxfev=self._maxfev, bounds=gaus_param_bounds,
+                                           ftol=self._tol, xtol=self._tol, gtol=self._tol,
+                                           check_finite=False, jac=jacobian_gauss)
         except RuntimeError as excp:
             error_msg = ("Model fitting failed due to not finding "
                          "parameters in the peak component fit.")

fooof/tests/core/test_jacobians.py

@@ -0,0 +1,33 @@
+"""Tests for fooof.core.jacobians."""
+
+from fooof.core.jacobians import *
+
+###################################################################################################
+###################################################################################################
+
+def test_jacobian_gauss():
+
+    xs = np.arange(1, 100)
+    ctr, hgt, wid = 50, 5, 10
+
+    jacobian = jacobian_gauss(xs, ctr, hgt, wid)
+    assert isinstance(jacobian, np.ndarray)
+    assert jacobian.shape == (len(xs), 3)
+
+def test_jacobian_expo():
+
+    xs = np.arange(1, 100)
+    off, knee, exp = 10, 5, 2
+
+    jacobian = jacobian_expo(xs, off, knee, exp)
+    assert isinstance(jacobian, np.ndarray)
+    assert jacobian.shape == (len(xs), 3)
+
+def test_jacobian_expo_nk():
+
+    xs = np.arange(1, 100)
+    off, exp = 10, 2
+
+    jacobian = jacobian_expo_nk(xs, off, exp)
+    assert isinstance(jacobian, np.ndarray)
+    assert jacobian.shape == (len(xs), 2)

fooof/tests/objs/test_fit.py

@@ -391,7 +391,7 @@ def test_fooof_fit_failure():
 
     ## Induce a runtime error, and check it runs through
     tfm = FOOOF(verbose=False)
-    tfm._maxfev = 5
+    tfm._maxfev = 2
 
     tfm.fit(*gen_power_spectrum([3, 50], [50, 2], [10, 0.5, 2, 20, 0.3, 4]))
 
@@ -417,7 +417,7 @@ def test_fooof_debug():
     """Test FOOOF in debug mode, including with fit failures."""
 
     tfm = FOOOF(verbose=False)
-    tfm._maxfev = 5
+    tfm._maxfev = 2
 
     tfm.set_debug_mode(True)
     assert tfm._debug is True