Create lead_field_processing

calibrain/lead_field_processing

@@ -0,0 +1,394 @@
+from logging import warning
+import os
+import mne
+import numpy as np
+from mne.io.constants import FIFF
+
+from calibrain.utils import get_data_path
+
+
+def _canonicalize_basis_columns(Q: np.ndarray) -> np.ndarray:
+    """
+    Stabilize basis column signs so that the entry with largest absolute value
+    in each column is nonnegative.
+
+    This removes arbitrary SVD sign flips across runs/platforms while leaving
+    the spanned subspace unchanged.
+
+    Parameters
+    ----------
+    Q : array, shape (n, k)
+        Basis matrix with k orthonormal columns, where n is the number of
+        orientations and k is the reduced dimensionality, e.g. 2 for MEG
+        free orientation.
+
+    Returns
+    -------
+    Q_out : array, shape (n, k)
+        Basis matrix with signs stabilized.
+    """
+    if Q.ndim != 2:
+        raise ValueError(f"Expected 2D basis matrix. Got {Q.shape}")
+
+    Q_out = np.array(Q, dtype=float, copy=True)
+
+    for j in range(Q_out.shape[1]):
+        idx = int(np.argmax(np.abs(Q_out[:, j])))
+        if Q_out[idx, j] < 0:
+            Q_out[:, j] *= -1.0
+
+    return Q_out
+
+
+def _extract_cortical_normals(fwd) -> np.ndarray:
+    """
+    Extract cortical-normal information for source simulation.
+
+    For fixed-orientation forward solutions, the cortical-normal direction is
+    already used implicitly by the fixed leadfield. We nevertheless save the
+    normal vectors for traceability.
+
+    For free-orientation forward solutions after
+
+        mne.convert_forward_solution(fwd, surf_ori=True, force_fixed=False)
+
+    the source components are represented in a local surface-oriented basis.
+    In this local basis, the cortical-normal direction is the third component:
+
+        [0, 0, 1].
+
+    Returns
+    -------
+    cortical_normals : array, shape (n_sources, 3)
+        Cortical-normal direction represented in the same source-component
+        coordinate system as the saved leadfield.
+    """
+    n_sources = fwd["nsource"]
+
+    if fwd["source_ori"] == FIFF.FIFFV_MNE_FREE_ORI:
+        cortical_normals = np.zeros((n_sources, 3), dtype=float)
+        cortical_normals[:, 2] = 1.0
+        return cortical_normals
+
+    source_nn = fwd.get("source_nn", None)
+
+    if source_nn is not None and source_nn.shape == (n_sources, 3):
+        cortical_normals = np.asarray(source_nn, dtype=float).copy()
+    else:
+        normals_list = []
+
+        for src in fwd["src"]:
+            if "nn" not in src or "vertno" not in src:
+                raise ValueError(
+                    "Could not extract cortical normals from the forward solution."
+                )
+            normals_list.append(src["nn"][src["vertno"]])
+
+        cortical_normals = np.vstack(normals_list)
+
+    norms = np.linalg.norm(cortical_normals, axis=1, keepdims=True)
+
+    if np.any(norms < 1e-12):
+        raise ValueError("Some cortical normal vectors have near-zero norm.")
+
+    cortical_normals = cortical_normals / norms
+
+    return cortical_normals
+
+
+def extract_leadfield(
+    fwd_datapath,
+    subject: str,
+    coil_name: str,
+    orientation_type: str,
+    save_path: str,
+):
+    """
+    Extract the leadfield from a forward solution file for a given subject and
+    channel type. The leadfield is saved as a .npz file.
+
+    For MEG data, channels are filtered by unit type:
+        - magnetometers: T
+        - gradiometers: T/m
+
+    For EEG data, channels are filtered by unit type:
+        - EEG: V
+
+    Parameters
+    ----------
+    fwd_datapath : str or Path
+        Path to the directory containing the forward solution files.
+
+    subject : str
+        Subject identifier, used to load the specific forward solution file.
+
+    coil_name : str
+        Channel type. Options are:
+            - 'mag'
+            - 'grad'
+            - 'eeg'
+
+    orientation_type : str
+        Orientation type of the leadfield. Options are:
+            - 'fixed'
+            - 'free'
+
+    save_path : str or Path
+        Path to save the extracted leadfield .npz file.
+
+    Notes
+    -----
+    Q_basis is:
+        - None for fixed orientation
+        - identity for free EEG, as an API-consistent placeholder in the retained 3D basis
+        - local 3x2 SVD basis for free MEG
+
+    cortical_normals is saved for source simulation:
+        - fixed orientation: normal is implicit in the fixed leadfield
+        - free orientation with surf_ori=True: normal direction is [0, 0, 1]
+          in the local source-component basis
+    """
+    fwd_path = f"{fwd_datapath}/{subject}-fwd.fif"
+    print(f"Loading forward solution from {fwd_path}")
+
+    fwd = mne.read_forward_solution(fwd_path, verbose="error")
+    print(f"Number of sources: {fwd['nsource']}")
+    print(f"Number of channels: {len(fwd['info']['chs'])}")
+
+    if orientation_type == "fixed":
+        if fwd["source_ori"] == FIFF.FIFFV_MNE_FIXED_ORI:
+            print("Forward solution orientation is already fixed")
+        elif fwd["source_ori"] == FIFF.FIFFV_MNE_FREE_ORI:
+            print("Forward solution orientation is free, converting to fixed orientation")
+
+        fwd = mne.convert_forward_solution(
+            fwd,
+            surf_ori=True,
+            force_fixed=True,
+        )
+
+        assert (
+            fwd["source_ori"] == FIFF.FIFFV_MNE_FIXED_ORI
+        ), "Failed to convert to fixed orientation"
+
+    elif orientation_type == "free":
+        if fwd["source_ori"] == FIFF.FIFFV_MNE_FREE_ORI:
+            print("Forward solution orientation is already free")
+        elif fwd["source_ori"] == FIFF.FIFFV_MNE_FIXED_ORI:
+            print("Forward solution orientation is fixed, converting to free orientation")
+
+        fwd = mne.convert_forward_solution(
+            fwd,
+            surf_ori=True,
+            force_fixed=False,
+        )
+
+        assert (
+            fwd["source_ori"] == FIFF.FIFFV_MNE_FREE_ORI
+        ), "Failed to convert to free orientation"
+
+    else:
+        raise ValueError(f"Unknown orientation type: {orientation_type}")
+
+    cortical_normals = _extract_cortical_normals(fwd)
+
+    # Filter channels by type
+    if coil_name == "mag":
+        mag_channels = [
+            ch["ch_name"]
+            for ch in fwd["info"]["chs"]
+            if ch["unit"] == FIFF.FIFF_UNIT_T
+        ]
+        coil_type = FIFF.FIFFV_COIL_VV_MAG_T1
+
+    elif coil_name == "grad":
+        mag_channels = [
+            ch["ch_name"]
+            for ch in fwd["info"]["chs"]
+            if ch["unit"] == FIFF.FIFF_UNIT_T_M
+        ]
+        coil_type = FIFF.FIFFV_COIL_VV_PLANAR_T1
+
+    elif coil_name == "eeg":
+        mag_channels = [
+            ch["ch_name"]
+            for ch in fwd["info"]["chs"]
+            if ch["unit"] == FIFF.FIFF_UNIT_V
+        ]
+        coil_type = FIFF.FIFFV_COIL_EEG
+
+    else:
+        raise ValueError(f"Unknown channel type: {coil_name}")
+
+    fwd = fwd.pick_channels(mag_channels)
+
+    if orientation_type == "fixed":
+        Q_basis = None
+
+    elif orientation_type == "free":
+        L = fwd["sol"]["data"]  # (M, 3N)
+        N = fwd["nsource"]
+        M = L.shape[0]
+
+        L_block = L.reshape(M, N, 3)  # (M, N, 3)
+
+        if coil_name in ["eeg"]:
+            Q_basis = np.repeat(
+                np.eye(3)[None, :, :],
+                fwd["nsource"],
+                axis=0,
+            )  # (N, 3, 3)
+
+            fwd["sol"]["data"] = L_block
+
+        elif coil_name in ["mag", "grad"]:
+            # Free MEG: reduce each local block Mx3 -> Mx2 by local SVD.
+            Q_basis = np.zeros((fwd["nsource"], 3, 2))
+            L_block_final = np.zeros(
+                (len(fwd["info"]["chs"]), fwd["nsource"], 2)
+            )
+
+            for i in range(fwd["nsource"]):
+                Li = L_block[:, i, :]  # (M, 3)
+
+                _, _, Vt = np.linalg.svd(Li, full_matrices=False)
+
+                Qi = _canonicalize_basis_columns(Vt[:2, :].T)  # (3, 2)
+
+                Q_basis[i] = Qi
+                L_block_final[:, i, :] = Li @ Qi  # (M, 2)
+
+            # Kept for debugging/compatibility, but not saved as the main leadfield.
+            L_flat = L_block_final.reshape(M, -1)  # (M, 2N), not used
+
+            fwd["sol"]["data"] = L_block_final
+
+    chs = fwd["info"]["chs"]
+    first = chs[0]
+
+    kind = first.get("kind")
+    unit = first.get("unit")
+    unit_mul = first.get("unit_mul")
+    coil_type = first.get("coil_type")
+
+    # Check for consistency across all channels.
+    for ch in chs[1:]:
+        if (
+            kind is not None
+            and ch.get("kind") is not None
+            and ch["kind"] != kind
+        ):
+            warning("Info contains mixed channel kinds; using the first one (%s).", kind)
+            break
+
+        if (
+            unit is not None
+            and ch.get("unit") is not None
+            and ch["unit"] != unit
+        ):
+            warning("Info contains mixed channel units; using the first one (%s).", unit)
+            break
+
+        if (
+            unit_mul is not None
+            and ch.get("unit_mul") is not None
+            and ch["unit_mul"] != unit_mul
+        ):
+            warning(
+                "Info contains mixed channel unit multipliers; using the first one (%s).",
+                unit_mul,
+            )
+            break
+
+        if (
+            coil_type is not None
+            and ch.get("coil_type") is not None
+            and ch["coil_type"] != coil_type
+        ):
+            warning("Info contains mixed coil types; using the first one (%s).", coil_type)
+            break
+
+    # Update working directory info to fwd['info']
+    with fwd["info"]._unlock():
+        fwd["info"]["working_dir"] = str(fwd_datapath)
+
+    print(f"Number of channels after picking: {len(fwd['info']['chs'])}")
+
+    if fwd["source_ori"] == FIFF.FIFFV_MNE_FIXED_ORI:
+        print(
+            f"Leadfield matrix shape (fixed orientation): "
+            f"{fwd['sol']['data'].shape} (channels, sources)"
+        )
+
+    elif fwd["source_ori"] == FIFF.FIFFV_MNE_FREE_ORI:
+        print(
+            f"Leadfield matrix shape (free orientation): "
+            f"{fwd['sol']['data'].shape}"
+        )
+
+    # Save the leadfield as npz file
+    first_ch = fwd["info"]["chs"][0]
+
+    leadfield_data = {
+        "subject": subject,
+        "leadfield": fwd["sol"]["data"],
+        "source_ori": fwd["source_ori"],
+        "src_coords": fwd["source_rr"],
+        "cortical_normals": cortical_normals,
+        "info": fwd["info"],
+        "orientation_type": orientation_type,
+        "coil_type": coil_type,
+        "sensor_kind": first_ch.get("kind"),
+        "sensor_units": first_ch.get("unit"),
+        "sensor_unitmult": first_ch.get("unit_mul"),
+        "Q_basis": Q_basis,
+    }
+
+    file_name = f"{subject}_{orientation_type}_leadfield.npz"
+
+    if save_path is None:
+        save_path = get_data_path() / "leadfield" / file_name
+    else:
+        save_path = save_path / file_name
+
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+    np.savez(save_path, **leadfield_data)
+
+    print(f"Leadfield saved to {save_path}")
+
+
+# full sources in fwd
+# fwd_datapath = get_data_path() / 'fwd'
+# save_path = get_data_path() / 'leadfield'
+
+# reduced sources in fwd
+fwd_datapath = get_data_path() / "1284src_fwd"
+save_path = get_data_path() / "1284src_leadfield"
+
+subjects_map = {
+    "eeg": ["fsaverage"],
+    "mag": ["CC120166", "CC120264", "CC120309", "CC120313"],
+}
+
+for coil_name, subjects in subjects_map.items():
+    for subject in subjects:
+        print(
+            f"\nExtracting leadfield (fixed & free) for subject: "
+            f"{subject}, coil_name: {coil_name}"
+        )
+
+        extract_leadfield(
+            fwd_datapath,
+            subject=subject,
+            coil_name=coil_name,
+            orientation_type="fixed",
+            save_path=save_path,
+        )
+
+        extract_leadfield(
+            fwd_datapath,
+            subject=subject,
+            coil_name=coil_name,
+            orientation_type="free",
+            save_path=save_path,
+        )