ENH: Add Forward-based projection reconstruction

Author: BabaSanfour

doc/changes/dev/newfeature.rst

@@ -0,0 +1 @@
+Projection reconstruction can now use an explicit :class:`~mne.Forward` model and integer spatial rank, by `Hamza Abdelhedi`_.

mne/_fiff/proj.py

@@ -13,6 +13,7 @@
 from ..fixes import _safe_svd
 from ..utils import (
     _check_option,
+    _ensure_int,
     _validate_type,
     fill_doc,
     logger,
@@ -561,7 +562,15 @@ def plot_projs_topomap(
         )
         return fig
 
-    def reconstruct_proj(self, *, projs=None, mode="accurate", origin="auto"):
+    def reconstruct_proj(
+        self,
+        *,
+        projs=None,
+        mode="accurate",
+        origin="auto",
+        forward=None,
+        rank=None,
+    ):
         """Apply SSP projectors and reconstruct the resulting signal in sensor space.
 
         Operates in place.
@@ -574,18 +583,46 @@ def reconstruct_proj(self, *, projs=None, mode="accurate", origin="auto"):
             ``None``, all projectors attached to the instance are used.
         mode : str
             Either ``'accurate'`` or ``'fast'``, determines the quality of the
-            Legendre polynomial expansion used for reconstruction.
+            Legendre polynomial expansion used for geometry-based
+            reconstruction. Ignored when ``forward`` is provided.
         origin : array-like, shape (3,) | str
             Origin of the sphere in the head coordinate frame and in meters.
             Can be ``'auto'`` (default), which means a head-digitization-based
-            origin fit.
+            origin fit. Used for geometry-based reconstruction and ignored when
+            ``forward`` is provided.
+        forward : instance of Forward | None
+            Forward model used to construct the reconstruction field mapping.
+            If ``None`` (default), use the geometry-based field mapping model.
+            If provided, ``rank`` must also be specified.
+        rank : int | None
+            Number of spatial modes to retain when reconstructing with
+            ``forward``. This is a sensor-space reconstruction rank, not a
+            number of sources or dipoles. Must be provided together with
+            ``forward``.
+
+        Notes
+        -----
+        When ``forward`` is provided, the reconstruction uses the sensor-space
+        field covariance formed from the Forward gain matrix. ``rank`` specifies
+        the number of sensor-space modes retained in the reconstruction.
 
         Returns
         -------
         self : same type as the input data
             The modified instance.
         """
-        from ..forward import _map_meg_or_eeg_channels
+        from ..forward import Forward, _map_meg_or_eeg_channels
+
+        if forward is None:
+            if rank is not None:
+                raise ValueError("rank can only be used when forward is provided")
+        else:
+            _validate_type(forward, Forward, "forward")
+            if rank is None:
+                raise ValueError("rank must be provided when forward is provided")
+            rank = _ensure_int(rank, "rank")
+            if rank <= 0:
+                raise ValueError(f"rank must be positive, got {rank}")
 
         if projs is None:
             if len(self.info["projs"]) == 0:
@@ -626,7 +663,12 @@ def reconstruct_proj(self, *, projs=None, mode="accurate", origin="auto"):
                         make_eeg_average_ref_proj(info_to, verbose=False)
                     ]
             mapping = _map_meg_or_eeg_channels(
-                info_from, info_to, mode=mode, origin=origin
+                info_from,
+                info_to,
+                mode=mode,
+                origin=origin,
+                forward=forward,
+                rank=rank,
             )
             self.data[..., picks, :] = np.matmul(mapping, self.data[..., picks, :])
         return self

mne/forward/_field_interpolation.py

@@ -12,7 +12,7 @@
 
 from .._fiff.constants import FIFF
 from .._fiff.meas_info import _simplify_info
-from .._fiff.pick import pick_info, pick_types
+from .._fiff.pick import pick_channels_forward, pick_info, pick_types
 from .._fiff.proj import _has_eeg_average_ref_proj, make_projector
 from ..bem import _check_origin
 from ..cov import make_ad_hoc_cov
@@ -21,22 +21,34 @@
 from ..fixes import _safe_svd
 from ..surface import get_head_surf, get_meg_helmet_surf
 from ..transforms import _find_trans, transform_surface_to
-from ..utils import _check_fname, _check_option, _pl, _reg_pinv, logger, verbose
+from ..utils import (
+    _check_fname,
+    _check_option,
+    _pl,
+    _reg_pinv,
+    logger,
+    verbose,
+)
 from ._lead_dots import _do_cross_dots, _do_self_dots, _do_surface_dots, _get_legen_fun
 from ._make_forward import _create_eeg_els, _create_meg_coils, _read_coil_defs
 
 
 def _setup_dots(mode, info, coils, ch_type):
     """Set up dot products."""
     int_rad = 0.06
