Repository: sappelhoff/pyprep (python)

Issue:
EOG get dropped
I suppose most people would want to use PyPrep as an early preprocessing pipeline to get ride of line noise and get a robust average reference (at least, that is my understanding from the documentation). To do that, it seems that PyPrep should make available a raw object that has been cleaned and re-referenced. My understanding is that it does, as the .raw attribute of the PrepPipeline class. However, in this Raw object, some channels (like EOG channels) are dropped by PyPrep. I don't think that they should be, since they are likely to be necessary for the following artifact rejection steps that will be run after PyPrep.

Pull request description:
Adding a full_raw property to PrepPipeline
Fix #33.

Patch:
diff --git a/docs/whats_new.rst b/docs/whats_new.rst
index 059a681..f33abe4 100644
--- a/docs/whats_new.rst
+++ b/docs/whats_new.rst
@@ -19,10 +19,12 @@ People who contributed to this software across releases (in alphabetical order)
 
 * `Aamna Lawrence`_
 * `Adam Li`_
+* `Christian Oreilly`_
 * `Stefan Appelhoff`_
 * `Victor Xiang`_
 * `Yorguin Mantilla`_
 
+
 .. _current:
 
 Current
@@ -30,6 +32,10 @@ Current
 
 Changelog
 ~~~~~~~~~
+- Channel types  are now available from a new ``ch_types_all`` attribute, and non-EEG channel names are now available from a new ``ch_names_non_eeg`` attribute from :class:`pyprep.PrepPipeline`, by `Yorguin Mantilla`_ (`#34 <https://github.com/sappelhoff/pyprep/pull/34>`_)
+- Renaming of ``ch_names`` attribute of ``PrepPipeline`` to ``ch_names_all``, by `Yorguin Mantilla`_ (`#34 <https://github.com/sappelhoff/pyprep/pull/34>`_)
+- It's now possible to pass ``'eeg'`` to ``ref_chs`` and ``reref_chs`` parameters to select only eeg channels for referencing, by `Yorguin Mantilla`_ (`#34 <https://github.com/sappelhoff/pyprep/pull/34>`_)
+- Prep will retain the non eeg channels through the ``raw`` attribute. The eeg-only and non-eeg parts will be in raw_eeg and raw_non_eeg respectively. See the ``raw`` attribute, by `Christian Oreilly`_ (`#34 <https://github.com/sappelhoff/pyprep/pull/34>`_)
 - When a ransac call needs more memory than available, pyprep will now automatically switch to a slower but less memory-consuming version of ransac. See the ``channel_wise``parameter, by `Yorguin Mantilla`_ (`#32 <https://github.com/sappelhoff/pyprep/pull/32>`_)
 - It's now possible to pass an empty list for the ``line_freqs`` param in ``PrepPipeline`` to skip the line noise removal, by `Yorguin Mantilla`_ (`#29 <https://github.com/sappelhoff/pyprep/pull/29>`_)
 - The three main classes ``PrepPipeline``, ``NoisyChannels``, and ``Reference`` now have a ``random_state`` parameter to set a seed that gets passed on to all their internal methods and class calls, by `Stefan Appelhoff`_ (`#31 <https://github.com/sappelhoff/pyprep/pull/31>`_)
@@ -108,5 +114,6 @@ Changelog
 .. _Stefan Appelhoff: http://stefanappelhoff.com/
 .. _Aamna Lawrence: https://github.com/AamnaLawrence
 .. _Adam Li: https://github.com/adam2392/
+.. _Christian Oreilly: https://github.com/christian-oreilly
 .. _Victor Xiang: https://github.com/Nick3151
 .. _Yorguin Mantilla: https://github.com/yjmantilla
diff --git a/examples/run_full_prep.py b/examples/run_full_prep.py
index cfe920a..33f1c4d 100644
--- a/examples/run_full_prep.py
+++ b/examples/run_full_prep.py
@@ -72,9 +72,6 @@ montage_kind = "standard_1005"
 montage = mne.channels.make_standard_montage(montage_kind)
 
 # Extract some info
