initial commit

mne/preprocessing/stim.py

@@ -0,0 +1,176 @@
+# Authors: The MNE-Python contributors.
+# License: BSD-3-Clause
+# Copyright the MNE-Python contributors.
+
+import numpy as np
+from scipy.interpolate import interp1d
+from scipy.signal.windows import hann
+
+from .._fiff.pick import _picks_to_idx
+from ..epochs import BaseEpochs
+from ..event import find_events
+from ..evoked import Evoked
+from ..io import BaseRaw
+from ..utils import _check_option, _check_preload, _validate_type, fill_doc
+
+
+def _get_window(start, end):
+    """Return window which has length as much as parameter start - end."""
+    window = 1 - np.r_[hann(4)[:2], np.ones(np.abs(end - start) - 4), hann(4)[-2:]].T
+    return window
+
+
+def _fix_artifact(
+    data, window, picks, first_samp, last_samp, base_tmin, base_tmax, mode
+):
+    """Modify original data by using parameter data."""
+    if mode == "linear":
+        x = np.array([first_samp, last_samp])
+        f = interp1d(x, data[:, (first_samp, last_samp)][picks])
+        xnew = np.arange(first_samp, last_samp)
+        interp_data = f(xnew)
+        data[picks, first_samp:last_samp] = interp_data
+    if mode == "window":
+        data[picks, first_samp:last_samp] = (
+            data[picks, first_samp:last_samp] * window[np.newaxis, :]
+        )
+    if mode == "constant":
+        data[picks, first_samp:last_samp] = data[picks, base_tmin:base_tmax].mean(
+            axis=1
+        )[:, None]
+
+
+@fill_doc
+def fix_stim_artifact(
+    inst,
+    events=None,
+    event_id=None,
+    tmin=0.0,
+    tmax=0.01,
+    *,
+    baseline=None,
+    mode="linear",
+    stim_channel=None,
+    picks=None,
+):
+    """Eliminate stimulation's artifacts from instance.
+
+    .. note:: This function operates in-place, consider passing
+              ``inst.copy()`` if this is not desired.
+
+    Parameters
+    ----------
+    inst : instance of Raw or Epochs or Evoked
+        The data.
+    events : array, shape (n_events, 3)
+        The list of events. Required only when inst is Raw.
+    event_id : int
+        The id of the events generating the stimulation artifacts.
+        If None, read all events. Required only when inst is Raw.
+    tmin : float
+        Start time of the interpolation window in seconds.
+    tmax : float
+        End time of the interpolation window in seconds.
+    baseline : None | tuple, shape (2,)
+        The baseline to use when ``mode='constant'``, in which case it
+        must be non-None.
+
+        .. versionadded:: 1.8
+    mode : 'linear' | 'window' | 'constant'
+        Way to fill the artifacted time interval.
+
+        ``"linear"``
+            Does linear interpolation.
+        ``"window"``
+            Applies a ``(1 - hanning)`` window.
+        ``"constant"``
+            Uses baseline average. baseline parameter must be provided.
+
+        .. versionchanged:: 1.8
+           Added the ``"constant"`` mode.
+    stim_channel : str | None
+        Stim channel to use.
+    %(picks_all_data)s
+
+    Returns
+    -------
+    inst : instance of Raw or Evoked or Epochs
+        Instance with modified data.
+    """
+    _check_option("mode", mode, ["linear", "window", "constant"])
+    s_start = int(np.ceil(inst.info["sfreq"] * tmin))
+    s_end = int(np.ceil(inst.info["sfreq"] * tmax))
+    if mode == "constant":
+        _validate_type(
+            baseline, (tuple, list), "baseline", extra="when mode='constant'"
+        )
+        _check_option("len(baseline)", len(baseline), [2])
+        for bi, b in enumerate(baseline):
+            _validate_type(
+                b, "numeric", f"baseline[{bi}]", extra="when mode='constant'"
+            )
+        b_start = int(np.ceil(inst.info["sfreq"] * baseline[0]))
+        b_end = int(np.ceil(inst.info["sfreq"] * baseline[1]))
+    else:
+        b_start = b_end = np.nan
+    if (mode == "window") and (s_end - s_start) < 4:
+        raise ValueError(
+            'Time range is too short. Use a larger interval or set mode to "linear".'
+        )
+    window = None
+    if mode == "window":
+        window = _get_window(s_start, s_end)
+
+    picks = _picks_to_idx(inst.info, picks, "data", exclude=())
+
+    _check_preload(inst, "fix_stim_artifact")
+    if isinstance(inst, BaseRaw):
+        if events is None:
+            events = find_events(inst, stim_channel=stim_channel)
+        if len(events) == 0:
+            raise ValueError("No events are found")
+        if event_id is None:
+            events_sel = np.arange(len(events))
+        else:
+            events_sel = events[:, 2] == event_id
+        event_start = events[events_sel, 0]
+        data = inst._data
+        for event_idx in event_start:
+            first_samp = int(event_idx) - inst.first_samp + s_start
+            last_samp = int(event_idx) - inst.first_samp + s_end
+            base_t1 = int(event_idx) - inst.first_samp + b_start
+            base_t2 = int(event_idx) - inst.first_samp + b_end
+            _fix_artifact(
+                data, window, picks, first_samp, last_samp, base_t1, base_t2, mode
+            )
+    elif isinstance(inst, BaseEpochs):
+        if inst.reject is not None:
+            raise RuntimeError(
+                "Reject is already applied. Use reject=None in the constructor."
+            )
+        e_start = int(np.ceil(inst.info["sfreq"] * inst.tmin))
+        first_samp = s_start - e_start
+        last_samp = s_end - e_start
+        data = inst._data
+        base_t1 = b_start - e_start
+        base_t2 = b_end - e_start
+        for epoch in data:
+            _fix_artifact(
+                epoch, window, picks, first_samp, last_samp, base_t1, base_t2, mode
+            )
+
+    elif isinstance(inst, Evoked):
+        first_samp = s_start - inst.first
+        last_samp = s_end - inst.first
+        data = inst.data
+        base_t1 = b_start - inst.first
+        base_t2 = b_end - inst.first
+
+        _fix_artifact(
+            data, window, picks, first_samp, last_samp, base_t1, base_t2, mode
+        )
+
+    else:
+        raise TypeError(f"Not a Raw or Epochs or Evoked (got {type(inst)}).")
+
+    return inst