-    noise = make_ad_hoc_cov(info, dict(mag=20e-15, grad=5e-13, eeg=1e-6))
+    noise = _make_field_mapping_noise(info)
     # "fast" uses a coarser (n_coeff=50) Legendre series than "accurate" (n_coeff=100)
     n_coeff = 50 if mode == "fast" else 100
     leg_fun, n_fact = _get_legen_fun(ch_type, False, n_coeff)
     return int_rad, noise, leg_fun, n_fact
 
 
-def _compute_mapping_matrix(fmd, info):
+def _make_field_mapping_noise(info):
+    """Create the ad hoc noise covariance used for field mapping."""
+    return make_ad_hoc_cov(info, dict(mag=20e-15, grad=5e-13, eeg=1e-6))
+
+
+def _compute_mapping_matrix(fmd, info, *, rank=None):
     """Do the hairy computations."""
     logger.info("    Preparing the mapping matrix...")
     # assemble a projector and apply it to the data
@@ -54,7 +66,9 @@ def _compute_mapping_matrix(fmd, info):
 
     # SVD is numerically better than the eigenvalue composition even if
     # mat is supposed to be symmetric and positive definite
-    if fmd.get("pinv_method", "tsvd") == "tsvd":
+    if rank is not None:
+        inv, _, _ = _reg_pinv(whitened_dots, reg=0, rank=rank)
+    elif fmd.get("pinv_method", "tsvd") == "tsvd":
         inv, fmd["nest"] = _pinv_trunc(whitened_dots, fmd["miss"])
     else:
         assert fmd["pinv_method"] == "tikhonov", fmd["pinv_method"]
@@ -110,7 +124,9 @@ def _pinv_tikhonov(x, reg):
     return inv, n
 
 
-def _map_meg_or_eeg_channels(info_from, info_to, mode, *, origin, miss=None):
+def _map_meg_or_eeg_channels(
+    info_from, info_to, mode, *, origin, miss=None, forward=None, rank=None
+):
     """Find mapping from one set of channels to another.
 
     Parameters
@@ -133,8 +149,6 @@ def _map_meg_or_eeg_channels(info_from, info_to, mode, *, origin, miss=None):
     mapping : array, shape (n_to, n_from)
         A mapping matrix.
     """
-    assert origin is not None  # should be assured elsewhere
-
     # no need to apply trans because both from and to coils are in device
     # coordinates
     info_kinds = set(ch["kind"] for ch in info_to["chs"])
@@ -150,6 +164,26 @@ def _map_meg_or_eeg_channels(info_from, info_to, mode, *, origin, miss=None):
     )
     kind = "eeg" if info_kinds[0] == FIFF.FIFFV_EEG_CH else "meg"
 
+    if forward is not None:
+        forward = pick_channels_forward(
+            forward, include=info_from["ch_names"], ordered=True
+        )
+        assert forward["sol"]["row_names"] == info_from["ch_names"]
+        lead_field = forward["sol"]["data"]
+        # Form the sensor-space field covariance from the Forward gain matrix.
+        # As with any Gram representation, very weak modes can be numerically unstable.
+        dots = lead_field @ lead_field.T
+        fmd = dict(
+            kind=kind,
+            ch_names=info_from["ch_names"],
+            noise=_make_field_mapping_noise(info_from),
+            self_dots=dots,
+            surface_dots=dots,
+        )
+        return _compute_mapping_matrix(fmd, info_from, rank=rank)
+
+    assert origin is not None  # should be assured elsewhere
+
     #
     # Step 1. Prepare the coil definitions
     #

mne/tests/test_proj.py

@@ -19,10 +19,14 @@
     compute_raw_covariance,
     convert_forward_solution,
     create_info,
+    make_forward_solution,
+    make_sphere_model,
+    pick_channels_forward,
     pick_types,
     read_events,
     read_forward_solution,
     read_source_estimate,
+    read_source_spaces,
     sensitivity_map,
 )
 from mne._fiff.proj import (
@@ -209,6 +213,101 @@ def test_reconstruct_proj(raw_orig, events):
         )
 
 
