Add support for dictionary-type ref_channels in set_eeg_reference()

doc/changes/devel/12366.newfeature.rst

@@ -0,0 +1 @@
+Add support for `dict` type argument ``ref_channels`` to :func:`mne.set_eeg_reference`, to allow flexible re-referencing (e.g. ``raw.set_eeg_reference(ref_channels={'A1': ['A2', 'A3']})`` will set the new A1 data to be ``A1 - (A2 + A3)/2``), by :newcontrib:`Alex Lepauvre` and `Qian Chu`_

doc/changes/names.inc

@@ -22,6 +22,8 @@
 
 .. _Alex Kiefer: https://home.alexk101.dev
 
+.. _Alex Lepauvre: https://github.com/AlexLepauvre
+
 .. _Alex Rockhill: https://github.com/alexrockhill/
 
 .. _Alexander Rudiuk: https://github.com/ARudiuk

mne/_fiff/reference.py

@@ -87,6 +87,71 @@ def _check_before_reference(inst, ref_from, ref_to, ch_type):
     return ref_to
 
 
+def _check_before_dict_reference(inst, ref_dict):
+    """Prepare instance for dict-based referencing."""
+    # Check to see that data is preloaded
+    _check_preload(inst, "Applying a reference")
+
+    # Promote all values to list-like. This simplifies our logic and also helps catch
+    # self-referencing cases like `{"Cz": ["Cz"]}`
+    _refdict = {k: [v] if isinstance(v, str) else list(v) for k, v in ref_dict.items()}
+
+    # Check that keys are strings and values are lists-of-strings
+    key_types = {type(k) for k in _refdict}
+    value_types = {type(v) for val in _refdict.values() for v in val}
+    for elem_name, elem in dict(key=key_types, value=value_types).items():
+        if bad_elem := elem - {str}:
+            raise TypeError(
+                f"{elem_name.capitalize()}s in the ref_channels dict must be strings. "
+                f"Your dict has {elem_name}s of type "
+                f'{", ".join(map(lambda x: x.__name__, bad_elem))}.'
+            )
+
+    # Check that keys are valid channels and values are lists-of-valid-channels
+    ch_set = set(inst.ch_names)
+    bad_ch_set = set(inst.info["bads"])
+    keys = set(_refdict)
+    values = set(sum(_refdict.values(), []))
+    for elem_name, elem in dict(key=keys, value=values).items():
+        if bad_elem := elem - ch_set:
+            raise ValueError(
+                f'ref_channels dict contains invalid {elem_name}(s) '
+                f'({", ".join(bad_elem)}) '
+                "that are not names of channels in the instance."
+            )
+        # Check that values are not bad channels
+        if bad_elem := elem.intersection(bad_ch_set):
+            warn(
+                f"ref_channels dict contains {elem_name}(s) "
+                f"({', '.join(bad_elem)}) "
+                "that are marked as bad channels."
+            )
+
+    # Check for self-referencing
+    self_ref = [[k] == v for k, v in _refdict.items()]
+    if any(self_ref):
+        which = np.array(list(_refdict))[np.nonzero(self_ref)]
+        for ch in which:
+            warn(f"Channel {ch} is self-referenced, which will nullify the channel.")
+
+    # Check that channel types match. First unpack list-like vals into separate items:
+    pairs = [(k, v) for k in _refdict for v in _refdict[k]]
+    ch_type_map = dict(zip(inst.ch_names, inst.get_channel_types()))
+    mismatch = [ch_type_map[k] != ch_type_map[v] for k, v in pairs]
+    if any(mismatch):
+        mismatch_pairs = np.array(pairs)[mismatch]
+        for k, v in mismatch_pairs:
+            warn(
+                f"Channel {k} ({ch_type_map[k]}) is referenced to channel {v} which is "
+                f"a different channel type ({ch_type_map[v]})."
+            )
+
+    # convert channel names to indices
+    keys_ix = pick_channels(inst.ch_names, list(_refdict), ordered=True)
+    vals_ix = (pick_channels(inst.ch_names, v, ordered=True) for v in _refdict.values())
+    return dict(zip(keys_ix, vals_ix))
+
+
 def _apply_reference(inst, ref_from, ref_to=None, forward=None, ch_type="auto"):
     """Apply a custom EEG referencing scheme."""
     ref_to = _check_before_reference(inst, ref_from, ref_to, ch_type)
@@ -128,6 +193,22 @@ def _apply_reference(inst, ref_from, ref_to=None, forward=None, ch_type="auto"):
     return inst, ref_data
 
 