-eeg_index = mne.pick_types(raw.info, eeg=True, eog=False, meg=False)
-ch_names = raw.info["ch_names"]
-ch_names_eeg = list(np.asarray(ch_names)[eeg_index])
 sample_rate = raw.info["sfreq"]
 
 # Make a copy of the data
@@ -84,18 +81,18 @@ raw_copy = raw.copy()
 # Set PREP parameters and run PREP
 # --------------------------------
 #
-# Notes: we keep all the default parameter settings as described in the PREP
-# paper except one, the fraction of bad time windows
-# (we change it from 0.01 to 0.1), because the EEG data is 60s long, which
-# means it gots only 60 time windows. We think the algorithm will be too
-# sensitive if using the default setting.
+# Notes: We keep all the default parameter settings as described in the PREP,
+# though for this case we think that the default fraction of bad time windows
+# being 0.01 is too sensitive since it gots only 60 time windows (the EEG data
+# is 60s long). As a result this example returns a lot of interpolated channels.
 
 # Fit prep
 prep_params = {
-    "ref_chs": ch_names_eeg,
-    "reref_chs": ch_names_eeg,
+    "ref_chs": "eeg",
+    "reref_chs": "eeg",
     "line_freqs": np.arange(60, sample_rate / 2, 60),
 }
+
 prep = PrepPipeline(raw_copy, prep_params, montage)
 prep.fit()
 
diff --git a/pyprep/prep_pipeline.py b/pyprep/prep_pipeline.py
index 2c18426..91f1115 100644
--- a/pyprep/prep_pipeline.py
+++ b/pyprep/prep_pipeline.py
@@ -17,22 +17,20 @@ class PrepPipeline:
     Parameters
     ----------
     raw : mne.raw
-        The data.
+        The data. Channel types must be correctly assigned (e.g.,
+        ocular channels are assigned the type 'eog').
     prep_params : dict
         Parameters of PREP which include at least the following keys:
-
-        - ref_chs : list
-            - A list of channel names to be used for rereferencing
-              [default: all channels]
-
-        - reref_chs : list
+        - ref_chs : list | 'eeg'
+            - A list of channel names to be used for rereferencing.
+               Can be a str 'eeg' to use all EEG channels.
+        - reref_chs : list | 'eeg'
             - A list of channel names to be used for line-noise removed, and
-              referenced [default: all channels]
-
+              referenced. Can be a str 'eeg' to use all EEG channels.
         - line_freqs : array_like
             - list of floats indicating frequencies to be removed.
-              Can be an empty list to skip this step.
-
+              For example for 60Hz you may specify ``np.arange(60, sfreq / 2, 60)``.
+              Specify an empty list to skip the line noise removal step.
     montage : DigMontage
         Digital montage of EEG data.
     ransac : boolean
@@ -43,6 +41,33 @@ class PrepPipeline:
         If None, the seed will be obtained from the operating system
         (see RandomState for details). Default is None.
 
+    Attributes
+    ----------
+    raw : mne.raw
+        The data including eeg and non eeg channels. It is unprocessed if
+        accessed before the fit method, processed if accessed after a
+        succesful fit method.
+    raw_eeg : mne.raw
+        The only-eeg part of the data. It is unprocessed if accessed before
+        the fit method, processed if accessed after a succesful fit method.
+    raw_non_eeg : mne.raw | None
+        The non-eeg part of the data. It is not processed when calling
+        the fit method. If the input was only EEG it will be None.
+    noisy_channels_original : list
+        bad channels before robust reference
+    bad_before_interpolation : list
+        bad channels after robust reference but before interpolation
+    EEG_before_interpolation : ndarray
+        EEG data in uV before the interpolation
+    reference_before_interpolation : ndarray
+        Reference signal in uV before interpolation.
+    reference_after_interpolation : ndarray
+        Reference signal in uV after interpolation.
+    interpolated_channels : list
+        Names of the interpolated channels.
+    still_noisy_channels : list
+        Names of the noisy channels after interpolation.
+
     References
     ----------
     .. [1] Bigdely-Shamlo, N., Mullen, T., Kothe, C., Su, K. M., Robbins, K. A.
