Aluode/PerceptionLabPortable

1# Authors: The MNE-Python contributors.2# License: BSD-3-Clause3# Copyright the MNE-Python contributors.4 5import numpy as np6from scipy.interpolate import interp1d7from scipy.signal.windows import hann8 9from .._fiff.pick import _picks_to_idx10from ..epochs import BaseEpochs11from ..event import find_events12from ..evoked import Evoked13from ..io import BaseRaw14from ..utils import _check_option, _check_preload, _validate_type, fill_doc15 16 17def _get_window(start, end):18    """Return window which has length as much as parameter start - end."""19    window = 1 - np.r_[hann(4)[:2], np.ones(np.abs(end - start) - 4), hann(4)[-2:]].T20    return window21 22 23def _fix_artifact(24    data, window, picks, first_samp, last_samp, base_tmin, base_tmax, mode25):26    """Modify original data by using parameter data."""27    if mode == "linear":28        x = np.array([first_samp, last_samp])29        f = interp1d(x, data[:, (first_samp, last_samp)][picks])30        xnew = np.arange(first_samp, last_samp)31        interp_data = f(xnew)32        data[picks, first_samp:last_samp] = interp_data33    if mode == "window":34        data[picks, first_samp:last_samp] = (35            data[picks, first_samp:last_samp] * window[np.newaxis, :]36        )37    if mode == "constant":38        data[picks, first_samp:last_samp] = data[picks, base_tmin:base_tmax].mean(39            axis=140        )[:, None]41 42 43@fill_doc44def fix_stim_artifact(45    inst,46    events=None,47    event_id=None,48    tmin=0.0,49    tmax=0.01,50    *,51    baseline=None,52    mode="linear",53    stim_channel=None,54    picks=None,55):56    """Eliminate stimulation's artifacts from instance.57 58    .. note:: This function operates in-place, consider passing59              ``inst.copy()`` if this is not desired.60 61    Parameters62    ----------63    inst : instance of Raw or Epochs or Evoked64        The data.65    events : array, shape (n_events, 3)66        The list of events. Required only when inst is Raw.67    event_id : int68        The id of the events generating the stimulation artifacts.69        If None, read all events. Required only when inst is Raw.70    tmin : float71        Start time of the interpolation window in seconds.72    tmax : float73        End time of the interpolation window in seconds.74    baseline : None | tuple, shape (2,)75        The baseline to use when ``mode='constant'``, in which case it76        must be non-None.77 78        .. versionadded:: 1.879    mode : 'linear' | 'window' | 'constant'80        Way to fill the artifacted time interval.81 82        ``"linear"``83            Does linear interpolation.84        ``"window"``85            Applies a ``(1 - hanning)`` window.86        ``"constant"``87            Uses baseline average. baseline parameter must be provided.88 89        .. versionchanged:: 1.890           Added the ``"constant"`` mode.91    stim_channel : str | None92        Stim channel to use.93    %(picks_all_data)s94 95    Returns96    -------97    inst : instance of Raw or Evoked or Epochs98        Instance with modified data.99    """100    _check_option("mode", mode, ["linear", "window", "constant"])101    s_start = int(np.ceil(inst.info["sfreq"] * tmin))102    s_end = int(np.ceil(inst.info["sfreq"] * tmax))103    if mode == "constant":104        _validate_type(105            baseline, (tuple, list), "baseline", extra="when mode='constant'"106        )107        _check_option("len(baseline)", len(baseline), [2])108        for bi, b in enumerate(baseline):109            _validate_type(110                b, "numeric", f"baseline[{bi}]", extra="when mode='constant'"111            )112        b_start = int(np.ceil(inst.info["sfreq"] * baseline[0]))113        b_end = int(np.ceil(inst.info["sfreq"] * baseline[1]))114    else:115        b_start = b_end = np.nan116    if (mode == "window") and (s_end - s_start) < 4:117        raise ValueError(118            'Time range is too short. Use a larger interval or set mode to "linear".'119        )120    window = None121    if mode == "window":122        window = _get_window(s_start, s_end)123 124    picks = _picks_to_idx(inst.info, picks, "data", exclude=())125 126    _check_preload(inst, "fix_stim_artifact")127    if isinstance(inst, BaseRaw):128        if events is None:129            events = find_events(inst, stim_channel=stim_channel)130        if len(events) == 0:131            raise ValueError("No events are found")132        if event_id is None:133            events_sel = np.arange(len(events))134        else:135            events_sel = events[:, 2] == event_id136        event_start = events[events_sel, 0]137        data = inst._data138        for event_idx in event_start:139            first_samp = int(event_idx) - inst.first_samp + s_start140            last_samp = int(event_idx) - inst.first_samp + s_end141            base_t1 = int(event_idx) - inst.first_samp + b_start142            base_t2 = int(event_idx) - inst.first_samp + b_end143            _fix_artifact(144                data, window, picks, first_samp, last_samp, base_t1, base_t2, mode145            )146    elif isinstance(inst, BaseEpochs):147        if inst.reject is not None:148            raise RuntimeError(149                "Reject is already applied. Use reject=None in the constructor."150            )151        e_start = int(np.ceil(inst.info["sfreq"] * inst.tmin))152        first_samp = s_start - e_start153        last_samp = s_end - e_start154        data = inst._data155        base_t1 = b_start - e_start156        base_t2 = b_end - e_start157        for epoch in data:158            _fix_artifact(159                epoch, window, picks, first_samp, last_samp, base_t1, base_t2, mode160            )161 162    elif isinstance(inst, Evoked):163        first_samp = s_start - inst.first164        last_samp = s_end - inst.first165        data = inst.data166        base_t1 = b_start - inst.first167        base_t2 = b_end - inst.first168 169        _fix_artifact(170            data, window, picks, first_samp, last_samp, base_t1, base_t2, mode171        )172 173    else:174        raise TypeError(f"Not a Raw or Epochs or Evoked (got {type(inst)}).")175 176    return inst177