ENH: Allow selecting projectors for reconstruction

Author: BabaSanfour

mne/_fiff/proj.py

@@ -515,19 +515,61 @@ def plot_projs_topomap(
         )
         return fig
 
-    def _reconstruct_proj(self, mode="accurate", origin="auto"):
+    def _reconstruct_proj(self, projs=None, mode="accurate", origin="auto"):
         from ..forward import _map_meg_or_eeg_channels
 
-        if len(self.info["projs"]) == 0:
-            return self
-        self.apply_proj()
+        if projs is None:
+            if len(self.info["projs"]) == 0:
+                return self
+            self.apply_proj()
+            info = self.info
+        else:
+            if isinstance(projs, Projection):
+                projs = [projs]
+            selected_projs = _check_projs(projs)
+            if len(selected_projs) == 0:
+                return self
+            projector, nproj, _ = make_projector(
+                selected_projs, self.info["ch_names"], self.info["bads"]
+            )
+            if nproj == 0:
+                return self
+            self.data[:] = np.matmul(projector, self.data)
+            active_eeg_ref_projs = [
+                proj
+                for proj in self.info["projs"]
+                if proj["active"] and _is_eeg_average_ref_proj(proj)
+            ]
+            attached_selected = [
+                attached
+                for attached in self.info["projs"]
+                if any(
+                    _proj_equal(attached, selected, check_active=False)
+                    for selected in selected_projs
+                )
+            ]
+            activate_proj(attached_selected, copy=False, verbose=False)
+            mapping_projs = _check_projs([*selected_projs, *active_eeg_ref_projs])
+            mapping_projs = _uniquify_projs(
+                mapping_projs, check_active=False, sort=False
+            )
+            mapping_projs = activate_proj(mapping_projs, copy=False, verbose=False)
+            info = self.info.copy()
+            with info._unlock():
+                info["projs"] = mapping_projs
         for kind in ("meg", "eeg"):
             kwargs = dict(meg=False)
             kwargs[kind] = True
-            picks = pick_types(self.info, **kwargs)
+            picks = pick_types(info, **kwargs)
             if len(picks) == 0:
                 continue
-            info_from = pick_info(self.info, picks)
+            info_from = pick_info(info, picks)
+            if projs is not None:
+                _, nproj, _ = make_projector(
+                    selected_projs, info_from["ch_names"], info_from["bads"]
+                )
+                if nproj == 0:
+                    continue
             info_to = info_from.copy()
             with info_to._unlock():
                 info_to["projs"] = []