+@pytest.fixture(scope="module")
+def eeg_forward():
+    """Create a small local EEG forward for projection reconstruction tests."""
+    raw = read_raw_fif(raw_fname, preload=False, verbose=False).pick(picks="eeg")
+    raw.pick(raw.ch_names[:8])
+    src = read_source_spaces(base_dir / "small-src.fif.gz", verbose=False)
+    sphere = make_sphere_model(verbose=False)
+    return make_forward_solution(
+        raw.info,
+        trans=None,
+        src=src,
+        bem=sphere,
+        meg=False,
+        eeg=True,
+        mindist=0.0,
+        verbose=False,
+    )
+
+
+def _direct_forward_reconstruction(evoked, forward, rank):
+    """Compute the independent direct lead-field reconstruction."""
+    projector = make_projector(evoked.info["projs"], evoked.ch_names)[0]
+    forward = pick_channels_forward(forward, include=evoked.ch_names, ordered=True)
+    lead_field = forward["sol"]["data"]
+    u, s, vh = np.linalg.svd(projector @ lead_field, full_matrices=False)
+    mapping = lead_field @ (vh[:rank].T / s[:rank]) @ u[:, :rank].T
+    if _has_eeg_average_ref_proj(evoked.info):
+        mapping -= mapping.mean(axis=0)
+    return mapping @ (projector @ evoked.data)
+
+
+def test_reconstruct_proj_forward(raw_orig, eeg_forward):
+    """Test Forward reconstruction and channel handling."""
+    raw = raw_orig.copy().pick(picks="eeg")
+    raw.pick(raw.ch_names[:8])
+    evoked = EvokedArray(raw.get_data()[:, :10], raw.info, tmin=0.0)
+    evoked.add_proj(
+        _make_test_proj(
+            evoked.ch_names,
+            np.arange(1.0, len(evoked.ch_names) + 1.0),
+            "Forward reconstruction",
+        ),
+        verbose=False,
+    )
+    rank = 3
+    for average_ref in (False, True):
+        this_evoked = evoked.copy()
+        if average_ref:
+            this_evoked.set_eeg_reference(projection=True)
+        expected = _direct_forward_reconstruction(this_evoked, eeg_forward, rank)
+        got = this_evoked.copy().reconstruct_proj(forward=eeg_forward, rank=rank).data
+        assert_allclose(got, expected, rtol=1e-10, atol=1e-12)
+
+    reordered = pick_channels_forward(
+        eeg_forward, include=eeg_forward.ch_names[::-1], ordered=True
+    )
+    got = evoked.copy().reconstruct_proj(forward=eeg_forward, rank=rank).data
+    got_reordered = evoked.copy().reconstruct_proj(forward=reordered, rank=rank).data
+    assert_allclose(got, got_reordered, rtol=1e-10, atol=1e-12)
+
+    evoked_bad = evoked.copy()
+    evoked_bad.info["bads"] = [evoked_bad.ch_names[0]]
+    got_bad = evoked_bad.reconstruct_proj(forward=eeg_forward, rank=rank).data
+    assert_allclose(got_bad[0], evoked_bad.data[0])
+
+
+def test_reconstruct_proj_forward_validation(eeg_forward):
+    """Test validation of the explicit Forward reconstruction arguments."""
+    info = create_info(eeg_forward["info"]["ch_names"], 100.0, "eeg")
+    evoked = EvokedArray(np.zeros((len(info["ch_names"]), 1)), info, tmin=0.0)
+    evoked.add_proj(
+        _make_test_proj(
+            evoked.ch_names,
+            np.ones(len(evoked.ch_names)),
+            "Forward reconstruction",
+        ),
+        verbose=False,
+    )
+    with pytest.raises(ValueError, match="rank can only be used"):
+        evoked.copy().reconstruct_proj(rank=1)
+    with pytest.raises(ValueError, match="rank must be provided"):
+        evoked.copy().reconstruct_proj(forward=eeg_forward)
+    with pytest.raises(TypeError, match="forward must be an instance of Forward"):
+        evoked.copy().reconstruct_proj(forward=[], rank=1)
+    for rank in (0, -1):
+        with pytest.raises(ValueError, match="rank must be positive"):
+            evoked.copy().reconstruct_proj(forward=eeg_forward, rank=rank)
+    for rank in (1.5, True):
+        with pytest.raises(TypeError, match="rank must be an int"):
+            evoked.copy().reconstruct_proj(forward=eeg_forward, rank=rank)
+    rank = len(evoked.ch_names) + 1
+    with pytest.raises(ValueError, match="Invalid value for the rank parameter"):
+        evoked.copy().reconstruct_proj(forward=eeg_forward, rank=rank)
+
+
 @pytest.mark.parametrize("kind", ["raw", "epochs", "evoked"])
 def test_apply_proj_default(kind):
     """Test that ``projs=None`` preserves legacy behavior."""