Fix gfp computation

pycrostates/preprocessing/extract_gfp_peaks.py

@@ -20,17 +20,25 @@
     _check_reject_by_annotation,
     _check_tmin_tmax,
     _check_type,
+    _check_value,
 )
 from ..utils._docs import fill_doc
 from ..utils._logs import logger, verbose
 
 if TYPE_CHECKING:
-    from typing import Optional, Union
+    from typing import Callable, Optional, Union
 
-    from .._typing import Picks, ScalarFloatArray
+    from .._typing import Picks
     from ..io import ChData
 
 
+_GFP_FUNC: dict[str, Callable] = {
+    "eeg": lambda x: np.std(x, axis=0),
+    "grad": lambda x: np.sqrt(np.mean(x**2, axis=0)),
+    "mag": lambda x: np.sqrt(np.mean(x**2, axis=0)),
+}
+
+
 @fill_doc
 @verbose
 def extract_gfp_peaks(
@@ -103,6 +111,11 @@ def extract_gfp_peaks(
     picks_all = _picks_to_idx(inst.info, inst.ch_names, none="all", exclude="bads")
     _check_picks_uniqueness(inst.info, picks)
 
+    # gfp function
+    ch_type = inst.get_channel_types(picks, unique=True)[0]
+    _check_value(ch_type, _GFP_FUNC, "ch_type")
+    gfp_function = _GFP_FUNC[ch_type]
+
     # set kwargs for .get_data()
     kwargs = dict(tmin=tmin, tmax=tmax)
     if isinstance(inst, BaseRaw):
@@ -113,7 +126,8 @@ def extract_gfp_peaks(
         # retrieve data array on which we look for GFP peaks
         data = inst.get_data(picks=picks, **kwargs)
         # retrieve indices of GFP peaks
-        ind_peaks = _extract_gfp_peaks(data, min_peak_distance)
+        gfp = gfp_function(data)
+        ind_peaks, _ = find_peaks(gfp, distance=min_peak_distance)
         # retrieve the peaks data
         if return_all:
             del data  # free up memory
@@ -124,7 +138,8 @@ def extract_gfp_peaks(
         for k in range(len(inst)):  # pylint: disable=consider-using-enumerate
             data = inst[k].get_data(picks=picks, **kwargs)[0, :, :]
             # data is 2D, of shape (n_channels, n_samples)
-            ind_peaks = _extract_gfp_peaks(data, min_peak_distance)
+            gfp = gfp_function(data)
+            ind_peaks, _ = find_peaks(gfp, distance=min_peak_distance)
             if return_all:
                 del data  # free up memory
                 data = inst[k].get_data(picks=picks_all, **kwargs)[0, :, :]
@@ -134,6 +149,7 @@ def extract_gfp_peaks(
     n_samples = inst.times.size
     if isinstance(inst, BaseEpochs):
         n_samples *= len(inst)
+
     logger.info(
         "%s GFP peaks extracted out of %s samples (%.2f%% of the original data).",
         peaks.shape[1],
@@ -143,27 +159,3 @@ def extract_gfp_peaks(
 
     info = pick_info(inst.info, picks_all if return_all else picks)
     return ChData(peaks, info)
-
-
-def _extract_gfp_peaks(
-    data: ScalarFloatArray, min_peak_distance: int = 2
-) -> ScalarFloatArray:
-    """Extract GFP peaks from input data.
-
-    Parameters
-    ----------
-    data : array of shape (n_channels, n_samples)
-        The data to extract GFP peaks from.
-    min_peak_distance : int
-        Required minimal horizontal distance (>= 1) in samples between neighboring
-        peaks. Smaller peaks are removed first until the condition is fulfilled for all
-        remaining peaks. Default to 2.
-
-    Returns
-    -------
-    peaks : array of shape (n_picks,)
-        The indices when peaks occur.
-    """
-    gfp = np.std(data, axis=0)
-    peaks, _ = find_peaks(gfp, distance=min_peak_distance)
-    return peaks

pycrostates/preprocessing/tests/test_extract_gfp_peaks.py

@@ -18,24 +18,39 @@
 dir_ = testing.data_path() / "MEG" / "sample"
 fname_raw_testing = dir_ / "sample_audvis_trunc_raw.fif"
 raw = mne.io.read_raw_fif(fname_raw_testing, preload=False)
-raw.pick("eeg").load_data()
 epochs = mne.make_fixed_length_epochs(raw, duration=2, preload=True)
 
 
 @pytest.mark.parametrize("inst", (raw, epochs))
 def test_extract_gfp(inst, caplog):
     """Test valid arguments for extract_gfp_peaks."""
+    # eeg
+    ch_data = extract_gfp_peaks(inst, picks="eeg")
+    assert isinstance(ch_data, ChData)
+    picks = _picks_to_idx(inst.info, "eeg")
+    assert ch_data.info == pick_info(inst.info, picks)
+
+    # grad
+    ch_data = extract_gfp_peaks(inst, picks="grad")
+    assert isinstance(ch_data, ChData)
+    picks = _picks_to_idx(inst.info, "grad")
+    assert ch_data.info == pick_info(inst.info, picks)
+
+    # mag
+    ch_data = extract_gfp_peaks(inst, picks="mag")
+    assert isinstance(ch_data, ChData)
+    picks = _picks_to_idx(inst.info, "mag")
+    assert ch_data.info == pick_info(inst.info, picks)
+
+    # default(eeg)
     ch_data = extract_gfp_peaks(inst)
     assert isinstance(ch_data, ChData)
-    picks = _picks_to_idx(inst.info, None)
+    picks = _picks_to_idx(inst.info, "eeg")
     assert ch_data.info == pick_info(inst.info, picks)
-    assert "GFP peaks extracted" in caplog.text
 
     # with min_peak_distance
     ch_data2 = extract_gfp_peaks(inst, min_peak_distance=10)
     assert isinstance(ch_data2, ChData)
-    picks = _picks_to_idx(inst.info, None)
-    assert ch_data2.info == pick_info(inst.info, picks)
     assert ch_data != ch_data2
 
     # with picks