#!/usr/bin/env python """Parcel pruning orchestrator for BrainRL. This is the *only* place atlas parcels are loaded, scored, and pruned. It runs once before training and writes a frozen candidate manifest to ``configs/parcel_candidates.json``. The OpenEnv environment then consumes that manifest and never touches NIfTI files at episode time. Pipeline stages: 1. Load atlas parcels (real Schaefer-200 via nilearn, or a deterministic synthetic fallback so the demo always runs). 2. Intersect parcels with the brain mask (voxel counts). 3. Compute cheap parcel scores: - voxel_count_quality (sigmoidal on n_voxels) - variance_score (deterministic per-parcel proxy) - semantic_prior (label/network heuristic) - cached_encoding_score (optional CSV from prior fits) 4. Apply ``min_voxels_per_parcel`` floor. 5. Rank by a weighted ``prune_score`` and keep the top ``max_candidates``. 6. Save a JSON manifest used by the OpenEnv environment. Example: python prepare_parcels.py \ --atlas schaefer200 \ --max-candidates 200 \ --min-voxels 20 \ --prompt-top-k 30 \ --selection-budget 20 \ --output configs/parcel_candidates.json """ from __future__ import annotations import argparse import csv import json import math import random import re from dataclasses import dataclass, field from pathlib import Path from typing import Any PROJECT_ROOT = Path(__file__).resolve().parent DEFAULT_OUTPUT = PROJECT_ROOT / "configs" / "parcel_candidates.json" # Network membership for Schaefer-200 7Networks. Each entry maps a network name # to (LH parcel count, RH parcel count). Counts approximate the canonical # Schaefer 2018 split and total exactly 200 across both hemispheres. SCHAEFER200_7NET_COUNTS: dict[str, tuple[int, int]] = { "Vis": (14, 14), "SomMot": (16, 16), "DorsAttn": (12, 11), "SalVentAttn": (11, 11), "Limbic": (4, 5), "Cont": (12, 14), "Default": (31, 29), } # Sub-region tokens used to make synthetic labels look anatomically plausible # and to give the semantic-prior heuristic something to grip. NETWORK_SUBREGIONS: dict[str, list[str]] = { "Vis": ["Striate", "ExStr", "ExStrSup", "ExStrInf"], "SomMot": ["S1", "M1", "Aud", "ParOper"], "DorsAttn": ["Post", "FEF", "PrCv", "PostC"], "SalVentAttn": ["ParOper", "Med", "FrOperIns", "PFCl"], "Limbic": ["TempPole", "OFC"], "Cont": ["Par", "Temp", "PFCl", "PFCmp", "Cing", "pCun"], "Default": ["Temp", "TempPar", "PFCv", "PFCm", "PFCdPFCm", "pCunPCC", "IPL", "PHC"], } # --------------------------------------------------------------------------- # Atlas loading # --------------------------------------------------------------------------- @dataclass class RawParcel: """Pre-pruning parcel record.""" label_index: int label: str hemisphere: str network: str sub_region: str n_voxels: int notes: str = "" def base_region_id(self) -> str: return f"parcel_{self.label_index:03d}" def _network_for_index(idx_within_hemi: int, counts: dict[str, tuple[int, int]], hemi_idx: int) -> tuple[str, int]: """Map a 1-based parcel index inside a hemisphere to (network, slot).""" cursor = 0 for network, (lh, rh) in counts.items(): n = lh if hemi_idx == 0 else rh if idx_within_hemi <= cursor + n: return network, idx_within_hemi - cursor cursor += n return "Default", idx_within_hemi def load_synthetic_schaefer200(seed: int = 42) -> list[RawParcel]: """Build a deterministic Schaefer-200-style parcel set without downloads.""" rng = random.Random(seed) parcels: list[RawParcel] = [] label_index = 0 for hemi_idx, hemi in enumerate(("LH", "RH")): per_hemi = sum(lh if hemi_idx == 0 else rh for lh, rh in SCHAEFER200_7NET_COUNTS.values()) for i in range(1, per_hemi + 1): label_index += 1 network, slot = _network_for_index(i, SCHAEFER200_7NET_COUNTS, hemi_idx) sub_options = NETWORK_SUBREGIONS[network] sub_region = sub_options[(slot - 1) % len(sub_options)] sub_idx = ((slot - 1) // len(sub_options)) + 1 label = f"7Networks_{hemi}_{network}_{sub_region}_{sub_idx}" n_voxels = max(8, int(rng.gauss(330, 120))) parcels.append( RawParcel( label_index=label_index, label=label, hemisphere="left" if hemi == "LH" else "right", network=network, sub_region=sub_region, n_voxels=n_voxels, ) ) return parcels def load_nilearn_schaefer(n_rois: int = 200, yeo_networks: int = 7) -> list[RawParcel]: """Fetch the real Schaefer-2018 atlas via nilearn and count voxels.""" try: import numpy as np from nilearn import datasets, image except ImportError as exc: # pragma: no cover raise SystemExit( "Real-atlas mode requires nilearn. Install with `pip install -e .[atlas]` " "or rerun with --atlas-source synthetic." ) from exc atlas = datasets.fetch_atlas_schaefer_2018( n_rois=n_rois, yeo_networks=yeo_networks, resolution_mm=2, ) labels_img = image.load_img(atlas.maps) label_data = np.asarray(labels_img.dataobj).astype(int) raw_labels: list[str] = [] for entry in atlas.labels: if isinstance(entry, bytes): raw_labels.append(entry.decode("utf-8")) else: raw_labels.append(str(entry)) parcels: list[RawParcel] = [] for label_index in range(1, n_rois + 1): if label_index - 1 >= len(raw_labels): break label = raw_labels[label_index - 1] n_voxels = int((label_data == label_index).sum()) hemi_match = re.search(r"_(LH|RH)_", label) hemisphere = "left" if hemi_match and hemi_match.group(1) == "LH" else "right" network_match = re.search( r"_(Vis|SomMot|DorsAttn|SalVentAttn|Limbic|Cont|Default)_", label ) network = network_match.group(1) if network_match else "Default" sub_match = re.search(rf"_{network}_([A-Za-z]+)", label) sub_region = sub_match.group(1) if sub_match else "Unknown" parcels.append( RawParcel( label_index=label_index, label=label, hemisphere=hemisphere, network=network, sub_region=sub_region, n_voxels=n_voxels, notes="schaefer2018_7networks", ) ) return parcels # --------------------------------------------------------------------------- # Scoring # --------------------------------------------------------------------------- # Higher = more relevant to auditory/language fMRI prediction. These priors are # intentionally small, transparent heuristics; the RL agent is supposed to # discover better orderings, not be handed the answer. NETWORK_LANGUAGE_PRIORS: dict[str, float] = { "Default": 0.78, "Cont": 0.66, "SalVentAttn": 0.62, "Limbic": 0.55, "DorsAttn": 0.40, "SomMot": 0.45, "Vis": 0.20, } LANGUAGE_KEYWORDS: dict[str, float] = { "Aud": 0.30, "Temp": 0.20, "TempPar": 0.22, "TempPole": 0.18, "IPL": 0.12, "PFCl": 0.10, "FrOperIns": 0.14, "ParOper": 0.10, "PFCv": 0.08, } def semantic_prior_for(parcel: RawParcel) -> float: base = NETWORK_LANGUAGE_PRIORS.get(parcel.network, 0.4) label = parcel.label bonus = 0.0 for keyword, value in LANGUAGE_KEYWORDS.items(): if keyword in label: bonus = max(bonus, value) if parcel.hemisphere == "left": bonus += 0.05 return float(min(0.99, base + bonus)) def voxel_count_quality(n_voxels: int) -> float: """Soft sigmoid: ~0 at 10 voxels, ~0.5 at ~150, ~1 at very large parcels.""" return float(1.0 / (1.0 + math.exp(-(n_voxels - 150) / 80.0))) def deterministic_variance_score(parcel: RawParcel, seed: int) -> float: """Stable per-parcel pseudo-variance proxy (no NIfTI required).""" rng = random.Random(f"{seed}:{parcel.label}") base = 0.45 if parcel.network in {"Default", "Cont", "SalVentAttn"} else 0.30 return float(max(0.0, min(1.0, base + rng.uniform(-0.10, 0.20)))) def base_r2_estimate(parcel: RawParcel, semantic: float, variance: float) -> float: """Cheap predicted-encoding score used as a hint, not a target.""" raw = 0.020 + 0.060 * semantic + 0.020 * variance return float(max(0.005, min(0.20, raw))) # --------------------------------------------------------------------------- # Cached encoding scores (optional) # --------------------------------------------------------------------------- def load_cached_scores(paths: list[str]) -> dict[str, float]: """Read optional encoding scores keyed by region label or parcel id.""" region_keys = ("roi_name", "region_id", "target", "parcel", "label") score_keys = ("mean_r2", "mean_corr", "score", "r2") scores: dict[str, float] = {} for raw_path in paths: path = Path(raw_path).expanduser() if not path.exists() or path.suffix.lower() != ".csv": continue with path.open("r", encoding="utf-8", newline="") as handle: for row in csv.DictReader(handle): key = None for region_key in region_keys: candidate = row.get(region_key) if candidate: key = str(candidate).strip() break if not key: continue value: float | None = None for score_key in score_keys: raw = row.get(score_key) if raw in (None, ""): continue try: parsed = float(raw) except ValueError: continue if score_key == "mean_corr": parsed = max(0.0, parsed) ** 2 value = max(0.0, min(0.25, parsed)) break if value is not None: scores[key] = value return scores # --------------------------------------------------------------------------- # Pruning # --------------------------------------------------------------------------- @dataclass class PruneWeights: activity: float = 0.40 semantic: float = 0.30 encoding: float = 0.20 voxel_quality: float = 0.10 @dataclass class ScoredParcel: parcel: RawParcel semantic_prior: float variance_score: float voxel_quality: float cached_encoding: float base_r2: float prune_score: float cost: float redundancy_group: str network_priority_kept: bool = False extras: dict[str, Any] = field(default_factory=dict) def redundancy_group_for(parcel: RawParcel) -> str: label = parcel.label if "Aud" in label or "Temp" in label or "TempPar" in label or "TempPole" in label: return "auditory_temporal" if "PFCl" in label or "FrOperIns" in label or "PFCv" in label: return "inferior_frontal" if "Vis" in label: return "visual" if "SomMot" in label or "S1" in label or "M1" in label: return "somatomotor" return "association" def cost_for(parcel: RawParcel) -> float: if parcel.network in {"Vis", "SomMot"}: return 1.05 if parcel.network == "Limbic": return 1.15 return 1.0 def score_parcel( parcel: RawParcel, *, weights: PruneWeights, cached_scores: dict[str, float], seed: int, ) -> ScoredParcel: semantic = semantic_prior_for(parcel) variance = deterministic_variance_score(parcel, seed) voxq = voxel_count_quality(parcel.n_voxels) encoding_score = 0.0 for key in (parcel.label, parcel.base_region_id(), parcel.sub_region): if key in cached_scores: encoding_score = cached_scores[key] break base_r2 = base_r2_estimate(parcel, semantic, variance) if encoding_score > 0: base_r2 = float(max(base_r2, encoding_score)) prune_score = ( weights.activity * variance + weights.semantic * semantic + weights.encoding * (encoding_score / 0.20 if encoding_score > 0 else semantic * 0.5) + weights.voxel_quality * voxq ) return ScoredParcel( parcel=parcel, semantic_prior=semantic, variance_score=variance, voxel_quality=voxq, cached_encoding=encoding_score, base_r2=base_r2, prune_score=float(prune_score), cost=cost_for(parcel), redundancy_group=redundancy_group_for(parcel), ) def prune_parcels( parcels: list[RawParcel], *, max_candidates: int, min_voxels: int, weights: PruneWeights, cached_scores: dict[str, float], seed: int, keep_language_floor: int = 30, ) -> list[ScoredParcel]: scored = [ score_parcel(p, weights=weights, cached_scores=cached_scores, seed=seed) for p in parcels ] eligible = [s for s in scored if s.parcel.n_voxels >= min_voxels] if not eligible: eligible = scored eligible.sort(key=lambda s: s.prune_score, reverse=True) # Reserve a floor for clearly language-relevant parcels even if their # prune score is borderline. Keeps the easy-curriculum signal alive. language_pool = sorted( (s for s in eligible if s.redundancy_group in {"auditory_temporal", "inferior_frontal"}), key=lambda s: s.prune_score, reverse=True, ) floor = min(keep_language_floor, len(language_pool), max_candidates) forced = {id(s) for s in language_pool[:floor]} kept: list[ScoredParcel] = [] for s in eligible: if id(s) in forced: s.network_priority_kept = True kept.append(s) if len(kept) >= max_candidates: break if len(kept) < max_candidates: for s in eligible: if id(s) in forced: continue kept.append(s) if len(kept) >= max_candidates: break kept.sort(key=lambda s: s.prune_score, reverse=True) return kept # --------------------------------------------------------------------------- # Manifest writing # --------------------------------------------------------------------------- def manifest_payload( *, atlas: str, selection_budget: int, prompt_top_k: int, cost_penalty: float, weights: PruneWeights, pruned: list[ScoredParcel], source_notes: str, ) -> dict[str, Any]: candidates = [] for idx, scored in enumerate(pruned): parcel = scored.parcel candidates.append( { "region_id": parcel.base_region_id(), "atlas": atlas, "label": parcel.label, "hemisphere": parcel.hemisphere, "network": parcel.network, "sub_region": parcel.sub_region, "n_voxels": int(parcel.n_voxels), "semantic_prior": float(round(scored.semantic_prior, 4)), "variance_score": float(round(scored.variance_score, 4)), "voxel_quality": float(round(scored.voxel_quality, 4)), "cached_encoding": float(round(scored.cached_encoding, 4)), "base_r2": float(round(scored.base_r2, 4)), "cost": float(round(scored.cost, 4)), "prune_score": float(round(scored.prune_score, 4)), "redundancy_group": scored.redundancy_group, "rank": idx, "kept_by_floor": bool(scored.network_priority_kept), } ) return { "atlas": atlas, "candidate_mode": "atlas_parcels", "max_candidates": len(candidates), "selection_budget": int(selection_budget), "prompt_top_k": int(prompt_top_k), "cost_penalty": float(cost_penalty), "prune_weights": { "activity": weights.activity, "semantic": weights.semantic, "encoding": weights.encoding, "voxel_quality": weights.voxel_quality, }, "source_notes": source_notes, "candidates": candidates, } # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( description="Load an atlas, prune parcels, and freeze the candidate manifest." ) parser.add_argument( "--atlas", type=str, default="schaefer200", help="Atlas identifier; only schaefer200 (7 networks) is implemented.", ) parser.add_argument( "--atlas-source", choices=("auto", "nilearn", "synthetic"), default="auto", help="Where to load parcels from. 'auto' tries nilearn then falls back.", ) parser.add_argument("--max-candidates", type=int, default=200) parser.add_argument("--min-voxels", type=int, default=20) parser.add_argument("--prompt-top-k", type=int, default=30) parser.add_argument("--selection-budget", type=int, default=20) parser.add_argument("--cost-penalty", type=float, default=0.002) parser.add_argument("--seed", type=int, default=42) parser.add_argument( "--output", type=str, default=str(DEFAULT_OUTPUT), help="Where to write the candidate manifest JSON.", ) parser.add_argument( "--cached-scores", action="append", default=None, help="Optional CSV with cached encoding scores; can be repeated.", ) parser.add_argument("--w-activity", type=float, default=0.40) parser.add_argument("--w-semantic", type=float, default=0.30) parser.add_argument("--w-encoding", type=float, default=0.20) parser.add_argument("--w-voxel-quality", type=float, default=0.10) parser.add_argument( "--keep-language-floor", type=int, default=30, help="Force-keep this many top auditory/language parcels.", ) parser.add_argument( "--quiet", action="store_true", help="Suppress per-stage logs.", ) return parser def _log(quiet: bool, message: str) -> None: if not quiet: print(message) def main() -> None: args = build_parser().parse_args() if args.atlas != "schaefer200": raise SystemExit( f"Unsupported atlas={args.atlas}. Only 'schaefer200' is implemented." ) weights = PruneWeights( activity=float(args.w_activity), semantic=float(args.w_semantic), encoding=float(args.w_encoding), voxel_quality=float(args.w_voxel_quality), ) cached_paths = list(args.cached_scores or []) cached_scores = load_cached_scores(cached_paths) _log(args.quiet, f"[1/6] Loaded {len(cached_scores)} cached encoding score rows.") parcels: list[RawParcel] = [] source_notes = "" if args.atlas_source in ("auto", "nilearn"): try: parcels = load_nilearn_schaefer() source_notes = "nilearn:schaefer2018_7networks_n200_2mm" _log(args.quiet, f"[2/6] Loaded {len(parcels)} parcels from nilearn Schaefer-2018.") except SystemExit: if args.atlas_source == "nilearn": raise parcels = [] except Exception as exc: # noqa: BLE001 if args.atlas_source == "nilearn": raise SystemExit(f"nilearn atlas load failed: {exc}") from exc _log(args.quiet, f"[2/6] nilearn atlas unavailable ({exc}); falling back to synthetic.") parcels = [] if not parcels: parcels = load_synthetic_schaefer200(seed=int(args.seed)) source_notes = "synthetic:schaefer200_7networks" _log(args.quiet, f"[2/6] Loaded {len(parcels)} synthetic Schaefer-200-style parcels.") _log(args.quiet, f"[3/6] Computed voxel-count metadata for {len(parcels)} parcels.") _log(args.quiet, "[4/6] Scoring parcels with weighted prune formula.") pruned = prune_parcels( parcels, max_candidates=int(args.max_candidates), min_voxels=int(args.min_voxels), weights=weights, cached_scores=cached_scores, seed=int(args.seed), keep_language_floor=int(args.keep_language_floor), ) _log(args.quiet, f"[5/6] Pruned to {len(pruned)} candidates (min_voxels>={args.min_voxels}).") payload = manifest_payload( atlas=args.atlas, selection_budget=int(args.selection_budget), prompt_top_k=int(args.prompt_top_k), cost_penalty=float(args.cost_penalty), weights=weights, pruned=pruned, source_notes=source_notes, ) output_path = Path(args.output).resolve() output_path.parent.mkdir(parents=True, exist_ok=True) with output_path.open("w", encoding="utf-8") as handle: json.dump(payload, handle, indent=2) handle.write("\n") if not args.quiet: top_preview = ", ".join( f"{c['region_id']}({c['network']})" for c in payload["candidates"][:5] ) print(f"[6/6] Wrote {len(payload['candidates'])} parcel candidates -> {output_path}") print(f" atlas={args.atlas} budget={args.selection_budget} " f"prompt_top_k={args.prompt_top_k} source={source_notes}") print(f" top: {top_preview}") if __name__ == "__main__": main()