[WIP] Flexible dtype results class

fooof/core/funcs.py

@@ -7,22 +7,30 @@
     - They are left available for easy swapping back in, if desired.
 """
 
+from inspect import isfunction
+
 import numpy as np
 
 from fooof.core.errors import InconsistentDataError
 
 ###################################################################################################
 ###################################################################################################
 
-def gaussian_function(xs, *params):
+def gaussian_function(xs, cf, pw, bw, *params):
     """Gaussian fitting function.
 
     Parameters
     ----------
     xs : 1d array
         Input x-axis values.
+    cf : float
+        The center of the gaussian.
+    pw : float
+        The height of the gaussian.
+    bw : float
+        The width of the gaussian.
     *params : float
-        Parameters that define gaussian function.
+        Additional centers, heights, and widths.
 
     Returns
     -------
@@ -32,16 +40,17 @@ def gaussian_function(xs, *params):
 
     ys = np.zeros_like(xs)
 
+    params = [cf, pw, bw, *params]
+
     for ii in range(0, len(params), 3):
 
         ctr, hgt, wid = params[ii:ii+3]
-
         ys = ys + hgt * np.exp(-(xs-ctr)**2 / (2*wid**2))
 
     return ys
 
 
-def expo_function(xs, *params):
+def expo_function(xs, offset, knee, exp):
     """Exponential fitting function, for fitting aperiodic component with a 'knee'.
 
     NOTE: this function requires linear frequency (not log).
@@ -50,26 +59,32 @@ def expo_function(xs, *params):
     ----------
     xs : 1d array
         Input x-axis values.
-    *params : float
-        Parameters (offset, knee, exp) that define Lorentzian function:
-        y = 10^offset * (1/(knee + x^exp))
+    offset : float
+        The y-intercept of the fit.
+    knee : float
+        The bend in the fit.
+    exp : float
+        The exponential slope of the fit.
 
     Returns
     -------
     ys : 1d array
         Output values for exponential function.
+
+    Notes
+    -----
+    Parameters (offset, knee, exp) that define Lorentzian function:
+    y = 10^offset * (1/(knee + x^exp))
     """
 
     ys = np.zeros_like(xs)
 
-    offset, knee, exp = params
-
     ys = ys + offset - np.log10(knee + xs**exp)
 
     return ys
 
 
-def expo_nk_function(xs, *params):
+def expo_nk_function(xs, offset, exp):
     """Exponential fitting function, for fitting aperiodic component without a 'knee'.
 
     NOTE: this function requires linear frequency (not log).
@@ -78,34 +93,40 @@ def expo_nk_function(xs, *params):
     ----------
     xs : 1d array
         Input x-axis values.
-    *params : float
-        Parameters (offset, exp) that define Lorentzian function:
-        y = 10^off * (1/(x^exp))
+    offset : float
+        The y-intercept of the fit.
+    exp : float
+        The exponential slope of the fit.
 
     Returns
     -------
     ys : 1d array
         Output values for exponential function, without a knee.
+
+    Notes
+    -----
+    Parameters (offset, exp) that define Lorentzian function:
+    y = 10^off * (1/(x^exp))
     """
 
     ys = np.zeros_like(xs)
 
-    offset, exp = params
-
     ys = ys + offset - np.log10(xs**exp)
 
     return ys
 
 
-def linear_function(xs, *params):
+def linear_function(xs, offset, slope):
     """Linear fitting function.
 
     Parameters
     ----------
     xs : 1d array
         Input x-axis values.
-    *params : float
-        Parameters that define linear function.
+    offset : float
+        The y-intercept of the fit.
+    slope : float
+        The slope of the fit.
 
     Returns
     -------
@@ -115,22 +136,24 @@ def linear_function(xs, *params):
 
     ys = np.zeros_like(xs)
 
-    offset, slope = params
-
     ys = ys + offset + (xs*slope)
 
     return ys
 
 
-def quadratic_function(xs, *params):
+def quadratic_function(xs, offset, slope, curve):
     """Quadratic fitting function.
 
     Parameters
     ----------
     xs : 1d array
         Input x-axis values.
