#!/usr/bin/env python """Fit and evaluate A1 alignment caches for target masks. This script consumes Phase 6 artifacts from `run_a1_bootstrap.py` and computes ridge-regression performance for: - Protocol C: cross-subject subject-holdout on a shared canonical stimulus It supports three target-mask modes: - `run_top10`: per-run-condition top-10% ISC masks from Swati outputs - `run_top25`: per-run-condition top-25% ISC masks from Swati outputs - `core_roi`: legacy evaluation with the 7 core language ROIs """ from __future__ import annotations import argparse import json import os from pathlib import Path from typing import Any import nibabel as nib from nibabel.processing import resample_from_to import numpy as np import pandas as pd from a1_pipeline.io_utils import ensure_directory, write_json from a1_pipeline.participant_runs import ( load_participant_run_map, resolve_subject_actual_run, resolve_subject_canonical_run, ) from a1_pipeline.targets import CORE_ROI_NAMES, load_core_roi_masks TARGET_MASK_MODE_CHOICES: tuple[str, ...] = ("run_top10", "run_top25", "core_roi") RUN_MASK_PERCENT_BY_MODE: dict[str, int] = { "run_top10": 10, "run_top25": 25, } SWATI_CONDITION_BY_CANONICAL_RUN: dict[int, str] = { 1: "single_f", 2: "single_m", 3: "mixed_f", 4: "mixed_m", } def _resolve_num_fit_workers(requested: str, n_layers: int) -> tuple[int, int]: """Resolve --num-fit-workers to (n_workers, blas_threads_per_worker). 'auto' picks min(n_layers, max(1, cpu_count // 2)) so each worker still has multiple BLAS threads available for matmul. """ cpu_count = os.cpu_count() or 4 token = str(requested).strip().lower() if token in {"", "auto"}: n_workers = max(1, min(int(n_layers), max(1, cpu_count // 2))) else: try: value = int(token) except ValueError as exc: raise ValueError(f"Invalid --num-fit-workers={requested!r}") from exc n_workers = max(1, min(value, int(n_layers))) blas_threads = max(1, cpu_count // n_workers) return n_workers, blas_threads def _parse_csv_int_list(raw: str) -> list[int]: tokens = [token.strip() for token in raw.split(",") if token.strip()] return [int(token) for token in tokens] def _parse_subjects(raw: str) -> list[str] | None: value = raw.strip().lower() if value == "all": return None tokens = [token.strip() for token in raw.split(",") if token.strip()] return sorted(set(tokens)) def _parse_layers(raw: str) -> list[int] | None: value = raw.strip().lower() if value == "all": return None layers = sorted(set(_parse_csv_int_list(raw))) if any(layer < 0 for layer in layers): raise ValueError("Layer indices must be non-negative") return layers def _parse_protocols(raw: str) -> set[str]: values = {token.strip().upper() for token in raw.split(",") if token.strip()} if not values: raise ValueError("At least one protocol is required") allowed = {"C"} unknown = values.difference(allowed) if unknown: raise ValueError(f"Unsupported protocol(s): {sorted(unknown)}") return values def _parse_subject_field(raw: Any) -> list[str]: text = str(raw).strip() if text == "" or text.lower() == "nan": return [] return [token.strip() for token in text.split(",") if token.strip()] def _build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description="Fit and evaluate A1 cached alignment inputs") parser.add_argument( "--bootstrap-output-dir", type=str, default=None, help="Output directory from run_a1_bootstrap.py (required)", ) parser.add_argument( "--model-slug", type=str, default="Qwen_Qwen3-0.6B", help="Model slug in alignment_regressor_summary.csv", ) parser.add_argument( "--model-id", type=str, default=None, help="Optional human-readable model ID for reporting", ) parser.add_argument( "--alpha", type=float, default=1e3, help="Ridge regularization coefficient", ) parser.add_argument( "--subjects", type=str, default="all", help="Comma-separated subject IDs or 'all'", ) parser.add_argument( "--layer-indices", type=str, default="all", help="Comma-separated layer indices or 'all'", ) parser.add_argument( "--protocols", type=str, default="C", help="Protocols to run: C (cross-subject shared-space).", ) parser.add_argument( "--target-mask-mode", type=str, default="run_top10", choices=sorted(TARGET_MASK_MODE_CHOICES), help=( "Target mask family to evaluate. 'run_top10' selects the ISC top-10%% mask for the " "canonical stimulus behind each actual run, 'run_top25' keeps the 25%% mask option, " "and 'core_roi' keeps the legacy 7-ROI evaluation." ), ) parser.add_argument( "--run-mask-dir", "--run-top10-mask-dir", "--run-top25-mask-dir", dest="run_mask_dir", type=str, default=None, help=( "Directory containing Swati ISC run-conditioned masks. Can point either to the Swati " "output root or directly to its isc_group folder." ), ) parser.add_argument( "--roi-mask-dir", type=str, default=None, help="Legacy core ROI mask directory, used only when --target-mask-mode core_roi.", ) parser.add_argument( "--output-dir", type=str, default=None, help="Fit result output directory (default: /fit_results/)", ) parser.add_argument( "--participant-run-info", type=str, default=str((Path(__file__).resolve().parent / "assets" / "participant_run_info.json")), help=( "JSON file mapping each subject run to canonical stimulus condition " "(default: code/assets/participant_run_info.json)" ), ) parser.add_argument( "--num-fit-workers", type=str, default="auto", help=( "Number of CPU workers to evaluate layers in parallel for protocol C. " "'auto' uses min(n_layers, cpu_count // 2). Use 1 to force serial." ), ) return parser def _resolve_core_paths(bootstrap_output_dir: Path) -> dict[str, Path]: csv_dir = bootstrap_output_dir / "csv" def _pick_csv(name: str) -> Path: csv_candidate = csv_dir / name if csv_candidate.exists(): return csv_candidate return bootstrap_output_dir / name return { "bootstrap_summary": bootstrap_output_dir / "bootstrap_summary.json", "analysis_mask": bootstrap_output_dir / "analysis_mask.nii.gz", "bold_summary": _pick_csv("alignment_bold_summary.csv"), "regressor_summary": _pick_csv("alignment_regressor_summary.csv"), "protocol_c": _pick_csv("protocol_c_cross_subject_folds.csv"), } def _check_required_files(path_map: dict[str, Path], protocols: set[str]) -> None: required_names = { "bootstrap_summary", "analysis_mask", "bold_summary", "regressor_summary", } if "C" in protocols: required_names.add("protocol_c") missing = [name for name in sorted(required_names) if not path_map[name].exists()] if missing: lines = [f"{name}: {path_map[name]}" for name in missing] raise FileNotFoundError("Missing required files:\n" + "\n".join(lines)) def _resolve_core_roi_mask_dir(args: argparse.Namespace, bootstrap_summary: dict[str, Any]) -> Path: if args.roi_mask_dir: return Path(args.roi_mask_dir).resolve() if "roi_mask_dir" in bootstrap_summary: return Path(str(bootstrap_summary["roi_mask_dir"])).resolve() return (Path(__file__).resolve().parent / "assets" / "roi_masks") def _canonical_run_target_name(canonical_run: int, top_percent: int) -> str: if canonical_run not in SWATI_CONDITION_BY_CANONICAL_RUN: raise KeyError( f"Unsupported canonical run for target mask selection: {canonical_run}. " f"Allowed runs: {sorted(SWATI_CONDITION_BY_CANONICAL_RUN)}" ) return f"{SWATI_CONDITION_BY_CANONICAL_RUN[canonical_run]}_top{int(top_percent)}" def _analysis_mask_components( analysis_mask_path: Path, ) -> tuple[nib.Nifti1Image, np.ndarray, np.ndarray]: analysis_img = nib.load(str(analysis_mask_path)) analysis_mask_bool = analysis_img.get_fdata() > 0.5 analysis_flat = analysis_mask_bool.ravel(order="C") analysis_flat_indices = np.flatnonzero(analysis_flat) return analysis_img, analysis_mask_bool, analysis_flat_indices def _mask_to_vector_indices( mask_bool: np.ndarray, analysis_mask_bool: np.ndarray, analysis_flat_indices: np.ndarray, mask_name: str, ) -> np.ndarray: mask_in_analysis = analysis_mask_bool & mask_bool mask_flat_indices = np.flatnonzero(mask_in_analysis.ravel(order="C")) if mask_flat_indices.size == 0: raise ValueError(f"Target mask {mask_name} has zero voxels inside analysis mask") return np.searchsorted(analysis_flat_indices, mask_flat_indices).astype(np.int64) def _build_core_roi_index_maps( analysis_mask_path: Path, roi_mask_dir: Path, ) -> dict[int, dict[str, np.ndarray]]: analysis_img, analysis_mask_bool, analysis_flat_indices = _analysis_mask_components( analysis_mask_path=analysis_mask_path, ) core_roi_masks, _ = load_core_roi_masks(roi_mask_dir=roi_mask_dir, reference_img=analysis_img) roi_index_map: dict[str, np.ndarray] = {} for roi_name in CORE_ROI_NAMES: roi_bool = core_roi_masks[roi_name] roi_index_map[roi_name] = _mask_to_vector_indices( mask_bool=roi_bool, analysis_mask_bool=analysis_mask_bool, analysis_flat_indices=analysis_flat_indices, mask_name=roi_name, ) return { canonical_run: dict(roi_index_map) for canonical_run in sorted(SWATI_CONDITION_BY_CANONICAL_RUN) } def _normalize_run_mask_dir(path: Path, top_percent: int) -> Path | None: candidates = [path, path / "isc_group"] required_files = [ f"{condition}_isc_top{int(top_percent)}.nii.gz" for condition in SWATI_CONDITION_BY_CANONICAL_RUN.values() ] for candidate in candidates: if not candidate.exists() or not candidate.is_dir(): continue if all((candidate / filename).exists() for filename in required_files): return candidate.resolve() return None def _resolve_run_mask_dir(args: argparse.Namespace, top_percent: int) -> Path: script_path = Path(__file__).resolve() workspace_root = script_path.parents[3] candidate_roots: list[Path] = [] if args.run_mask_dir: candidate_roots.append(Path(args.run_mask_dir).expanduser().resolve()) candidate_roots.extend( [ workspace_root / "data" / "isc_group", workspace_root / "data", workspace_root / "swati" / "TEAM-9" / "output", workspace_root / "swati" / "TEAM-9" / "output" / "isc_group", ] ) for candidate in candidate_roots: normalized = _normalize_run_mask_dir(candidate, top_percent=top_percent) if normalized is not None: return normalized searched = "\n".join(str(path) for path in candidate_roots) raise FileNotFoundError( f"Could not locate Swati run-top{int(top_percent)} masks. Checked:\n" + searched ) def _build_run_mask_index_maps( analysis_mask_path: Path, run_mask_dir: Path, top_percent: int, ) -> dict[int, dict[str, np.ndarray]]: analysis_img, analysis_mask_bool, analysis_flat_indices = _analysis_mask_components( analysis_mask_path=analysis_mask_path, ) index_maps: dict[int, dict[str, np.ndarray]] = {} for canonical_run, condition in sorted(SWATI_CONDITION_BY_CANONICAL_RUN.items()): mask_name = _canonical_run_target_name(canonical_run, top_percent=top_percent) mask_path = run_mask_dir / f"{condition}_isc_top{int(top_percent)}.nii.gz" if not mask_path.exists(): raise FileNotFoundError( f"Missing run-top{int(top_percent)} mask for {condition}: {mask_path}" ) mask_img = nib.load(str(mask_path)) resampled = resample_from_to( mask_img, (analysis_img.shape, analysis_img.affine), order=0, ) mask_bool = resampled.get_fdata() > 0.0 index_maps[canonical_run] = { mask_name: _mask_to_vector_indices( mask_bool=mask_bool, analysis_mask_bool=analysis_mask_bool, analysis_flat_indices=analysis_flat_indices, mask_name=mask_name, ) } return index_maps def _build_bold_path_map(bold_summary_df: pd.DataFrame) -> dict[tuple[str, int], Path]: required_columns = {"subject", "run", "bold_z_path"} missing = required_columns.difference(bold_summary_df.columns) if missing: raise ValueError(f"alignment_bold_summary.csv missing columns: {sorted(missing)}") duplicate_check = bold_summary_df.groupby(["subject", "run"]).size() duplicates = duplicate_check[duplicate_check > 1] if not duplicates.empty: raise ValueError( "Expected unique bold_z_path per (subject,run). Found duplicates for: " + ", ".join([f"({subject},{run})" for subject, run in duplicates.index.tolist()]) ) mapping: dict[tuple[str, int], Path] = {} for row in bold_summary_df.itertuples(index=False): subject = str(getattr(row, "subject")) run = int(getattr(row, "run")) path = Path(str(getattr(row, "bold_z_path"))) mapping[(subject, run)] = path return mapping def _build_regressor_path_map( regressor_summary_df: pd.DataFrame, model_slug: str, ) -> tuple[dict[tuple[int, int], Path], str]: required_columns = {"model_id", "model_slug", "run", "layer_idx", "regressor_z_path"} missing = required_columns.difference(regressor_summary_df.columns) if missing: raise ValueError(f"alignment_regressor_summary.csv missing columns: {sorted(missing)}") model_df = regressor_summary_df[regressor_summary_df["model_slug"] == model_slug].copy() if model_df.empty: available = sorted(set(str(value) for value in regressor_summary_df["model_slug"].tolist())) raise ValueError( f"No regressors found for model_slug={model_slug}. Available slugs: {available}" ) selected_model_id = str(model_df.iloc[0]["model_id"]) duplicate_check = model_df.groupby(["run", "layer_idx"]).size() duplicates = duplicate_check[duplicate_check > 1] if not duplicates.empty: details = ", ".join([f"(run={run},layer={layer})" for run, layer in duplicates.index.tolist()]) raise ValueError( "Multiple regressors found for run/layer pairs. " "This fitter expects one regressor per (run,layer) for the chosen model. " f"Conflicts: {details}" ) mapping: dict[tuple[int, int], Path] = {} for row in model_df.itertuples(index=False): run = int(getattr(row, "run")) layer = int(getattr(row, "layer_idx")) path = Path(str(getattr(row, "regressor_z_path"))) mapping[(run, layer)] = path return mapping, selected_model_id def _load_subject_roi_runs( subject: str, runs: list[int], bold_path_map: dict[tuple[str, int], Path], participant_run_map: dict[str, dict[int, int]], target_index_map_by_canonical_run: dict[int, dict[str, np.ndarray]], ) -> dict[int, dict[str, np.ndarray]]: out: dict[int, dict[str, np.ndarray]] = {} for run in runs: key = (subject, run) if key not in bold_path_map: raise KeyError(f"Missing BOLD cache for subject={subject}, run={run}") bold_path = bold_path_map[key] bold_mmap = np.load(bold_path, mmap_mode="r") canonical_run = resolve_subject_canonical_run( participant_run_map=participant_run_map, subject=subject, run=run, ) if canonical_run not in target_index_map_by_canonical_run: raise KeyError( f"Missing target mask indices for canonical_run={canonical_run}, subject={subject}, run={run}" ) roi_run: dict[str, np.ndarray] = {} for roi_name, roi_indices in target_index_map_by_canonical_run[canonical_run].items(): roi_run[roi_name] = np.asarray(bold_mmap[:, roi_indices], dtype=np.float32) out[run] = roi_run return out def _ridge_projection_matrix(x_train: np.ndarray, alpha: float) -> np.ndarray: if alpha <= 0: raise ValueError("alpha must be positive") x_train64 = np.asarray(x_train, dtype=np.float64) xtx = x_train64.T @ x_train64 reg = np.eye(xtx.shape[0], dtype=np.float64) * float(alpha) return np.linalg.solve(xtx + reg, x_train64.T) def _safe_corr_per_voxel(y_true: np.ndarray, y_pred: np.ndarray, eps: float = 1e-12) -> np.ndarray: yt = np.asarray(y_true, dtype=np.float64) yp = np.asarray(y_pred, dtype=np.float64) yt_centered = yt - np.mean(yt, axis=0, keepdims=True) yp_centered = yp - np.mean(yp, axis=0, keepdims=True) numerator = np.sum(yt_centered * yp_centered, axis=0) denom = np.sqrt(np.sum(yt_centered**2, axis=0) * np.sum(yp_centered**2, axis=0)) corr = np.full(denom.shape, np.nan, dtype=np.float64) valid = denom > eps corr[valid] = numerator[valid] / denom[valid] return corr def _safe_r2_per_voxel(y_true: np.ndarray, y_pred: np.ndarray, eps: float = 1e-12) -> np.ndarray: yt = np.asarray(y_true, dtype=np.float64) yp = np.asarray(y_pred, dtype=np.float64) sse = np.sum((yt - yp) ** 2, axis=0) yt_mean = np.mean(yt, axis=0, keepdims=True) sst = np.sum((yt - yt_mean) ** 2, axis=0) r2 = np.full(sst.shape, np.nan, dtype=np.float64) valid = sst > eps r2[valid] = 1.0 - (sse[valid] / sst[valid]) return r2 def _two_v_two_accuracy(y_true: np.ndarray, y_pred: np.ndarray, eps: float = 1e-12) -> tuple[float, int]: yt = np.asarray(y_true, dtype=np.float64) yp = np.asarray(y_pred, dtype=np.float64) if yt.shape != yp.shape: raise ValueError( "2v2 accuracy expects y_true and y_pred to share shape, got " f"{yt.shape} and {yp.shape}" ) if yt.ndim != 2: raise ValueError(f"2v2 accuracy expects 2D matrices, got ndim={yt.ndim}") if yt.shape[0] < 2 or yt.shape[1] == 0: return float("nan"), 0 yt_centered = yt - np.mean(yt, axis=1, keepdims=True) yp_centered = yp - np.mean(yp, axis=1, keepdims=True) yt_norm = np.linalg.norm(yt_centered, axis=1) yp_norm = np.linalg.norm(yp_centered, axis=1) valid_rows = (yt_norm > eps) & (yp_norm > eps) if int(np.sum(valid_rows)) < 2: return float("nan"), 0 yt_unit = yt_centered[valid_rows] / yt_norm[valid_rows, None] yp_unit = yp_centered[valid_rows] / yp_norm[valid_rows, None] similarity = yt_unit @ yp_unit.T diagonal = np.diag(similarity) # Margin > 0 means matched assignment beats swapped assignment for a pair. pair_margin = diagonal[:, None] + diagonal[None, :] - similarity - similarity.T pair_idx = np.triu_indices(pair_margin.shape[0], k=1) margins = pair_margin[pair_idx] n_pairs = int(margins.size) if n_pairs == 0: return float("nan"), 0 wins = float(np.sum(margins > eps)) ties = float(np.sum(np.abs(margins) <= eps)) accuracy = (wins + 0.5 * ties) / float(n_pairs) return float(accuracy), n_pairs def _score_matrix(y_true: np.ndarray, y_pred: np.ndarray) -> dict[str, float]: corr = _safe_corr_per_voxel(y_true=y_true, y_pred=y_pred) r2 = _safe_r2_per_voxel(y_true=y_true, y_pred=y_pred) two_v_two_accuracy, n_2v2_pairs = _two_v_two_accuracy(y_true=y_true, y_pred=y_pred) finite_corr = corr[np.isfinite(corr)] finite_r2 = r2[np.isfinite(r2)] return { "mean_corr": float(np.mean(finite_corr)) if finite_corr.size > 0 else float("nan"), "median_corr": float(np.median(finite_corr)) if finite_corr.size > 0 else float("nan"), "mean_r2": float(np.mean(finite_r2)) if finite_r2.size > 0 else float("nan"), "two_v_two_accuracy": float(two_v_two_accuracy), "n_2v2_pairs": int(n_2v2_pairs), "n_voxels_scored": int(finite_corr.size), } def _collect_protocol_c_subjects(cross_subject_df: pd.DataFrame) -> list[str]: if cross_subject_df.empty: return [] subjects: set[str] = set() for row in cross_subject_df.itertuples(index=False): subjects.update(_parse_subject_field(getattr(row, "train_subjects"))) subjects.update(_parse_subject_field(getattr(row, "test_subjects"))) return sorted(subjects) def _evaluate_protocol_c( layers: list[int], alpha: float, cross_subject_df: pd.DataFrame, bold_path_map: dict[tuple[str, int], Path], regressor_path_map: dict[tuple[int, int], Path], target_index_map_by_canonical_run: dict[int, dict[str, np.ndarray]], participant_run_map: dict[str, dict[int, int]], model_slug: str, model_id: str, n_workers: int = 1, blas_threads_per_worker: int = 0, ) -> pd.DataFrame: if cross_subject_df.empty: return pd.DataFrame() all_subjects = _collect_protocol_c_subjects(cross_subject_df) canonical_runs = sorted({int(value) for value in cross_subject_df["canonical_run"].tolist()}) required_actual_runs: dict[str, list[int]] = {} for subject in all_subjects: actual_runs = sorted( { resolve_subject_actual_run( participant_run_map=participant_run_map, subject=subject, canonical_run=canonical_run, ) for canonical_run in canonical_runs } ) required_actual_runs[subject] = actual_runs subject_roi_cache: dict[str, dict[int, dict[str, np.ndarray]]] = {} for subject in all_subjects: subject_roi_cache[subject] = _load_subject_roi_runs( subject=subject, runs=required_actual_runs[subject], bold_path_map=bold_path_map, participant_run_map=participant_run_map, target_index_map_by_canonical_run=target_index_map_by_canonical_run, ) if blas_threads_per_worker <= 0: blas_threads_per_worker = max(1, (os.cpu_count() or 4) // max(1, n_workers)) def _layer_task(layer_idx: int) -> list[dict[str, Any]]: return _evaluate_protocol_c_layer( layer_idx=int(layer_idx), alpha=alpha, cross_subject_df=cross_subject_df, regressor_path_map=regressor_path_map, canonical_runs=canonical_runs, subject_roi_cache=subject_roi_cache, target_index_map_by_canonical_run=target_index_map_by_canonical_run, participant_run_map=participant_run_map, model_slug=model_slug, model_id=model_id, blas_threads=blas_threads_per_worker, ) if n_workers > 1 and len(layers) > 1: from joblib import Parallel, delayed print( f"[fit] Protocol C: parallel layers across {n_workers} threads " f"(layers={len(layers)}, blas_threads/worker={blas_threads_per_worker})", flush=True, ) results = Parallel(n_jobs=int(n_workers), prefer="threads")( delayed(_layer_task)(layer_idx) for layer_idx in layers ) else: results = [_layer_task(layer_idx) for layer_idx in layers] rows: list[dict[str, Any]] = [] for layer_rows in results: rows.extend(layer_rows) df = pd.DataFrame(rows) if not df.empty: df = df.sort_values(["layer_idx", "fold_id", "subject", "roi_name"]).reset_index(drop=True) return df def _evaluate_protocol_c_noise_ceiling( cross_subject_df: pd.DataFrame, bold_path_map: dict[tuple[str, int], Path], target_index_map_by_canonical_run: dict[int, dict[str, np.ndarray]], participant_run_map: dict[str, dict[int, int]], ) -> pd.DataFrame: if cross_subject_df.empty: return pd.DataFrame() all_subjects = _collect_protocol_c_subjects(cross_subject_df) canonical_runs = sorted({int(value) for value in cross_subject_df["canonical_run"].tolist()}) required_actual_runs: dict[str, list[int]] = {} for subject in all_subjects: actual_runs = sorted( { resolve_subject_actual_run( participant_run_map=participant_run_map, subject=subject, canonical_run=canonical_run, ) for canonical_run in canonical_runs } ) required_actual_runs[subject] = actual_runs subject_roi_cache: dict[str, dict[int, dict[str, np.ndarray]]] = {} for subject in all_subjects: subject_roi_cache[subject] = _load_subject_roi_runs( subject=subject, runs=required_actual_runs[subject], bold_path_map=bold_path_map, participant_run_map=participant_run_map, target_index_map_by_canonical_run=target_index_map_by_canonical_run, ) rows: list[dict[str, Any]] = [] for fold_row in cross_subject_df.itertuples(index=False): fold_id = str(getattr(fold_row, "fold_id")) canonical_run = int(getattr(fold_row, "canonical_run")) condition_label = str(getattr(fold_row, "condition_label")) train_subjects = _parse_subject_field(getattr(fold_row, "train_subjects")) test_subjects = _parse_subject_field(getattr(fold_row, "test_subjects")) if not train_subjects or not test_subjects: raise ValueError(f"Protocol C fold {fold_id} must have non-empty train and test subjects") for roi_name in target_index_map_by_canonical_run[canonical_run].keys(): y_train_blocks: list[np.ndarray] = [] for subject in train_subjects: actual_run = resolve_subject_actual_run( participant_run_map=participant_run_map, subject=subject, canonical_run=canonical_run, ) y_train_subject = subject_roi_cache[subject][actual_run][roi_name] y_train_blocks.append(y_train_subject) y_train_stack = np.stack(y_train_blocks, axis=0) y_pred = np.mean(y_train_stack, axis=0, dtype=np.float64) for subject in test_subjects: actual_run = resolve_subject_actual_run( participant_run_map=participant_run_map, subject=subject, canonical_run=canonical_run, ) y_test = subject_roi_cache[subject][actual_run][roi_name] if y_test.shape != y_pred.shape: raise ValueError( "Noise ceiling test TR/voxel mismatch for " f"subject={subject}, actual_run={actual_run}, canonical_run={canonical_run}, " f"roi={roi_name}. y_pred={y_pred.shape}, y_test={y_test.shape}" ) scores = _score_matrix(y_true=y_test, y_pred=y_pred) rows.append( { "protocol": "C_cross_subject_noise_ceiling", "subject": subject, "run": int(actual_run), "canonical_run": int(canonical_run), "condition_label": condition_label, "fold_id": fold_id, "train_subjects": ",".join(train_subjects), "test_subjects": ",".join(test_subjects), "n_train_subjects": int(len(train_subjects)), "n_test_subjects": int(len(test_subjects)), "roi_name": roi_name, "n_train_tr": int(y_pred.shape[0] * len(train_subjects)), "n_test_tr": int(y_test.shape[0]), "n_voxels_roi": int(y_test.shape[1]), **scores, } ) df = pd.DataFrame(rows) if not df.empty: df = df.sort_values(["fold_id", "subject", "roi_name"]).reset_index(drop=True) return df def _evaluate_protocol_c_layer( layer_idx: int, alpha: float, cross_subject_df: pd.DataFrame, regressor_path_map: dict[tuple[int, int], Path], canonical_runs: list[int], subject_roi_cache: dict[str, dict[int, dict[str, np.ndarray]]], target_index_map_by_canonical_run: dict[int, dict[str, np.ndarray]], participant_run_map: dict[str, dict[int, int]], model_slug: str, model_id: str, blas_threads: int, ) -> list[dict[str, Any]]: """Evaluate Protocol C for a single layer. Returns list of result rows.""" from threadpoolctl import threadpool_limits rows: list[dict[str, Any]] = [] with threadpool_limits(limits=int(blas_threads)): x_by_canonical_run: dict[int, np.ndarray] = {} for canonical_run in canonical_runs: key = (canonical_run, layer_idx) if key not in regressor_path_map: raise KeyError( "Missing regressor cache for " f"canonical_run={canonical_run}, layer={layer_idx}, model={model_slug}" ) x_by_canonical_run[canonical_run] = np.asarray(np.load(regressor_path_map[key]), dtype=np.float32) projector_cache: dict[tuple[int, int], tuple[np.ndarray, int]] = {} for fold_row in cross_subject_df.itertuples(index=False): fold_id = str(getattr(fold_row, "fold_id")) canonical_run = int(getattr(fold_row, "canonical_run")) condition_label = str(getattr(fold_row, "condition_label")) train_subjects = _parse_subject_field(getattr(fold_row, "train_subjects")) test_subjects = _parse_subject_field(getattr(fold_row, "test_subjects")) if not train_subjects or not test_subjects: raise ValueError(f"Protocol C fold {fold_id} must have non-empty train and test subjects") x_test = x_by_canonical_run[canonical_run] projector_key = (canonical_run, len(train_subjects)) if projector_key not in projector_cache: x_train = np.vstack([x_test for _ in train_subjects]) projector_cache[projector_key] = ( _ridge_projection_matrix(x_train=x_train, alpha=alpha), int(x_train.shape[0]), ) projector, n_train_tr = projector_cache[projector_key] for roi_name in target_index_map_by_canonical_run[canonical_run].keys(): y_train_blocks: list[np.ndarray] = [] for subject in train_subjects: actual_run = resolve_subject_actual_run( participant_run_map=participant_run_map, subject=subject, canonical_run=canonical_run, ) y_train_subject = subject_roi_cache[subject][actual_run][roi_name] if y_train_subject.shape[0] != x_test.shape[0]: raise ValueError( "Protocol C train TR mismatch between regressors and BOLD for " f"subject={subject}, actual_run={actual_run}, canonical_run={canonical_run}, " f"layer={layer_idx}, roi={roi_name}. " f"x_train_tr_per_subject={x_test.shape[0]}, y_train_tr={y_train_subject.shape[0]}" ) y_train_blocks.append(y_train_subject) y_train = np.vstack(y_train_blocks) if y_train.shape[0] != n_train_tr: raise ValueError( "Protocol C pooled train TR mismatch between regressors and BOLD for " f"fold={fold_id}, layer={layer_idx}, roi={roi_name}. " f"x_train_tr={n_train_tr}, y_train_tr={y_train.shape[0]}" ) weights = projector @ np.asarray(y_train, dtype=np.float64) y_pred = np.asarray(x_test, dtype=np.float64) @ weights for subject in test_subjects: actual_run = resolve_subject_actual_run( participant_run_map=participant_run_map, subject=subject, canonical_run=canonical_run, ) y_test = subject_roi_cache[subject][actual_run][roi_name] if y_test.shape[0] != x_test.shape[0]: raise ValueError( "Protocol C test TR mismatch between regressors and BOLD for " f"subject={subject}, actual_run={actual_run}, canonical_run={canonical_run}, " f"layer={layer_idx}, roi={roi_name}. " f"x_test_tr={x_test.shape[0]}, y_test_tr={y_test.shape[0]}" ) scores = _score_matrix(y_true=y_test, y_pred=y_pred) rows.append( { "protocol": "C_cross_subject_shared_space", "model_id": model_id, "model_slug": model_slug, "subject": subject, "run": int(actual_run), "canonical_run": int(canonical_run), "condition_label": condition_label, "fold_id": fold_id, "train_subjects": ",".join(train_subjects), "test_subjects": ",".join(test_subjects), "n_train_subjects": int(len(train_subjects)), "n_test_subjects": int(len(test_subjects)), "layer_idx": int(layer_idx), "roi_name": roi_name, "alpha": float(alpha), "n_train_tr": int(n_train_tr), "n_test_tr": int(x_test.shape[0]), "n_voxels_roi": int(y_test.shape[1]), **scores, } ) return rows def _summarize_layers(scores_df: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]: if scores_df.empty: return pd.DataFrame(), pd.DataFrame() layer_summary = ( scores_df.groupby(["protocol", "model_id", "model_slug", "layer_idx", "roi_name"], as_index=False) .agg( mean_corr=("mean_corr", "mean"), std_corr=("mean_corr", "std"), mean_r2=("mean_r2", "mean"), mean_2v2_accuracy=("two_v_two_accuracy", "mean"), std_2v2_accuracy=("two_v_two_accuracy", "std"), n_records=("mean_corr", "size"), ) .sort_values(["protocol", "roi_name", "layer_idx"]) .reset_index(drop=True) ) best_rows: list[pd.Series] = [] for (_, roi_name), group_df in layer_summary.groupby(["protocol", "roi_name"]): idx = int(group_df["mean_corr"].idxmax()) best_rows.append(layer_summary.loc[idx]) best_layer_summary = pd.DataFrame(best_rows).reset_index(drop=True) best_layer_summary = best_layer_summary.sort_values(["protocol", "roi_name"]).reset_index(drop=True) return layer_summary, best_layer_summary def _summarize_noise_ceiling(scores_df: pd.DataFrame) -> pd.DataFrame: if scores_df.empty: return pd.DataFrame() return ( scores_df.groupby(["protocol", "roi_name"], as_index=False) .agg( mean_corr=("mean_corr", "mean"), std_corr=("mean_corr", "std"), mean_r2=("mean_r2", "mean"), mean_2v2_accuracy=("two_v_two_accuracy", "mean"), std_2v2_accuracy=("two_v_two_accuracy", "std"), n_records=("mean_corr", "size"), ) .sort_values(["protocol", "roi_name"]) .reset_index(drop=True) ) def _compare_best_layer_to_noise_ceiling( best_layer_df: pd.DataFrame, noise_ceiling_summary_df: pd.DataFrame, ) -> pd.DataFrame: if best_layer_df.empty or noise_ceiling_summary_df.empty: return pd.DataFrame() ceiling_df = noise_ceiling_summary_df.rename( columns={ "mean_corr": "noise_ceiling_mean_corr", "std_corr": "noise_ceiling_std_corr", "mean_r2": "noise_ceiling_mean_r2", "mean_2v2_accuracy": "noise_ceiling_mean_2v2_accuracy", "std_2v2_accuracy": "noise_ceiling_std_2v2_accuracy", "n_records": "noise_ceiling_n_records", } ) merged = best_layer_df.merge(ceiling_df, on="roi_name", how="left") if merged.empty: return merged for metric in ["mean_corr", "mean_r2", "mean_2v2_accuracy"]: denom = np.asarray(merged[f"noise_ceiling_{metric}"], dtype=np.float64) numer = np.asarray(merged[metric], dtype=np.float64) ratio = np.full(denom.shape, np.nan, dtype=np.float64) valid = np.isfinite(denom) & (np.abs(denom) > 1e-12) ratio[valid] = numer[valid] / denom[valid] merged[f"fraction_of_noise_ceiling_{metric}"] = ratio return merged def main() -> None: parser = _build_parser() args = parser.parse_args() bootstrap_output_dir = Path(args.bootstrap_output_dir).resolve() protocols = _parse_protocols(args.protocols) path_map = _resolve_core_paths(bootstrap_output_dir=bootstrap_output_dir) _check_required_files(path_map=path_map, protocols=protocols) with path_map["bootstrap_summary"].open("r", encoding="utf-8") as handle: bootstrap_summary = json.load(handle) target_mask_mode = str(args.target_mask_mode).strip().lower() if target_mask_mode not in TARGET_MASK_MODE_CHOICES: raise ValueError( f"Unsupported --target-mask-mode={target_mask_mode!r}. " f"Allowed: {sorted(TARGET_MASK_MODE_CHOICES)}" ) roi_mask_dir: Path | None = None run_mask_dir: Path | None = None target_top_percent: int | None = None if target_mask_mode == "core_roi": roi_mask_dir = _resolve_core_roi_mask_dir(args=args, bootstrap_summary=bootstrap_summary) if not roi_mask_dir.exists(): raise FileNotFoundError(f"ROI mask dir not found: {roi_mask_dir}") else: target_top_percent = RUN_MASK_PERCENT_BY_MODE[target_mask_mode] run_mask_dir = _resolve_run_mask_dir(args=args, top_percent=target_top_percent) output_dir = ( Path(args.output_dir).resolve() if args.output_dir else bootstrap_output_dir / "fit_results" / str(args.model_slug) ) ensure_directory(output_dir) bold_summary_df = pd.read_csv(path_map["bold_summary"]) regressor_summary_df = pd.read_csv(path_map["regressor_summary"]) cross_subject_df = ( pd.read_csv(path_map["protocol_c"]) if path_map["protocol_c"].exists() else pd.DataFrame() ) requested_subjects = _parse_subjects(args.subjects) if requested_subjects is not None and "C" in protocols: raise ValueError("Protocol C currently requires --subjects all so fold membership stays valid") if requested_subjects is not None: subject_set = set(requested_subjects) bold_summary_df = bold_summary_df[bold_summary_df["subject"].isin(subject_set)].copy() if cross_subject_df.empty and "C" in protocols: raise ValueError("Protocol C was requested but no cross-subject rows were available") bold_path_map = _build_bold_path_map(bold_summary_df=bold_summary_df) regressor_path_map, detected_model_id = _build_regressor_path_map( regressor_summary_df=regressor_summary_df, model_slug=str(args.model_slug), ) model_id = str(args.model_id) if args.model_id else detected_model_id available_layers = sorted({layer for (_, layer) in regressor_path_map.keys()}) requested_layers = _parse_layers(args.layer_indices) if requested_layers is None: layers = available_layers else: layer_set = set(available_layers) missing_layers = [layer for layer in requested_layers if layer not in layer_set] if missing_layers: raise ValueError( f"Requested layers are unavailable for model {args.model_slug}: {missing_layers}. " f"Available: {available_layers}" ) layers = requested_layers if not layers: raise ValueError("No layers selected for fitting") protocol_c_subjects = _collect_protocol_c_subjects(cross_subject_df) if "C" in protocols else [] subjects = sorted(set(protocol_c_subjects)) if not subjects: raise ValueError("No subjects available after filtering") participant_run_info_path = Path(args.participant_run_info).resolve() required_runs = sorted({int(value) for value in bold_summary_df["run"].tolist()}) if not required_runs: raise ValueError("No required runs detected for participant run mapping") participant_run_map = load_participant_run_map( participant_run_info_path=participant_run_info_path, subjects=subjects, runs=required_runs, ) if target_mask_mode == "core_roi": target_index_map_by_canonical_run = _build_core_roi_index_maps( analysis_mask_path=path_map["analysis_mask"], roi_mask_dir=roi_mask_dir, ) target_mask_source_dir = roi_mask_dir else: target_index_map_by_canonical_run = _build_run_mask_index_maps( analysis_mask_path=path_map["analysis_mask"], run_mask_dir=run_mask_dir, top_percent=target_top_percent, ) target_mask_source_dir = run_mask_dir canonical_runs_used = sorted( { canonical_run for subject in subjects for canonical_run in participant_run_map[subject].values() } ) protocol_c_df = pd.DataFrame() noise_ceiling_df = pd.DataFrame() n_fit_workers, blas_threads_per_worker = _resolve_num_fit_workers( requested=str(args.num_fit_workers), n_layers=int(len(layers)), ) if n_fit_workers > 1: print( f"[fit] Layer-parallel fitting enabled: workers={n_fit_workers}, " f"blas_threads/worker={blas_threads_per_worker}, layers={len(layers)}, " f"cpu_count={os.cpu_count()}", flush=True, ) if "C" in protocols: protocol_c_df = _evaluate_protocol_c( layers=layers, alpha=float(args.alpha), cross_subject_df=cross_subject_df, bold_path_map=bold_path_map, regressor_path_map=regressor_path_map, target_index_map_by_canonical_run=target_index_map_by_canonical_run, participant_run_map=participant_run_map, model_slug=str(args.model_slug), model_id=model_id, n_workers=n_fit_workers, blas_threads_per_worker=blas_threads_per_worker, ) protocol_c_df.to_csv(output_dir / "protocol_c_core_roi_scores.csv", index=False) noise_ceiling_df = _evaluate_protocol_c_noise_ceiling( cross_subject_df=cross_subject_df, bold_path_map=bold_path_map, target_index_map_by_canonical_run=target_index_map_by_canonical_run, participant_run_map=participant_run_map, ) noise_ceiling_df.to_csv(output_dir / "noise_ceiling_protocol_c_core_roi_scores.csv", index=False) combined_df = protocol_c_df.copy() combined_df.to_csv(output_dir / "core_roi_scores_all.csv", index=False) layer_summary_df, best_layer_df = _summarize_layers(scores_df=combined_df) layer_summary_df.to_csv(output_dir / "core_roi_layer_summary.csv", index=False) best_layer_df.to_csv(output_dir / "core_roi_best_layer_summary.csv", index=False) noise_ceiling_summary_df = _summarize_noise_ceiling(scores_df=noise_ceiling_df) noise_ceiling_summary_df.to_csv(output_dir / "noise_ceiling_core_roi_summary.csv", index=False) best_vs_noise_ceiling_df = _compare_best_layer_to_noise_ceiling( best_layer_df=best_layer_df, noise_ceiling_summary_df=noise_ceiling_summary_df, ) best_vs_noise_ceiling_df.to_csv(output_dir / "best_layer_vs_noise_ceiling.csv", index=False) run_summary = { "bootstrap_output_dir": str(bootstrap_output_dir), "output_dir": str(output_dir), "model_slug": str(args.model_slug), "model_id": model_id, "alpha": float(args.alpha), "subjects": subjects, "n_subjects": int(len(subjects)), "layers": [int(value) for value in layers], "n_layers": int(len(layers)), "protocols": sorted(protocols), "target_mask_mode": target_mask_mode, "target_top_percent": int(target_top_percent) if target_top_percent is not None else None, "target_mask_source_dir": str(target_mask_source_dir), "roi_mask_dir": str(roi_mask_dir) if roi_mask_dir is not None else None, "run_mask_dir": str(run_mask_dir) if run_mask_dir is not None else None, "run_top10_mask_dir": str(run_mask_dir) if target_top_percent == 10 and run_mask_dir is not None else None, "run_top25_mask_dir": str(run_mask_dir) if target_top_percent == 25 and run_mask_dir is not None else None, "participant_run_info_path": str(participant_run_info_path), "required_runs": [int(value) for value in required_runs], "canonical_runs_used": [int(value) for value in canonical_runs_used], "target_names_by_canonical_run": { str(canonical_run): list(target_index_map_by_canonical_run[canonical_run].keys()) for canonical_run in sorted(target_index_map_by_canonical_run) }, "n_protocol_c_rows": int(len(protocol_c_df)), "n_noise_ceiling_rows": int(len(noise_ceiling_df)), "n_total_rows": int(len(combined_df)), } write_json(output_dir / "fit_run_summary.json", run_summary) print("=" * 72) print("A1 fit complete") print(f"Model slug: {args.model_slug}") print(f"Model id: {model_id}") print(f"Subjects: {len(subjects)}") print(f"Layers: {len(layers)}") print(f"Alpha: {float(args.alpha)}") print(f"Target mask mode: {target_mask_mode}") print(f"Target mask source: {target_mask_source_dir}") print(f"Protocol C rows: {len(protocol_c_df)}") print(f"Noise ceiling rows: {len(noise_ceiling_df)}") print(f"Output directory: {output_dir}") print("=" * 72) if __name__ == "__main__": main()