+def _apply_dict_reference(inst, ref_dict):
+    """Apply a dict-based custom EEG referencing scheme."""
+    # this converts all keys to channel indices and all values to arrays of ch. indices:
+    ref_dict = _check_before_dict_reference(inst, ref_dict)
+
+    data = inst._data
+    orig_data = data.copy()
+    for ref_to, ref_from in ref_dict.items():
+        ref_data = orig_data[..., ref_from, :].mean(-2, keepdims=True)
+        data[..., [ref_to], :] -= ref_data
+
+    with inst.info._unlock():
+        inst.info["custom_ref_applied"] = FIFF.FIFFV_MNE_CUSTOM_REF_ON
+    return inst, None
+
+
 @fill_doc
 def add_reference_channels(inst, ref_channels, copy=True):
     """Add reference channels to data that consists of all zeros.
@@ -319,18 +400,23 @@ def set_eeg_reference(
     Returns
     -------
     inst : instance of Raw | Epochs | Evoked
-        Data with EEG channels re-referenced. If ``ref_channels='average'`` and
+        Data with EEG channels re-referenced. If ``ref_channels="average"`` and
         ``projection=True`` a projection will be added instead of directly
         re-referencing the data.
     ref_data : array
         Array of reference data subtracted from EEG channels. This will be
-        ``None`` if ``projection=True`` or ``ref_channels='REST'``.
+        ``None`` if ``projection=True``, or if ``ref_channels`` is ``"REST"`` or a
+        :class:`dict`.
     %(set_eeg_reference_see_also_notes)s
     """
     from ..forward import Forward
 
     _check_can_reref(inst)
 
+    if isinstance(ref_channels, dict):
+        logger.info("Applying a custom dict-based reference.")
+        return _apply_dict_reference(inst, ref_channels)
+
     ch_type = _get_ch_type(inst, ch_type)
 
     if projection:  # average reference projector

mne/_fiff/tests/test_reference.py

@@ -364,6 +364,160 @@ def test_set_eeg_reference_rest():
     assert 0.995 < exp_var <= 1
 
 
+@testing.requires_testing_data
+@pytest.mark.parametrize("inst_type", ["raw", "epochs"])
+@pytest.mark.parametrize(
+    "ref_channels, expectation",
+    [
+        (
+            {2: "EEG 001"},
+            pytest.raises(
+                TypeError,
+                match="Keys in the ref_channels dict must be strings. "
+                "Your dict has keys of type int.",
+            ),
+        ),
+        (
+            {"EEG 001": (1, 2)},
+            pytest.raises(
+                TypeError,
+                match="Values in the ref_channels dict must be strings. "
+                "Your dict has values of type int.",
+            ),
+        ),
+        (
+            {"EEG 001": [1, 2]},
+            pytest.raises(
+                TypeError,
+                match="Values in the ref_channels dict must be strings. "
+                "Your dict has values of type int.",
+            ),
+        ),
+        (
+            {"EEG 999": "EEG 001"},
+            pytest.raises(
+                ValueError,
+                match=r"ref_channels dict contains invalid key\(s\) \(EEG 999\) "
+                "that are not names of channels in the instance.",
+            ),
+        ),
+        (
+            {"EEG 001": "EEG 999"},
+            pytest.raises(
+                ValueError,
+                match=r"ref_channels dict contains invalid value\(s\) \(EEG 999\) "
+                "that are not names of channels in the instance.",
+            ),
+        ),
+        (
+            {"EEG 001": "EEG 057"},
+            pytest.warns(
+                RuntimeWarning,
+                match=r"ref_channels dict contains value\(s\) \(EEG 057\) "
+                "that are marked as bad channels.",
+            ),
+        ),
+        (
+            {"EEG 001": "STI 001"},
+            pytest.warns(
+                RuntimeWarning,
+                match=(
+                    r"Channel EEG 001 \(eeg\) is referenced to channel "
+                    r"STI 001 which is a different channel type \(stim\)."
+                ),
+            ),
+        ),
+        (
+            {"EEG 001": "EEG 001"},
+            pytest.warns(
+                RuntimeWarning,
+                match=(
+                    "Channel EEG 001 is self-referenced, "
+                    "which will nullify the channel."
+                ),
+            ),
+        ),
+        (
+            {"EEG 001": "EEG 002", "EEG 002": "EEG 003", "EEG 003": "EEG 005"},
+            nullcontext(),
+        ),
+        (
+            {
+                "EEG 001": ["EEG 002", "EEG 003"],
+                "EEG 002": "EEG 003",
+                "EEG 003": "EEG 005",
+            },
+            nullcontext(),
+        ),
+    ],
+)
+def test_set_eeg_reference_dict(ref_channels, inst_type, expectation):
+    """Test setting dict-based reference."""
+    if inst_type == "raw":
+        inst = read_raw_fif(fif_fname).crop(0, 1).pick(picks=["eeg", "stim"])
+    # Test re-referencing Epochs object
+    elif inst_type == "epochs":
+        raw = read_raw_fif(fif_fname, preload=False)
+        events = read_events(eve_fname)
+        inst = Epochs(
+            raw,
+            events=events,
+            event_id=1,
+            tmin=-0.2,
+            tmax=0.5,
+            preload=False,
+        )
+    with pytest.raises(
+        RuntimeError,
+        match="By default, MNE does not load data.*Applying a reference requires.*",
+    ):
+        inst.set_eeg_reference(ref_channels=ref_channels)
+    inst.load_data()
+    inst.info["bads"] = ["EEG 057"]
+    with expectation:
+        reref, _ = set_eeg_reference(inst.copy(), ref_channels, copy=False)
+
+    if isinstance(expectation, nullcontext):
+        # Check that the custom_ref_applied is set correctly:
+        assert reref.info["custom_ref_applied"] == FIFF.FIFFV_MNE_CUSTOM_REF_ON
+
+        # Get raw data
+        _data = inst._data
+
+        # Get that channels that were and weren't re-referenced:
+        ch_raw = pick_channels(
+            inst.ch_names,
+            [ch for ch in inst.ch_names if ch not in list(ref_channels.keys())],
+        )
+        ch_reref = pick_channels(inst.ch_names, list(ref_channels.keys()), ordered=True)
+
+        # Check that the non re-reference channels are untouched:
+        assert_allclose(
+            _data[..., ch_raw, :], reref._data[..., ch_raw, :], 1e-6, atol=1e-15
+        )
+
+        # Compute the reference data:
+        ref_data = []
+        for val in ref_channels.values():
+            if isinstance(val, str):
+                val = [val]  # pick_channels expects a list
+            ref_data.append(
+                _data[..., pick_channels(inst.ch_names, val, ordered=True), :].mean(
+                    -2, keepdims=True
+                )
+            )
+        if inst_type == "epochs":
+            ref_data = np.concatenate(ref_data, axis=1)
+        else:
+            ref_data = np.squeeze(np.array(ref_data))
+        assert_allclose(
+            _data[..., ch_reref, :],
+            reref._data[..., ch_reref, :] + ref_data,
+            1e-6,
+            atol=1e-15,
+        )
+
+
 @testing.requires_testing_data
 @pytest.mark.parametrize("inst_type", ("raw", "epochs", "evoked"))
 def test_set_bipolar_reference(inst_type):

