| |
| """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: <bootstrap-output-dir>/fit_results/<model-slug>)", |
| ) |
| 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) |
|
|
| |
| 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() |
|
|
