simnibs_analyze._pipeline_io

  1"""
  2SimNIBS pipeline I/O module.
  3Centralises all input/output operations: file discovery, NIfTI image loading,
  4CSV and YAML configuration reading/writing.
  5"""
  6
  7from __future__ import annotations
  8
  9from pathlib import Path
 10from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 11
 12import nibabel as nib
 13import numpy as np
 14import pandas as pd
 15
 16from ._config import PipelineConfig
 17from ._logging import get_logger
 18
 19logger = get_logger(__name__)
 20
 21SPACE_MNI = "mni"
 22SPACE_NATIVE = "native"
 23
 24
 25# ---------------------------------------------------------------------------
 26# Output naming context
 27# ---------------------------------------------------------------------------
 28
 29
 30class OutputContext:
 31    """
 32    Groups all parameters needed to build pipeline output filenames.
 33    Instantiate in the main loop (run.py) and pass to save / get_*_path functions.
 34
 35    Parameters
 36    ----------
 37    subject : str
 38        Subject identifier (e.g. '0008').
 39    condition : str
 40        ROI / stimulation condition name (e.g. 'fef', 'ips-left').
 41        Used in preprocessed filenames and mask names.
 42    mode : str
 43        'simulation' or 'optimization'.
 44    space : str
 45        'mni' or 'native'.
 46    roi_name : str
 47        ROI key as defined in config.yaml → target_generation.rois.
 48        Usually identical to *condition*; useful to distinguish the config key
 49        from the filename.
 50    base_name : str or None
 51        Stem of the source e-field file (e.g. 'sub-0008_scalar_MNI_magnE').
 52        Required for get_preproc_paths.  Set after find_efield_files.
 53    results_dir : Path or None
 54        Root results directory (config.paths.results_dir).
 55    simnibs_output : Path or None
 56        Root SimNIBS output directory (config.paths.simnibs_output).
 57    """
 58
 59    def __init__(
 60        self,
 61        subject: str = "",
 62        condition: str = "",
 63        mode: str = "",
 64        space: str = SPACE_MNI,
 65        roi_name: str = "",
 66        base_name: str = "",
 67        results_dir: Optional[Path] = None,
 68        simnibs_output: Optional[Path] = None,
 69    ) -> None:
 70        self.subject = subject
 71        self.condition = condition
 72        self.mode = mode
 73        self.space = space
 74        self.roi_name = roi_name or condition  # default: same as condition
 75        self.base_name = base_name
 76        self.results_dir = results_dir
 77        self.simnibs_output = simnibs_output
 78
 79
 80def normalize_space(space: str) -> str:
 81    """Normalize and validate the working space value."""
 82    normalized = str(space).lower().strip()
 83    if normalized not in {SPACE_MNI, SPACE_NATIVE}:
 84        raise ValueError(
 85            f"Invalid 'space' parameter: {space}. Allowed values: {SPACE_MNI}, {SPACE_NATIVE}"
 86        )
 87    return normalized
 88
 89
 90def space_tag(space: str) -> str:
 91    """Return tagged space label used in output paths and file names."""
 92    return f"space-{normalize_space(space)}"
 93
 94
 95def get_preps_root(paths_config) -> Path:
 96    """Retourne la racine des segmentations charm (m2m_*)."""
 97    if paths_config.simnibs_preps is not None:
 98        return Path(paths_config.simnibs_preps)
 99    return Path(paths_config.simnibs_output)