@@ -53,20 +78,50 @@ class PrepPipeline:
 
     def __init__(self, raw, prep_params, montage, ransac=True, random_state=None):
         """Initialize PREP class."""
-        self.raw = raw.copy()
-        self.ch_names = self.raw.ch_names
-        self.raw.set_montage(montage)
-        self.raw.pick_types(eeg=True, eog=False, meg=False)
-        self.ch_names_eeg = self.raw.ch_names
-        self.EEG_raw = self.raw.get_data() * 1e6
+        self.raw_eeg = raw.copy()
+
+        # split eeg and non eeg channels
+        self.ch_names_all = raw.ch_names.copy()
+        self.ch_types_all = raw.get_channel_types()
+        self.ch_names_eeg = [
+            self.ch_names_all[i]
+            for i in range(len(self.ch_names_all))
+            if self.ch_types_all[i] == "eeg"
+        ]
+        self.ch_names_non_eeg = list(set(self.ch_names_all) - set(self.ch_names_eeg))
+        self.raw_eeg.pick_channels(self.ch_names_eeg)
+        if self.ch_names_non_eeg == []:
+            self.raw_non_eeg = None
+        else:
+            self.raw_non_eeg = raw.copy()
+            self.raw_non_eeg.pick_channels(self.ch_names_non_eeg)
+
+        self.raw_eeg.set_montage(montage)
+        # raw_non_eeg may not be compatible with the montage
+        # so it is not set for that object
+
+        self.EEG_raw = self.raw_eeg.get_data() * 1e6
         self.prep_params = prep_params
-        self.sfreq = self.raw.info["sfreq"]
+        if self.prep_params["ref_chs"] == "eeg":
+            self.prep_params["ref_chs"] = self.ch_names_eeg
+        if self.prep_params["reref_chs"] == "eeg":
+            self.prep_params["reref_chs"] = self.ch_names_eeg
+        self.sfreq = self.raw_eeg.info["sfreq"]
         self.ransac = ransac
         self.random_state = check_random_state(random_state)
 
+    @property
+    def raw(self):
+        """Return a version of self.raw_eeg that includes the non-eeg channels."""
+        full_raw = self.raw_eeg.copy()
+        if self.raw_non_eeg is None:
+            return full_raw
+        else:
+            return full_raw.add_channels([self.raw_non_eeg])
+
     def fit(self):
         """Run the whole PREP pipeline."""
-        noisy_detector = NoisyChannels(self.raw, random_state=self.random_state)
+        noisy_detector = NoisyChannels(self.raw_eeg, random_state=self.random_state)
         noisy_detector.find_bad_by_nan_flat()
         # unusable_channels = _union(
         #     noisy_detector.bad_by_nan, noisy_detector.bad_by_flat
@@ -89,17 +144,17 @@ class PrepPipeline:
 
             # Add Trend back
             self.EEG = self.EEG_raw - self.EEG_new + self.EEG_clean
-            self.raw._data = self.EEG * 1e-6
+            self.raw_eeg._data = self.EEG * 1e-6
 
         # Step 3: Referencing
         reference = Reference(
-            self.raw,
+            self.raw_eeg,
             self.prep_params,
             ransac=self.ransac,
             random_state=self.random_state,
         )
         reference.perform_reference()
-        self.raw = reference.raw
+        self.raw_eeg = reference.raw
         self.noisy_channels_original = reference.noisy_channels_original
         self.bad_before_interpolation = reference.bad_before_interpolation
         self.EEG_before_interpolation = reference.EEG_before_interpolation
diff --git a/pyprep/utils.py b/pyprep/utils.py
index 33e4123..eda8185 100644
--- a/pyprep/utils.py
+++ b/pyprep/utils.py
@@ -4,6 +4,7 @@ from cmath import sqrt
 
 import numpy as np
 import scipy.interpolate
+import mne
 
 
 def _union(list1, list2):
@@ -86,3 +87,66 @@ def split_list(mylist, chunk_size):
     return [
         mylist[offs : offs + chunk_size] for offs in range(0, len(mylist), chunk_size)
     ]
