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

@@ -155,6 +155,46 @@ 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."""
+    # Check that ref_dict channel exist in instance:
+    ref_dict_channels = list(
+        set(
+            list(ref_dict.keys())
+            + [item for sublist in ref_dict.values() for item in sublist]
+        )
+    )
+    assert all([ref_ch in inst.ch_names for ref_ch in ref_dict_channels]), (
+        "The custom referencing dictionary "
+        "contains channels which are not in the "
+        "instance!"
+    )
+    # Raise a warning if one of the reference channel is bad:
+    if any([ch in ref_dict_channels for ch in inst.info["bads"]]):
+        msg = "Channels in the reference scheme are marked as bad!"
+        _on_missing("warns", msg)
+
+    # Copy the data instance to re-reference:
+    ref_to_data = inst.copy()._data
+    if len(ref_dict) > 0:
+        # Loop through each channel to re-reference:
+        for ch in ref_dict.keys():
+            assert len(ref_dict[ch]) > 0, f"No channel to re-reference ch-{ch}"
+            # Get indices of the channels to use as reference
+            ref_from = pick_channels(inst.ch_names, ref_dict[ch], ordered=True)
+            # Get indice of channel to re.reference:
+            ref_to = pick_channels(inst.ch_names, ch, ordered=True)
+            # Compute the reference data:
+            ref_data = inst._data[..., ref_from, :].mean(-2, keepdims=True)
+            # Subtract the reference data to the channel to re-reference:
+            ref_to_data[..., ref_to, :] -= ref_data
+    # Add the data back to the instance:
+    inst._data = ref_to_data
+    # Set that custom reference was applied:
+    inst.info["custom_ref_applied"] = FIFF.FIFFV_MNE_CUSTOM_REF_ON
+    return inst
+
+
 @fill_doc
 def add_reference_channels(inst, ref_channels, copy=True):
     """Add reference channels to data that consists of all zeros.
@@ -430,7 +470,10 @@ def set_eeg_reference(
             "reference."
         )
 
-    return _apply_reference(inst, ref_channels, ch_sel, forward, ch_type=ch_type)
+    if isinstance(ref_channels, dict):
+        return _apply_dict_reference(inst, ref_channels, ch_type=ch_type)
+    else:
+        return _apply_reference(inst, ref_channels, ch_sel, forward, ch_type=ch_type)
 
 
 def _get_ch_type(inst, ch_type):

mne/utils/docs.py

@@ -3660,13 +3660,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 channels to be referenced to (a list of)
+      names of channels to use as reference. This is the most flexible
+      re-referencing approaching. 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']}.
     - An empty list, in which case MNE will not attempt any re-referencing of
       the data
 """