ENH: Add Forward-based projection reconstruction (#14235)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

Author: BabaSanfour

doc/changes/dev/14235.newfeature.rst

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

mne/_fiff/proj.py

@@ -13,6 +13,7 @@
 from ..fixes import _safe_svd
 from ..utils import (
     _check_option,
+    _check_rank,
     _validate_type,
     fill_doc,
     logger,
@@ -561,7 +562,17 @@ def plot_projs_topomap(
         )
         return fig
 
-    def reconstruct_proj(self, *, projs=None, mode="accurate", origin="auto"):
+    @verbose
+    def reconstruct_proj(
+        self,
+        *,
+        projs=None,
+        mode="accurate",
+        origin="auto",
+        forward=None,
+        rank=None,
+        verbose=None,
+    ):
         """Apply SSP projectors and reconstruct the resulting signal in sensor space.
 
         Operates in place.
@@ -574,18 +585,48 @@ 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.
+        %(rank)s
+
+            Only used when ``forward`` is provided, where the default ``None``
+            estimates the rank of the projected field covariance. ``'full'`` is
+            not supported, as that covariance is rank-deficient after
+            projection. Without ``forward``, ``rank`` must be ``None``, and the
+            geometry-based field mapping uses its own internal truncation
+            rather than a rank estimated from the data.
+        %(verbose)s
 
         Returns
         -------
         self : same type as the input data
             The modified instance.
+
+        Notes
+        -----
+        When ``forward`` is provided, ``rank`` controls the sensor-space rank
+        of the projected Forward field covariance.
         """
-        from ..forward import _map_meg_or_eeg_channels
+        from ..forward import Forward, _map_meg_or_eeg_channels
+
+        if forward is not None:
+            _validate_type(forward, Forward, "forward")
+        rank = _check_rank(rank)
+        if forward is None and rank is not None:
+            raise ValueError("rank can only be used when forward is provided")
+        if forward is not None and rank == "full":
+            raise ValueError(
+                "rank='full' is incompatible with Forward reconstruction "
+                "after projection; use rank='info', None, or an explicit rank dict"
+            )
 
         if projs is None:
             if len(self.info["projs"]) == 0:
@@ -626,7 +667,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,31 +12,43 @@
 
 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
+from ..cov import Covariance, make_ad_hoc_cov
 from ..epochs import BaseEpochs, EpochsArray
 from ..evoked import Evoked, EvokedArray
-from ..fixes import _safe_svd
+from ..rank import _compute_rank_int
 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
@@ -52,12 +64,26 @@ def _compute_mapping_matrix(fmd, info):
     whitener = np.diag(1.0 / np.sqrt(noise_cov["data"].ravel()))
     whitened_dots = np.dot(whitener.T, np.dot(proj_dots, whitener))
 
-    # SVD is numerically better than the eigenvalue composition even if
-    # mat is supposed to be symmetric and positive definite
+    # whitened_dots is symmetric and positive semi-definite, so _reg_pinv (which
+    # requires square Hermitian input) can do the truncated pseudoinversion
     if fmd.get("pinv_method", "tsvd") == "tsvd":
-        inv, fmd["nest"] = _pinv_trunc(whitened_dots, fmd["miss"])
+        n = len(whitened_dots)
+        if rank is None:
+            # truncate at most "miss" fraction of the singular value energy
+            s = np.linalg.svd(whitened_dots, compute_uv=False, hermitian=True)
+            varexp = np.cumsum(s)
+            varexp /= varexp[-1]
+            rank = np.where(varexp >= 1.0 - fmd["miss"])[0][0] + 1
+            logger.info(
+                f"    Truncating at {rank}/{n} components to omit less than "
+                f"{fmd['miss']:g} ({1.0 - varexp[rank - 1]:0.2g})"
+            )
+        else:
+            logger.info(f"    Truncating at {rank}/{n} components")
+        inv, _, fmd["nest"] = _reg_pinv(whitened_dots, reg=0, rank=rank)
     else:
         assert fmd["pinv_method"] == "tikhonov", fmd["pinv_method"]
+        assert rank is None, rank  # only the tsvd path supports an explicit rank
         inv, fmd["nest"] = _pinv_tikhonov(whitened_dots, fmd["miss"])
 
     # Sandwich with the whitener
@@ -81,26 +107,6 @@ def _compute_mapping_matrix(fmd, info):
     return mapping_mat
 
 
-def _pinv_trunc(x, miss):
-    """Compute pseudoinverse, truncating at most "miss" fraction of varexp."""
-    u, s, v = _safe_svd(x, full_matrices=False)
-
-    # Eigenvalue truncation
-    varexp = np.cumsum(s)
-    varexp /= varexp[-1]
-    n = np.where(varexp >= (1.0 - miss))[0][0] + 1
-    logger.info(
-        "    Truncating at %d/%d components to omit less than %g (%0.2g)",
-        n,
-        len(s),
-        miss,
-        1.0 - varexp[n - 1],
-    )
-    s = 1.0 / s[:n]
-    inv = ((u[:, :n] * s) @ v[:n]).T
-    return inv, n
-
-
 def _pinv_tikhonov(x, reg):
     # _reg_pinv requires square Hermitian, which we have here
     inv, _, n = _reg_pinv(x, reg=reg, rank=None)
@@ -110,7 +116,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
@@ -127,14 +135,18 @@ def _map_meg_or_eeg_channels(info_from, info_to, mode, *, origin, miss=None):
         Origin of the sphere in the head coordinate frame and in meters.
         Can be ``'auto'``, which means a head-digitization-based origin
         fit.
+    forward : instance of Forward | None
+        Forward model used instead of geometry-based field interpolation.
+    rank : None | 'info' | dict
+        Rank specification for Forward-based reconstruction, where ``None``
+        estimates it from the projected field covariance. Must be ``None`` for
+        the geometry-based path, which uses its own truncation instead.
 
     Returns
     -------
     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 +162,34 @@ 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
+        field_cov = Covariance(
+            dots,
+            info_from["ch_names"],
+            info_from["bads"],
+            info_from["projs"],
+            nfree=1,
+        )
+        rank_int = _compute_rank_int(field_cov, rank=rank, info=info_from)
+        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_int)
+
+    assert origin is not None  # should be assured elsewhere
+
     #
     # Step 1. Prepare the coil definitions
     #