#!/usr/bin/env python3 """rDCIM Direct Propagation Visualizer. Uses the rDCIM (regression Dynamic Causal Modeling) effective connectivity matrix directly to simulate information propagation between brain regions. This is an alternative to the MDN flow-based visualization: instead of following a continuous vector field, we use the discrete ROI-to-ROI connectivity matrix to model how perturbations in one region propagate through the brain network. Usage: python examples/rdcim_propagation.py python examples/rdcim_propagation.py --global-state "someone feeling anxious" python examples/rdcim_propagation.py --perturb 42 --perturbation "sudden fear response" Key bindings: Click select ROI to perturb p perturb selected ROI (prompts in console) P (shift+p) propagate perturbation through network s initialize brain states from global state r reset all states +/- increase/decrease connection threshold t toggle connection labels Escape quit """ import argparse import json import os import sys from pathlib import Path import numpy as np import vtk # Add parent to path for imports sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) from src.brain_state import BrainStateDB PROJECT_ROOT = Path(__file__).resolve().parent.parent DATA_DIR = PROJECT_ROOT / "data" DEFAULT_RDCIM = DATA_DIR / "sch400_rDCM_A.npy" DEFAULT_CENTROIDS = DATA_DIR / "schaefer400_centroids_MNI.npy" DEFAULT_STATE_FILE = DATA_DIR / "brain_states_rdcim.json" # Schaefer 400 network names (7-network parcellation) NETWORK_COLORS = { "Vis": (0.55, 0.0, 0.75), # purple "SomMot": (0.0, 0.4, 0.8), # blue "DorsAttn": (0.0, 0.7, 0.0), # green "SalVentAttn": (0.8, 0.0, 0.5), # magenta "Limbic": (0.9, 0.7, 0.0), # gold "Cont": (0.9, 0.4, 0.0), # orange "Default": (0.8, 0.2, 0.2), # red } def _ensure_ssl(): try: import certifi cert_file = certifi.where() cur = os.environ.get("SSL_CERT_FILE", "") if not cur or not os.path.isfile(cur): os.environ["SSL_CERT_FILE"] = cert_file cur2 = os.environ.get("REQUESTS_CA_BUNDLE", "") if not cur2 or not os.path.isfile(cur2): os.environ["REQUESTS_CA_BUNDLE"] = cert_file except ImportError: pass def load_rdcim(path: Path) -> np.ndarray: """Load rDCIM effective connectivity matrix.""" if not path.exists(): raise FileNotFoundError(f"Missing rDCIM file: {path}") A = np.load(path) A = np.asarray(A, dtype=np.float32) if A.ndim != 2 or A.shape[0] != A.shape[1]: raise RuntimeError(f"rDCIM must be square, got {A.shape}") print(f"[rDCIM] loaded {A.shape[0]}x{A.shape[1]} matrix") return A def load_centroids(path: Path) -> np.ndarray: """Load ROI centroid positions in MNI space.""" if not path.exists(): raise FileNotFoundError(f"Missing centroids: {path}") P = np.load(path).astype(np.float32) if P.ndim != 2 or P.shape[1] not in (2, 3): raise RuntimeError(f"Centroids must be Nx3, got {P.shape}") if P.shape[1] == 2: P = np.c_[P, np.zeros((P.shape[0], 1), np.float32)] print(f"[centroids] loaded {P.shape[0]} ROIs") return P def get_roi_names(R: int) -> list[str]: """Generate Schaefer-style ROI names.""" names = [] for i in range(R): hemi = "LH" if i < R // 2 else "RH" names.append(f"{hemi}_ROI_{i+1:03d}") return names def map_rois_to_regions(centroids: np.ndarray, roi_names: list[str]) -> list[str]: """Map each ROI centroid to its nearest Allen atlas brain region. Uses a voxel label grid (same as mesh_overlay) for robust mapping. Returns updated ROI names like 'LH_ROI_042 (precentral gyrus)'. """ from src.mesh_overlay import FlowMeshOverlay import vtk alignment_file = DATA_DIR / "brain_alignment.json" mesh_dir = DATA_DIR / "meshes" # Dummy renderer/window (no display needed) ren = vtk.vtkRenderer() win = vtk.vtkRenderWindow() win.SetOffScreenRendering(1) win.AddRenderer(ren) try: overlay = FlowMeshOverlay(ren=ren, win=win, mesh_dir=mesh_dir, alignment_file=alignment_file) grid_cache = DATA_DIR / "label_grid_cache.npz" if not overlay.load_label_grid(grid_cache): overlay.build_label_grid() overlay.save_label_grid(grid_cache) except Exception as e: print(f"[roi-map] Could not build label grid: {e}") return roi_names mapped_names = [] matched = 0 for i, (name, pos) in enumerate(zip(roi_names, centroids)): key = overlay.get_region_at_point(pos) if key is None: key = overlay.find_nearest_region(pos, search_radius=4) if key is not None: region_name = overlay.get_region_name(key) hemi = overlay.get_hemisphere_label(key, pos) hemi_tag = f", {hemi}" if hemi else "" # Add relative position within the region center = overlay.get_mesh_center(key) bounds = overlay.get_mesh_bounds(key) pos_tag = "" if center is not None and bounds is not None: diff = pos - center extent = np.array([bounds[1]-bounds[0], bounds[3]-bounds[2], bounds[5]-bounds[4]]) extent = np.maximum(extent, 1e-6) rel = diff / (extent * 0.5) parts = [] if abs(rel[2]) > 0.3: parts.append("dorsal" if rel[2] > 0 else "ventral") if abs(rel[0]) > 0.3: parts.append("lateral" if abs(rel[0]) > 0.5 else "medial") if abs(rel[1]) > 0.3: parts.append("anterior" if rel[1] > 0 else "posterior") if parts: pos_tag = f", near {' '.join(parts)} section" else: pos_tag = ", near center" mapped_names.append(f"{name} ({region_name}{hemi_tag}{pos_tag})") matched += 1 else: # Still add hemisphere based on x-coordinate for unmapped ROIs hemi_label = "left hemisphere" if pos[0] < 0 else "right hemisphere" mapped_names.append(f"{name} ({hemi_label})") print(f"[roi-map] Mapped {matched}/{len(roi_names)} ROIs to Allen regions") return mapped_names def get_strongest_connections(A: np.ndarray, source_idx: int, top_k: int = 20, depth: int = 2) -> list[dict]: """Get strongest outgoing connections from source, with multi-hop propagation. Args: A: connectivity matrix (R x R) source_idx: source ROI index top_k: max connections per level depth: propagation depth Returns: List of {target_idx, weight, depth, path} """ R = A.shape[0] Ac = A.copy() np.fill_diagonal(Ac, 0.0) results = [] visited = {source_idx} current_level = [source_idx] for d in range(1, depth + 1): next_level = [] for src in current_level: weights = Ac[src] abs_w = np.abs(weights) # Get top-k strongest if len(abs_w) > top_k: top_idx = np.argpartition(abs_w, -top_k)[-top_k:] else: top_idx = np.arange(len(abs_w)) top_idx = top_idx[abs_w[top_idx] > 0] for ti in top_idx: if ti in visited: continue results.append({ "target_idx": int(ti), "weight": float(weights[ti]), "abs_weight": float(abs_w[ti]), "depth": d, "source_idx": int(src), }) visited.add(int(ti)) next_level.append(int(ti)) current_level = next_level if not current_level: break # Sort by absolute weight results.sort(key=lambda x: x["abs_weight"], reverse=True) return results class RDCIMVisualizer: """Interactive VTK visualizer for rDCIM connectivity.""" def __init__(self, A: np.ndarray, centroids: np.ndarray, roi_names: list[str], brain_state_db: BrainStateDB): self.A = A self.centroids = centroids self.R = A.shape[0] self.roi_names = roi_names self.brain_state_db = brain_state_db self.selected_roi = None self.connection_threshold = 0.02 # fraction of max |A| self.show_labels = False self._connection_actors = [] self._label_actors = [] self._roi_actors = [] self._highlight_actor = None self._state_label_actors = [] # 3D text labels showing states next to ROIs self._propagation_line_actors = [] # animated connection lines during propagation self._perturb_proposals = None # list of 4 proposals when ready self._perturb_region = None # region name being perturbed self._perturb_waiting = False # True while waiting for user to press 1-5 self._anim_steps = [] # propagation steps for looped animation self._anim_source_idx = None # source ROI for animation self._anim_frame = 0 # current animation frame self._anim_playing = False # True while loop animation is active self.propagation_depth = 2 # user-adjustable depth for Shift+P self.propagation_top_k = 10 # connections per level # Compute stats Ac = A.copy() np.fill_diagonal(Ac, 0.0) self.max_weight = float(np.abs(Ac).max()) if Ac.size else 1.0 # Setup VTK self.ren = vtk.vtkRenderer() self.ren.SetBackground(0.05, 0.05, 0.1) self.win = vtk.vtkRenderWindow() self.win.AddRenderer(self.ren) self.win.SetSize(1400, 900) self.win.SetWindowName("rDCIM Propagation Visualizer") self._create_roi_spheres() self._create_text_overlay() self.picker = vtk.vtkCellPicker() self.picker.SetTolerance(0.005) def _create_roi_spheres(self): """Create sphere actors for each ROI.""" # Normalize centroids to reasonable scale center = self.centroids.mean(axis=0) scale = max(1.0, float(np.abs(self.centroids - center).max())) for i in range(self.R): sphere = vtk.vtkSphereSource() pos = self.centroids[i] sphere.SetCenter(float(pos[0]), float(pos[1]), float(pos[2])) sphere.SetRadius(1.5) sphere.SetThetaResolution(12) sphere.SetPhiResolution(12) mapper = vtk.vtkPolyDataMapper() mapper.SetInputConnection(sphere.GetOutputPort()) actor = vtk.vtkActor() actor.SetMapper(mapper) # Color by network (based on position heuristic) r, g, b = 0.6, 0.6, 0.6 actor.GetProperty().SetColor(r, g, b) actor.GetProperty().SetOpacity(0.7) self.ren.AddActor(actor) self._roi_actors.append(actor) self.ren.ResetCamera() def _create_text_overlay(self): """Create info text actors.""" self._info_actor = vtk.vtkTextActor() self._info_actor.SetInput("Click ROI to select | p: perturb | " "Shift+P: propagate | s: init states") tp = self._info_actor.GetTextProperty() tp.SetColor(1, 1, 1) tp.SetFontSize(13) tp.SetFontFamilyToCourier() self._info_actor.SetPosition(10, 10) self.ren.AddActor(self._info_actor) # Selected ROI info / perturbation options (top-center yellow) self._sel_actor = vtk.vtkTextActor() self._sel_actor.SetInput("") tp2 = self._sel_actor.GetTextProperty() tp2.SetColor(1.0, 0.9, 0.3) tp2.SetFontSize(14) tp2.SetFontFamilyToCourier() tp2.SetJustificationToCentered() tp2.SetVerticalJustificationToTop() self._sel_actor.GetPositionCoordinate().SetCoordinateSystemToNormalizedDisplay() self._sel_actor.GetPositionCoordinate().SetValue(0.5, 0.97) self.ren.AddActor(self._sel_actor) # Propagation summary (top-left, green) self._summary_actor = vtk.vtkTextActor() self._summary_actor.SetInput("") tp3 = self._summary_actor.GetTextProperty() tp3.SetColor(0.8, 1.0, 0.8) tp3.SetFontSize(11) tp3.SetFontFamilyToCourier() tp3.SetJustificationToLeft() tp3.SetVerticalJustificationToTop() self._summary_actor.GetPositionCoordinate().SetCoordinateSystemToNormalizedDisplay() self._summary_actor.GetPositionCoordinate().SetValue(0.01, 0.97) self._summary_actor.VisibilityOff() self.ren.AddActor(self._summary_actor) # Information flow story (top-right, light blue) self._story_actor = vtk.vtkTextActor() self._story_actor.SetInput("") tp4 = self._story_actor.GetTextProperty() tp4.SetColor(0.7, 0.9, 1.0) tp4.SetFontSize(11) tp4.SetFontFamilyToCourier() tp4.SetJustificationToRight() tp4.SetVerticalJustificationToTop() self._story_actor.GetPositionCoordinate().SetCoordinateSystemToNormalizedDisplay() self._story_actor.GetPositionCoordinate().SetValue(0.99, 0.97) self._story_actor.VisibilityOff() self.ren.AddActor(self._story_actor) def _show_connections(self, source_idx: int, depth: int = 2, top_k: int = 15): """Show connections from selected ROI. Always draws direct connections (depth=1) from the source ROI first, then adds multi-hop connections from intermediate ROIs. """ # Clear old connections for act in self._connection_actors: self.ren.RemoveActor(act) self._connection_actors.clear() connections = get_strongest_connections(self.A, source_idx, top_k=top_k, depth=depth) if not connections: return # Separate by depth to ensure direct connections are always shown depth1 = [c for c in connections if c["depth"] == 1] depth2 = [c for c in connections if c["depth"] > 1] # Show up to 10 direct + 5 multi-hop shown = depth1[:10] + depth2[:5] if not shown: return max_w = max(c["abs_weight"] for c in shown) src_pos = self.centroids[source_idx] info_lines = [f"Selected: {self.roi_names[source_idx]}"] state = self.brain_state_db.get(self.roi_names[source_idx]) if state: info_lines.append(f"State: {state[:80]}") info_lines.append(f"Connections ({len(depth1)} direct, {len(depth2)} multi-hop):") for c in shown: ti = c["target_idx"] w = c["weight"] # For depth-1: line goes from selected ROI to target # For depth-2+: line goes from intermediate source to target from_idx = c["source_idx"] from_pos = self.centroids[from_idx] to_pos = self.centroids[ti] # Create line line = vtk.vtkLineSource() line.SetPoint1(float(from_pos[0]), float(from_pos[1]), float(from_pos[2])) line.SetPoint2(float(to_pos[0]), float(to_pos[1]), float(to_pos[2])) mapper = vtk.vtkPolyDataMapper() mapper.SetInputConnection(line.GetOutputPort()) actor = vtk.vtkActor() actor.SetMapper(mapper) # Color: green for positive, red for negative; dimmer for multi-hop norm_w = c["abs_weight"] / max(max_w, 1e-9) is_multihop = c["depth"] > 1 if w > 0: if is_multihop: actor.GetProperty().SetColor(0.4, 0.7, 0.4) # dimmer green else: actor.GetProperty().SetColor(0.2, 0.9, 0.2) else: if is_multihop: actor.GetProperty().SetColor(0.7, 0.4, 0.4) # dimmer red else: actor.GetProperty().SetColor(0.9, 0.2, 0.2) line_width = max(1.0, 4.0 * norm_w) if not is_multihop else max(1.0, 2.5 * norm_w) actor.GetProperty().SetLineWidth(line_width) actor.GetProperty().SetOpacity(max(0.3, norm_w * (0.7 if is_multihop else 1.0))) self.ren.AddActor(actor) self._connection_actors.append(actor) # Highlight target ROI if ti < len(self._roi_actors): if w > 0: self._roi_actors[ti].GetProperty().SetColor(0.2, 0.9, 0.2) else: self._roi_actors[ti].GetProperty().SetColor(0.9, 0.2, 0.2) opacity = min(1.0, 0.5 + norm_w * 0.5) if not is_multihop else min(0.9, 0.3 + norm_w * 0.4) self._roi_actors[ti].GetProperty().SetOpacity(opacity) depth_tag = f"d{c['depth']}" if is_multihop else " " sign_tag = "inh" if w < 0 else "exc" info_lines.append(f" {depth_tag}: {self.roi_names[ti]} " f"w={w:+.4f} ({sign_tag})") self._sel_actor.SetInput("\n".join(info_lines)) self.win.Render() def _clear_state_labels(self): """Remove all 3D state text labels and propagation lines.""" for act in self._state_label_actors: self.ren.RemoveActor(act) self._state_label_actors.clear() for act in self._propagation_line_actors: self.ren.RemoveActor(act) self._propagation_line_actors.clear() @staticmethod def _word_wrap(text: str, width: int = 45, max_lines: int = 20) -> str: """Word-wrap text to fit in a VTK text actor.""" lines = [] for line in text.split("\n"): while len(line) > width: brk = line.rfind(" ", 0, width) if brk <= 0: brk = width lines.append(line[:brk]) line = line[brk:].lstrip() lines.append(line) return "\n".join(lines[:max_lines]) def _add_state_label(self, roi_idx: int, text: str, color=(1.0, 0.4, 0.4)): """Add a small 3D text label centered on an ROI showing its updated state.""" pos = self.centroids[roi_idx] actor = vtk.vtkBillboardTextActor3D() # Truncate to keep labels compact short = text[:60] + "..." if len(text) > 60 else text actor.SetInput(short) # Position directly at the ROI center — billboard text will face camera actor.SetPosition(float(pos[0]), float(pos[1]), float(pos[2])) tp = actor.GetTextProperty() tp.SetColor(*color) tp.SetFontSize(10) tp.SetFontFamilyToCourier() tp.SetBold(True) tp.SetJustificationToCentered() tp.SetVerticalJustificationToCentered() self.ren.AddActor(actor) self._state_label_actors.append(actor) # Highlight the ROI sphere in red if roi_idx < len(self._roi_actors): self._roi_actors[roi_idx].GetProperty().SetColor(*color) self._roi_actors[roi_idx].GetProperty().SetOpacity(1.0) def _add_propagation_line(self, source_idx: int, target_idx: int, color=(1.0, 0.3, 0.3)): """Draw a line from source to target ROI during animated propagation.""" p1 = self.centroids[source_idx] p2 = self.centroids[target_idx] pts = vtk.vtkPoints() pts.InsertNextPoint(float(p1[0]), float(p1[1]), float(p1[2])) pts.InsertNextPoint(float(p2[0]), float(p2[1]), float(p2[2])) line = vtk.vtkLine() line.GetPointIds().SetId(0, 0) line.GetPointIds().SetId(1, 1) cells = vtk.vtkCellArray() cells.InsertNextCell(line) pd = vtk.vtkPolyData() pd.SetPoints(pts) pd.SetLines(cells) mapper = vtk.vtkPolyDataMapper() mapper.SetInputData(pd) actor = vtk.vtkActor() actor.SetMapper(mapper) actor.GetProperty().SetColor(*color) actor.GetProperty().SetLineWidth(3.0) actor.GetProperty().SetOpacity(0.8) self.ren.AddActor(actor) self._propagation_line_actors.append(actor) def _select_roi(self, idx: int): """Select an ROI.""" # Reset old selection if self.selected_roi is not None and self.selected_roi < len(self._roi_actors): self._roi_actors[self.selected_roi].GetProperty().SetColor(0.6, 0.6, 0.6) self._roi_actors[self.selected_roi].GetProperty().SetOpacity(0.7) # Reset all ROIs for act in self._roi_actors: act.GetProperty().SetColor(0.6, 0.6, 0.6) act.GetProperty().SetOpacity(0.7) self.selected_roi = idx if idx < len(self._roi_actors): self._roi_actors[idx].GetProperty().SetColor(1.0, 1.0, 0.0) self._roi_actors[idx].GetProperty().SetOpacity(1.0) self._show_connections(idx, depth=self.propagation_depth, top_k=self.propagation_top_k) print(f"[select] ROI {idx}: {self.roi_names[idx]} " f"(depth={self.propagation_depth}, top_k={self.propagation_top_k})") def run(self): """Start the interactive visualizer.""" iren = vtk.vtkRenderWindowInteractor() iren.SetRenderWindow(self.win) style = vtk.vtkInteractorStyleTrackballCamera() iren.SetInteractorStyle(style) def on_click(obj, ev): x, y = obj.GetEventPosition() if self.picker.Pick(x, y, 0, self.ren) <= 0: return # Find nearest ROI px, py, pz = self.picker.GetPickPosition() pos = np.array([px, py, pz], np.float32) dists = np.linalg.norm(self.centroids - pos, axis=1) nearest = int(np.argmin(dists)) if dists[nearest] < 10.0: # threshold self._anim_playing = False self._summary_actor.VisibilityOff() self._story_actor.VisibilityOff() self._clear_state_labels() self._select_roi(nearest) def on_key(obj, ev): key = obj.GetKeySym() key_lower = key.lower() if key else "" shift = bool(obj.GetShiftKey()) if key_lower == "escape": obj.TerminateApp() elif key_lower == "p" and not shift: # Perturb selected ROI (synchronous — VTK freezes during GPT) self._anim_playing = False if self.selected_roi is None: print("[perturb] Select an ROI first (click).") return name = self.roi_names[self.selected_roi] current_state = self.brain_state_db.get(name) self._sel_actor.SetInput(f"Generating perturbation options\nfor {name}...\n(window will freeze briefly)") self.win.Render() print(f"\n[perturb] Region: {name}") if current_state: print(f" Current state: {current_state}") print(f" Fetching perturbation options from GPT...") sys.stdout.flush() try: proposals = self.brain_state_db.propose_perturbations(name) except Exception as e: print(f"[perturb] GPT proposal failed: {e}") proposals = [ "Heightened activity in this region", "Suppressed activity in this region", "Shift to an alternative processing mode", "Disrupted connectivity with downstream regions", ] # Show options in VTK and console lines = [f"Perturbation options for:", name[:60], ""] for i, p in enumerate(proposals, 1): lines.append(f" {i}. {p[:55]}") lines.append(f" 5. (Custom — type in console)") lines.append("") lines.append("Press 1-5 in VTK window") self._sel_actor.SetInput("\n".join(lines)) self.win.Render() print(f"\n Perturbation options for {name}:") for i, p in enumerate(proposals, 1): print(f" {i}. {p}") print(f" 5. (Write your own)") print(f" >>> Press 1-5 in the VTK window <<<\n") sys.stdout.flush() self._perturb_proposals = proposals self._perturb_region = name self._perturb_waiting = True elif key_lower in ("1", "2", "3", "4") and self._perturb_waiting: # User picked a perturbation option (synchronous) idx = int(key_lower) - 1 desc = self._perturb_proposals[idx] self._perturb_waiting = False self._perturb_proposals = None name = self._perturb_region print(f"[perturb] Applying option {key_lower}: {desc}") self._sel_actor.SetInput(f"Applying perturbation\nto {name}...\n(window will freeze briefly)") self.win.Render() try: new_state = self.brain_state_db.alter_region_state( name, desc, skip_validation=True) print(f"[perturb] {name} -> {new_state}") self._sel_actor.SetInput(f"Perturbed: {name}\n{new_state[:80]}") except Exception as e: print(f"[perturb] Error: {e}") self._sel_actor.SetInput(f"Perturbation failed: {e}") self.win.Render() elif key_lower == "5" and self._perturb_waiting: # Custom perturbation — synchronous console input self._perturb_waiting = False self._perturb_proposals = None name = self._perturb_region self._sel_actor.SetInput(f"Type perturbation in console\nthen press Enter") self.win.Render() print(f" >>> Type your custom perturbation in the console <<<") sys.stdout.flush() desc = input(" Enter your perturbation: ").strip() if not desc: print("[perturb] Cancelled.") return print(f"[perturb] Applying: {desc}") self._sel_actor.SetInput(f"Applying perturbation\nto {name}...") self.win.Render() try: new_state = self.brain_state_db.alter_region_state( name, desc, skip_validation=True) print(f"[perturb] {name} -> {new_state}") self._sel_actor.SetInput(f"Perturbed: {name}\n{new_state[:80]}") except Exception as e: print(f"[perturb] Error: {e}") self._sel_actor.SetInput(f"Perturbation failed: {e}") self.win.Render() elif key_lower == "p" and shift: # Propagate through network (synchronous) if self.selected_roi is None: print("[propagate] Select an ROI first.") return if not self.brain_state_db.has_states(): print("[propagate] Initialize states first (s).") return source_name = self.roi_names[self.selected_roi] source_idx = self.selected_roi connections = get_strongest_connections( self.A, source_idx, top_k=self.propagation_top_k, depth=self.propagation_depth ) self._sel_actor.SetInput(f"Propagating from {source_name}...\n(window will freeze briefly)") self._clear_state_labels() self._summary_actor.VisibilityOff() self._story_actor.VisibilityOff() self.win.Render() print(f"[propagate] From {source_name} to " f"{len(connections)} regions (graph-based)...") sys.stdout.flush() # Build name->index map and target->source map for visualization name_to_idx = {n: i for i, n in enumerate(self.roi_names)} # Map target_name -> source_idx from the connections list target_source_map = {} for c in connections: tname = self.roi_names[c["target_idx"]] target_source_map[tname] = c["source_idx"] # Collect propagation steps for looped animation self._anim_steps = [] # list of (from_idx, to_idx, state_text, depth) def _render_update(msg): """Called per-region during propagation.""" print(f" {msg}") sys.stdout.flush() if ": " in msg: parts = msg.strip().split(": ", 1) target_name = parts[0].strip() new_state = parts[1].strip() ti = name_to_idx.get(target_name) if ti is not None: # Draw line from the actual source (not always the original) from_idx = target_source_map.get(target_name, source_idx) # Find depth for color depth = 1 for c in connections: if c["target_idx"] == ti: depth = c.get("depth", 1) break color = (1.0, 0.3, 0.3) if depth == 1 else (1.0, 0.6, 0.2) self._add_propagation_line(from_idx, ti, color=color) self._add_state_label(ti, new_state, color=color) self.win.Render() # Store for animation loop self._anim_steps.append((from_idx, ti, new_state, depth)) try: for c in connections: c["target_name"] = self.roi_names[c["target_idx"]] updates = self.brain_state_db.propagate_through_graph( source_name, connections, A=self.A, roi_names=self.roi_names, callback=_render_update ) if updates: summary = self.brain_state_db.summarize_changes( updates, source_name ) story = self.brain_state_db.generate_flow_story( updates, source_name, connections=connections ) print(f"\n--- PROPAGATION SUMMARY ---") print(summary) print(f"\n--- INFORMATION FLOW STORY ---") print(story) print("--- END ---\n") # Show summary in top-left self._summary_actor.SetInput( self._word_wrap(f"PROPAGATION SUMMARY\n{summary}", 45, 20)) self._summary_actor.VisibilityOn() # Show story in top-right self._story_actor.SetInput( self._word_wrap(f"INFORMATION FLOW\n{story}", 45, 20)) self._story_actor.VisibilityOn() # Clear yellow status text self._sel_actor.SetInput("") # Start looped animation self._anim_source_idx = source_idx self._anim_frame = 0 self._anim_playing = True else: print("[propagate] No changes.") except Exception as e: print(f"[propagate] Error: {e}") import traceback traceback.print_exc() self.win.Render() elif key_lower == "s": # Initialize states (synchronous) region_names = self.roi_names self._sel_actor.SetInput("Enter brain state in console...") self.win.Render() print("[states] Enter global brain state (type in console):") sys.stdout.flush() gs = input(" > ").strip() self._sel_actor.SetInput(f"Initializing states...\n(window will freeze)") self.win.Render() try: result = self.brain_state_db.initialize_from_global( gs, region_names, callback=lambda m: print(f" {m}") ) print(f"[states] {result}") self._sel_actor.SetInput("States initialized.") except Exception as e: print(f"[states] Error: {e}") self._sel_actor.SetInput(f"State init failed: {e}") self.win.Render() elif key_lower == "r": self.brain_state_db.clear() print("[states] Cleared all states.") elif key_lower in ("plus", "equal"): self.connection_threshold = min(0.5, self.connection_threshold * 1.5) print(f"[threshold] {self.connection_threshold:.3f}") if self.selected_roi is not None: self._show_connections(self.selected_roi, depth=self.propagation_depth, top_k=self.propagation_top_k) elif key_lower in ("minus", "underscore"): self.connection_threshold = max(0.001, self.connection_threshold / 1.5) print(f"[threshold] {self.connection_threshold:.3f}") if self.selected_roi is not None: self._show_connections(self.selected_roi, depth=self.propagation_depth, top_k=self.propagation_top_k) elif key_lower == "bracketright": # ] = increase propagation depth self.propagation_depth = min(6, self.propagation_depth + 1) print(f"[depth] propagation depth = {self.propagation_depth}") self._sel_actor.SetInput( f"Propagation depth: {self.propagation_depth}\n" f"(top_k={self.propagation_top_k} per level)") if self.selected_roi is not None: self._show_connections(self.selected_roi, depth=self.propagation_depth, top_k=self.propagation_top_k) self.win.Render() elif key_lower == "bracketleft": # [ = decrease propagation depth self.propagation_depth = max(1, self.propagation_depth - 1) print(f"[depth] propagation depth = {self.propagation_depth}") self._sel_actor.SetInput( f"Propagation depth: {self.propagation_depth}\n" f"(top_k={self.propagation_top_k} per level)") if self.selected_roi is not None: self._show_connections(self.selected_roi, depth=self.propagation_depth, top_k=self.propagation_top_k) self.win.Render() elif key_lower == "period": # . = increase top_k (more connections per level) self.propagation_top_k = min(50, self.propagation_top_k + 5) print(f"[top_k] connections per level = {self.propagation_top_k}") if self.selected_roi is not None: self._show_connections(self.selected_roi, depth=self.propagation_depth, top_k=self.propagation_top_k) self.win.Render() elif key_lower == "comma": # , = decrease top_k self.propagation_top_k = max(3, self.propagation_top_k - 5) print(f"[top_k] connections per level = {self.propagation_top_k}") if self.selected_roi is not None: self._show_connections(self.selected_roi, depth=self.propagation_depth, top_k=self.propagation_top_k) self.win.Render() def on_anim_timer(_o, _e): """Looped propagation animation — show signal traveling step by step.""" if not self._anim_playing or not self._anim_steps: return frame = self._anim_frame steps = self._anim_steps total_frames = len(steps) + 6 # steps + pause frames at end if frame == 0: # Clear previous animation visuals self._clear_state_labels() # Show source ROI highlighted if self._anim_source_idx < len(self._roi_actors): self._roi_actors[self._anim_source_idx].GetProperty().SetColor(1.0, 1.0, 0.0) self._roi_actors[self._anim_source_idx].GetProperty().SetOpacity(1.0) if frame < len(steps): from_idx, to_idx, state_text, depth = steps[frame] color = (1.0, 0.3, 0.3) if depth == 1 else (1.0, 0.6, 0.2) self._add_propagation_line(from_idx, to_idx, color=color) self._add_state_label(to_idx, state_text, color=color) # Highlight target ROI if to_idx < len(self._roi_actors): self._roi_actors[to_idx].GetProperty().SetColor(*color) self._roi_actors[to_idx].GetProperty().SetOpacity(1.0) self._anim_frame = frame + 1 if self._anim_frame >= total_frames: # Reset ROI colors and loop for act in self._roi_actors: act.GetProperty().SetColor(0.6, 0.6, 0.6) act.GetProperty().SetOpacity(0.7) self._anim_frame = 0 self.win.Render() iren.Initialize() iren.AddObserver("LeftButtonPressEvent", on_click) iren.AddObserver("KeyPressEvent", on_key) iren.AddObserver("TimerEvent", on_anim_timer) iren.CreateRepeatingTimer(500) # animation tick every 500ms self.win.Render() print("\n=== rDCIM Propagation Visualizer ===") print(f" ROIs: {self.R}") print(f" Max |weight|: {self.max_weight:.4f}") print(f" Click ROI to select | p perturb | Shift+P propagate") print(f" s init states | r reset | +/- threshold") print(f" [/] depth ({self.propagation_depth}) | ,/. top_k ({self.propagation_top_k}) | Esc quit") print() iren.Start() def main(): _ensure_ssl() try: from dotenv import load_dotenv load_dotenv(PROJECT_ROOT / ".env") except ImportError: pass ap = argparse.ArgumentParser(description="rDCIM Propagation Visualizer") ap.add_argument("--rdcim", type=Path, default=DEFAULT_RDCIM, help="rDCIM matrix .npy file") ap.add_argument("--centroids", type=Path, default=DEFAULT_CENTROIDS, help="ROI centroid positions .npy file") ap.add_argument("--global-state", type=str, default="", help="Initialize brain states with this global state") ap.add_argument("--perturb", type=int, default=None, help="ROI index to perturb on startup") ap.add_argument("--perturbation", type=str, default="", help="Perturbation description") ap.add_argument("--hq", action="store_true", help="Use high-quality GPT model (gpt-5.4) instead of gpt-5.4-mini") ap.add_argument("--debug", action="store_true", help="Print full LLM prompts to console before each call") ap.add_argument("--extra-parcellation", type=Path, default=None, help="Path to NIfTI parcellation file for finer subregion labels") ap.add_argument("--extra-parcellation-labels", type=Path, default=None, help="JSON label map for extra parcellation") args = ap.parse_args() model = "gpt-5.4" if args.hq else "gpt-5.4-mini" print(f"[config] LLM model: {model}") A = load_rdcim(args.rdcim) centroids = load_centroids(args.centroids) if A.shape[0] != centroids.shape[0]: print(f"[warning] matrix ({A.shape[0]}) != centroids ({centroids.shape[0]})") n = min(A.shape[0], centroids.shape[0]) A = A[:n, :n] centroids = centroids[:n] roi_names = get_roi_names(A.shape[0]) # Map ROI names to actual Allen atlas regions print("[init] Mapping ROIs to anatomical regions...") roi_names = map_rois_to_regions(centroids, roi_names) brain_db = BrainStateDB(DEFAULT_STATE_FILE, model=model, debug=args.debug) if args.global_state: print(f"[init] Initializing brain states: {args.global_state}") brain_db.initialize_from_global(args.global_state, roi_names) viz = RDCIMVisualizer(A, centroids, roi_names, brain_db) if args.perturb is not None and args.perturbation: if 0 <= args.perturb < len(roi_names): brain_db.alter_region_state(roi_names[args.perturb], args.perturbation) viz.run() if __name__ == "__main__": main()