update set check modes

Author: TomDonoghue

fooof/objs/fit.py

@@ -45,9 +45,11 @@
 _debug : bool
     Whether the object is set in debug mode.
     This should be controlled by using the `set_debug_mode` method.
-_check_data : bool
-    Whether to check added data for NaN or Inf values, and fail out if present.
-    This should be controlled by using the `set_check_data_mode` method.
+_check_data, _check_freqs : bool
+    Whether to check added inputs for incorrect inputs, failing if present.
+    Frequency data is checked for linear spacing.
+    Power values are checked for data for NaN or Inf values.
+    These modes default to True, and can be controlled with the `set_check_modes` method.
 
 Code Notes
 ----------
@@ -723,6 +725,24 @@ def set_debug_mode(self, debug):
         self._debug = debug
 
 
+    def set_check_modes(self, check_freqs=None, check_data=None):
+        """Set check modes, which controls if an error is raised based on check on the inputs.
+
+        Parameters
+        ----------
+        check_freqs : bool, optional
+            Whether to run in check freqs mode, which checks the frequency data.
+        check_data : bool, optional
+            Whether to run in check data mode, which checks the power spectrum values data.
+        """
+
+        if check_freqs is not None:
+            self._check_freqs = check_freqs
+        if check_data is not None:
+            self._check_data = check_data
+
+
+    # This kept for backwards compatibility, but to be removed in 2.0 in favor of `set_check_modes`
     def set_check_data_mode(self, check_data):
         """Set check data mode, which controls if an error is raised if NaN or Inf data are added.
 
@@ -732,7 +752,7 @@ def set_check_data_mode(self, check_data):
             Whether to run in check data mode.
         """
 
-        self._check_data = check_data
+        self.set_check_modes(check_data=check_data)
 
 
     def to_df(self, peak_org):

fooof/tests/objs/test_fit.py

@@ -414,17 +414,24 @@ def test_fooof_debug():
     with raises(FitError):
         tfm.fit(*gen_power_spectrum([3, 50], [50, 2], [10, 0.5, 2, 20, 0.3, 4]))
 
-def test_fooof_check_data():
-    """Test FOOOF in with check data mode turned off, including with NaN data."""
+def test_fooof_set_check_modes(tfm):
+    """Test changing check_modes using set_check_modes, and that checks get turned off.
+    Note that testing for checks raising errors happens in test_fooof_checks.`"""
 
     tfm = FOOOF(verbose=False)
 
-    tfm.set_check_data_mode(False)
+    tfm.set_check_modes(False, False)
+    assert tfm._check_freqs is False
     assert tfm._check_data is False
 
-    # Add data, with check data turned off
-    #   In check data mode, adding data with NaN should run
-    freqs = gen_freqs([3, 50], 0.5)
+    # Add bad frequency data, with check freqs turned off
+    freqs = np.array([1, 2, 4])
+    powers = np.array([1, 2, 3])
+    tfm.add_data(freqs, powers)
+    assert tfm.has_data
+
+    # Add bad power values data, with check data turned off
+    freqs = gen_freqs([3, 30], 1)
     powers = np.ones_like(freqs) * np.nan
     tfm.add_data(freqs, powers)
     assert tfm.has_data
@@ -433,6 +440,11 @@ def test_fooof_check_data():
     tfm.fit()
     assert not tfm.has_model
 
+    # Reset check modes to true
+    tfm.set_check_modes(True, True)
+    assert tfm._check_freqs is True
+    assert tfm._check_data is True
+
 def test_fooof_to_df(tfm, tbands, skip_if_no_pandas):
 
     df1 = tfm.to_df(2)