"""
Brain-alignment data: Phi matrices, voxel coordinates, DynaDiff loader.
All data is loaded once at module import time from --phi-dir / --dynadiff-dir.
Public flags HAS_PHI and HAS_DYNADIFF tell panels what is available.
"""
import base64
import io
import os
import sys
import threading
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
from .args import args
# ---------- Nilearn surface rendering (TribeV2-style) ----------
_NILEARN_AVAILABLE = False
_fsavg5 = None # cached fsaverage5 surface data
_fsavg5_tree = None # cached KDTree over pial-left coords
_fsavg5_pials = None # cached (pial_left_xyz, pial_right_xyz)
try:
import nibabel as _nib
from nilearn.datasets import fetch_surf_fsaverage as _fetch_surf_fsaverage
from nilearn.plotting import plot_surf_stat_map as _plot_surf_stat_map
from scipy.spatial import cKDTree as _cKDTree
_NILEARN_AVAILABLE = True
except ImportError:
pass
def _get_fsavg5():
global _fsavg5, _fsavg5_pials, _fsavg5_tree
if _fsavg5 is None:
_fsavg5 = _fetch_surf_fsaverage('fsaverage5')
pl = _nib.load(_fsavg5['pial_left']).darrays[0].data
pr = _nib.load(_fsavg5['pial_right']).darrays[0].data
_fsavg5_pials = (pl, pr)
_fsavg5_tree = None
return _fsavg5
_surf_interp = None # cached per-hemisphere (mask, idxs, weights) for fast IDW
_surf_interp_lock = threading.Lock()
def _ensure_surf_interp(coords: np.ndarray, max_dist: float = 12.0, k: int = 5):
"""Build and cache the KDTree + IDW weights for voxel-to-surface projection."""
global _surf_interp, _fsavg5_tree
if _surf_interp is not None:
return
with _surf_interp_lock:
if _surf_interp is not None:
return
_get_fsavg5()
_fsavg5_tree = _cKDTree(coords)
result = []
for pial_xyz in _fsavg5_pials:
dists, idxs = _fsavg5_tree.query(pial_xyz, k=k, workers=-1)
mask = dists[:, 0] <= max_dist
d = np.where(dists[mask] == 0, 1e-10, dists[mask])
w = 1.0 / d
w /= w.sum(axis=1, keepdims=True)
result.append((mask, idxs[mask], w, pial_xyz.shape[0]))
_surf_interp = result
def _voxels_to_surface(values: np.ndarray, coords: np.ndarray,
max_dist: float = 12.0, k: int = 5):
"""Interpolate voxel values onto fsaverage5 pial vertices via KDTree IDW.
Returns (tex_left, tex_right) each shape (10242,) with NaN for vertices
farther than max_dist mm from any voxel.
"""
_ensure_surf_interp(coords, max_dist, k)
textures = []
for mask, idxs, weights, n_verts in _surf_interp:
tex = np.full(n_verts, np.nan, dtype=np.float32)
if mask.any():
tex[mask] = (weights * values[idxs]).sum(axis=1)
textures.append(tex)
return textures[0], textures[1]
def _render_brain_surface_b64(values: np.ndarray, title: str = '',
compact: bool = False, cbar_label: str = '',
figsize=(12, 3.5), dpi=80) -> str | None:
"""Render voxel values on fsaverage5 cortical surface.
Uses KDTree IDW to project values onto pial vertices, then renders with
nilearn's plot_surf_stat_map on the inflated fsaverage5 mesh.
compact=True → single left-posterior view; False → 4-view (lat+med, both hemis).
Returns base64 PNG or None if nilearn unavailable.
"""
if not _NILEARN_AVAILABLE or _voxel_coords is None:
return None
fs = _get_fsavg5()
tex_l, tex_r = _voxels_to_surface(values, _voxel_coords)
vmax = float(np.nanpercentile(np.abs(values), 98)) or 1e-6
kwargs = dict(cmap='RdBu_r', colorbar=False, vmin=-vmax, vmax=vmax,
bg_on_data=True)
_VIEWS_FULL = [
(tex_l, 'infl_left', 'sulc_left', 'left', (0, -135)),
(tex_l, 'infl_left', 'sulc_left', 'left', (0, 0)),
(tex_r, 'infl_right', 'sulc_right', 'right', (0, 180)),
(tex_r, 'infl_right', 'sulc_right', 'right', (0, -45)),
]
if compact:
fig, ax = plt.subplots(1, 1, figsize=(3.5, 2.8),
subplot_kw={'projection': '3d'},
facecolor='#f8f8f8')
_plot_surf_stat_map(surf_mesh=fs['infl_left'], stat_map=tex_l,
bg_map=fs['sulc_left'], hemi='left', view=(0, -135),
axes=ax, figure=fig, **kwargs)
ax.set_box_aspect(None, zoom=1.4)
else:
fig, axes = plt.subplots(
1, 4, figsize=figsize, facecolor='#f8f8f8',
subplot_kw={'projection': '3d'},
gridspec_kw={'wspace': -0.1, 'hspace': 0},
)
for ax, (tex, infl_k, sulc_k, hemi, view) in zip(axes, _VIEWS_FULL):
_plot_surf_stat_map(surf_mesh=fs[infl_k], stat_map=tex,
bg_map=fs[sulc_k], hemi=hemi, view=view,
axes=ax, figure=fig, **kwargs)
ax.set_box_aspect(None, zoom=1.4)
sm = plt.cm.ScalarMappable(cmap='RdBu_r',
norm=plt.Normalize(vmin=-vmax, vmax=vmax))
sm.set_array([])
cbar_ax = fig.add_axes([0.92, 0.2, 0.015, 0.6])
cbar = fig.colorbar(sm, cax=cbar_ax)
if cbar_label:
cbar.set_label(cbar_label, fontsize=9)
if title:
fig.suptitle(title, fontsize=10)
buf = io.BytesIO()
fig.savefig(buf, format='png', dpi=dpi, bbox_inches='tight',
facecolor='#f8f8f8')
plt.close(fig)
return base64.b64encode(buf.getvalue()).decode('utf-8')
# ---------- Phi (brain alignment) ----------
_phi_cv = None # (C, V) concept-by-voxel matrix, memory-mapped
_phi_c = None # (C,) per-concept cortical leverage scores
_voxel_coords = None # (V, 3) MNI voxel coordinates
_voxel_to_vertex = None # (V,) fsaverage vertex → voxel map (surface-space phi only)
_N_VOXELS_DD = 15724
_N_VERTS_FSAVG = 37984
def _pick_best_file(candidates: list, model_key: str, search_dir: str) -> str | None:
"""Prefer model_key substring match; fall back to largest file."""
if not candidates:
return None
if model_key:
matched = [f for f in candidates if model_key in f.lower()]
if matched:
return sorted(matched)[0]
print(f"[Phi] WARNING: --phi-model '{model_key}' matched nothing in {candidates}; "
"falling back to largest file")
return max(candidates, key=lambda f: os.path.getsize(os.path.join(search_dir, f)))
if args.phi_dir and os.path.isdir(args.phi_dir):
_pdir = args.phi_dir
_model_key = (args.phi_model or "").lower()
# Phi_cv matrix
_phi_mat_files = [f for f in os.listdir(_pdir)
if f.lower().startswith('phi_cv') and f.endswith('.npy')]
_phi_pick = _pick_best_file(_phi_mat_files, _model_key, _pdir)
if _phi_pick:
_phi_path = os.path.join(_pdir, _phi_pick)
_phi_cv = np.load(_phi_path, mmap_mode='r')
print(f"[Phi] Loaded {_phi_pick}: shape {_phi_cv.shape}, dtype {_phi_cv.dtype}")
if _phi_cv.shape[1] == _N_VERTS_FSAVG:
_v2v_path = os.path.join(_pdir, 'voxel_to_vertex_map.npy')
if os.path.exists(_v2v_path):
_voxel_to_vertex = np.load(_v2v_path)
print(f"[Phi] Surface-space phi; loaded voxel_to_vertex_map: "
f"{_voxel_to_vertex.shape}")
else:
print("[Phi] WARNING: surface-space phi but voxel_to_vertex_map.npy not found")
elif _phi_cv.shape[1] == _N_VOXELS_DD:
print("[Phi] Voxel-space phi detected.")
else:
print(f"[Phi] WARNING: unexpected phi dimension {_phi_cv.shape[1]}")
else:
print(f"[Phi] WARNING: no Phi_cv_*.npy found in {_pdir}")
# phi_c leverage scores
_phi_c_files = [f for f in os.listdir(_pdir)
if f.lower().startswith('phi_c')
and not f.lower().startswith('phi_cv')
and f.endswith('.npy')]
_phi_c_pick = _pick_best_file(_phi_c_files, _model_key, _pdir)
if _phi_c_pick:
_phi_c = np.load(os.path.join(_pdir, _phi_c_pick))
print(f"[Phi] Leverage scores {_phi_c_pick}: shape {_phi_c.shape}, "
f"range [{_phi_c.min():.4f}, {_phi_c.max():.4f}]")
else:
print(f"[Phi] No phi_c_*.npy found in {_pdir} — leverage scores unavailable")
# Voxel coordinates
_coords_path = os.path.join(_pdir, 'voxel_coords.npy')
if os.path.exists(_coords_path):
_voxel_coords = np.load(_coords_path)
print(f"[Phi] Voxel coordinates: {_voxel_coords.shape}")
else:
print("[Phi] voxel_coords.npy not found — cortical scatter unavailable")
HAS_PHI = _phi_cv is not None
# ---------- DynaDiff ----------
_dd_loader = None
HAS_DYNADIFF = False
_scripts_dir = os.path.dirname(os.path.abspath(__file__)) + '/..'
if args.dynadiff_modal_url:
# ── Modal HTTP mode (no local GPU needed) ────────────────────────────────
if not HAS_PHI:
print("[DynaDiff] WARNING: --phi-dir not set; steering panel requires Phi data. "
"Disabling.")
else:
try:
sys.path.insert(0, _scripts_dir)
from dynadiff_loader import HTTPDynaDiffLoader
_token = (args.dynadiff_modal_token
or os.environ.get("DYNADIFF_MODAL_TOKEN", ""))
_dd_loader = HTTPDynaDiffLoader(
url=args.dynadiff_modal_url,
token=_token,
)
_dd_loader.start()
HAS_DYNADIFF = True
print(f"[DynaDiff] Modal endpoint: {args.dynadiff_modal_url}")
except Exception as err:
print(f"[DynaDiff] WARNING: Could not init Modal loader ({err}). "
"Steering panel will be disabled.")
elif args.dynadiff_dir and os.path.isdir(args.dynadiff_dir):
# ── In-process mode (original, requires local GPU + dynadiff repo) ───────
if not HAS_PHI:
print("[DynaDiff] WARNING: --phi-dir not set; steering panel requires Phi data. "
"Disabling.")
else:
try:
sys.path.insert(0, _scripts_dir)
from dynadiff_loader import get_loader
_h5 = args.dynadiff_h5
if not os.path.isabs(_h5):
_h5 = os.path.join(args.dynadiff_dir, _h5)
_dd_loader = get_loader(
dynadiff_dir = args.dynadiff_dir,
checkpoint = args.dynadiff_checkpoint,
h5_path = _h5,
nsd_thumb_dir = args.brain_thumbnails,
subject_idx = 0,
)
HAS_DYNADIFF = True
print(f"[DynaDiff] In-process loader ready "
f"(checkpoint: {args.dynadiff_checkpoint})")
except Exception as err:
print(f"[DynaDiff] WARNING: Could not start loader ({err}). "
"Steering panel will be disabled.")
# ---------- Per-feature helpers ----------
def phi_cv_shape() -> tuple | None:
"""Return (_phi_cv.shape[0], _phi_cv.shape[1]) or None if not loaded."""
return _phi_cv.shape if _phi_cv is not None else None
def phi_c_for_feat(feat: int) -> float | None:
"""Cortical leverage score for a feature, or None."""
if _phi_c is None or feat >= len(_phi_c):
return None
return float(_phi_c[feat])
def phi_voxel_row(feat: int) -> np.ndarray | None:
"""Return the phi row in voxel space (15724,) float32, or None."""
if _phi_cv is None or feat >= _phi_cv.shape[0]:
return None
row = np.array(_phi_cv[feat], dtype=np.float32)
if _voxel_to_vertex is not None:
return row[_voxel_to_vertex]
return row
def phi_c_vals(indices) -> list:
"""Return phi_c leverage values for a list of feature indices (0.0 when unavailable)."""
if _phi_c is None:
return [0.0] * len(indices)
return [float(_phi_c[i]) if i < len(_phi_c) else 0.0 for i in indices]
def feat_display_name(feat: int | None) -> str:
"""Best-effort display name for DynaDiff feature table."""
if feat is None:
return 'unknown'
from .state import active_ds
ds = active_ds()
return ds['feature_names'].get(feat) or f'feat {feat}'
def dynadiff_request(sample_idx: int, steerings: list, seed: int) -> dict:
"""Run DynaDiff reconstruction. Raises RuntimeError if model not ready."""
status, err = _dd_loader.status
if status == 'loading':
raise RuntimeError('DynaDiff model still loading — try again shortly')
if status == 'error':
raise RuntimeError(f'DynaDiff model load failed: {err}')
return _dd_loader.reconstruct(sample_idx, steerings, seed)
# ---------- Rendering helpers ----------
def _render_phi_map_b64_compact(feat: int, figsize=(3.5, 2.8), dpi=70) -> str | None:
"""Single left-lateral surface view of phi, small enough for a steering card."""
from .state import active_ds
cached = active_ds().get('phi_map_cache', {}).get(feat)
if cached is not None:
return cached
phi_vox = phi_voxel_row(feat)
if phi_vox is None:
return None
b64 = _render_brain_surface_b64(phi_vox, compact=True, dpi=dpi)
if b64 is not None:
return b64
# Fallback: axial scatter
if _voxel_coords is None:
return None
vmax = float(np.abs(phi_vox).max()) or 1e-6
fig, ax = plt.subplots(1, 1, figsize=figsize, facecolor='#f8f8f8')
ax.scatter(_voxel_coords[:, 0], _voxel_coords[:, 1],
c=phi_vox, cmap='RdBu_r', s=3, alpha=0.8,
vmin=-vmax, vmax=vmax, rasterized=True, marker='s')
ax.set_aspect('equal'); ax.set_xticks([]); ax.set_yticks([])
ax.set_facecolor('#f8f8f8')
fig.tight_layout(pad=0.2)
buf = io.BytesIO()
fig.savefig(buf, format='png', dpi=dpi, bbox_inches='tight', facecolor='#f8f8f8')
plt.close(fig)
return base64.b64encode(buf.getvalue()).decode('utf-8')
def _render_cortical_profile_b64(feat: int) -> str | None:
"""Base64 PNG of cortical profile on fsaverage5 surface (TribeV2-style)."""
from .state import active_ds
cached = active_ds().get('cortical_profile_cache', {}).get(feat)
if cached is not None:
return cached
phi_vox = phi_voxel_row(feat)
if phi_vox is None:
return None
phi_c_val = phi_c_for_feat(feat)
phi_c_str = f' (φ_c = {phi_c_val:.4f})' if phi_c_val is not None else ''
title_str = f'Cortical Profile — Feature {feat}{phi_c_str}'
b64 = _render_brain_surface_b64(phi_vox, title=title_str,
cbar_label='Φ weight', dpi=90)
if b64 is not None:
return b64
# Fallback: 2-view axial/coronal scatter
if _voxel_coords is None:
return None
vmax = float(np.abs(phi_vox).max()) or 1e-6
fig, axes = plt.subplots(1, 2, figsize=(10, 4.0), facecolor='#f8f8f8')
for ax, (t, xi, yi) in zip(axes, [("Axial (x–y)", 0, 1),
("Coronal (x–z)", 0, 2)]):
sc = ax.scatter(_voxel_coords[:, xi], _voxel_coords[:, yi],
c=phi_vox, cmap='RdBu_r', s=4, alpha=0.75,
vmin=-vmax, vmax=vmax, rasterized=True, marker='s')
ax.set_title(t, fontsize=10); ax.set_aspect('equal')
ax.set_xticks([]); ax.set_yticks([]); ax.set_facecolor('#f8f8f8')
fig.subplots_adjust(right=0.88, top=0.88)
cbar_ax = fig.add_axes([0.91, 0.15, 0.02, 0.65])
fig.colorbar(sc, cax=cbar_ax).set_label('Φ weight', fontsize=9)
fig.suptitle(title_str, fontsize=11)
buf = io.BytesIO()
fig.savefig(buf, format='png', dpi=90, bbox_inches='tight', facecolor='#f8f8f8')
plt.close(fig)
return base64.b64encode(buf.getvalue()).decode('utf-8')
def get_dd_fmri(sample_idx: int) -> np.ndarray | None:
"""Return raw fMRI (N_VOXELS,) for a DynaDiff sample index, or None."""
if _dd_loader is None:
return None
try:
return _dd_loader.get_fmri(sample_idx)
except Exception:
return None
def apply_steering_fmri(fmri: np.ndarray, steerings: list) -> np.ndarray:
"""Apply steering perturbations to fMRI in-place (numpy).
steerings: list of (phi_voxel np.ndarray, lam float, threshold float)
"""
if _dd_loader is None:
return fmri
beta_std = _dd_loader.beta_std
if beta_std is None:
return fmri
result = fmri.copy()
for phi_voxel, lam, thr in steerings:
if phi_voxel is None:
continue
phi_max = float(np.abs(phi_voxel).max())
if phi_max < 1e-12:
continue
scale = beta_std / phi_max
if thr < 1.0:
cutoff = float(np.percentile(np.abs(phi_voxel), 100.0 * (1.0 - thr)))
mask = np.abs(phi_voxel) >= cutoff
else:
mask = np.ones(len(phi_voxel), dtype=bool)
perturb = lam * scale * phi_voxel
perturb[~mask] = 0.0
result += perturb
return result
def render_fmri_brain_compact_b64(fmri_voxels: np.ndarray,
title: str = '') -> str | None:
"""Compact left-lateral surface view of fMRI voxel activity, returns base64 PNG."""
if fmri_voxels is None or _voxel_coords is None:
return None
while fmri_voxels.ndim > 1:
fmri_voxels = fmri_voxels.mean(axis=-1)
b64 = _render_brain_surface_b64(fmri_voxels, title=title, compact=True, dpi=70)
if b64 is not None:
return b64
# Fallback: axial scatter
vmax = float(np.abs(fmri_voxels).max()) or 1e-6
fig, ax = plt.subplots(1, 1, figsize=(3.5, 2.8), facecolor='#f8f8f8')
ax.scatter(_voxel_coords[:, 0], _voxel_coords[:, 1],
c=fmri_voxels, cmap='RdBu_r', s=3, alpha=0.8,
vmin=-vmax, vmax=vmax, rasterized=True, marker='s')
ax.set_aspect('equal'); ax.set_xticks([]); ax.set_yticks([])
ax.set_facecolor('#f8f8f8')
if title:
ax.set_title(title, fontsize=9)
fig.tight_layout(pad=0.2)
buf = io.BytesIO()
fig.savefig(buf, format='png', dpi=70, bbox_inches='tight', facecolor='#f8f8f8')
plt.close(fig)
return base64.b64encode(buf.getvalue()).decode('utf-8')
def render_cortical_profile(feat: int) -> str:
"""Two-view scatter of phi voxel weights as an inline PNG HTML block."""
b64 = _render_cortical_profile_b64(feat)
if b64 is None:
return ""
return (
'Cortical Profile (Φ)
'
f'
'
)