@@ -1082,9 +1124,7 @@ def _has_eeg_average_ref_proj(
         return False
     found_names = list()
     for proj in projs:
-        if proj["kind"] == FIFF.FIFFV_PROJ_ITEM_EEG_AVREF or re.match(
-            "^Average .* reference$", proj["desc"]
-        ):
+        if _is_eeg_average_ref_proj(proj):
             if not check_active or proj["active"]:
                 found_names.extend(proj["data"]["col_names"])
     # If some are missing we have a problem (keep order for the message,
@@ -1097,6 +1137,13 @@ def _has_eeg_average_ref_proj(
     return True
 
 
+def _is_eeg_average_ref_proj(proj):
+    """Check whether a projector is an average EEG reference."""
+    return proj["kind"] == FIFF.FIFFV_PROJ_ITEM_EEG_AVREF or bool(
+        re.match("^Average .* reference$", proj["desc"])
+    )
+
+
 def _needs_eeg_average_ref_proj(info):
     """Determine if the EEG needs an average EEG reference.
 

mne/tests/test_proj.py

@@ -261,6 +261,113 @@ def test_compute_proj_epochs(tmp_path):
         write_proj(fname, ["foo"], overwrite=True)
 
 
+@pytest.mark.parametrize(
+    ("ch_type", "picks", "proj_kwargs"),
+    [
+        ("meg", [0, 1, 2, 3, 4, 6, 7], dict(n_grad=2, n_mag=2, n_eeg=0)),
+        ("eeg", np.arange(340, 360), dict(n_grad=0, n_mag=0, n_eeg=2)),
+    ],
+)
+def test_reconstruct_proj_selection(ch_type, picks, proj_kwargs):
+    """Test selecting projectors for reconstruction."""
+    raw = read_raw_fif(raw_fname)
+    raw.add_proj([], remove_existing=True)
+    events = read_events(event_fname)
+    evoked = Epochs(
+        raw,
+        events[:5],
+        1,
+        -0.1,
+        0.1,
+        picks=picks,
+        decim=10,
+        preload=True,
+        verbose="error",
+    ).average()
+    if ch_type == "eeg":
+        evoked.set_eeg_reference(projection=True)
+        evoked_car_active = evoked.copy().apply_proj()
+    evoked_proj = evoked.copy().apply_proj().crop(None, 0)
+    projs = compute_proj_evoked(evoked_proj, **proj_kwargs)
+    assert len(projs) >= 2
+    evoked.add_proj(projs)
+    original_projs = cp.deepcopy(evoked.info["projs"])
+
+    # The default is unchanged, including activating all attached projectors.
+    default = evoked.copy()._reconstruct_proj()
+    default_none = evoked.copy()._reconstruct_proj(projs=None)
+    assert_allclose(default_none.data, default.data)
+    assert all(proj["active"] for proj in default.info["projs"])
+
+    # A scalar projector and a list of projectors use exactly the requested
+    # artifact directions, ignoring unrelated projectors attached to Info.
+    selected = list(projs[:2])
+    reconstructed = []
+    for proj in selected:
+        actual = evoked.copy()._reconstruct_proj(projs=proj)
+        expected = evoked.copy().del_proj().add_proj(proj)._reconstruct_proj()
+        assert_allclose(actual.data, expected.data)
+        assert [p["active"] for p in actual.info["projs"]] == [
+            p["desc"] == proj["desc"] for p in original_projs
+        ]
+        reconstructed.append(actual.data)
+    assert np.linalg.norm(reconstructed[0] - reconstructed[1]) > 1e-3 * np.linalg.norm(
+        reconstructed[0]
+    )
+
+    actual = evoked.copy()._reconstruct_proj(projs=selected)
+    expected = evoked.copy().del_proj().add_proj(selected)._reconstruct_proj()
+    assert_allclose(actual.data, expected.data)
+    selected_descs = {proj["desc"] for proj in selected}
+    assert [p["active"] for p in actual.info["projs"]] == [
+        p["desc"] in selected_descs for p in original_projs
+    ]
+    assert np.linalg.norm(actual.data - reconstructed[0]) > 1e-3 * np.linalg.norm(
+        reconstructed[0]
+    )
+
+    if ch_type == "eeg":
+        # An inactive average-reference projector is not part of the explicit
+        # selection. An already-active one remains part of the mapping state.
+        expected = evoked_car_active.copy().add_proj(projs[0])._reconstruct_proj()
+        evoked_car_active.add_proj(projs)
+        actual = evoked_car_active._reconstruct_proj(projs=projs[0])
+        assert_allclose(actual.data, expected.data)
+        assert_allclose(actual.data.mean(axis=0), 0.0, atol=1e-20)
+        assert [p["active"] for p in actual.info["projs"]] == [True, True, False]
+
+
+def test_reconstruct_proj_selection_mixed():
+    """Test that explicit reconstruction leaves unaffected channel types alone."""
+    raw = read_raw_fif(raw_fname)
+    raw.add_proj([], remove_existing=True)
+    events = read_events(event_fname)
+    picks = [0, 1, 2, 3, 4, 6, 7, *range(340, 360)]
+    evoked = Epochs(
+        raw,
+        events[:5],
+        1,
+        -0.1,
+        0.1,
+        picks=picks,
+        decim=10,
+        preload=True,
+        verbose="error",
+    ).average()
+    projs = compute_proj_evoked(
+        evoked.copy().pick("meg").crop(None, 0), n_grad=2, n_mag=2, n_eeg=0
+    )
+    evoked.add_proj(projs)
+    eeg_picks = pick_types(evoked.info, meg=False, eeg=True)
+    meg_picks = pick_types(evoked.info, meg=True, eeg=False)
+    data_before = evoked.data.copy()
+    evoked._reconstruct_proj(projs=projs[0])
+    assert_allclose(evoked.data[eeg_picks], data_before[eeg_picks], rtol=0, atol=0)
+    assert not np.array_equal(evoked.data[meg_picks], data_before[meg_picks])
+    assert evoked.info["projs"][0]["active"]
+    assert not any(proj["active"] for proj in evoked.info["projs"][1:])
+
+
 @pytest.mark.slowtest
 def test_compute_proj_raw(tmp_path):
     """Test SSP computation on raw."""