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,63 @@ 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")
+
+    def check_value_str(inst, value, key_ch_type):
+        # Check that value is in ch_names
+        assert (
+            value in inst.ch_names
+        ), f"Channel {value} in ref_channels is not in the instance"
+        # If value is a bad channel, issue a warning
+        if value in inst.info["bads"]:
+            msg = f"Channel {value} in ref_channels is marked as bad!"
+            _on_missing("warn", msg)
+        # If key and value are of different channel types, issue a warning
+        value_pick = pick_channels(inst.ch_names, [value], ordered=True)
+        value_ch_type = inst.get_channel_types(picks=value_pick)[0]
+        if key_ch_type != value_ch_type:
+            msg = (
+                f"Channel {key} is of type {DEFAULTS['titles'][key_ch_type]}, "
+                f"but reference channel {value} is of type "
+                f"{DEFAULTS['titles'][value_ch_type]}."
+            )
+            _on_missing("warn", msg)
+
+    for key, value in ref_dict.items():
+        # Check that keys are strings
+        assert isinstance(key, str), (
+            "Keys in dict-type ref_channels must be strings. You provided "
+            f"{type(key)}."
+        )
+        # Check that keys are in ch_names
+        assert (
+            key in inst.ch_names
+        ), f"Channel {key} in ref_channels is not in the instance"
+        key_pick = pick_channels(inst.ch_names, [key], ordered=True)
+        key_ch_type = inst.get_channel_types(picks=key_pick)[0]
+        # Check values
+        if isinstance(value, str):
+            check_value_str(inst, value, key_ch_type)
+        elif isinstance(value, list):
+            for val in value:
+                # Check that values are strings
+                assert isinstance(val, str), (
+                    "Values in dict-type ref_channels must be strings or "
+                    "lists of strings. You provided a list of "
+                    f"{type(val)}"
+                )
+                check_value_str(inst, val, key_ch_type)
+        else:
+            raise ValueError(
+                "Values in dict-type ref_channels must be strings or "
+                "lists of strings. You provided "
+                f"{type(value)}"
+            )
+
+
 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 +185,26 @@ 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_before_dict_reference(inst, ref_dict)
+
+    orig_data = inst.copy()._data
+    data = inst._data
+    if len(ref_dict) > 0:
+        for key, value in ref_dict.items():
+            if isinstance(value, str):
+                value = [value]  # pick_channels expects a list
+            ref_from = pick_channels(inst.ch_names, value, ordered=True)
+            ref_to = pick_channels(inst.ch_names, [key], ordered=True)
+            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, data
+
+
 @fill_doc
 def add_reference_channels(inst, ref_channels, copy=True):
     """Add reference channels to data that consists of all zeros.
@@ -331,6 +408,10 @@ def set_eeg_reference(
 
     _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,80 @@ def test_set_eeg_reference_rest():
     assert 0.995 < exp_var <= 1
 
 
+@testing.requires_testing_data
+@pytest.mark.parametrize(
+    "ref_channels, expectation",
+    [
+        (
+            {2: "EEG 001"},
+            pytest.raises(
+                AssertionError, match=f"Keys in dict-type.*You provided {int}"
+            ),
+        ),
+        (
+            {"EEG 001": (1, 2)},
+            pytest.raises(
+                ValueError, match=f"Values in dict-type.*You provided {type((1,2))}"
+            ),
+        ),
+        (
+            {"EEG 001": [1, 2]},
+            pytest.raises(
+                AssertionError,
+                match="Values in dict-type.*You provided a list of <class 'int'>",
+            ),
+        ),
+        (
+            {"EEG 999": "EEG 001"},
+            pytest.raises(
+                AssertionError,
+                match="Channel EEG 999 in ref_channels is not in the instance",
+            ),
+        ),
+        (
+            {"EEG 001": "EEG 999"},
+            pytest.raises(
+                AssertionError,
+                match="Channel EEG 999 in ref_channels is not in the instance",
+            ),
+        ),
+        (
+            {"EEG 001": "EEG 057"},
+            pytest.warns(
+                RuntimeWarning,
+                match="Channel EEG 057 in ref_channels is marked as bad!",
+            ),
+        ),
+        (
+            {"EEG 001": "STI 001"},
+            pytest.warns(
+                RuntimeWarning,
+                match="Channel EEG 001 is of type EEG, but reference channel STI 001 is of type Stimulus.",
+            ),
+        ),
+        (
+            {"EEG 001": "EEG 002", "EEG 002": "EEG 002", "EEG 003": "EEG 005"},
+            nullcontext(),
+        ),
+    ],
+)
+def test_set_eeg_reference_dict(ref_channels, expectation):
+    """Test setting dict-based reference."""
+    raw = read_raw_fif(fif_fname).crop(0, 1).pick(picks=["eeg", "stim"])
+    with pytest.raises(
+        RuntimeError,
+        match="By default, MNE does not load data.*Applying a reference requires.*",
+    ):
+        raw.set_eeg_reference(ref_channels=ref_channels)
+    raw.load_data()
+    raw.info["bads"] = ["EEG 057"]
+    with expectation:
+        raw.set_eeg_reference(ref_channels=ref_channels)
+
+    # Data testing goes below here @Alex
+    # if expectation == nullcontext():
+
+
 @testing.requires_testing_data
 @pytest.mark.parametrize("inst_type", ("raw", "epochs", "evoked"))
 def test_set_bipolar_reference(inst_type):

mne/utils/docs.py

@@ -3665,13 +3665,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
 """
@@ -3999,6 +4006,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.