"""Candidate loading for the BrainRL OpenEnv environment. Two candidate modes are supported: * ``atlas_parcels`` (default) reads the frozen manifest produced by ``prepare_parcels.py`` at ``configs/parcel_candidates.json``. This is the hackathon-target setup with ~200 Schaefer-style parcels. * ``roi_priors`` reads the small ``configs/region_priors.json`` fixture and is kept as the easy-curriculum smoke-test path. The loader never touches NIfTI files; pruning happens once in ``prepare_parcels.py`` and the environment treats the resulting manifest as read-only at episode time. """ from __future__ import annotations import json import os import random from dataclasses import dataclass from pathlib import Path from typing import Any PROJECT_ROOT = Path(__file__).resolve().parent DEFAULT_CONFIG_DIR = PROJECT_ROOT / "configs" DEFAULT_CONFIG_PATH = DEFAULT_CONFIG_DIR / "subset_config.yaml" DEFAULT_PARCEL_MANIFEST = DEFAULT_CONFIG_DIR / "parcel_candidates.json" DEFAULT_PRIOR_PATH = DEFAULT_CONFIG_DIR / "region_priors.json" @dataclass(frozen=True) class RegionCandidate: """Compact metadata + scoring parameters for one candidate region/parcel.""" region_id: str label: str hemisphere: str semantic_prior: float base_r2: float cost: float redundancy_group: str notes: str = "" atlas: str = "" network: str = "" sub_region: str = "" n_voxels: int = 0 prune_score: float = 0.0 def as_dict(self) -> dict[str, Any]: return { "region_id": self.region_id, "label": self.label, "hemisphere": self.hemisphere, "semantic_prior": float(self.semantic_prior), "base_r2": float(self.base_r2), "cost": float(self.cost), "redundancy_group": self.redundancy_group, "notes": self.notes, "atlas": self.atlas, "network": self.network, "sub_region": self.sub_region, "n_voxels": int(self.n_voxels), "prune_score": float(self.prune_score), } def prompt_dict(self) -> dict[str, Any]: """Compact view used by the prompt builder (no internal hints).""" return { "region_id": self.region_id, "label": self.label, "hemisphere": self.hemisphere, "network": self.network, "semantic_prior": float(round(self.semantic_prior, 4)), "base_r2_hint": float(round(self.base_r2, 4)), "cost": float(round(self.cost, 4)), "n_voxels": int(self.n_voxels), } @dataclass(frozen=True) class BrainSubset: """Loaded candidate set + task parameters.""" dataset_name: str candidates: list[RegionCandidate] selection_budget: int cost_penalty: float source: str candidate_mode: str atlas: str prompt_top_k: int @property def n_regions(self) -> int: return len(self.candidates) # --------------------------------------------------------------------------- # Config / I/O helpers # --------------------------------------------------------------------------- def _read_json(path: Path) -> dict[str, Any]: with path.open("r", encoding="utf-8") as handle: return json.load(handle) def config_dir() -> Path: """Resolve config directory, allowing HF Dataset bundles to override it.""" override = os.getenv("BRAINRL_CONFIG_DIR") return Path(override).expanduser() if override else DEFAULT_CONFIG_DIR def default_config_path() -> Path: return config_dir() / "subset_config.yaml" def default_parcel_manifest() -> Path: return config_dir() / "parcel_candidates.json" def default_prior_path() -> Path: return config_dir() / "region_priors.json" def _resolve_config_path(raw_path: str | Path, *, base_dir: Path) -> Path: path = Path(raw_path).expanduser() if path.is_absolute(): return path candidates = ( base_dir / path, PROJECT_ROOT / path, base_dir.parent / path, ) for candidate in candidates: if candidate.exists(): return candidate return base_dir / path def _parse_simple_yaml(path: Path) -> dict[str, Any]: """Tiny YAML reader so we don't drag PyYAML into the runtime deps.""" if not path.exists(): return {} config: dict[str, Any] = {} current_list_key: str | None = None with path.open("r", encoding="utf-8") as handle: for raw_line in handle: line = raw_line.split("#", 1)[0].rstrip() if not line.strip(): continue if line.startswith(" - ") and current_list_key: config.setdefault(current_list_key, []).append(line[4:].strip()) continue if ":" not in line: continue key, raw_value = line.split(":", 1) key = key.strip() value = raw_value.strip() if value == "": config[key] = [] current_list_key = key continue current_list_key = None if value.lower() in {"true", "false"}: config[key] = value.lower() == "true" else: try: config[key] = int(value) except ValueError: try: config[key] = float(value) except ValueError: config[key] = value.strip('"').strip("'") return config # --------------------------------------------------------------------------- # Atlas-parcel mode # --------------------------------------------------------------------------- def _candidates_from_parcel_manifest(payload: dict[str, Any]) -> list[RegionCandidate]: parcels = payload.get("candidates", []) candidates: list[RegionCandidate] = [] for entry in parcels: candidates.append( RegionCandidate( region_id=str(entry["region_id"]), label=str(entry.get("label", entry["region_id"])), hemisphere=str(entry.get("hemisphere", "unknown")), semantic_prior=float(entry.get("semantic_prior", 0.5)), base_r2=float(entry.get("base_r2", 0.04)), cost=float(entry.get("cost", 1.0)), redundancy_group=str(entry.get("redundancy_group", "association")), notes=str(entry.get("notes", entry.get("network", ""))), atlas=str(entry.get("atlas", payload.get("atlas", ""))), network=str(entry.get("network", "")), sub_region=str(entry.get("sub_region", "")), n_voxels=int(entry.get("n_voxels", 0)), prune_score=float(entry.get("prune_score", 0.0)), ) ) return candidates def _load_atlas_parcels(config: dict[str, Any]) -> tuple[BrainSubset, str]: base_dir = Path(str(config.get("_config_dir", config_dir()))).expanduser() manifest_path = _resolve_config_path( config.get("parcel_manifest_path", str(default_parcel_manifest())), base_dir=base_dir, ) if not manifest_path.exists(): raise FileNotFoundError( f"Parcel manifest not found at {manifest_path}. " "Run `python prepare_parcels.py` (see Makefile target `prepare`)." ) payload = _read_json(manifest_path) candidates = _candidates_from_parcel_manifest(payload) max_candidates = int(config.get("max_candidates", payload.get("max_candidates", len(candidates)))) candidates = candidates[: max(1, max_candidates)] selection_budget = int( config.get("selection_budget", payload.get("selection_budget", min(20, len(candidates)))) ) selection_budget = max(1, min(selection_budget, len(candidates))) prompt_top_k = int(config.get("prompt_top_k", payload.get("prompt_top_k", 30))) prompt_top_k = max(1, min(prompt_top_k, len(candidates))) cost_penalty = float(config.get("cost_penalty", payload.get("cost_penalty", 0.002))) atlas = str(payload.get("atlas", config.get("atlas_name", "schaefer200"))) subset = BrainSubset( dataset_name=str(config.get("dataset_name", "le_petit_prince_atlas_parcels")), candidates=candidates, selection_budget=selection_budget, cost_penalty=cost_penalty, source=f"parcel_manifest:{manifest_path.name}", candidate_mode="atlas_parcels", atlas=atlas, prompt_top_k=prompt_top_k, ) return subset, str(manifest_path) # --------------------------------------------------------------------------- # ROI-priors (legacy easy curriculum) # --------------------------------------------------------------------------- def _redundancy_group_for_roi(region_id: str) -> str: if "STS" in region_id or region_id == "TP": return "auditory_temporal" if region_id.startswith("BA"): return "inferior_frontal" return "association" def _load_roi_priors(config: dict[str, Any]) -> tuple[BrainSubset, str]: base_dir = Path(str(config.get("_config_dir", config_dir()))).expanduser() prior_path = _resolve_config_path( config.get("prior_path", str(default_prior_path())), base_dir=base_dir, ) if not prior_path.exists(): raise FileNotFoundError(f"ROI priors file not found at {prior_path}.") payload = _read_json(prior_path) priors = payload.get("priors", []) rng = random.Random(int(config.get("seed", 42))) candidates: list[RegionCandidate] = [] for prior in priors: region_id = str(prior["region_id"]) semantic_prior = float(prior.get("semantic_prior", 0.5)) jitter = rng.uniform(-0.01, 0.015) base_r2 = max(0.005, min(0.2, 0.025 + 0.09 * semantic_prior + jitter)) candidates.append( RegionCandidate( region_id=region_id, label=str(prior.get("label", region_id)), hemisphere=str(prior.get("hemisphere", "left")), semantic_prior=semantic_prior, base_r2=base_r2, cost=float(prior.get("cost", 1.0)), redundancy_group=_redundancy_group_for_roi(region_id), notes=str(prior.get("notes", "")), atlas="language_rois", network="language", n_voxels=0, prune_score=semantic_prior, ) ) candidates.sort(key=lambda item: item.semantic_prior, reverse=True) max_regions = int(config.get("max_regions", len(candidates))) candidates = candidates[: max(1, max_regions)] requested_budget = int(config.get("selection_budget", min(5, len(candidates)))) selection_budget = max(1, min(requested_budget, len(candidates))) subset = BrainSubset( dataset_name=str(config.get("dataset_name", "le_petit_prince_small_roi")), candidates=candidates, selection_budget=selection_budget, cost_penalty=float(config.get("cost_penalty", 0.002)), source=f"roi_priors:{prior_path.name}", candidate_mode="roi_priors", atlas="language_rois", prompt_top_k=int(config.get("prompt_top_k", len(candidates))), ) return subset, str(prior_path) # --------------------------------------------------------------------------- # Public API # --------------------------------------------------------------------------- def load_brain_subset(config_path: Path | str | None = None) -> BrainSubset: """Load a candidate set based on subset_config.yaml. Falls back to roi_priors mode if the parcel manifest is missing so the repository remains runnable without first running ``prepare_parcels.py``. """ resolved_config_path = Path(config_path).expanduser() if config_path else default_config_path() config = _parse_simple_yaml(resolved_config_path) config["_config_dir"] = str(resolved_config_path.parent) mode = str(config.get("candidate_mode", "atlas_parcels")).strip().lower() if mode == "atlas_parcels": try: subset, _ = _load_atlas_parcels(config) return subset except FileNotFoundError: # Graceful fallback so smoke tests work without a prepared manifest. subset, _ = _load_roi_priors(config) return subset if mode == "roi_priors": subset, _ = _load_roi_priors(config) return subset raise ValueError( f"Unknown candidate_mode={mode!r}. Expected 'atlas_parcels' or 'roi_priors'." ) def candidate_table(subset: BrainSubset) -> list[dict[str, Any]]: """Return JSON-serializable candidate metadata.""" return [candidate.as_dict() for candidate in subset.candidates]