+
+
+def make_random_mne_object(
+    ch_names,
+    ch_types,
+    times,
+    sfreq,
+    n_freq_comps=5,
+    freq_range=[10, 60],
+    scale=1e-6,
+    RNG=np.random.RandomState(1337),
+):
+    """Make a random MNE object to use for testing.
+
+    Parameters
+    ----------
+    ch_names : list
+        names of channels
+    ch_types : list
+        types of channels
+    times : 1d ndarray
+        Time vector to use.
+    sfreq : float
+        Sampling frequency associated with the time vector.
+    n_freq_comps : int
+        Number of signal components summed to make a signal.
+    freq_range : list, len==2
+        Signals will contain freqs from this range.
+    scale : float
+        Scaling factor applied to the signal in volts. For example 1e-6 to
+        get microvolts.
+    RNG : np.random.RandomState
+        Random state seed.
+
+    Returns
+    -------
+    raw : mne raw object
+        The mne object for performing the tests.
+    n_freq_comps : int
+        Number of random frequency components to introduce.
+    freq_range : tuple
+        The low (`freq_range[0]`) and high (`freq_range[1]`) endpoints of
+        a frequency range from which random draws will be made to
+        introduce frequency components in the test data.
+    """
+    n_chans = len(ch_names)
+    signal_len = times.shape[0]
+    # Make a random signal
+    signal = np.zeros((n_chans, signal_len))
+    low = freq_range[0]
+    high = freq_range[1]
+    for chan in range(n_chans):
+        # Each channel signal is a sum of random freq sine waves
+        for freq_i in range(n_freq_comps):
+            freq = RNG.randint(low, high, signal_len)
+            signal[chan, :] += np.sin(2 * np.pi * times * freq)
+
+    signal *= scale  # scale
+
+    # Make mne object
+    info = mne.create_info(ch_names=ch_names, sfreq=sfreq, ch_types=ch_types)
+    raw = mne.io.RawArray(signal, info)
+    return raw, n_freq_comps, freq_range


Test patch:
diff --git a/tests/test_noisy.py b/tests/test_noisy.py
index 57295ce..dea3443 100644
--- a/tests/test_noisy.py
+++ b/tests/test_noisy.py
@@ -5,88 +5,42 @@ import numpy as np
 import pytest
 
 from pyprep.noisy import Noisydata, find_bad_epochs
+from pyprep.utils import make_random_mne_object
 
 RNG = np.random.RandomState(1337)
 
 
-def make_random_mne_object(
-    sfreq=1000.0, t_secs=600, n_freq_comps=5, freq_range=[10, 60]
-):
-    """Make a random MNE object to use for testing.
-
-    Parameters
-    ----------
-    sfreq : float
-        Sampling frequency in Hz.
-
-    t_secs : int
-        Recording length in seconds.
-
-    n_freq_comps : int
-        Number of signal components summed to make a signal.
-
-    freq_range : list, len==2
-        Signals will contain freqs from this range.
-
-    Returns
-    -------
-    raw : mne raw object
-        The mne object for performing the tests.
-
-    n_freq_comps : int
-
-    freq_range : list, len==2
-
-    """
-    t = np.arange(0, t_secs, 1.0 / sfreq)
-    signal_len = t.shape[0]
-    ch_names = [
-        "Fpz",
-        "AFz",
-        "Fz",
-        "FCz",
-        "Cz",
-        "CPz",
-        "Pz",
-        "POz",
-        "Oz",
-        "C1",
-        "C2",
-        "C3",
-        "C4",
-        "C5",
-        "C6",
-    ]
-    ch_types = ["eeg" for chn in ch_names]
-    n_chans = len(ch_names)
-
-    # Make a random signal
-    signal = np.zeros((n_chans, signal_len))
-    low = freq_range[0]
-    high = freq_range[1]
-    for chan in range(n_chans):
-        # Each channel signal is a sum of random freq sine waves
-        for freq_i in range(n_freq_comps):
-            freq = RNG.randint(low, high, signal_len)
-            signal[chan, :] += np.sin(2 * np.pi * t * freq)
-
-    signal *= 1e-6  # scale to Volts
-
-    # Make mne object
-    info = mne.create_info(ch_names=ch_names, sfreq=sfreq, ch_types=ch_types)
-    raw = mne.io.RawArray(signal, info)
-    return raw, n_freq_comps, freq_range
-
-
 # We make new random mne objects until we have one without inherent bad chans
 # This is required so that we can then selectively insert noise in the tests.
 found_good_test_object = False
