Better handling of flat data in numba function and spindles detect (#86)

docs/changelog.rst

@@ -21,6 +21,7 @@ c. Added a new dataset containing 8 hours of ECG data. The dataset is in compres
 
 a. When using an MNE.Raw object, conversion of the data from Volts to micro-Volts is now performed within MNE. `PR 70 <https://github.com/raphaelvallat/yasa/pull/70>`_
 b. Use `black <https://black.readthedocs.io/en/stable/>`_ code formatting.
+c. Better handling of flat data in :py:func:`yasa.spindles_detect`. The function previously returned a division by zero error if part of the data was flat. See `issue 85 <https://github.com/raphaelvallat/yasa/issues/85>`_
 
 **Dependencies**
 

yasa/numba.py

@@ -27,6 +27,8 @@ def _corr(x, y):
     r_d1 = np.sqrt(xm2s)
     r_d2 = np.sqrt(ym2s)
     r_den = r_d1 * r_d2
+    if r_den == 0:
+        return np.nan
     return r_num / r_den
 
 
@@ -69,6 +71,8 @@ def _slope_lstsq(x, y):
         sy += y[j]
     den = n_times * sx2 - (sx**2)
     num = n_times * sxy - sx * sy
+    if den == 0:
+        return np.nan
     return num / den
 
 

yasa/tests/test_detection.py

@@ -91,6 +91,12 @@ def test_spindles_detect(self):
         # Test with hypnogram
         spindles_detect(data, sf, hypno=np.ones(data.size))
 
+        # Test with 1-sec of flat data -- we should still have 2 detected spindles
+        data_flat = data.copy()
+        data_flat[100:200] = 1
+        sp = spindles_detect(data_flat, sf).summary()
+        assert sp.shape[0] == 2
+
         # Single channel with Isolation Forest + hypnogram
         sp = spindles_detect(data_full[1, :], sf, hypno=hypno_full, remove_outliers=True)
 

yasa/tests/test_numba.py

@@ -10,15 +10,19 @@ def test_numba(self):
         """Test numba functions"""
         x = np.asarray([4, 5, 7, 8, 5, 6], dtype=np.float64)
         y = np.asarray([1, 5, 4, 6, 8, 5], dtype=np.float64)
+        y_flat = np.ones_like(x)
 
         np.testing.assert_almost_equal(_corr(x, y), np.corrcoef(x, y)[0][1])
+        assert np.isnan(_corr(x, y_flat))
         assert _covar(x, y) == np.cov(x, y)[0][1]
         assert _rms(x) == np.sqrt(np.mean(np.square(x)))
 
         # Least square slope and detrending
         y = np.arange(30) + 3 * np.random.random(30)
         times = np.arange(y.size, dtype=np.float64)
         slope = _slope_lstsq(times, y)
+        assert np.isnan(_slope_lstsq(y_flat, x))
+        assert np.array_equal(_detrend(x, y_flat), np.zeros_like(y_flat))
         np.testing.assert_array_almost_equal(_detrend(times, y), detrend(y, type="linear"))
         X = times[..., np.newaxis]
         X = np.column_stack((np.ones(X.shape[0]), X))