add compute_presence helper function

Author: TomDonoghue

specparam/tests/utils/test_data.py

@@ -48,6 +48,23 @@ def _dispersion_callable(data):
     assert isinstance(out4, np.ndarray)
     assert np.array_equal(out4, out2)
 
+def test_compute_presence():
+
+    data1_full = np.array([0, 1, 2, 3, 4])
+    data1_nan = np.array([0, np.nan, 2, 3, np.nan])
+    assert compute_presence(data1_full) == 1.0
+    assert compute_presence(data1_nan) == 0.6
+    assert compute_presence(data1_nan, output='percent') == 60.0
+
+    data2_full = np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]])
+    data2_nan = np.array([[0, np.nan, 2, 3, np.nan], [np.nan, 6, 7, 8, np.nan]])
+    assert np.array_equal(compute_presence(data2_full), np.array([1.0, 1.0, 1.0, 1.0, 1.0]))
+    assert np.array_equal(compute_presence(data2_nan), np.array([0.5, 0.5, 1.0, 1.0, 0.0]))
+    assert np.array_equal(compute_presence(data2_nan, output='percent'),
+                          np.array([50.0, 50.0, 100.0, 100.0, 0.0]))
+    assert compute_presence(data2_full, average=True) == 1.0
+    assert compute_presence(data2_nan, average=True) == 0.6
+
 def test_trim_spectrum():
 
     f_in = np.array([0., 1., 2., 3., 4., 5.])

specparam/utils/data.py

@@ -84,6 +84,46 @@ def compute_dispersion(data, dispersion='std'):
     return dispersion_data
 
 
+def compute_presence(data, average=False, output='ratio'):
+    """Compute data presence (as number of non-NaN values) from an array of data.
+
+    Parameters
+    ----------
+    data : 1d or 2d array
+        Data array to check presence of.
+    average : bool, optional, default: False
+        Whether to average across . Only used for 2d array inputs.
+        If False, for 2d array, the output is an array matching the length of the 0th dimension of the input.
+        If True, for 2d arrays, the output is a single value averaged across the whole array.
+    output : {'ratio', 'percent'}
+        Representation for the output:
+            'ratio' - ratio value, between 0.0, 1.0.
+            'percent' - percent value, betweeon 0-100%.
+
+    Returns
+    -------
+    presence : float or array of float
+        The computed presence in the given array.
+    """
+
+    assert output in ['ratio', 'percent'], 'Setting for output type not understood.'
+
+    if data.ndim == 1:
+        presence = sum(~np.isnan(data)) / len(data)
+
+    elif data.ndim == 2:
+        if average:
+            presence = compute_presence(data.flatten())
+        else:
+            n_events, n_windows = data.shape
+            presence = np.sum(~np.isnan(data), 0) / (np.ones(n_windows) * n_events)
+
+    if output == 'percent':
+        presence *= 100
+
+    return presence
+
+
 def trim_spectrum(freqs, power_spectra, f_range):
     """Extract a frequency range from power spectra.