mne/utils/docs.py

@@ -3693,13 +3693,20 @@ def _reflow_param_docstring(docstring, has_first_line=True, width=75):
 """
 
 docdict["ref_channels_set_eeg_reference"] = """
-ref_channels : list of str | str
+ref_channels : list of str | str | dict
     Can be:
 
-    - The name(s) of the channel(s) used to construct the reference.
+    - The name(s) of the channel(s) used to construct the reference for
+      every channel of ``ch_type``.
     - ``'average'`` to apply an average reference (default)
     - ``'REST'`` to use the Reference Electrode Standardization Technique
       infinity reference :footcite:`Yao2001`.
+    - A dictionary mapping names of data channels to (lists of) names of
+      reference channels. For example, {'A1': 'A3'} would replace the
+      data in channel 'A1' with the difference between 'A1' and 'A3'. To take
+      the average of multiple channels as reference, supply a list of channel
+      names as the dictionary value, e.g. {'A1': ['A2', 'A3']} would replace
+      channel A1 with ``A1 - mean(A2, A3)``.
     - An empty list, in which case MNE will not attempt any re-referencing of
       the data
 """
@@ -4027,6 +4034,16 @@ def _reflow_param_docstring(docstring, has_first_line=True, width=75):
     The given EEG electrodes are referenced to a point at infinity using the
     lead fields in ``forward``, which helps standardize the signals.
 
+- Different references for different channels
+    Set ``ref_channels`` to a dictionary mapping source channel names (str)
+    to the reference channel names (str or list of str). Unlike the other
+    approaches where the same reference is applied globally, you can set
+    different references for different channels with this method. For example,
+    to re-reference channel 'A1' to 'A2' and 'B1' to the average of 'B2' and
+    'B3', set ``ref_channels={'A1': 'A2', 'B1': ['B2', 'B3']}``. Warnings are
+    issued when a bad channel is used as a reference or when a mapping involves
+    channels of different types.
+
 1. If a reference is requested that is not the average reference, this
    function removes any pre-existing average reference projections.