add get data & model component methods to fit object

Author: TomDonoghue

fooof/objs/fit.py

@@ -63,6 +63,7 @@
 from numpy.linalg import LinAlgError
 from scipy.optimize import curve_fit
 
+from fooof.core.utils import unlog
 from fooof.core.items import OBJ_DESC
 from fooof.core.info import get_indices
 from fooof.core.io import save_fm, load_json
@@ -596,6 +597,95 @@ def get_meta_data(self):
                              for key in OBJ_DESC['meta_data']})
 
 
+    def get_data_component(self, component='full', space='log'):
+        """Get a data component.
+
+        Parameters
+        ----------
+        component : {'full', 'aperiodic', 'peak'}
+            Which data component to return.
+                'full' - full power spectrum
+                'aperiodic' - isolated aperiodic data component
+                'peak' - isolated peak data component
+        space : {'log', 'linear'}
+            Which space to return the data component in.
+                'log' - returns in log10 space.
+                'linear' - returns in linear space.
+
+        Returns
+        -------
+        output : 1d array
+            Specified data component, in specified spacing.
+
+        Notes
+        -----
+        The 'space' parameter doesn't just define the spacing of the data component
+        values, but rather defines the space of the additive data definiton such that
+        `power_spectrum = aperiodic_component + peak_component`.
+        With space set as 'log', this combination holds in log space.
+        With space set as 'linear', this combination holds in linear space.
+        """
+
+        assert space in ['linear', 'log'], "Input for 'space' invalid."
+
+        if component == 'full':
+            output = self.power_spectrum if space == 'log' else unlog(self.power_spectrum)
+        elif component == 'aperiodic':
+            output = self._spectrum_peak_rm if space == 'log' else \
+                unlog(self.power_spectrum) / unlog(self._peak_fit)
+        elif component == 'peak':
+            output = self._spectrum_flat if space == 'log' else \
+                unlog(self.power_spectrum) - unlog(self._ap_fit)
+        else:
+            raise ValueError('Input for component invalid.')
+
+        return output
+
+
+    def get_model_component(self, component='full', space='log'):
+        """Get a model component.
+
+        Parameters
+        ----------
+        component : {'full', 'aperiodic', 'peak'}
+            Which model component to return.
+                'full' - full model
+                'aperiodic' - isolated aperiodic model component
+                'peak' - isolated peak model component
+        space : {'log', 'linear'}
+            Which space to return the model component in.
+                'log' - returns in log10 space.
+                'linear' - returns in linear space.
+
+        Returns
+        -------
+        output : 1d array
+            Specified model component, in specified spacing.
+
+        Notes
+        -----
+        The 'space' parameter doesn't just define the spacing of the model component
+        values, but rather defines the space of the additive model such that
+        `model = aperiodic_component + peak_component`.
+        With space set as 'log', this combination holds in log space.
+        With space set as 'lienar', this combination holds in linear space.
+        """
+
+        assert space in ['linear', 'log'], "Input for 'space' invalid."
+
+        if component == 'full':
+            output = self.fooofed_spectrum_ if space == 'log' else unlog(self.fooofed_spectrum_)
+        elif component == 'aperiodic':
+            output = self._ap_fit if space == 'log' else unlog(self._ap_fit)
+        elif component == 'peak':
+            output = self._peak_fit if space == 'log' else \
+                unlog(self.fooofed_spectrum_) - unlog(self._ap_fit)
+        else:
+            raise ValueError('Input for component invalid.')
+
+        return output
+
+
     def get_params(self, name, col=None):
         """Return model fit parameters for specified feature(s).
 

fooof/tests/objs/test_fit.py

@@ -311,6 +311,17 @@ def test_obj_gets(tfm):
     results = tfm.get_results()
     assert isinstance(results, FOOOFResults)
 
+def test_get_components(tfm):
+
+    # Make sure test object has been fit
+    tfm.fit()
+
+    # Test get data & model components
+    for comp in ['full', 'aperiodic', 'peak']:
+        for space in ['log', 'linear']:
+            assert isinstance(tfm.get_data_component(comp, space), np.ndarray)
+            assert isinstance(tfm.get_model_component(comp, space), np.ndarray)
+
 def test_get_params(tfm):
     """Test the get_params method."""