MRG+1: fix for set_bipolar_reference (#5780)

* fix for set_bipolar_reference

* updated code with suggestions in set_bipolar_reference and test

* test update

* add drop_ref parameter to set_bipolar_reference

* mmagnuski nitpicks

doc/whats_new.rst

@@ -19,6 +19,8 @@ Current
 Changelog
 ~~~~~~~~~
 
+- Add ``drop_refs=True`` parameter to :func:`set_bipolar_reference` to drop/keep anode and cathode channels after applying the reference by `Cristóbal Moënne-Loccoz`_. 
+
 - Add 448-labels subdivided aparc cortical parcellation by `Denis Engemann`_ and `Sheraz Khan`_
 
 - Add keyboard shortcuts to nativate volume source estimates in time using (shift+)left/right arrow keys by `Mainak Jas`_
@@ -28,6 +30,8 @@ Changelog
 Bug
 ~~~
 
+- Fix :func:`set_bipolar_reference` in the case of generating all bipolar combinations and also in the case of repeated channels in both lists (anode and cathode) by `Cristóbal Moënne-Loccoz`_
+
 - Fix missing code for computing the median when ``method='median'`` in :meth:`mne.Epochs.average` by `Cristóbal Moënne-Loccoz`_
 
 - Fix CTF helmet plotting in :func:`mne.viz.plot_evoked_field` by `Eric Larson`_

mne/io/reference.py

@@ -398,7 +398,7 @@ def set_eeg_reference(inst, ref_channels='average', copy=True,
 
 @verbose
 def set_bipolar_reference(inst, anode, cathode, ch_name=None, ch_info=None,
-                          copy=True, verbose=None):
+                          drop_refs=True, copy=True, verbose=None):
     """Re-reference selected channels using a bipolar referencing scheme.
 
     A bipolar reference takes the difference between two channels (the anode
@@ -430,6 +430,8 @@ def set_bipolar_reference(inst, anode, cathode, ch_name=None, ch_info=None,
         This parameter can be used to supply a dictionary (or a dictionary for
         each bipolar channel) containing channel information to merge in,
         overwriting the default values. Defaults to None.
+    drop_refs : bool
+        Whether to drop the anode/cathode channels from the instance.
     copy : bool
         Whether to operate on a copy of the data (True) or modify it in-place
         (False). Defaults to True.
@@ -508,11 +510,10 @@ def set_bipolar_reference(inst, anode, cathode, ch_name=None, ch_info=None,
     if copy:
         inst = inst.copy()
 
-    rem_ca = list(cathode)
     for i, (an, ca, name, chs) in enumerate(
             zip(anode, cathode, ch_name, new_chs)):
-        if an in anode[i + 1:]:
-            # Make a copy of anode if it's still needed later
+        if an in anode[i + 1:] or an in cathode[i + 1:] or not drop_refs:
+            # Make a copy of the channel if it's still needed later
             # otherwise it's modified inplace
             _copy_channel(inst, an, 'TMP')
             an = 'TMP'
@@ -522,11 +523,10 @@ def set_bipolar_reference(inst, anode, cathode, ch_name=None, ch_info=None,
         inst.info['chs'][an_idx]['ch_name'] = name
         logger.info('Bipolar channel added as "%s".' % name)
         inst.info._update_redundant()
-        if an in rem_ca:
-            idx = rem_ca.index(an)
-            del rem_ca[idx]
 
-    # Drop remaining cathode channels
-    inst.drop_channels(rem_ca)
+    # Drop remaining channels.
+    if drop_refs:
+        drop_channels = (set(anode) | set(cathode)) & set(inst.ch_names)
+        inst.drop_channels(drop_channels)
 
     return inst

mne/io/tests/test_reference.py

@@ -4,6 +4,7 @@
 #
 # License: BSD (3-clause)
 
+import itertools
 import os.path as op
 import numpy as np
 
@@ -12,9 +13,9 @@
 
 from mne import (pick_channels, pick_types, Epochs, read_events,
                  set_eeg_reference, set_bipolar_reference,
-                 add_reference_channels)
+                 add_reference_channels, create_info)
 from mne.epochs import BaseEpochs
-from mne.io import read_raw_fif
+from mne.io import RawArray, read_raw_fif
 from mne.io.constants import FIFF
 from mne.io.proj import _has_eeg_average_ref_proj, Projection
 from mne.io.reference import _apply_reference
@@ -473,4 +474,48 @@ def test_add_reference():
     pytest.raises(ValueError, add_reference_channels, raw, 1)
 
 
+def test_bipolar_combinations():
+    """Test bipolar channel generation."""
+    ch_names = ['CH' + str(ni + 1) for ni in range(10)]
+    info = create_info(
+        ch_names=ch_names, sfreq=1000., ch_types=['eeg'] * len(ch_names))
+    raw_data = np.random.randn(len(ch_names), 1000)
+    raw = RawArray(raw_data, info)
+
+    def _check_bipolar(raw_test, ch_a, ch_b):
+        picks = [raw_test.ch_names.index(ch_a + '-' + ch_b)]
+        get_data_res = raw_test.get_data(picks=picks)[0, :]
+        manual_a = raw_data[ch_names.index(ch_a), :]
+        manual_b = raw_data[ch_names.index(ch_b), :]
+        assert_array_equal(get_data_res, manual_a - manual_b)
+
+    # test classic EOG/ECG bipolar reference (only two channels per pair).
+    raw_test = set_bipolar_reference(raw, ['CH2'], ['CH1'], copy=True)
+    _check_bipolar(raw_test, 'CH2', 'CH1')
+
+    # test all combinations.
+    a_channels, b_channels = zip(*itertools.combinations(ch_names, 2))
+    a_channels, b_channels = list(a_channels), list(b_channels)
+    raw_test = set_bipolar_reference(raw, a_channels, b_channels, copy=True)
+    for ch_a, ch_b in zip(a_channels, b_channels):
+        _check_bipolar(raw_test, ch_a, ch_b)
+    # check if reference channels have been dropped.
+    assert (len(raw_test.ch_names) == len(a_channels))
+
+    raw_test = set_bipolar_reference(
+        raw, a_channels, b_channels, drop_refs=False, copy=True)
+    # check if reference channels have been kept correctly.
+    assert (len(raw_test.ch_names) == len(a_channels) + len(ch_names))
+    for idx, ch_label in enumerate(ch_names):
+        manual_ch = raw_data[idx, :]
+        assert_array_equal(
+            raw_test._data[raw_test.ch_names.index(ch_label), :], manual_ch)
+
+    # test bipolars with a channel in both list (anode & cathode).
+    raw_test = set_bipolar_reference(
+        raw, ['CH2', 'CH1'], ['CH1', 'CH2'], copy=True)
+    _check_bipolar(raw_test, 'CH2', 'CH1')
+    _check_bipolar(raw_test, 'CH1', 'CH2')
+
+
 run_tests_if_main()