"""Subject/condition-aware data splits for BrainRL. The Le Petit Prince derivative ships with ``participant_run_info.json`` that maps each subject's four runs to one of four conditions (``single_m``, ``single_f``, ``mixed_m``, ``mixed_f``). For the hackathon we typically train on a single condition (e.g. ``single_m``) and split subjects into train/test groups so generalization claims are honest. This module is intentionally tiny: pure stdlib, no dependency on the environment, so it can be used both at preparation time and at episode time. """ from __future__ import annotations import json import os import re from dataclasses import dataclass from pathlib import Path from typing import Iterable PROJECT_ROOT = Path(__file__).resolve().parent DEFAULT_PARTICIPANT_INFO = PROJECT_ROOT / "configs" / "participant_run_info.json" LEGACY_PARTICIPANT_INFO = PROJECT_ROOT.parent / "derivatives" / "participant_run_info.json" VALID_CONDITIONS: tuple[str, ...] = ("single_m", "single_f", "mixed_m", "mixed_f") @dataclass(frozen=True) class SubjectRunPair: """One concrete (subject, run, condition) tuple usable as an episode.""" subject_id: str run_id: str condition: str def episode_seed(self, base_seed: int = 0) -> int: h = abs(hash((self.subject_id, self.run_id, self.condition, int(base_seed)))) return h % (2**31 - 1) def as_dict(self) -> dict[str, str]: return { "subject_id": self.subject_id, "run_id": self.run_id, "condition": self.condition, } # --------------------------------------------------------------------------- # Loading + filtering # --------------------------------------------------------------------------- def load_participant_run_info(path: str | Path | None = None) -> dict[str, dict[str, str]]: """Load the raw subject -> run -> condition mapping.""" if path: info_path = Path(path) elif os.getenv("BRAINRL_CONFIG_DIR"): info_path = Path(os.environ["BRAINRL_CONFIG_DIR"]) / "participant_run_info.json" else: info_path = DEFAULT_PARTICIPANT_INFO info_path = info_path.expanduser() if path is None and not info_path.exists() and LEGACY_PARTICIPANT_INFO.exists(): info_path = LEGACY_PARTICIPANT_INFO if not info_path.exists(): raise FileNotFoundError( f"participant_run_info.json not found at {info_path}. " "Pass --participant-info to point at the correct file." ) with info_path.open("r", encoding="utf-8") as handle: payload = json.load(handle) return { str(subject): {str(run): str(cond) for run, cond in runs.items()} for subject, runs in payload.items() } def pairs_for_condition( info: dict[str, dict[str, str]], condition: str, ) -> list[SubjectRunPair]: """Return every (subject, run) tuple matching ``condition``.""" if condition not in VALID_CONDITIONS: raise ValueError( f"Unknown condition={condition!r}. Expected one of {VALID_CONDITIONS}." ) pairs: list[SubjectRunPair] = [] for subject, runs in info.items(): for run_id, cond in runs.items(): if cond == condition: pairs.append( SubjectRunPair(subject_id=subject, run_id=run_id, condition=cond) ) pairs.sort(key=lambda p: (p.subject_id, p.run_id)) return pairs # --------------------------------------------------------------------------- # Subject parsing # --------------------------------------------------------------------------- _SUB_RE = re.compile(r"^sub-(\d+)$") def _normalize_subject(token: str) -> str: token = token.strip() if not token: return "" if token.isdigit(): return f"sub-{int(token):02d}" match = _SUB_RE.match(token) if match: return f"sub-{int(match.group(1)):02d}" return token def parse_subject_spec(spec: str | None, available: Iterable[str]) -> list[str]: """Parse a CLI subject spec into a sorted list of subject ids. Accepts: * ``"sub-01:sub-20"`` – inclusive range * ``"sub-01,sub-05,sub-09"`` – explicit comma list * ``"01:20"`` / ``"01,05,09"`` – shorthand without ``sub-`` prefix * ``None`` / ``""`` / ``"all"`` – every available subject """ available_set = sorted({s for s in available}) if not spec or spec.lower() == "all": return list(available_set) if ":" in spec: start_raw, end_raw = spec.split(":", 1) start = _normalize_subject(start_raw) end = _normalize_subject(end_raw) ordered = sorted(available_set) try: start_idx = ordered.index(start) end_idx = ordered.index(end) except ValueError as exc: raise ValueError( f"Subject range {spec!r} does not match available subjects: {ordered}" ) from exc if start_idx > end_idx: start_idx, end_idx = end_idx, start_idx return ordered[start_idx : end_idx + 1] items = [_normalize_subject(t) for t in spec.split(",") if t.strip()] missing = [s for s in items if s not in available_set] if missing: raise ValueError( f"Subjects {missing} not in available list {available_set}." ) return sorted(items) def parse_exclude_spec(spec: str | None) -> set[str]: """Parse an exclusion spec into a set of normalized subject ids. Same syntax as ``parse_subject_spec`` but does not require the subjects to exist in any participant list - excluding a missing id is a no-op so that scripts stay robust if the corrupted-subject list drifts. """ if not spec: return set() if spec.lower() == "none": return set() out: set[str] = set() if ":" in spec: start_raw, end_raw = spec.split(":", 1) start = _normalize_subject(start_raw) end = _normalize_subject(end_raw) try: start_idx = int(start.split("-")[1]) end_idx = int(end.split("-")[1]) except (IndexError, ValueError) as exc: raise ValueError(f"Bad exclusion range {spec!r}.") from exc if start_idx > end_idx: start_idx, end_idx = end_idx, start_idx for i in range(start_idx, end_idx + 1): out.add(f"sub-{i:02d}") return out for token in spec.split(","): token = token.strip() if not token: continue out.add(_normalize_subject(token)) return out # --------------------------------------------------------------------------- # Splits # --------------------------------------------------------------------------- @dataclass(frozen=True) class ConditionSplit: condition: str train_pairs: list[SubjectRunPair] test_pairs: list[SubjectRunPair] train_subjects: list[str] test_subjects: list[str] excluded_subjects: list[str] def pairs_for(self, split: str) -> list[SubjectRunPair]: split = split.lower() if split == "train": return list(self.train_pairs) if split == "test": return list(self.test_pairs) if split == "all": return list(self.train_pairs) + list(self.test_pairs) raise ValueError(f"Unknown split={split!r}; expected train/test/all.") def summary(self) -> dict[str, object]: return { "condition": self.condition, "n_train_subjects": len(self.train_subjects), "n_test_subjects": len(self.test_subjects), "n_train_pairs": len(self.train_pairs), "n_test_pairs": len(self.test_pairs), "train_subjects": self.train_subjects, "test_subjects": self.test_subjects, "excluded_subjects": self.excluded_subjects, } def build_condition_split( *, condition: str, participant_info_path: str | Path | None = None, train_subjects_spec: str | None = None, test_subjects_spec: str | None = None, exclude_subjects: str | Iterable[str] | None = None, train_frac: float = 0.75, ) -> ConditionSplit: """High-level helper used by all CLI scripts. ``exclude_subjects`` can be a comma list / range / iterable of subject ids (e.g. ``"sub-03,sub-18"``). Excluded subjects are dropped before any train/test logic and never appear in either pair list, so corrupted or held-out subjects can be skipped consistently across the whole stack. """ info = load_participant_run_info(participant_info_path) all_pairs = pairs_for_condition(info, condition) if not all_pairs: raise ValueError(f"No (subject, run) pairs found for condition={condition!r}.") if isinstance(exclude_subjects, str) or exclude_subjects is None: excluded = parse_exclude_spec(exclude_subjects if isinstance(exclude_subjects, str) else None) else: excluded = {_normalize_subject(s) for s in exclude_subjects if s} if excluded: all_pairs = [p for p in all_pairs if p.subject_id not in excluded] if not all_pairs: raise ValueError( f"All subjects for condition={condition!r} were excluded by " f"exclude_subjects={sorted(excluded)}." ) available_subjects = sorted({p.subject_id for p in all_pairs}) if train_subjects_spec or test_subjects_spec: train_subjects = [ s for s in parse_subject_spec(train_subjects_spec, available_subjects) if s not in excluded ] if test_subjects_spec: test_subjects = [ s for s in parse_subject_spec(test_subjects_spec, available_subjects) if s not in excluded ] else: test_subjects = [s for s in available_subjects if s not in set(train_subjects)] if not train_subjects: train_subjects = available_subjects[: max(1, int(len(available_subjects) * train_frac))] if not test_subjects: test_subjects = [s for s in available_subjects if s not in set(train_subjects)] else: cutoff = max(1, int(round(len(available_subjects) * train_frac))) train_subjects = available_subjects[:cutoff] test_subjects = available_subjects[cutoff:] train_set = set(train_subjects) test_set = set(test_subjects) train_pairs = [p for p in all_pairs if p.subject_id in train_set] test_pairs = [p for p in all_pairs if p.subject_id in test_set] return ConditionSplit( condition=condition, train_pairs=train_pairs, test_pairs=test_pairs, train_subjects=sorted(train_subjects), test_subjects=sorted(test_subjects), excluded_subjects=sorted(excluded), )