100
101
102def get_simu_root(paths_config) -> Path:
103    """Retourne la racine des simulations/optimisations."""
104    if paths_config.simnibs_simu is not None:
105        return Path(paths_config.simnibs_simu)
106    return Path(paths_config.simnibs_output)
107
108
109def get_subject_paths(simnibs_output_dir: Path, subject: str) -> Dict[str, Path]:
110    """Return canonical subject-level paths used across the pipeline."""
111    subject_dir = Path(simnibs_output_dir) / subject
112    return {
113        "subject_dir": subject_dir,
114        "m2m_dir": subject_dir / f"m2m_{subject}",
115        "subject_target_dir": subject_dir / "subject_target",
116    }
117
118
119def get_subject_paths_from_config(paths_config, subject: str) -> Dict[str, Path]:
120    """
121    Return canonical subject-level paths, handling split preps/simu directories.
122
123    - ``subject_dir``        : simu root / subject  (simulations live here)
124    - ``m2m_dir``            : preps root / subject / m2m_{subject}  (charm segmentation)
125    - ``subject_target_dir`` : simu root / subject / subject_target
126    """
127    preps_root = get_preps_root(paths_config)
128    simu_root = get_simu_root(paths_config)
129    simu_subject_dir = simu_root / subject
130    preps_subject_dir = preps_root / subject
131    return {
132        "subject_dir": simu_subject_dir,
133        "m2m_dir": preps_subject_dir / f"m2m_{subject}",
134        "subject_target_dir": simu_subject_dir / "subject_target",
135    }
136
137
138def get_analysis_dir(results_dir: Path, space: str) -> Path:
139    """Return the shared analysis output directory."""
140    return Path(results_dir) / "3-analysis"
141
142
143def get_features_csv_path(results_dir: Path, space: str) -> Path:
144    """Return the canonical features CSV path for a given space."""
145    return get_analysis_dir(results_dir, space) / f"all_features_{space_tag(space)}.csv"
146
147
148def get_inter_subject_summary_csv_path(results_dir: Path, space: str) -> Path:
149    """Return the inter-subject summary CSV path for a given space."""
150    return (
151        get_analysis_dir(results_dir, space)
152        / f"inter_subject_summary_{space_tag(space)}.csv"
153    )
154
155
156def get_intra_subject_diff_csv_path(
157    results_dir: Path, space: str, condition: str
158) -> Path:
159    """Return the intra-subject diff CSV path for a condition and space."""
160    return (
161        get_analysis_dir(results_dir, space)
162        / f"intra_subject_diff_{condition}_{space_tag(space)}.csv"
163    )
164
165
166def get_clusters_csv_path(results_dir: Path, space: str) -> Path:
167    """Return the clustering CSV path for a given space."""
168    return get_analysis_dir(results_dir, space) / f"clusters_{space_tag(space)}.csv"
169
170
171def load_config(config_path: Path) -> PipelineConfig:
172    """Load and validate the YAML configuration file via Pydantic models."""
173    from _config import load_and_validate
174
175    return load_and_validate(config_path)
176
177
178def find_raw_efield(
179    simnibs_output_dir: Path, subject: str, roi: str, mode: str
180) -> Optional[Path]:
181    """
182    Find the raw (unprocessed) e-field file in the SimNIBS output directory.
183
184    Parameters
185    ----------
186    simnibs_output_dir : Path
187        SimNIBS output directory.
188    subject : str
189        Subject ID.
190    roi : str
191        ROI name.
192    mode : str
193        Mode (simulation or optimization).
194
195    Returns
196    -------
197    Path or None
198        Path to the file, or None if not found.
199    """
200    subject_dir = simnibs_output_dir / subject
201
202    if not subject_dir.exists():
203        return None
204
205    if mode == "simulation":
206        base_dir = subject_dir / "simulations"
207    else:
208        base_dir = subject_dir / "optimizations"
209
210    if not base_dir.exists():
211        return None
212
213    pattern = f"{mode}_{mode}_{roi}_*"
214    matching_dirs = list(base_dir.glob(pattern))
215
216    if not matching_dirs:
217        return None
218
219    mode_dir = matching_dirs[0]
220
221    if mode == "optimization":
222        mni_volumes_dir = mode_dir / "simulation_with_optimal_montage" / "mni_volumes"
223    else:
224        mni_volumes_dir = mode_dir / "mni_volumes"
225
226    if not mni_volumes_dir.exists():
227        return None
228
229    efield_files = list(mni_volumes_dir.glob("*_scalar_MNI_magnE.nii.gz"))
230    return efield_files[0] if efield_files else None
231
232
233def find_simulation_dirs(
234    subject_dir: Path, condition: str, mode: str, folder_pattern: str | None = None
235) -> List[Path]:
236    """
237    Find all simulation/optimization directories for a given condition.
238    Handles hashes in folder names.
239
240    Parameters
241    ----------
242    subject_dir : Path
243        Subject directory (e.g. 001-CC).
244    condition : str
245        Stimulation condition (e.g. fef, ips-left).
246        Used as the search fragment if *folder_pattern* is not provided.
247    mode : str
248        Mode (simulation or optimization).
249    folder_pattern : str or None
250        Glob fragment to use instead of the condition name when searching
251        SimNIBS folders.  Useful when the ROI name differs from the folder name
252        (e.g. ROI 'ips-left' but folders named '…ips_left…').
253
254    Returns
255    -------
256    List[Path]
257        List of matching directories.
258    """
259    fragment = folder_pattern if folder_pattern is not None else condition
260    pattern = f"*{mode}_{fragment}*"
261
262    if mode == "simulation":
263        base_dir = subject_dir / "simulations"
264    else:
265        base_dir = subject_dir / "optimizations"
266
267    if not base_dir.exists():
268        logger.warning(f"{mode} directory not found: {base_dir}")
269        return []
270
271    found_dirs = list(base_dir.glob(pattern))
272
273    if not found_dirs:
274        logger.warning(f"No {mode} found for pattern: {pattern} in {base_dir}")
275
276    return found_dirs
277
278
279def find_efield_files(
280    simulation_dir: Path, mode: str, space: str = SPACE_MNI
281) -> List[Path]:
282    """
283    Find e-field files in a simulation/optimization directory.
284
285    Parameters
286    ----------
287    simulation_dir : Path
288        Simulation or optimization directory.
289    mode : str
290        Mode (simulation or optimization).
291    space : str
292        ``'mni'`` (default): ``*_scalar_MNI_magnE.nii.gz`` files in ``mni_volumes/``.
293        ``'native'``: ``*_scalar_magnE.nii.gz`` files in ``subject_volumes/``.
294
295    Returns
296    -------
297    List[Path]
298        List of e-field files found.
299    """
300    space = normalize_space(space)
301
302    if space == SPACE_NATIVE:
303        if mode == "optimization":
304            volumes_dir = (
305                simulation_dir / "simulation_with_optimal_montage" / "subject_volumes"
306            )
307        else:
308            volumes_dir = simulation_dir / "subject_volumes"
309        glob_pattern = "*_scalar_magnE.nii.gz"
310    else:
311        if mode == "optimization":
312            volumes_dir = (
313                simulation_dir / "simulation_with_optimal_montage" / "mni_volumes"
314            )
315        else:
316            volumes_dir = simulation_dir / "mni_volumes"
317        glob_pattern = "*_scalar_MNI_magnE.nii.gz"
318
319    if not volumes_dir.exists():
320        logger.warning(f"{space}_volumes directory not found: {volumes_dir}")
321        return []
322
323    efield_files = list(volumes_dir.glob(glob_pattern))
324
325    if not efield_files:
326        logger.warning(f"No e-field files found in {volumes_dir}")
327
328    return efield_files
329
330
331def get_t1_conform(
332    m2m_dir: Path,
333    filename: str = "segmentation/T1_bias_corrected.nii.gz",
334) -> Path:
335    """
336    Return the path to the T1 file inside ``m2m_dir``.
337
338    Parameters
339    ----------
340    m2m_dir : Path
341        ``m2m_<subject>`` directory produced by SimNIBS.
342    filename : str
343        Relative path to the T1 file (default: ``segmentation/T1_bias_corrected.nii.gz``).
344
345    Raises
346    ------
347    FileNotFoundError
348    """
349    path = Path(m2m_dir) / filename
350    if not path.exists():
351        raise FileNotFoundError(f"T1 not found: {path}")
352    return path
353
354
355def get_brainmask(
356    m2m_dir: Optional[Path] = None,
357    filename: str = "label_prep/tissue_labeling_upsampled.nii.gz",
358    space: str = SPACE_NATIVE,
359    mni_mask_path: Optional[Path] = None,
360) -> Path:
361    """
362    Return the path to the brain mask.
363
364    Parameters
365    ----------
366    m2m_dir : Path or None
367        ``m2m_<subject>`` directory produced by SimNIBS. Ignored when ``space='mni'``.
368    filename : str
369        Relative path to the mask file inside ``m2m_dir`` (subject space only).
370    space : str
371        ``'native'`` (default): mask inside ``m2m_dir``.
372        ``'mni'``: MNI mask passed via ``mni_mask_path`` (read from config).
373    mni_mask_path : Path or None
374        MNI mask path, required when ``space='mni'``.
375        Must come from ``config['paths']['mni_brain_mask']``.
376
377    Raises
378    ------
379    FileNotFoundError, ValueError
380    """
381    space = normalize_space(space)
382
383    if space == SPACE_MNI:
384        if mni_mask_path is None:
385            raise ValueError(
386                "mni_mask_path is required when space='mni' (config['paths']['mni_brain_mask'])"
387            )
388        path = Path(mni_mask_path)
389    else:
390        if m2m_dir is None:
391            raise ValueError("m2m_dir is required when space='native'")
392        path = Path(m2m_dir) / filename
393    if not path.exists():
394        raise FileNotFoundError(f"Brain mask not found: {path}")
395    return path
396
397
398def get_mni_tissues(m2m_dir: Path) -> Path:
399    """
400    Return the path to the tissue segmentation in MNI space.
401
402    Produced by SimNIBS at ``toMNI/final_tissues_MNI.nii.gz``.
403    Labels: 1=WM, 2=GM, 3=CSF, 4=Bone, 5=Scalp …
404    """
405    path = Path(m2m_dir) / "toMNI" / "final_tissues_MNI.nii.gz"
406    if not path.exists():
407        raise FileNotFoundError(f"MNI tissues not found: {path}")
408    return path
409
410
411def get_roi_mask_path(
412    simnibs_output_dir: Path,
413    condition: str,
414    space: str = SPACE_MNI,
415    subject: Optional[str] = None,
416) -> Path:
417    """
418    Return the ROI mask path for a given condition.
419
420    Parameters
421    ----------
422    simnibs_output_dir : Path
423        SimNIBS output directory.
424    condition : str
425        Stimulation condition.
426    space : str
427        ``'mni'`` (default): masks in ``mni_target/``.
428        ``'native'``: subject masks in ``<subject>/subject_target/``.
429    subject : str or None
430        Subject ID, required when ``space='native'``.
431
432    Returns
433    -------
434    Path
435        ROI mask path.
436
437    Raises
438    ------
439    FileNotFoundError
440        If the requested mask does not exist.
441    """
442    space = normalize_space(space)
443
444    if space == SPACE_NATIVE:
445        if not subject:
446            raise ValueError("subject is required when space='native'")
447        mask_path = (
448            simnibs_output_dir
449            / subject
450            / "subject_target"
451            / f"{condition}_mask_{space_tag(space)}.nii.gz"
452        )
453    else:
454        mask_path = (
455            simnibs_output_dir
456            / "mni_target"
457            / f"{condition}_mask_{space_tag(space)}.nii.gz"
458        )
459
460    if not mask_path.exists():
461        raise FileNotFoundError(f"ROI mask not found ({space} space): {mask_path}")
462
463    return mask_path
464
465
466def get_preproc_dir(sim_dir: Path, mode: str, space: str = SPACE_MNI) -> Path:
467    """
468    Return the preprocessing directory for a given simulation directory.
469
470    Parameters
471    ----------
472    sim_dir : Path
473        SimNIBS simulation directory.
474    mode : str
475        ``'simulation'`` or ``'optimization'``.
476    space : str
477        ``'mni'`` (default) or ``'native'``.
478    """
479    space = normalize_space(space)
480    volumes_dir = "subject_volumes" if space == SPACE_NATIVE else "mni_volumes"
481    if mode == "optimization":
482        return sim_dir / "simulation_with_optimal_montage" / volumes_dir
483    return sim_dir / volumes_dir
484
485
486def get_preproc_paths(preproc_dir: Path, base_name: str, roi_name: str) -> dict:
487    """
488    Return the intra- and extra-ROI preprocessed file paths.
489
490    Parameters
491    ----------
492    preproc_dir : Path
493        Preprocessing output directory (e.g. mni_volumes/).
494    base_name : str
495        Base name of the e-field file (without extension).
496    roi_name : str
497        ROI name used (e.g. 'fef', 'ips_left').
498
499    Returns
500    -------
501    dict with keys:
502        ``intra_masked``, ``intra_cleaned``,
503        ``extra_masked``,  ``extra_cleaned``
504    """
505    return {
506        "intra_masked": preproc_dir / f"{base_name}_{roi_name}_masked.nii.gz",
507        "intra_cleaned": preproc_dir / f"{base_name}_{roi_name}_cleaned.nii.gz",
508        "extra_masked": preproc_dir / f"{base_name}_extra_{roi_name}_masked.nii.gz",
509        "extra_cleaned": preproc_dir / f"{base_name}_extra_{roi_name}_cleaned.nii.gz",
510    }
511
512
513def load_nifti(path: Path) -> Tuple[np.ndarray, nib.Nifti1Image]:
514    """
515    Load a NIfTI file.
516
517    Parameters
518    ----------
519    path : Path
520        Path to the NIfTI file.
521
522    Returns
523    -------
524    data : np.ndarray
525        Volume data array.
526    img : nib.Nifti1Image
527        Full NIfTI image.
528    """
529    img = nib.load(str(path))
530    data = img.get_fdata()
531    return data, img
532
533
534def load_img(
535    path_or_img: Union[str, Path, nib.spatialimages.SpatialImage],
536) -> nib.spatialimages.SpatialImage:
537    """Load a NIfTI image from a path or return the image if already loaded."""
538    if isinstance(path_or_img, (str, Path)):
539        return nib.load(str(path_or_img))
540    if isinstance(path_or_img, nib.spatialimages.SpatialImage):
541        return path_or_img
542    raise TypeError(f"Expected path or nibabel image, got {type(path_or_img)}")
543
544
545def validate_binary(data: np.ndarray, name: str = "mask") -> None:
546    """Raise ValueError if *data* contains values other than 0 and 1."""
547    unique_values = np.unique(data)
548    if not np.all(np.isin(unique_values, [0, 1])):
549        raise ValueError(
550            f"{name} must be binary (contain only 0 and 1), "
551            f"but contains values: {unique_values}"
552        )
553
554
555def check_output(path: Path, if_exists: str = "overwrite") -> bool:
556    """Return True if the file should be written, False if it should be skipped.
557
558    Parameters
559    ----------
560    path : Path
561        Target output path.
562    if_exists : str
563        ``'overwrite'`` — always write (default).
564        ``'skip'``      — silently skip if file exists.
565        ``'error'``     — raise FileExistsError if file exists.
566    """
567    path = Path(path)
568    if path.exists():
569        if if_exists == "skip":
570            logger.info(f"Skip (already exists): {path.name}")
571            return False
572        elif if_exists == "error":
573            msg = f"Output already exists (if_exists='error'): {path}"
574            logger.error(msg)
575            raise FileExistsError(msg)
576    return True
577
578
579def save_nifti(
580    img: nib.spatialimages.SpatialImage, output_path: Path, if_exists: str = "overwrite"
581) -> None:
582    """Save a NIfTI image to disk, creating parent directories as needed."""
583    output_path = Path(output_path)
584    if not check_output(output_path, if_exists):
585        return
586    output_path.parent.mkdir(parents=True, exist_ok=True)
587    img.to_filename(str(output_path))
588
589
590def load_csvs(csv_paths: Iterable[Path]) -> pd.DataFrame:
591    """Load and concatenate multiple CSV files into a single DataFrame.
592
593    Parameters
594    ----------
595    csv_paths : Iterable[Path]
596        Iterable of paths to CSV files
597
598    Returns
599    -------
600    pd.DataFrame
601        Concatenated DataFrame from all CSV files
602    """
603    frames = [pd.read_csv(p) for p in csv_paths]
604    return pd.concat(frames, ignore_index=True)
605
606
607def save_rows(rows: List[Dict], out_csv: Path, if_exists: str = "overwrite") -> None:
608    """Save a list of row dicts to a CSV file, creating parent directories as needed.
609
610    Parameters
611    ----------
612    rows :
613        List of dicts, each representing one row (e.g. feature extraction output).
614    out_csv :
615        Destination CSV path.
616    if_exists :
617        ``'overwrite'`` (default), ``'skip'``, or ``'error'``.
618    """
619    out_csv = Path(out_csv)
620    if not check_output(out_csv, if_exists):
621        return
622    out_csv.parent.mkdir(parents=True, exist_ok=True)
623    pd.DataFrame(rows).to_csv(out_csv, index=False)
624
625
626def save_dataframe(
627    df: pd.DataFrame, out_path: Path, if_exists: str = "overwrite", **to_csv_kwargs
628) -> None:
629    """Save a DataFrame to CSV, creating parent directories as needed."""
630    out_path = Path(out_path)
631    if not check_output(out_path, if_exists):
632        return
633    out_path.parent.mkdir(parents=True, exist_ok=True)
634    df.to_csv(out_path, **to_csv_kwargs)
635
636
637def save_ants_image(img: Any, out_path: Path, if_exists: str = "overwrite") -> None:
638    """Save an ANTsPy image to disk, creating parent directories as needed."""
639    import ants
640
641    out_path = Path(out_path)
642    if not check_output(out_path, if_exists):
643        return
644    out_path.parent.mkdir(parents=True, exist_ok=True)
645    ants.image_write(img, str(out_path))
646
647
648def save_figure(out_path: Path, if_exists: str = "overwrite", **savefig_kwargs) -> bool:
649    """Save the current matplotlib figure to disk.
650
651    Returns True if the figure was saved, False if skipped.
652    Closes the figure in both cases.
653    """
654    import matplotlib.pyplot as plt
655
656    out_path = Path(out_path)
657    if not check_output(out_path, if_exists):
658        logger.debug(f"  skip figure (exists): {out_path.name}")
659        plt.close()
660        return False
661    out_path.parent.mkdir(parents=True, exist_ok=True)
662    plt.savefig(out_path, **savefig_kwargs)
663    plt.close()
664    return True