Source code for pyavs.source.forward

"""
Forward modeling for pyAVS package.

This module provides functions for creating forward models, BEM models,
and handling coregistration for source reconstruction.
"""

import os
import mne
import numpy as np
from pathlib import Path
from typing import List, Optional, Tuple, Dict, Any, Union

from ..layout import bids_stem, get_layout
from ..utils.validation import validate_subject_id
from ..utils.logging import get_logger

logger = get_logger('source.forward')


[docs] def create_bem_model(subject: str, subjects_dir: str, conductivity: Tuple[float, float, float] = (0.3, 0.006, 0.3), ico: Optional[int] = 4, verbose: bool = True) -> mne.bem.ConductorModel: """ Create BEM (Boundary Element Method) model for source reconstruction. Parameters ---------- subject : str Subject name in FreeSurfer subjects directory subjects_dir : str Path to FreeSurfer subjects directory conductivity : tuple of float, optional Conductivity values for (brain, skull, scalp) (default: (0.3, 0.006, 0.3)) ico : int, optional Icosahedral subdivision number (default: 4) verbose : bool, optional Whether to print progress information (default: True) Returns ------- mne.bem.ConductorModel BEM conductor model """ if verbose: logger.info(f"Creating BEM model for subject {subject}") # Check if BEM surfaces exist bem_dir = os.path.join(subjects_dir, subject, 'bem') # Create BEM surfaces if they don't exist surfaces_needed = ['inner_skull', 'outer_skull', 'outer_skin'] for surface in surfaces_needed: surf_file = os.path.join(bem_dir, f'{subject}-{surface}.surf') if not os.path.exists(surf_file): if verbose: logger.info(f"BEM surface {surface} not found, creating...") # This would typically require FreeSurfer to be run # For now, we'll check if it exists and warn if not logger.warning(f"BEM surface {surface} not found at {surf_file}") logger.info("Run FreeSurfer watershed algorithm to create BEM surfaces") try: # Create BEM model model = mne.make_bem_model( subject=subject, ico=ico, conductivity=conductivity, subjects_dir=subjects_dir, verbose=verbose ) if verbose: logger.info("BEM model created successfully") return model except Exception as e: logger.error(f"Error creating BEM model: {e}") raise
[docs] def create_bem_solution(bem_model: mne.bem.ConductorModel, verbose: bool = True) -> mne.bem.ConductorModel: """ Create BEM solution from BEM model. Parameters ---------- bem_model : mne.bem.ConductorModel BEM model verbose : bool, optional Whether to print progress information (default: True) Returns ------- mne.bem.ConductorModel BEM solution """ if verbose: logger.info("Computing BEM solution...") try: bem_solution = mne.make_bem_solution(bem_model, verbose=verbose) if verbose: logger.info("BEM solution computed successfully") return bem_solution except Exception as e: logger.error(f"Error computing BEM solution: {e}") raise
[docs] def create_source_space(subject: str, subjects_dir: str, spacing: str = 'ico4', surface: str = 'white', add_dist: bool = True, verbose: bool = True) -> mne.SourceSpaces: """ Create cortical source space. Parameters ---------- subject : str Subject name in FreeSurfer subjects directory subjects_dir : str Path to FreeSurfer subjects directory spacing : str, optional Spacing between sources (default: 'ico4') surface : str, optional Surface to use (default: 'white') add_dist : bool, optional Whether to add distance information (default: True) verbose : bool, optional Whether to print progress information (default: True) Returns ------- mne.SourceSpaces Source space """ if verbose: logger.info(f"Creating source space for subject {subject}") try: src = mne.setup_source_space( subject=subject, spacing=spacing, surface=surface, subjects_dir=subjects_dir, add_dist=add_dist, verbose=verbose ) if verbose: logger.info(f"Source space created with {src[0]['nuse']} + {src[1]['nuse']} sources") return src except Exception as e: logger.error(f"Error creating source space: {e}") raise
[docs] def create_forward_model(raw: mne.io.Raw, trans: Union[str, mne.transforms.Transform], src: mne.SourceSpaces, bem_solution: mne.bem.ConductorModel, meg: bool = True, eeg: bool = False, mindist: float = 5.0, ignore_ref: bool = True, verbose: bool = True) -> mne.Forward: """ Create forward model for source reconstruction. Parameters ---------- raw : mne.io.Raw MEG/EEG raw data (used for sensor information) trans : str or mne.transforms.Transform Transformation from head to MRI coordinates src : mne.SourceSpaces Source space bem_solution : mne.bem.ConductorModel BEM solution meg : bool, optional Whether to include MEG channels (default: True) eeg : bool, optional Whether to include EEG channels (default: False) mindist : float, optional Minimum distance between sources and inner skull (default: 5.0) ignore_ref : bool, optional Whether to ignore reference channels (default: True) verbose : bool, optional Whether to print progress information (default: True) Returns ------- mne.Forward Forward solution """ if verbose: logger.info("Creating forward model...") try: fwd = mne.make_forward_solution( raw.info, trans=trans, src=src, bem=bem_solution, meg=meg, eeg=eeg, mindist=mindist, ignore_ref=ignore_ref, verbose=verbose ) if verbose: logger.info(f"Forward model created with {fwd['nsource']} sources") return fwd except Exception as e: logger.error(f"Error creating forward model: {e}") raise
[docs] def setup_coregistration(subject: str, subjects_dir: str, raw: mne.io.Raw, fiducials: str = 'auto', verbose: bool = True) -> mne.transforms.Transform: """ Set up coregistration between MEG and MRI coordinate systems. Parameters ---------- subject : str Subject name in FreeSurfer subjects directory subjects_dir : str Path to FreeSurfer subjects directory raw : mne.io.Raw MEG raw data fiducials : str, optional How to handle fiducials (default: 'auto') verbose : bool, optional Whether to print progress information (default: True) Returns ------- mne.transforms.Transform Head-to-MRI transformation """ if verbose: logger.info(f"Setting up coregistration for subject {subject}") # Check if transformation file exists. The release ships it under # mri/transforms/, mirroring FreeSurfer's own convention. trans_file = os.path.join(subjects_dir, subject, 'mri', 'transforms', f'{subject}-trans.fif') if os.path.exists(trans_file): if verbose: logger.info(f"Loading existing transformation: {trans_file}") trans = mne.read_trans(trans_file) else: if verbose: logger.info("No existing transformation found") logger.info("Manual coregistration required using mne.gui.coregistration()") logger.info("or automatic coregistration with mne.coreg.fit_matched_points()") # For automated processing, we might try to use fiducials try: # This is a simplified approach - in practice, would need proper fiducial setup if fiducials == 'auto': # Attempt automatic coregistration based on head shape from mne.coreg import fit_matched_points # This would require digitization points and head surface # For now, create identity transform as placeholder trans = mne.transforms.Transform('head', 'mri', np.eye(4)) if verbose: logger.warning("Using identity transformation - manual coregistration recommended") except Exception as e: if verbose: logger.warning(f"Automatic coregistration failed: {e}") # Create identity transformation as fallback trans = mne.transforms.Transform('head', 'mri', np.eye(4)) if verbose: logger.warning("Using identity transformation - manual coregistration required") return trans
[docs] def check_forward_model(fwd: mne.Forward, verbose: bool = True) -> Dict[str, Any]: """ Check forward model for potential issues. Parameters ---------- fwd : mne.Forward Forward solution verbose : bool, optional Whether to print check results (default: True) Returns ------- dict Dictionary with check results """ checks = { 'n_sources': fwd['nsource'], 'n_channels': fwd['nchan'], 'coord_frame': fwd['coord_frame'], 'has_meg': 'meg' in fwd, 'has_eeg': 'eeg' in fwd, 'is_free_orientation': fwd['surf_ori'] == mne.io.constants.FIFF.FIFFV_MNE_FREE_ORI, 'issues': [] } # Check for common issues if fwd['nsource'] < 1000: checks['issues'].append("Low number of sources - check source space") if fwd['nchan'] < 100: checks['issues'].append("Low number of channels - check channel selection") # Check condition number try: G = fwd['sol']['data'] cond_num = np.linalg.cond(G @ G.T) checks['condition_number'] = cond_num if cond_num > 1e12: checks['issues'].append("High condition number - check coregistration") except Exception: checks['condition_number'] = None checks['issues'].append("Could not compute condition number") if verbose: logger.info("Forward model check:") logger.info(f" Sources: {checks['n_sources']}") logger.info(f" Channels: {checks['n_channels']}") logger.info(f" MEG: {checks['has_meg']}") logger.info(f" EEG: {checks['has_eeg']}") if checks['condition_number'] is not None: logger.info(f" Condition number: {checks['condition_number']:.2e}") if checks['issues']: logger.warning(" Issues found:") for issue in checks['issues']: logger.warning(f" - {issue}") else: logger.info(" No issues found") return checks
def _derivatives_forward_path(subject_id: int, session: int, data_path: Optional[str] = None) -> Path: """Path of a pyAVS-computed forward inside the derivatives tree. ``{derivatives}/sub-{id:02d}/ses-{sess:02d}/source/ sub-{id:02d}_ses-{sess:02d}_task-avs_fwd.fif`` """ layout = get_layout(data_path) return (layout.deriv_dir(subject_id, session, 'source') / f"{bids_stem(subject_id, session)}_fwd.fif")
[docs] def save_forward_model(fwd: mne.Forward, subject_id: int, session: int, data_path: Optional[str] = None, overwrite: bool = True) -> str: """ Save forward model to derivatives directory. Parameters ---------- fwd : mne.Forward Forward solution subject_id : int Subject ID session : int Session number data_path : str, optional Path to data directory. If None, uses configured data path overwrite : bool, optional Whether to overwrite existing files (default: True) Returns ------- str Path to saved forward model """ validate_subject_id(subject_id) fwd_path = _derivatives_forward_path(subject_id, session, data_path) fwd_path.parent.mkdir(parents=True, exist_ok=True) mne.write_forward_solution(str(fwd_path), fwd, overwrite=overwrite) logger.info(f"Saved forward model to: {fwd_path}") return str(fwd_path)
[docs] def load_forward_model(subject_id: int, session: Optional[int] = None, data_path: Optional[str] = None, verbose: bool = True) -> mne.Forward: """ Load a forward model from the dataset. Two locations are searched, in order: 1. **pyAVS derivatives** (only when ``session`` is given) — a forward recomputed with :func:`save_forward_model`: ``{derivatives}/sub-{id:02d}/ses-{sess:02d}/source/ sub-{id:02d}_ses-{sess:02d}_task-avs_fwd.fif`` 2. **The shipped forward** — one per subject, session-independent: ``{root}/derivatives/freesurfer/sub-{id:02d}/bem/sub-{id:02d}-fwd.fif`` The shipped forward is what the release provides; the derivatives path only exists if you recomputed one yourself, in which case it takes precedence. Parameters ---------- subject_id : int Subject ID session : int, optional Session number. When given, a per-session forward in the pyAVS derivatives tree is preferred over the shipped one. The shipped forward does not depend on session. data_path : str, optional Path to the ``avs-public`` root. If None, uses the configured data path. verbose : bool, optional Whether to print loading information (default: True) Returns ------- mne.Forward Forward solution """ validate_subject_id(subject_id) layout = get_layout(data_path) candidates = [] if session is not None: candidates.append(_derivatives_forward_path(subject_id, session, data_path)) candidates.append(layout.forward(subject_id)) for fwd_path in candidates: if fwd_path.exists(): break else: searched = "\n ".join(str(p) for p in candidates) raise FileNotFoundError( f"Forward model not found for subject {subject_id}. Searched:\n {searched}") if verbose: logger.info(f"Loading forward model from: {fwd_path}") return mne.read_forward_solution(str(fwd_path), verbose=verbose)