#!/usr/bin/env python """Serve a small local GUI for approving feature labels.""" from __future__ import annotations import argparse import csv import json import os import re import sys import threading import tempfile from datetime import datetime, timezone from http import HTTPStatus from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from io import BytesIO from pathlib import Path from typing import Any from urllib.parse import parse_qs, urlparse CONFIDENCE_VALUES = ("high", "medium", "low", "unclear") INTERPRETATION_TYPES = ( "form", "word", "phrase", "sentence", "long-range", "global", "position", "sophisticated", "unclear", ) SUBJECT = "UTS03" N_VERTICES = 20484 HEMISPHERES = ("lh", "rh") BRAIN_ENCODING = Path(__file__).resolve().parents[1] DEFAULT_SCRATCH_ROOT = Path(os.environ.get("BRAINENCODING_SCRATCH_ROOT", "/storage/scratch1/8/whuang409")) DEFAULT_PYCORTEX_DB = DEFAULT_SCRATCH_ROOT / "visionLM_data/ds003020/derivatives/pycortex-db" DEFAULT_FREESURFER_SUBJECTS = DEFAULT_SCRATCH_ROOT / "visionLM_data/ds003020/derivatives/freesurfer_subjdir" DEFAULT_FSAVERAGE5 = Path("/storage/project/r-aivanova7-0/shared/.env/freesurfer/subjects/fsaverage5") DEFAULT_NOISE_CEILING = ( Path("/storage/project/r-aivanova7-0/whuang409/loopedEncoding/results/noise_ceiling_fsaverage5") / "UTS03_wheretheressmoke_noise_ceiling_fsaverage5.npz" ) DEFAULT_ICA_PREDICTIVITY = ( BRAIN_ENCODING / "results/gemma-2-2b_layer13_lookback128_hook_resid_post_ica_lasso_fixed_a0.02207_ridge_sweep/predictivity/" "UTS03_gemma-2-2b_ica_lasso_fixed_lookback_128_hook_resid_post_layer_13_a0.02207_ridge_sweep_predictivity.npz" ) DEFAULT_SAE_PREDICTIVITY = ( BRAIN_ENCODING / "results/gemma-2-2b_layer13_lookback128_hook_resid_post_sae_dense_alpha_sweep/predictivity/" "UTS03_gemma-2-2b_sae_dense_lookback_128_hook_resid_post_layer_13_alpha_sweep_predictivity.npz" ) STOPWORDS = { "a", "an", "and", "are", "as", "at", "be", "by", "for", "from", "in", "is", "it", "of", "on", "or", "that", "the", "this", "to", "with", } SEMANTIC_GROUPS = ( ( "visual", "visual vision see saw seen look looks looking watch watching image picture photo camera light color dark bright face body scene object eye eyes draw drawing art mirror", ), ( "social", "social person people family mother father mom dad brother sister aunt child children boy girl woman man wife husband friend partner relationship married kinship", ), ( "mental", "mental mind think thought know knew believe realize understand remember memory worry hope fear afraid ashamed self identity confidence reflection inner", ), ( "speech", "speech say said tell told ask asked answer quote quoted dialogue conversation voice talk called utterance instruction command", ), ( "motion", "motion move moving walk run drive fly fall jump landing turn pull push go come arrive leave direction forward back toward", ), ( "time", "time date year years day night morning early late week month birthday age first start opening beginning after before", ), ( "number", "number count counted counting one two three four five seven ten twenty seventy eighty ninety hundred thousand million digit math", ), ( "body", "body hand hands arm leg head eye eyes mouth face skin blood hurt pain sick illness disease epilepsy doctor medical hospital", ), ( "place", "place location room house home school college university office stage car taxi boat store camp prison church city", ), ( "emotion", "emotion feel feeling felt love joy happy sad grief cry crying shame fear worried anxious angry regret relief", ), ( "religion", "religion religious church god soul souls prayer pray temple faith sin spiritual exorcism offering", ), ( "academic", "academic school college university class student teacher phd mit harvard princeton exam study science database actuarial", ), ( "object", "object thing stuff phone key keys card box door car bag clothes furniture couch table book paper photograph", ), ( "function", "function grammar article determiner preposition conjunction pronoun deictic demonstrative modal infinitive filler pause um uh and the for", ), ) SEMANTIC_EXPANSIONS: dict[str, set[str]] = {} SEMANTIC_GROUP_TOKENS: dict[str, set[str]] = {} for _group_name, _words in SEMANTIC_GROUPS: _tokens = {_group_name, *re.findall(r"[a-z0-9]+", _words.lower())} SEMANTIC_GROUP_TOKENS[_group_name] = _tokens for _token in _tokens: SEMANTIC_EXPANSIONS.setdefault(_token, set()).update(_tokens) def main() -> int: args = parse_args() packets = load_packets(args.packets) labels = load_labels(args.labels) highlights = load_highlights(args.highlights) server = LabelServer( (args.host, int(args.port)), Handler, packets=packets, packet_paths=args.packets, labels=labels, highlights=highlights, labels_path=args.labels, highlights_path=args.highlights, stage_a_root=args.stage_a_root, layer=args.layer, encoding_dir=args.encoding_dir, ica_contribution_dir=args.ica_contribution_dir, sae_coeff_dir=args.sae_coeff_dir, sae_erf_csv=args.sae_erf_csv, noise_ceiling_npz=args.noise_ceiling_npz, ica_predictivity_npz=args.ica_predictivity_npz, sae_predictivity_npz=args.sae_predictivity_npz, ncsnr_threshold=args.ncsnr_threshold, mask_reliable_vertices=args.mask_reliable_vertices, roi_dir=args.roi_dir, brainmap_cache=args.brainmap_cache, pycortex_db=args.pycortex_db, freesurfer_subjects=args.freesurfer_subjects, fsaverage5=args.fsaverage5, flatmap_height=args.flatmap_height, ) print(f"serving http://{args.host}:{args.port}", flush=True) print(f"packets={len(packets)} labels={args.labels}", flush=True) print(f"highlights={args.highlights}", flush=True) server.serve_forever() return 0 def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) root = Path(__file__).resolve().parents[1] label_root = root / "results/story_ica_labels_UTS03_layer13" encoding_root = root / "results/story_ica_encoding_UTS03_layer13/encoding/layer_13/UTS03" parser.add_argument("--packets", type=Path, nargs="+", default=[label_root / "ic_label_packets.jsonl"]) parser.add_argument("--labels", type=Path, default=label_root / "ic_labels.json") parser.add_argument("--highlights", type=Path, default=label_root / "ic_highlights.json") parser.add_argument("--stage-a-root", type=Path, default=root / "results/story_ica_stage_a_UTS03_layer13") parser.add_argument("--encoding-dir", type=Path, default=encoding_root) parser.add_argument("--ica-contribution-dir", type=Path, default=None) parser.add_argument("--sae-coeff-dir", type=Path, default=None) parser.add_argument("--sae-erf-csv", type=Path, default=None) parser.add_argument("--noise-ceiling-npz", type=Path, default=DEFAULT_NOISE_CEILING) parser.add_argument("--ica-predictivity-npz", type=Path, default=DEFAULT_ICA_PREDICTIVITY) parser.add_argument("--sae-predictivity-npz", type=Path, default=DEFAULT_SAE_PREDICTIVITY) parser.add_argument("--ncsnr-threshold", type=float, default=0.20) parser.add_argument("--mask-reliable-vertices", action=argparse.BooleanOptionalAction, default=True) parser.add_argument("--roi-dir", type=Path, default=root / "masks/requested_roi_groups_20484") parser.add_argument("--brainmap-cache", type=Path, default=label_root / "ic_brainmaps") parser.add_argument("--pycortex-db", type=Path, default=DEFAULT_PYCORTEX_DB) parser.add_argument("--freesurfer-subjects", type=Path, default=DEFAULT_FREESURFER_SUBJECTS) parser.add_argument("--fsaverage5", type=Path, default=DEFAULT_FSAVERAGE5) parser.add_argument("--flatmap-height", type=int, default=1024) parser.add_argument("--layer", type=int, default=13) parser.add_argument("--host", default="127.0.0.1") parser.add_argument("--port", type=int, default=8765) return parser.parse_args() def load_packets(paths: list[Path]) -> list[dict[str, Any]]: packets = [] for path in paths: if not path.exists(): continue with path.open(encoding="utf-8") as handle: for line in handle: if line.strip(): packet = json.loads(line) packet.setdefault("feature_source", "ica") packet.setdefault("feature_label", "IC" if packet["feature_source"] == "ica" else str(packet["feature_source"]).upper()) packet["_search_evidence_text"] = search_evidence_text(packet) packet["_search_evidence_tokens"] = semantic_tokens(packet["_search_evidence_text"]) packets.append(packet) return sorted( packets, key=lambda packet: ( 1 if packet.get("feature_source", "ica") == "ica" else 0, int(packet.get("encoding_summary", {}).get("selected_voxels") or 0), int(packet.get("encoding_summary", {}).get("activation_count") or 0), int(packet.get("encoding_summary", {}).get("selected_rows") or 0), -int(packet.get("component") or 0), 1 if packet.get("side") == "positive" else 0, ), reverse=True, ) def load_labels(path: Path) -> dict[str, Any]: if path.exists(): return json.loads(path.read_text(encoding="utf-8")) return {"metadata": {"created_utc": datetime.now(timezone.utc).isoformat()}, "labels": {}} def load_highlights(path: Path) -> dict[str, Any]: if path.exists(): payload = json.loads(path.read_text(encoding="utf-8")) else: payload = {"metadata": {"created_utc": datetime.now(timezone.utc).isoformat()}, "components": [], "features": []} payload.setdefault("metadata", {}) payload["components"] = sorted({int(component) for component in payload.get("components", [])}) legacy = {f"ica:{component}" for component in payload["components"]} payload["features"] = sorted({str(feature) for feature in payload.get("features", [])} | legacy) return payload class LabelServer(ThreadingHTTPServer): def __init__( self, *args: Any, packets: list[dict[str, Any]], packet_paths: list[Path], labels: dict[str, Any], highlights: dict[str, Any], labels_path: Path, highlights_path: Path, stage_a_root: Path, layer: int, encoding_dir: Path, ica_contribution_dir: Path | None, sae_coeff_dir: Path | None, sae_erf_csv: Path | None, noise_ceiling_npz: Path | None, ica_predictivity_npz: Path | None, sae_predictivity_npz: Path | None, ncsnr_threshold: float, mask_reliable_vertices: bool, roi_dir: Path, brainmap_cache: Path, pycortex_db: Path, freesurfer_subjects: Path, fsaverage5: Path, flatmap_height: int, **kwargs: Any, ) -> None: super().__init__(*args, **kwargs) self.packets = packets self.packet_paths = packet_paths self.packet_mtimes = self.current_packet_mtimes() self.labels = labels self.highlights = highlights self.labels_path = labels_path self.highlights_path = highlights_path self.labels_mtime = labels_path.stat().st_mtime if labels_path.exists() else 0.0 self.highlights_mtime = highlights_path.stat().st_mtime if highlights_path.exists() else 0.0 self.stage_a_root = stage_a_root self.layer = int(layer) self.encoding_dir = encoding_dir self.ica_contribution_dir = ica_contribution_dir self.sae_coeff_dir = sae_coeff_dir self.sae_erf_csv = sae_erf_csv self.noise_ceiling_npz = noise_ceiling_npz self.ica_predictivity_npz = ica_predictivity_npz self.sae_predictivity_npz = sae_predictivity_npz self.ncsnr_threshold = float(ncsnr_threshold) self.mask_reliable_vertices = bool(mask_reliable_vertices) self.roi_dir = roi_dir self.brainmap_cache = brainmap_cache self.pycortex_db = pycortex_db self.freesurfer_subjects = freesurfer_subjects self.fsaverage5 = fsaverage5 self.flatmap_height = int(flatmap_height) self._brainmap_context: ICBrainmapContext | None = None self._brainmap_lock = threading.Lock() self._highlight_lock = threading.Lock() def current_packet_mtimes(self) -> dict[str, float]: return {str(path): path.stat().st_mtime for path in self.packet_paths if path.exists()} def reload_packets_if_changed(self) -> None: mtimes = self.current_packet_mtimes() if mtimes != self.packet_mtimes: self.packets = load_packets(self.packet_paths) self.packet_mtimes = mtimes def reload_labels_if_changed(self) -> None: if not self.labels_path.exists(): return mtime = self.labels_path.stat().st_mtime if mtime > self.labels_mtime: self.labels = load_labels(self.labels_path) self.labels_mtime = mtime def reload_highlights_if_changed(self) -> None: if not self.highlights_path.exists(): return mtime = self.highlights_path.stat().st_mtime if mtime > self.highlights_mtime: self.highlights = load_highlights(self.highlights_path) self.highlights_mtime = mtime def brainmap_context(self) -> "ICBrainmapContext": if self._brainmap_context is not None: return self._brainmap_context with self._brainmap_lock: if self._brainmap_context is None: self._brainmap_context = ICBrainmapContext( layer=self.layer, encoding_dir=self.encoding_dir, ica_contribution_dir=self.ica_contribution_dir, sae_coeff_dir=self.sae_coeff_dir, noise_ceiling_npz=self.noise_ceiling_npz, ica_predictivity_npz=self.ica_predictivity_npz, sae_predictivity_npz=self.sae_predictivity_npz, ncsnr_threshold=self.ncsnr_threshold, mask_reliable_vertices=self.mask_reliable_vertices, stage_a_root=self.stage_a_root, roi_dir=self.roi_dir, output_dir=self.brainmap_cache, pycortex_db=self.pycortex_db, freesurfer_subjects=self.freesurfer_subjects, fsaverage5=self.fsaverage5, flatmap_height=self.flatmap_height, ) return self._brainmap_context class Handler(BaseHTTPRequestHandler): server: LabelServer def do_GET(self) -> None: # noqa: N802 parsed = urlparse(self.path) if parsed.path == "/": self.send_html(INDEX_HTML) return if parsed.path == "/api/state": query = parse_qs(parsed.query) self.server.reload_packets_if_changed() self.server.reload_labels_if_changed() self.server.reload_highlights_if_changed() self.send_json(self.state(query)) return if parsed.path == "/api/labels": self.server.reload_labels_if_changed() self.send_json(self.server.labels) return if parsed.path == "/api/highlights": self.server.reload_highlights_if_changed() self.send_json(self.server.highlights) return if parsed.path in {"/api/ic-brainmap", "/api/feature-brainmap"}: self.send_ic_brainmap(parse_qs(parsed.query)) return if parsed.path == "/api/erf-brainmap": self.send_erf_brainmap(parse_qs(parsed.query)) return self.send_error(HTTPStatus.NOT_FOUND) def do_POST(self) -> None: # noqa: N802 parsed = urlparse(self.path) if parsed.path == "/api/label": self.save_label() return if parsed.path == "/api/highlight": self.save_highlight() return self.send_error(HTTPStatus.NOT_FOUND) def read_json_payload(self) -> dict[str, Any]: length = int(self.headers.get("content-length") or 0) return json.loads(self.rfile.read(length).decode("utf-8")) def save_label(self) -> None: self.server.reload_labels_if_changed() payload = self.read_json_payload() packet_id = str(payload["packet_id"]) packet = next((item for item in self.server.packets if str(item.get("packet_id")) == packet_id), {}) item = { "packet_id": packet_id, "component": int(payload["component"]), "feature_source": str(payload.get("feature_source") or "ica"), "side": str(payload["side"]), "working_label": str(payload.get("working_label", "")).strip(), "confidence": normalize_choice(payload.get("confidence"), CONFIDENCE_VALUES), "interpretation_type": normalize_choice(payload.get("interpretation_type"), INTERPRETATION_TYPES), "approved": bool(payload.get("approved", False)), "notes": str(payload.get("notes", "")).strip(), "updated_utc": datetime.now(timezone.utc).isoformat(), } item["semantic_tags"] = normalize_semantic_tags(payload.get("semantic_tags")) or broad_semantic_tags(packet, item) self.server.labels.setdefault("labels", {})[packet_id] = item save_json_atomic(self.server.labels_path, self.server.labels) self.server.labels_mtime = self.server.labels_path.stat().st_mtime self.send_json({"ok": True, "label": item}) def save_highlight(self) -> None: self.server.reload_highlights_if_changed() payload = self.read_json_payload() component = int(payload["component"]) feature_source = str(payload.get("feature_source") or "ica") feature_key = make_feature_key(feature_source, component) highlighted = bool(payload.get("highlighted", False)) with self.server._highlight_lock: components = {int(value) for value in self.server.highlights.get("components", [])} features = {str(value) for value in self.server.highlights.get("features", [])} if highlighted: features.add(feature_key) if feature_source == "ica": components.add(component) else: features.discard(feature_key) if feature_source == "ica": components.discard(component) self.server.highlights["components"] = sorted(components) self.server.highlights["features"] = sorted(features) self.server.highlights.setdefault("metadata", {})["updated_utc"] = datetime.now(timezone.utc).isoformat() save_json_atomic(self.server.highlights_path, self.server.highlights) self.server.highlights_mtime = self.server.highlights_path.stat().st_mtime self.send_json({"ok": True, "component": component, "feature_source": feature_source, "highlighted": highlighted}) def state(self, query: dict[str, list[str]]) -> dict[str, Any]: packets = self.server.packets labels = self.server.labels.get("labels", {}) highlighted_components = {int(value) for value in self.server.highlights.get("components", [])} highlighted_features = {str(value) for value in self.server.highlights.get("features", [])} ica_erf = load_erf(self.server.stage_a_root, self.server.layer) sae_erf = load_sae_erf(self.server.sae_erf_csv, self.server.sae_coeff_dir, self.server.layer) q = (query.get("q", [""])[0] or "").lower().strip() status = query.get("status", ["all"])[0] side = query.get("side", ["all"])[0] source = query.get("source", ["all"])[0] interpretation_type = query.get("interpretation_type", ["all"])[0] semantic_tag_filter = query.get("semantic_tag", ["all"])[0] highlight_filter = query.get("highlight_filter", ["all"])[0] sort_by = query.get("sort_by", ["selected_voxels"])[0] sort_then = query.get("sort_then", ["erf"])[0] sort_dir = query.get("sort_dir", ["desc"])[0] limit = min(500, max(1, int(query.get("limit", ["100"])[0] or 100))) offset = max(0, int(query.get("offset", ["0"])[0] or 0)) filtered = [] for packet in packets: packet_id = packet["packet_id"] label = labels.get(packet_id, {}) component = int(packet["component"]) packet_source = str(packet.get("feature_source") or "ica") feature_key = make_feature_key(packet_source, component) is_highlighted = feature_key in highlighted_features or ( packet_source == "ica" and component in highlighted_components ) if source != "all" and packet_source != source: continue if side != "all" and packet.get("side") != side: continue if status == "approved" and not label.get("approved"): continue if status == "unapproved" and label.get("approved"): continue if interpretation_type != "all" and label.get("interpretation_type", "unclear") != interpretation_type: continue if highlight_filter == "highlighted" and not is_highlighted: continue if highlight_filter == "unhighlighted" and is_highlighted: continue label = dict(label) label["semantic_tags"] = normalize_semantic_tags(label.get("semantic_tags")) or broad_semantic_tags(packet, label) if semantic_tag_filter != "all" and semantic_tag_filter not in label["semantic_tags"]: continue search_score = semantic_search_score(q, packet, label) if q else 0.0 if q and search_score <= 0.0: continue packet_copy = {key: value for key, value in packet.items() if not key.startswith("_")} packet_copy["label"] = label packet_copy["highlighted"] = is_highlighted packet_copy["feature_key"] = feature_key packet_copy["search_score"] = search_score if packet_source == "ica" and component in ica_erf: packet_copy["erf"] = ica_erf[component] elif packet_source == "sae" and component in sae_erf: packet_copy["erf"] = sae_erf[component] filtered.append(packet_copy) if q and sort_by != "search_score": sort_then = sort_by sort_by = "search_score" filtered = sort_packets(filtered, sort_by=sort_by, sort_then=sort_then, sort_dir=sort_dir) page = filtered[offset : offset + limit] return { "packets": page, "total_packets": len(packets), "shown_packets": len(filtered), "returned_packets": len(page), "offset": offset, "limit": limit, "approved_count": sum(1 for value in labels.values() if value.get("approved")), "highlighted_count": len(highlighted_components), "confidence_values": list(CONFIDENCE_VALUES), "interpretation_types": list(INTERPRETATION_TYPES), "semantic_tags": sorted(SEMANTIC_GROUP_TOKENS), "sources": sorted({str(packet.get("feature_source") or "ica") for packet in packets}), "sort_values": ["search_score", "selected_voxels", "activation_count", "max_score", "erf", "confidence", "component"], } def send_json(self, payload: Any) -> None: raw = json.dumps(payload, ensure_ascii=False).encode("utf-8") self.send_response(HTTPStatus.OK) self.send_header("content-type", "application/json; charset=utf-8") self.send_header("content-length", str(len(raw))) self.end_headers() self.wfile.write(raw) def send_html(self, html: str) -> None: raw = html.encode("utf-8") self.send_response(HTTPStatus.OK) self.send_header("content-type", "text/html; charset=utf-8") self.send_header("content-length", str(len(raw))) self.end_headers() self.wfile.write(raw) def send_ic_brainmap(self, query: dict[str, list[str]]) -> None: try: component = int(query.get("component", [""])[0]) side = query.get("side", [""])[0] hemi = query.get("hemi", [""])[0] source = query.get("source", ["ica"])[0] or "ica" style = query.get("style", ["flatmap"])[0] or "flatmap" mask_reliable = query_bool( query, "mask_reliable", default=self.server.mask_reliable_vertices, ) mask_significant = query_bool( query, "mask_significant", default=mask_reliable, ) except ValueError: self.send_error(HTTPStatus.BAD_REQUEST, "Invalid feature brainmap query.") return if source not in {"ica", "sae"} or side not in {"positive", "negative"} or style not in {"flatmap", "litcoder"}: self.send_error(HTTPStatus.BAD_REQUEST, "Invalid feature brainmap source, side, or hemisphere.") return if style == "flatmap" and hemi not in HEMISPHERES: self.send_error(HTTPStatus.BAD_REQUEST, "Invalid feature brainmap hemisphere.") return try: raw = self.server.brainmap_context().render( source, component, side, hemi, style=style, mask_reliable_vertices=mask_reliable, mask_significant_vertices=mask_significant, ) except Exception as exc: # noqa: BLE001 self.send_error(HTTPStatus.INTERNAL_SERVER_ERROR, f"Could not render feature brainmap: {exc}") return self.send_response(HTTPStatus.OK) self.send_header("content-type", "image/png") self.send_header("content-length", str(len(raw))) self.send_header("cache-control", "no-store") self.end_headers() self.wfile.write(raw) def send_erf_brainmap(self, query: dict[str, list[str]]) -> None: hemi = query.get("hemi", [""])[0] style = query.get("style", ["flatmap"])[0] or "flatmap" if style not in {"flatmap", "litcoder"}: self.send_error(HTTPStatus.BAD_REQUEST, "Invalid ERF brainmap style.") return if style == "flatmap" and hemi not in HEMISPHERES: self.send_error(HTTPStatus.BAD_REQUEST, "Invalid ERF brainmap hemisphere.") return try: raw = self.server.brainmap_context().render_erf(hemi, style=style) except Exception as exc: # noqa: BLE001 self.send_error(HTTPStatus.INTERNAL_SERVER_ERROR, f"Could not render ERF brainmap: {exc}") return self.send_response(HTTPStatus.OK) self.send_header("content-type", "image/png") self.send_header("content-length", str(len(raw))) self.send_header("cache-control", "no-store") self.end_headers() self.wfile.write(raw) def log_message(self, format: str, *args: Any) -> None: return def normalize_choice(value: Any, allowed: tuple[str, ...]) -> str: text = str(value or "unclear") return text if text in allowed else "unclear" def make_feature_key(source: str, component: int) -> str: return f"{source}:{int(component)}" def semantic_search_score(query: str, packet: dict[str, Any], label: dict[str, Any]) -> float: query = query.lower().strip() if not query: return 0.0 label_text = search_label_text(packet, label) evidence_text = str(packet.get("_search_evidence_text") or search_evidence_text(packet)) all_text = f"{label_text} {evidence_text}".lower() if query in all_text: score = 30.0 if query in label_text.lower() else 18.0 else: score = 0.0 query_tokens = semantic_tokens(query) if not query_tokens: return score label_tokens = semantic_tokens(label_text) evidence_tokens = set(packet.get("_search_evidence_tokens") or semantic_tokens(evidence_text)) all_tokens = label_tokens | evidence_tokens expanded_terms = expanded_query_terms(query_tokens) tags = set(label.get("semantic_tags") or broad_semantic_tags(packet, label)) required = [token for token in query_tokens if token not in STOPWORDS] if required and not any(search_token_matches(token, label_tokens, evidence_tokens, tags) for token in required): return score for token in query_tokens: if token in tags: score += 25.0 if token in label_tokens: score += 8.0 elif token in evidence_tokens: score += 4.0 else: prefix_match = token_prefix_match(token, all_tokens) if prefix_match: score += 2.0 for token in expanded_terms - query_tokens: if token in label_tokens: score += 2.0 if len(query_tokens) > 1: covered = sum(1 for token in query_tokens if search_token_matches(token, label_tokens, evidence_tokens, tags)) score += 3.0 * covered / max(1, len(query_tokens)) return score def search_token_matches(token: str, label_tokens: set[str], evidence_tokens: set[str], tags: set[str]) -> bool: if token in tags or token in label_tokens or token in evidence_tokens: return True if token_prefix_match(token, label_tokens | evidence_tokens): return True return bool(SEMANTIC_EXPANSIONS.get(token, set()) & label_tokens) def broad_semantic_tags(packet: dict[str, Any], label: dict[str, Any], *, max_tags: int = 3) -> list[str]: label_text = " ".join(str(label.get(key) or "") for key in ("working_label", "notes")) label_tokens = semantic_tokens(label_text) evidence_tokens = set(packet.get("_search_evidence_tokens") or semantic_tokens(search_evidence_text(packet))) scored = [] for group, raw_tokens in SEMANTIC_GROUP_TOKENS.items(): group_tokens = {normalize_search_token(token) for token in raw_tokens} group_tokens = {token for token in group_tokens if token and token not in STOPWORDS} label_hits = group_tokens & label_tokens evidence_hits = group_tokens & evidence_tokens score = 3.0 * len(label_hits) + 0.35 * min(len(evidence_hits), 8) if group in label_tokens: score += 8.0 if score >= 2.5 and (label_hits or score >= 5.0): scored.append((score, group)) scored.sort(key=lambda item: (-item[0], item[1])) return [group for _score, group in scored[: max(0, int(max_tags))]] def normalize_semantic_tags(value: Any) -> list[str]: allowed = set(SEMANTIC_GROUP_TOKENS) if value is None: return [] if isinstance(value, str): raw_values = re.split(r"[,;\s]+", value) elif isinstance(value, (list, tuple, set)): raw_values = list(value) else: raw_values = [value] tags = [] for raw in raw_values: tag = normalize_search_token(str(raw)) if tag in allowed and tag not in tags: tags.append(tag) return tags def search_label_text(packet: dict[str, Any], label: dict[str, Any]) -> str: return " ".join( str(value or "") for value in ( packet.get("packet_id"), packet.get("feature_label"), packet.get("feature_source"), packet.get("side"), packet.get("component"), label.get("working_label"), label.get("notes"), " ".join(label.get("semantic_tags") or []), label.get("confidence"), label.get("interpretation_type"), ) ) def search_evidence_text(packet: dict[str, Any]) -> str: parts = [] for key in ("strong_examples", "near_zero_examples", "opposite_examples"): for example in packet.get(key) or []: parts.extend( str(example.get(field) or "") for field in ( "story", "target_text", "snippet", "sentence_text", "left_text", "right_text", "context_text", ) ) erf = packet.get("erf") or {} parts.extend(str(value or "") for value in erf.values()) return " ".join(parts) def semantic_tokens(text: str) -> set[str]: tokens = set() for raw in re.findall(r"[a-z0-9']+", text.lower()): token = normalize_search_token(raw) if token and token not in STOPWORDS: tokens.add(token) return tokens def expanded_query_terms(tokens: set[str]) -> set[str]: expanded = set(tokens) for token in tokens: expanded.update(SEMANTIC_EXPANSIONS.get(token, ())) return {normalize_search_token(token) for token in expanded if normalize_search_token(token)} def normalize_search_token(token: str) -> str: token = token.lower().strip("'") if token.endswith("'s"): token = token[:-2] if len(token) > 6 and token.endswith("ing"): token = token[:-3] elif len(token) > 5 and token.endswith("ed"): token = token[:-2] elif len(token) > 5 and token.endswith("es"): token = token[:-2] elif len(token) > 4 and token.endswith("s") and not token.endswith(("ss", "us", "is")): token = token[:-1] return token def token_prefix_match(query_token: str, tokens: set[str]) -> bool: if len(query_token) < 4: return False return any(token.startswith(query_token) or query_token.startswith(token) for token in tokens if len(token) >= 4) def save_json_atomic(path: Path, payload: dict[str, Any]) -> None: path.parent.mkdir(parents=True, exist_ok=True) with tempfile.NamedTemporaryFile("w", encoding="utf-8", dir=path.parent, delete=False) as handle: json.dump(payload, handle, indent=2, sort_keys=True) handle.write("\n") tmp = Path(handle.name) tmp.replace(path) def query_bool(query: dict[str, list[str]], key: str, *, default: bool) -> bool: value = query.get(key, [None])[0] if value is None or value == "": return bool(default) return str(value).strip().lower() not in {"0", "false", "no", "off"} def sort_packets(packets: list[dict[str, Any]], *, sort_by: str, sort_then: str, sort_dir: str) -> list[dict[str, Any]]: reverse = sort_dir != "asc" def packet_key(packet: dict[str, Any]) -> tuple[float, float, float, int, int]: return ( sort_value(packet, sort_by), sort_value(packet, sort_then), sort_value(packet, "selected_voxels"), -int(packet.get("component") or 0), 1 if packet.get("side") == "positive" else 0, ) pinned = [packet for packet in packets if packet.get("highlighted")] unpinned = [packet for packet in packets if not packet.get("highlighted")] return sorted(pinned, key=packet_key, reverse=reverse) + sorted(unpinned, key=packet_key, reverse=reverse) def sort_value(packet: dict[str, Any], key: str) -> float: if key == "selected_voxels": return float(packet.get("encoding_summary", {}).get("selected_voxels") or 0) if key == "activation_count": summary = packet.get("encoding_summary", {}) return float(summary.get("activation_count") or summary.get("selected_rows") or 0) if key == "max_score": summary = packet.get("encoding_summary", {}) return float(summary.get("max_coefficient") or summary.get("max_score") or 0) if key == "erf": erf = packet.get("erf") or {} try: return float(erf.get("mean_erf")) except (TypeError, ValueError): return -1.0 if key == "confidence": confidence = str(packet.get("label", {}).get("confidence") or "unclear") return {"high": 3.0, "medium": 2.0, "low": 1.0, "unclear": 0.0}.get(confidence, 0.0) if key == "search_score": return float(packet.get("search_score") or 0.0) if key == "component": return float(packet.get("component") or 0) return 0.0 def png_bytes(image: Any) -> bytes: buffer = BytesIO() image.save(buffer, format="PNG") return buffer.getvalue() def load_erf(stage_a_root: Path, layer: int) -> dict[int, dict[str, str]]: for path in ( stage_a_root / "plots" / f"layer_{layer:02d}_component_erf.csv", stage_a_root / "erf" / f"layer_{layer:02d}_component_erf.csv", ): rows = load_erf_csv(path) if rows: return rows return {} def load_sae_erf(path: Path | None, sae_coeff_dir: Path | None, layer: int) -> dict[int, dict[str, str]]: if path is None and sae_coeff_dir is not None: path = sae_coeff_dir.parent / "erf" / f"layer_{layer:02d}_sae_component_erf.csv" if path is None: return {} return load_erf_csv(path) def load_erf_csv(path: Path) -> dict[int, dict[str, str]]: if not path.exists(): return {} rows = {} with path.open(newline="", encoding="utf-8") as handle: for row in csv.DictReader(handle): rows[int(row["component"])] = dict(row) return rows class ICBrainmapContext: def __init__( self, *, layer: int, encoding_dir: Path, ica_contribution_dir: Path | None, sae_coeff_dir: Path | None, noise_ceiling_npz: Path | None, ica_predictivity_npz: Path | None, sae_predictivity_npz: Path | None, ncsnr_threshold: float, mask_reliable_vertices: bool, stage_a_root: Path, roi_dir: Path, output_dir: Path, pycortex_db: Path, freesurfer_subjects: Path, fsaverage5: Path, flatmap_height: int, ) -> None: scripts_dir = Path(__file__).resolve().parent if str(scripts_dir) not in sys.path: sys.path.insert(0, str(scripts_dir)) import numpy as np from PIL import Image, ImageDraw from plot_predictivity_brainmaps import ( HEMI_HALF, HEMI_LABEL, composite_rgba_on_white, flatmap_roi_masks, load_subject_flatmap, load_subject_roi_masks, native_to_flatmap, project_fsaverage5_to_native, ) from flatmap_plotting import ( colorize_positive, colorize_scalar, crop_box_for_half, half_mask, load_font, render_colorbar, render_positive_colorbar, text_size, ) from litcoder_style_plotting import litcoder_style_png_bytes from roi_overlay import draw_roi_legend, overlay_roi_borders self.np = np self.Image = Image self.ImageDraw = ImageDraw self.HEMI_HALF = HEMI_HALF self.HEMI_LABEL = HEMI_LABEL self.composite_rgba_on_white = composite_rgba_on_white self.native_to_flatmap = native_to_flatmap self.project_fsaverage5_to_native = project_fsaverage5_to_native self.colorize_positive = colorize_positive self.colorize_scalar = colorize_scalar self.crop_box_for_half = crop_box_for_half self.half_mask = half_mask self.load_font = load_font self.render_colorbar = render_colorbar self.render_positive_colorbar = render_positive_colorbar self.text_size = text_size self.draw_roi_legend = draw_roi_legend self.overlay_roi_borders = overlay_roi_borders self.litcoder_style_png_bytes = litcoder_style_png_bytes self.layer = int(layer) self.output_dir = output_dir self.fsaverage5 = fsaverage5 self.flatmap = load_subject_flatmap(SUBJECT, pycortex_db, freesurfer_subjects, fsaverage5, flatmap_height) roi_masks = load_subject_roi_masks(roi_dir, SUBJECT) flat_roi_masks = flatmap_roi_masks(roi_masks, self.flatmap) self.roi_groups = self.requested_roi_groups(flat_roi_masks) self.encoding_dir = encoding_dir self.ica_contribution_dir = ica_contribution_dir self.ica_predictivity_npz = ica_predictivity_npz self.sae_predictivity_npz = sae_predictivity_npz self.component_side_values, self.vmax = self.load_component_side_values(encoding_dir / "selected_ic_records.csv") self.ica_contribution_ids, self.ica_contribution_values, self.ica_contribution_vmax, self.ica_contribution_manifest = ( self.load_ica_contribution_maps(ica_contribution_dir, encoding_dir) ) self.ica_contribution_signature = self.ica_contribution_file_signature() self.sae_coeff_dir = sae_coeff_dir self.sae_feature_ids, self.sae_coefficients, self.sae_vmax, self.sae_coefficient_scale = self.load_sae_coefficients(sae_coeff_dir) self.ncsnr_threshold = float(ncsnr_threshold) self.reliability_mask = self.load_ncsnr_mask(noise_ceiling_npz, self.ncsnr_threshold) self.ica_significant_mask = self.load_positive_predictivity_mask(ica_predictivity_npz) self.ica_predictivity_mtime = self.file_mtime(ica_predictivity_npz) self.sae_significant_mask = self.load_positive_predictivity_mask(sae_predictivity_npz) self.sae_predictivity_mtime = self.file_mtime(sae_predictivity_npz) self.mask_reliable_vertices = bool(mask_reliable_vertices) self.erf_values, self.erf_vmax, self.erf_voxel_count = self.load_top5_erf_values( encoding_dir / "selected_ic_records.csv", self.component_erf_csv(stage_a_root, self.layer), ) self.flatmap_background_rgba = (226, 230, 235, 255) self.unselected_voxel_rgba = (0, 0, 0, 255) def file_mtime(self, path: Path | None) -> float: return path.stat().st_mtime if path is not None and path.exists() else 0.0 def requested_roi_groups(self, flat_roi_masks: dict[str, Any]) -> list[dict[str, object]]: specs = [ ("AC", "auditory cortex", (0, 220, 255), ("AC",)), ("Language", "temporal/parietal/frontal", (255, 150, 0), ("temporal_language", "parietal_language", "language_frontal")), ("High-level visual", "functional visual ROIs", (255, 35, 210), ("high_level_visual",)), ("Theory of mind", "approximate ToM network", (170, 105, 255), ("theory_of_mind",)), ("Multiple demand", "frontoparietal MD", (0, 180, 85), ("multiple_demand",)), ("Default mode", "Yeo default mode", (255, 60, 100), ("default_mode",)), ] groups = [] for label, description, color, keys in specs: present = [key for key in keys if key in flat_roi_masks] if not present: continue mask = self.np.zeros(self.flatmap.mask.shape, dtype=bool) for key in present: mask |= flat_roi_masks[key] if not self.np.any(mask): continue groups.append( { "label": label, "description": description, "color": color, "rois": present, "mask": mask, } ) return groups def load_component_side_values(self, path: Path) -> tuple[dict[tuple[int, str], dict[int, float]], float]: if not path.exists(): return {}, 1.0 values: dict[tuple[int, str], dict[int, float]] = {} all_values = [] with path.open(newline="", encoding="utf-8") as handle: for row in csv.DictReader(handle): coef = float(row["signed_coef_sum"]) if coef == 0: continue side = "positive" if coef > 0 else "negative" magnitude = abs(coef) component = int(row["ic"]) vertex = int(row["voxel_index"]) if not 0 <= vertex < N_VERTICES: continue component_values = values.setdefault((component, side), {}) component_values[vertex] = max(component_values.get(vertex, 0.0), magnitude) all_values.append(magnitude) if not all_values: return values, 1.0 vmax = float(self.np.percentile(self.np.asarray(all_values, dtype=self.np.float32), 99.0)) if not self.np.isfinite(vmax) or vmax <= 0: vmax = 1.0 return values, vmax def ica_contribution_candidates(self, path: Path | None, encoding_dir: Path) -> list[Path]: candidates = [] if path is not None: candidates.append(path) candidates.extend( [ encoding_dir / "signed_unique_contribution_maps", encoding_dir.parent / "signed_unique_contribution_maps", BRAIN_ENCODING / "results/gemma_pile10k_ica_labels_layer13_residpost/signed_unique_contribution_maps", BRAIN_ENCODING / "results/gemma_ica_labels_UTS03_layer13_residpost/signed_unique_contribution_maps", ] ) return candidates def ica_contribution_file_signature(self) -> tuple[tuple[str, float], ...]: out = [] for directory in self.ica_contribution_candidates(self.ica_contribution_dir, self.encoding_dir): for name in ( "component_ids.npy", "signed_unique_partial_correlation_thresholded.npy", "signed_unique_partial_correlation.npy", "manifest.json", ): path = directory / name if path.exists(): out.append((str(path), path.stat().st_mtime)) return tuple(out) def reload_ica_contribution_maps_if_changed(self) -> None: signature = self.ica_contribution_file_signature() if signature != self.ica_contribution_signature: self.ica_contribution_ids, self.ica_contribution_values, self.ica_contribution_vmax, self.ica_contribution_manifest = ( self.load_ica_contribution_maps(self.ica_contribution_dir, self.encoding_dir) ) self.ica_contribution_signature = signature def reload_predictivity_masks_if_changed(self) -> None: ica_mtime = self.file_mtime(self.ica_predictivity_npz) if ica_mtime != self.ica_predictivity_mtime: self.ica_significant_mask = self.load_positive_predictivity_mask(self.ica_predictivity_npz) self.ica_predictivity_mtime = ica_mtime sae_mtime = self.file_mtime(self.sae_predictivity_npz) if sae_mtime != self.sae_predictivity_mtime: self.sae_significant_mask = self.load_positive_predictivity_mask(self.sae_predictivity_npz) self.sae_predictivity_mtime = sae_mtime def load_ica_contribution_maps(self, path: Path | None, encoding_dir: Path) -> tuple[Any, Any, float, dict[str, Any]]: candidates = self.ica_contribution_candidates(path, encoding_dir) chosen = None values_path = None for candidate in candidates: ids_candidate = candidate / "component_ids.npy" thresholded_candidate = candidate / "signed_unique_partial_correlation_thresholded.npy" raw_candidate = candidate / "signed_unique_partial_correlation.npy" value_candidate = thresholded_candidate if thresholded_candidate.exists() else raw_candidate if ids_candidate.exists() and value_candidate.exists(): chosen = candidate values_path = value_candidate break if chosen is None or values_path is None: return None, None, 1.0, {} ids_path = chosen / "component_ids.npy" component_ids = self.np.load(ids_path, mmap_mode="r") values = self.np.load(values_path, mmap_mode="r") manifest = {} manifest_path = chosen / "manifest.json" if manifest_path.exists(): try: manifest = json.loads(manifest_path.read_text(encoding="utf-8")) except Exception: manifest = {} vmax = float(manifest.get("vmax_abs_p99_thresholded") or 1.0) if not self.np.isfinite(vmax) or vmax <= 0: nonzero = self.np.asarray(values[self.np.abs(values) > 0], dtype=self.np.float32) vmax = float(self.np.percentile(self.np.abs(nonzero), 99.0)) if nonzero.size else 1.0 if not self.np.isfinite(vmax) or vmax <= 0: vmax = 1.0 manifest["loaded_path"] = str(chosen) manifest["loaded_values"] = values_path.name return component_ids, values, vmax, manifest def load_sae_coefficients(self, path: Path | None) -> tuple[Any, Any, float, float]: if path is None: return None, None, 1.0, 1.0 ids_path = path / "feature_ids.npy" coef_path = path / "coef_magnitude.npy" quantized_path = path / "coef_magnitude_uint8.npy" manifest_path = path / "manifest.json" if not ids_path.exists(): return None, None, 1.0, 1.0 use_quantized = not coef_path.exists() and quantized_path.exists() if not coef_path.exists() and not use_quantized: return None, None, 1.0, 1.0 feature_ids = self.np.load(ids_path, mmap_mode="r") coefficients = self.np.load(quantized_path if use_quantized else coef_path, mmap_mode="r") vmax = 1.0 scale = 1.0 if manifest_path.exists(): try: manifest = json.loads(manifest_path.read_text(encoding="utf-8")) vmax = float(manifest.get("vmax") or manifest.get("scale_vmax") or vmax) if use_quantized: scale_denominator = float(manifest.get("scale_denominator") or 255.0) scale_vmax = float(manifest.get("scale_vmax") or vmax) scale = scale_vmax / scale_denominator if scale_denominator > 0 else 1.0 except Exception: vmax = 1.0 if not self.np.isfinite(vmax) or vmax <= 0: vmax = 1.0 return feature_ids, coefficients, vmax, scale def load_ncsnr_mask(self, path: Path | None, threshold: float) -> Any: if path is None or not path.exists(): return None payload = self.np.load(path, allow_pickle=False) if "ncsnr" in payload.files: ncsnr = self.np.asarray(payload["ncsnr"], dtype=self.np.float32) elif "signal_power" in payload.files and "noise_power" in payload.files: signal = self.np.asarray(payload["signal_power"], dtype=self.np.float32) noise = self.np.asarray(payload["noise_power"], dtype=self.np.float32) ncsnr = self.np.sqrt( self.np.divide( self.np.maximum(signal, 0.0), noise, out=self.np.zeros_like(signal, dtype=self.np.float32), where=noise > 0, ) ) else: return None if ncsnr.shape != (N_VERTICES,): raise ValueError(f"{path} NCSNR array has shape {ncsnr.shape}, expected {(N_VERTICES,)}") mask = self.np.isfinite(ncsnr) & (ncsnr > float(threshold)) if "coverage_mask" in payload.files: mask &= self.np.asarray(payload["coverage_mask"], dtype=bool) return mask def load_positive_predictivity_mask(self, path: Path | None) -> Any: if path is None or not path.exists(): return None payload = self.np.load(path, allow_pickle=False) if "significant_mask" not in payload.files or "correlations" not in payload.files: return None significant = self.np.asarray(payload["significant_mask"], dtype=bool) correlations = self.np.asarray(payload["correlations"], dtype=self.np.float32) if significant.shape != (N_VERTICES,) or correlations.shape != (N_VERTICES,): raise ValueError(f"{path} predictivity arrays have shape {significant.shape}/{correlations.shape}, expected {(N_VERTICES,)}") return significant & self.np.isfinite(correlations) & (correlations > 0.0) def load_component_erfs(self, path: Path) -> dict[int, float]: if not path.exists(): raise FileNotFoundError(f"Missing component ERF file: {path}") out: dict[int, float] = {} with path.open(newline="", encoding="utf-8") as handle: for row in csv.DictReader(handle): value = float(row["mean_erf"]) if self.np.isfinite(value): out[int(row["component"])] = value return out def component_erf_csv(self, stage_a_root: Path, layer: int) -> Path: candidates = [ stage_a_root / "plots" / f"layer_{int(layer):02d}_component_erf.csv", stage_a_root / "erf" / f"layer_{int(layer):02d}_component_erf.csv", ] return next((path for path in candidates if path.exists()), candidates[0]) def load_top5_erf_values(self, selected_path: Path, erf_path: Path) -> tuple[Any, float, int]: if not selected_path.exists() or not erf_path.exists(): return self.np.zeros(N_VERTICES, dtype=self.np.float32), 11.0, 0 component_erfs = self.load_component_erfs(erf_path) ranked_by_vertex: dict[int, list[tuple[float, float]]] = {} with selected_path.open(newline="", encoding="utf-8") as handle: for row in csv.DictReader(handle): vertex = int(row["voxel_index"]) if not 0 <= vertex < N_VERTICES: continue component = int(row["ic"]) erf = component_erfs.get(component) if erf is None: continue score = abs(float(row["signed_coef_sum"])) ranked_by_vertex.setdefault(vertex, []).append((score, erf)) values = self.np.zeros(N_VERTICES, dtype=self.np.float32) voxel_count = 0 for vertex, entries in ranked_by_vertex.items(): top = sorted(entries, key=lambda item: item[0], reverse=True)[:5] if not top: continue values[vertex] = float(self.np.mean([erf for _score, erf in top])) voxel_count += 1 return values, 11.0, voxel_count def render( self, source: str, component: int, side: str, hemi: str, *, style: str = "flatmap", mask_reliable_vertices: bool | None = None, mask_significant_vertices: bool | None = None, ) -> bytes: values = self.np.zeros(N_VERTICES, dtype=self.np.float32) metric_label = "ridge coefficient magnitude" scale_label = "ridge coefficient magnitude, shared p99 scale" signed_scale = False intrinsic_mask_labels = [] if source == "sae": self.reload_predictivity_masks_if_changed() if side != "positive": raise ValueError("SAE brainmaps are only defined for the positive side.") if self.sae_feature_ids is None or self.sae_coefficients is None: self.sae_feature_ids, self.sae_coefficients, self.sae_vmax, self.sae_coefficient_scale = self.load_sae_coefficients(self.sae_coeff_dir) if self.sae_feature_ids is None or self.sae_coefficients is None: raise FileNotFoundError("SAE coefficient maps are not available yet.") matches = self.np.flatnonzero(self.sae_feature_ids == int(component)) if matches.size == 0: raise KeyError(f"SAE feature {component} has no coefficient map.") values = self.np.asarray(self.sae_coefficients[int(matches[0])], dtype=self.np.float32) * float(self.sae_coefficient_scale) vmax = self.sae_vmax else: self.reload_ica_contribution_maps_if_changed() self.reload_predictivity_masks_if_changed() if self.ica_contribution_ids is not None and self.ica_contribution_values is not None: matches = self.np.flatnonzero(self.ica_contribution_ids == int(component)) if matches.size == 0: raise KeyError(f"ICA component {component} has no signed contribution map.") values = self.np.asarray(self.ica_contribution_values[int(matches[0])], dtype=self.np.float32) vmax = self.ica_contribution_vmax signed_scale = True metric_label = "signed unique partial r" scale_label = "signed IC-specific held-out contribution" if self.ica_contribution_manifest.get("loaded_values", "").endswith("_thresholded.npy"): alpha = self.ica_contribution_manifest.get("permutation_alpha") if alpha is not None: intrinsic_mask_labels.append(f"IC permutation p<={float(alpha):g}") else: intrinsic_mask_labels.append("IC permutation null") else: component_values = self.component_side_values.get((int(component), side), {}) for vertex, value in component_values.items(): values[vertex] = value vmax = self.vmax use_reliability_mask = self.mask_reliable_vertices if mask_reliable_vertices is None else bool(mask_reliable_vertices) use_significant_mask = use_reliability_mask if mask_significant_vertices is None else bool(mask_significant_vertices) mask_labels = list(intrinsic_mask_labels) if use_reliability_mask or use_significant_mask: values = self.np.asarray(values, dtype=self.np.float32).copy() if use_reliability_mask: if self.reliability_mask is not None: values[~self.reliability_mask] = 0.0 mask_labels.append(f"NCSNR>{self.ncsnr_threshold:g}") if use_significant_mask: significant_mask = self.sae_significant_mask if source == "sae" else self.ica_significant_mask if significant_mask is not None: values[~significant_mask] = 0.0 mask_labels.append("model-significant positive predictivity") n_selected = int(self.np.count_nonzero(values)) max_value = float(self.np.max(self.np.abs(values))) if signed_scale and values.size else float(self.np.max(values)) if values.size else 0.0 if style == "litcoder": return self.write_litcoder_image( values=values, source=source, component=int(component), side=side, n_selected=n_selected, max_value=max_value, vmax=vmax, metric_label=metric_label, mask_labels=mask_labels, signed_scale=signed_scale, ) if style != "flatmap": raise ValueError(f"Unknown brainmap style: {style}") native_values = self.project_fsaverage5_to_native(values, self.flatmap) flat_values = self.native_to_flatmap(native_values, self.flatmap) return self.write_image( flat_values=flat_values, source=source, component=int(component), side=side, hemi=hemi, n_selected=n_selected, max_value=max_value, vmax=vmax, metric_label=metric_label, scale_label=scale_label, mask_labels=mask_labels, mask_requested=use_reliability_mask or use_significant_mask, signed_scale=signed_scale, ) def render_erf(self, hemi: str, *, style: str = "flatmap") -> bytes: values = self.np.asarray(self.erf_values, dtype=self.np.float32).copy() if self.reliability_mask is not None: values[~self.reliability_mask] = 0.0 if self.ica_significant_mask is not None: values[~self.ica_significant_mask] = 0.0 if style == "litcoder": title = f"{SUBJECT} Gemma-2-2B L{self.layer} top-5 IC mean ERF; NCSNR>{self.ncsnr_threshold:g} + significant predictivity" return self.litcoder_style_png_bytes( values, title=title, fsaverage5=self.fsaverage5, positive=True, vmax=self.erf_vmax, cmap="viridis", ) if style != "flatmap": raise ValueError(f"Unknown ERF brainmap style: {style}") native_values = self.project_fsaverage5_to_native(values, self.flatmap) flat_values = self.native_to_flatmap(native_values, self.flatmap) return self.write_erf_image(flat_values=flat_values, hemi=hemi, n_selected=int(self.np.count_nonzero(values))) def write_litcoder_image( self, *, values: Any, source: str, component: int, side: str, n_selected: int, max_value: float, vmax: float, metric_label: str, mask_labels: list[str], signed_scale: bool, ) -> bytes: feature_name = "SAE" if source == "sae" else "ICA" mask_label = " + ".join(mask_labels) if mask_labels else "unmasked" title = ( f"{SUBJECT} Gemma-2-2B L{self.layer} {feature_name} {component} {side}; " f"{metric_label}; {mask_label}; {n_selected} nonzero; max {max_value:.3g}" ) return self.litcoder_style_png_bytes( values, title=title, fsaverage5=self.fsaverage5, positive=not signed_scale, vmax=vmax, cmap="magma" if not signed_scale else None, ) def write_image( self, *, flat_values: Any, source: str, component: int, side: str, hemi: str, n_selected: int, max_value: float, vmax: float, metric_label: str, scale_label: str, mask_labels: list[str], mask_requested: bool, signed_scale: bool, ) -> bytes: half = self.HEMI_HALF[hemi] panel_mask = self.half_mask(self.flatmap.mask, self.flatmap.split, half) crop = self.crop_box_for_half(self.flatmap.mask, self.flatmap.split, half, margin=0) crop_mask = panel_mask[crop] roi_groups = [group for group in self.roi_groups if self.np.any(self.np.asarray(group["mask"], dtype=bool) & panel_mask)] display_values = self.np.asarray(flat_values[crop], dtype=self.np.float32).copy() display_values[~self.np.isfinite(display_values)] = self.np.nan value_mask = crop_mask & self.np.isfinite(display_values) if signed_scale: value_mask &= self.np.abs(display_values) > 0.0 panel = self.colorize_scalar(display_values, value_mask, vmax, cmap_name="coolwarm") else: value_mask &= display_values > 0.0 panel = self.colorize_positive(display_values, value_mask, vmax, cmap_name="magma") panel = self.compose_flatmap_panel(panel, crop_mask) panel = self.overlay_roi_borders(panel, roi_groups, crop, crop_mask, width=1, smooth_iterations=1, dot_step=7).convert("RGB") pad = 18 title_h = 44 label_h = 34 colorbar_h = 52 title_font = self.load_font(24, bold=True) label_font = self.load_font(18, bold=True) small_font = self.load_font(15) roi_title_font = self.load_font(15, bold=True) probe = self.Image.new("RGB", (1, 1), "white") probe_draw = self.ImageDraw.Draw(probe) roi_legend_h = self.estimate_roi_legend_height(probe_draw, roi_groups, panel.width, small_font, roi_title_font) legend_h = colorbar_h + (roi_legend_h + 10 if roi_groups else 0) canvas_w = pad * 2 + panel.width canvas_h = pad * 2 + title_h + label_h + panel.height + legend_h canvas = self.Image.new("RGB", (canvas_w, canvas_h), "white") draw = self.ImageDraw.Draw(canvas) feature_name = "SAE" if source == "sae" else "ICA" title = f"{SUBJECT} {self.HEMI_LABEL[hemi]} {feature_name} {component} {side} map" draw.text((pad, pad), title, fill=(20, 20, 20), font=title_font) mask_label = " + ".join(mask_labels) if mask_labels else ("unmasked; mask unavailable" if mask_requested else "unmasked") max_prefix = "max |" if signed_scale else "max " max_suffix = "|" if signed_scale else "" label = f"Gemma-2-2B L{self.layer}; {mask_label}; {n_selected} nonzero voxels; {max_prefix}{metric_label}{max_suffix} {max_value:.3g}" label_w, _ = self.text_size(draw, label, label_font) y0 = pad + title_h draw.text((pad + max((panel.width - label_w) // 2, 0), y0), label, fill=(20, 20, 20), font=label_font) canvas.paste(panel, (pad, y0 + label_h)) legend_y = y0 + label_h + panel.height + 14 bar_w = min(720, canvas_w - pad * 2) bar = self.render_colorbar(bar_w, 18, cmap_name="coolwarm") if signed_scale else self.render_positive_colorbar(bar_w, 18, cmap_name="magma") canvas.paste(bar.convert("RGB"), (pad, legend_y)) draw.rectangle((pad, legend_y, pad + bar_w - 1, legend_y + 17), outline=(40, 40, 40), width=1) left = f"-{vmax:.3g}" if signed_scale else "0" draw.text((pad, legend_y + 22), left, fill=(20, 20, 20), font=small_font) right = f"+{vmax:.3g}" right_w, _ = self.text_size(draw, right, small_font) draw.text((pad + bar_w - right_w, legend_y + 22), right, fill=(20, 20, 20), font=small_font) mid = scale_label mid_w, _ = self.text_size(draw, mid, small_font) draw.text((pad + max((bar_w - mid_w) // 2, 0), legend_y + 22), mid, fill=(20, 20, 20), font=small_font) if roi_groups: self.draw_roi_legend(draw, roi_groups, pad, legend_y + 54, canvas_w - pad * 2, small_font, roi_title_font) return png_bytes(canvas) def write_erf_image(self, *, flat_values: Any, hemi: str, n_selected: int) -> bytes: half = self.HEMI_HALF[hemi] panel_mask = self.half_mask(self.flatmap.mask, self.flatmap.split, half) crop = self.crop_box_for_half(self.flatmap.mask, self.flatmap.split, half, margin=0) crop_mask = panel_mask[crop] roi_groups = [group for group in self.roi_groups if self.np.any(self.np.asarray(group["mask"], dtype=bool) & panel_mask)] display_values = self.np.asarray(flat_values[crop], dtype=self.np.float32).copy() display_values[~self.np.isfinite(display_values)] = self.np.nan value_mask = crop_mask & self.np.isfinite(display_values) & (display_values > 0.0) panel = self.colorize_positive(display_values, value_mask, self.erf_vmax, cmap_name="viridis") panel = self.compose_flatmap_panel(panel, crop_mask) panel = self.overlay_roi_borders(panel, roi_groups, crop, crop_mask, width=1, smooth_iterations=1, dot_step=7).convert("RGB") pad = 18 title_h = 44 label_h = 34 colorbar_h = 52 title_font = self.load_font(24, bold=True) label_font = self.load_font(18, bold=True) small_font = self.load_font(15) roi_title_font = self.load_font(15, bold=True) probe = self.Image.new("RGB", (1, 1), "white") probe_draw = self.ImageDraw.Draw(probe) roi_legend_h = self.estimate_roi_legend_height(probe_draw, roi_groups, panel.width, small_font, roi_title_font) legend_h = colorbar_h + (roi_legend_h + 10 if roi_groups else 0) canvas_w = pad * 2 + panel.width canvas_h = pad * 2 + title_h + label_h + panel.height + legend_h canvas = self.Image.new("RGB", (canvas_w, canvas_h), "white") draw = self.ImageDraw.Draw(canvas) title = f"{SUBJECT} {self.HEMI_LABEL[hemi]} top-5 IC mean ERF" draw.text((pad, pad), title, fill=(20, 20, 20), font=title_font) label = f"Gemma-2-2B layer {self.layer}; NCSNR>{self.ncsnr_threshold:g} + significant predictivity; {n_selected} voxels" label_w, _ = self.text_size(draw, label, label_font) y0 = pad + title_h draw.text((pad + max((panel.width - label_w) // 2, 0), y0), label, fill=(20, 20, 20), font=label_font) canvas.paste(panel, (pad, y0 + label_h)) legend_y = y0 + label_h + panel.height + 14 bar_w = min(720, canvas_w - pad * 2) bar = self.render_positive_colorbar(bar_w, 18, cmap_name="viridis") canvas.paste(bar.convert("RGB"), (pad, legend_y)) draw.rectangle((pad, legend_y, pad + bar_w - 1, legend_y + 17), outline=(40, 40, 40), width=1) draw.text((pad, legend_y + 22), "0", fill=(20, 20, 20), font=small_font) right = f"{self.erf_vmax:.0f}" right_w, _ = self.text_size(draw, right, small_font) draw.text((pad + bar_w - right_w, legend_y + 22), right, fill=(20, 20, 20), font=small_font) mid = "mean component ERF over top 5 selected ICs per voxel" mid_w, _ = self.text_size(draw, mid, small_font) draw.text((pad + max((bar_w - mid_w) // 2, 0), legend_y + 22), mid, fill=(20, 20, 20), font=small_font) if roi_groups: self.draw_roi_legend(draw, roi_groups, pad, legend_y + 54, canvas_w - pad * 2, small_font, roi_title_font) return png_bytes(canvas) def compose_flatmap_panel(self, layer: Any, crop_mask: Any) -> Any: base = self.Image.new("RGBA", layer.size, self.flatmap_background_rgba) rgba = self.np.asarray(base).copy() rgba[crop_mask] = self.unselected_voxel_rgba base = self.Image.fromarray(rgba, mode="RGBA") base.alpha_composite(layer.convert("RGBA")) return base def estimate_roi_legend_height(self, draw: Any, roi_groups: list[dict[str, object]], max_width: int, font: Any, title_font: Any) -> int: if not roi_groups: return 0 title = "ROI borders:" title_w, _ = self.text_size(draw, title, title_font) cursor_x = title_w + 18 cursor_y = 1 line_h = 24 for group in roi_groups: label = f"{group['label']}: {group['description']}" label_w, _ = self.text_size(draw, label, font) item_w = 34 + label_w + 22 if cursor_x + item_w > max_width and cursor_x > title_w + 18: cursor_x = 0 cursor_y += line_h cursor_x += item_w return cursor_y + line_h + 6 INDEX_HTML = r""" Feature Label Review
Feature Label Review

Select a feature.

""" if __name__ == "__main__": raise SystemExit(main())