-while not found_good_test_object:
-    raw, n_freq_comps, freq_range = make_random_mne_object()
-    nd = Noisydata(raw)
-    nd.find_all_bads(ransac=False)
-    if nd.get_bads() == []:
-        found_good_test_object = True
+t_secs = 600
+sfreq = 1000.0
+t = np.arange(0, t_secs, 1.0 / sfreq)
+ch_names = [
+    "Fpz",
+    "AFz",
+    "Fz",
+    "FCz",
+    "Cz",
+    "CPz",
+    "Pz",
+    "POz",
+    "Oz",
+    "C1",
+    "C2",
+    "C3",
+    "C4",
+    "C5",
+    "C6",
+]
+ch_types = ["eeg" for chn in ch_names]
+
+raw, n_freq_comps, freq_range = make_random_mne_object(
+    ch_names, ch_types, t, sfreq, RNG=RNG
+)
+nd = Noisydata(raw)
+nd.find_all_bads(ransac=False)
+assert nd.get_bads() == []
 
 # Make some arbitrary events sampled from the mid-section of raw.times
 n_events = 3
diff --git a/tests/test_prep_pipeline.py b/tests/test_prep_pipeline.py
index ca202fb..07a634c 100644
--- a/tests/test_prep_pipeline.py
+++ b/tests/test_prep_pipeline.py
@@ -6,6 +6,7 @@ import pytest
 import scipy.io as sio
 
 from pyprep.prep_pipeline import PrepPipeline
+from pyprep.utils import make_random_mne_object
 
 
 @pytest.mark.usefixtures("raw", "montage")
@@ -204,3 +205,45 @@ def test_prep_pipeline(raw, montage):
     # axs[4, 0].set_title('Interpolation', loc='left', fontsize=14)
     # axs[4, 0].set_ylabel('Channel Number', fontsize=14)
     axs[4, 1].set_xlabel("Time(s)", fontsize=14)
+
+
+def test_prep_pipeline_non_eeg(raw, montage):
+    """Test prep pipeline with non eeg channels."""
+    raw_copy = raw.copy()
+
+    # make arbitrary non eeg channels from the register
+    sfreq = raw_copy.info["sfreq"]  # Sampling frequency
+    times = np.array(list(range(raw_copy._data.shape[1])))
+    ch_names_non_eeg = ["misc" + str(i) for i in range(4)]
+    ch_types_non_eeg = ["misc" for i in range(4)]
+    raw_non_eeg, _, _ = make_random_mne_object(
+        ch_names_non_eeg,
+        ch_types_non_eeg,
+        times,
+        sfreq,
+        RNG=np.random.RandomState(1337),
+    )
+
+    raw_copy.add_channels([raw_non_eeg])
+
+    prep_params = {
+        "ref_chs": "eeg",
+        "reref_chs": "eeg",
+        "line_freqs": np.arange(60, sfreq / 2, 60),
+    }
+    prep = PrepPipeline(raw_copy, prep_params, montage, random_state=42)
+
+    prep.fit()
+
+    # correct non-eeg channels configured in init
+    assert set(prep.ch_names_non_eeg) == set(ch_names_non_eeg)
+    # original (all) channel names same as full_raw names
+    assert set(prep.raw.ch_names) == set(raw_copy.ch_names)
+    # names of raw (only eeg)  same as full names - non eeg names
+    assert set(prep.raw_eeg.ch_names) == set(raw_copy.ch_names) - set(ch_names_non_eeg)
+    # quantity of raw (only eeg) same as quantity of all - non eeg lists
+    assert prep.raw_eeg._data.shape[0] == len(raw_copy.ch_names) - len(
+        prep.ch_names_non_eeg
+    )
+    # quantity of channels in is the same as qty of full raw
+    assert raw_copy._data.shape[0] == prep.raw._data.shape[0]