-    *params : float
-        Parameters that define quadratic function.
+    offset : float
+        The y-intercept of the fit.
+    slope : float
+        The slope of the fit.
+    curve : float
+        The curve of the fit.
 
     Returns
     -------
@@ -140,8 +163,6 @@ def quadratic_function(xs, *params):
 
     ys = np.zeros_like(xs)
 
-    offset, slope, curve = params
-
     ys = ys + offset + (xs*slope) + ((xs**2)*curve)
 
     return ys
@@ -167,7 +188,9 @@ def get_pe_func(periodic_mode):
 
     """
 
-    if periodic_mode == 'gaussian':
+    if isfunction(periodic_mode):
+        pe_func = periodic_mode
+    elif periodic_mode == 'gaussian':
         pe_func = gaussian_function
     else:
         raise ValueError("Requested periodic mode not understood.")
@@ -194,7 +217,9 @@ def get_ap_func(aperiodic_mode):
         If the specified aperiodic mode label is not understood.
     """
 
-    if aperiodic_mode == 'fixed':
+    if isfunction(aperiodic_mode):
+        ap_func = aperiodic_mode
+    elif aperiodic_mode == 'fixed':
         ap_func = expo_nk_function
     elif aperiodic_mode == 'knee':
         ap_func = expo_function

fooof/data/__init__.py

@@ -1,3 +1,3 @@
 """Data sub-module for FOOOF."""
 
-from .data import FOOOFSettings, FOOOFMetaData, FOOOFResults, SimParams
+from .data import FOOOFSettings, FOOOFMetaData, FOOOFResults, SimParams, FitParams

fooof/data/data.py

@@ -8,7 +8,13 @@
     - this means no additional attributes can be defined (which is more memory efficient)
 """
 
-from collections import namedtuple
+from collections import namedtuple, OrderedDict
+
+import numpy as np
+
+from fooof.core.modutils import safe_import
+
+pd = safe_import('pandas')
 
 ###################################################################################################
 ###################################################################################################
@@ -78,8 +84,30 @@ class FOOOFResults(namedtuple('FOOOFResults', ['aperiodic_params', 'peak_params'
     -----
     This object is a data object, based on a NamedTuple, with immutable data attributes.
     """
+
     __slots__ = ()
 
+    def to_dict(self):
+
+        # Combined peak, aperiodic, and goodness of fit params
+        results_dict = OrderedDict()
+        results_dict.update(self.peak_params.to_dict())
+        results_dict.update(self.aperiodic_params.to_dict())
+        results_dict.update(OrderedDict(r_squared=self.r_squared, error=self.error))
+
+        return results_dict
+
+    def to_pandas(self):
+
+        if pd is None:
+            raise ValueError("Pandas is not installed.")
+
+        results_dict = self.to_dict()
+
+        results_df = pd.DataFrame(results_dict)
+
+        return results_df
+
 
 class SimParams(namedtuple('SimParams', ['aperiodic_params', 'periodic_params', 'nlv'])):
     """Parameters that define a simulated power spectrum.
@@ -98,3 +126,33 @@ class SimParams(namedtuple('SimParams', ['aperiodic_params', 'periodic_params',
     This object is a data object, based on a NamedTuple, with immutable data attributes.
     """
     __slots__ = ()
+
+
+class FitParams(np.ndarray):
+
+    def __new__(cls, params, labels):
+
+        return np.asarray(params).view(cls)
+
+    def __init__(self, params, labels):
+
+        self.params = params
+        self.labels = labels
+
+    def to_dict(self):
+
+        params = OrderedDict((k, v) for k, v in \
+            zip(self.labels, self.params.transpose()))
+
+        return params
+
+    def to_pandas(self):
+
+        if pd is None:
+            raise ValueError("Pandas is not installed.")
+
+        params = self.to_dict()
+
+        params = pd.DataFrame(params)
+
+        return params