Add pyvista renderer, precomputed phi map cache, and background-thread steering brain renders

scripts/explorer/brain.py

@@ -17,9 +17,10 @@ import numpy as np
 
 from .args import args
 
-# ---------- Nilearn surface rendering (TribeV2-style) ----------
+# ---------- Surface rendering ----------
 
 _NILEARN_AVAILABLE = False
+_PYVISTA_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)
@@ -27,12 +28,24 @@ _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
 
+try:
+    import pyvista as _pv
+    _pv.OFF_SCREEN = True
+    _PYVISTA_AVAILABLE = True
+except ImportError:
+    pass
+
+if not _PYVISTA_AVAILABLE:
+    try:
+        from nilearn.plotting import plot_surf_stat_map as _plot_surf_stat_map
+    except ImportError:
+        pass
+
 
 def _get_fsavg5():
     global _fsavg5, _fsavg5_pials, _fsavg5_tree
@@ -41,7 +54,7 @@ def _get_fsavg5():
         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   # built on first render, needs voxel coords
+        _fsavg5_tree  = None
     return _fsavg5
 
 
@@ -82,40 +95,171 @@ def _voxels_to_surface(values: np.ndarray, coords: np.ndarray,
     return textures[0], textures[1]
 
 
-def _render_brain_surface_b64(values: np.ndarray, title: str = '',
+# ---------- PyVista surface renderer ----------
+
+_pv_meshes = None  # cached (mesh_left, mesh_right, sulc_left, sulc_right)
+
+
+def _get_pv_meshes():
+    """Load fsaverage5 inflated meshes as pyvista PolyData (cached)."""
+    global _pv_meshes
+    if _pv_meshes is not None:
+        return _pv_meshes
+    fs = _get_fsavg5()
+    meshes = {}
+    for hemi in ('left', 'right'):
+        gii = _nib.load(fs[f'infl_{hemi}'])
+        coords = gii.darrays[0].data
+        faces_tri = gii.darrays[1].data.astype(np.int64)
+        n_faces = faces_tri.shape[0]
+        pv_faces = np.empty((n_faces, 4), dtype=np.int64)
+        pv_faces[:, 0] = 3
+        pv_faces[:, 1:] = faces_tri
+        mesh = _pv.PolyData(coords, pv_faces.ravel())
+        sulc = _nib.load(fs[f'sulc_{hemi}']).darrays[0].data
+        meshes[hemi] = (mesh, sulc)
+    _pv_meshes = meshes
+    return _pv_meshes
+
+
+# Camera positions: (position, focal_point, viewup)
+# Matched to the nilearn (elev, azim) views for visual cortex visibility
+_PV_VIEWS = {
+    'L_post_lat': {'position': (-100, -100, 30), 'focal_point': (0, 0, 0), 'viewup': (0, 0, 1)},
+    'L_medial':   {'position': (100, 0, 30),     'focal_point': (0, 0, 0), 'viewup': (0, 0, 1)},
+    'R_medial':   {'position': (-100, 0, 30),    'focal_point': (0, 0, 0), 'viewup': (0, 0, 1)},
+    'R_post_lat': {'position': (100, -100, 30),  'focal_point': (0, 0, 0), 'viewup': (0, 0, 1)},
+}
+
+
+def _render_pv_single(mesh, sulc, tex, vmax, cam, size=(350, 280)):
+    """Render a single hemisphere view with pyvista, return RGBA numpy array."""
+    p = _pv.Plotter(off_screen=True, window_size=size)
+    p.set_background('#f8f8f8')
+
+    # Sulcal depth as grey background shading
+    sulc_norm = np.clip(sulc, -2, 2)
+    sulc_grey = 0.55 - 0.15 * (sulc_norm / 2.0)
+    sulc_rgb = np.column_stack([sulc_grey, sulc_grey, sulc_grey])
+    bg_mesh = mesh.copy()
+    bg_mesh.point_data['sulc_rgb'] = sulc_rgb
+    p.add_mesh(bg_mesh, scalars='sulc_rgb', rgb=True,
+               lighting=True, opacity=1.0)
+
+    # Stat map overlay — only show non-NaN vertices
+    valid = ~np.isnan(tex)
+    if valid.any():
+        overlay = mesh.copy()
+        overlay.point_data['stat'] = np.where(valid, tex, 0.0)
+        # Opacity proportional to absolute value
+        abs_tex = np.abs(np.where(valid, tex, 0.0))
+        alpha = np.where(valid, np.clip(abs_tex / vmax, 0.15, 1.0), 0.0)
+        overlay.point_data['alpha'] = alpha
+        p.add_mesh(overlay, scalars='stat', cmap='RdBu_r',
+                   clim=[-vmax, vmax], opacity=alpha,
+                   show_scalar_bar=False, lighting=True)
+
+    p.camera_position = [cam['position'], cam['focal_point'], cam['viewup']]
+    img = p.screenshot(return_img=True)
+    p.close()
+    return img
+
+
+def _render_brain_pyvista_b64(values: np.ndarray, title: str = '',
+                               compact: bool = False,
+                               cbar_label: str = '') -> str | None:
+    """Render voxel values on fsaverage5 with pyvista (fast offscreen VTK)."""
+    if not _PYVISTA_AVAILABLE or not _NILEARN_AVAILABLE or _voxel_coords is None:
+        return None
+    try:
+        meshes = _get_pv_meshes()
+    except Exception:
+        return None
+
+    tex_l, tex_r = _voxels_to_surface(values, _voxel_coords)
+    vmax = float(np.nanpercentile(np.abs(values), 98)) or 1e-6
+
+    from PIL import Image as _PILImg
+
+    if compact:
+        mesh_l, sulc_l = meshes['left']
+        img = _render_pv_single(mesh_l, sulc_l, tex_l, vmax,
+                                _PV_VIEWS['L_post_lat'], size=(350, 280))
+        pil = _PILImg.fromarray(img)
+    else:
+        views = [
+            ('left',  tex_l, 'L_post_lat'),
+            ('left',  tex_l, 'L_medial'),
+            ('right', tex_r, 'R_medial'),
+            ('right', tex_r, 'R_post_lat'),
+        ]
+        panels = []
+        for hemi, tex, view_key in views:
+            mesh, sulc = meshes[hemi]
+            panel = _render_pv_single(mesh, sulc, tex, vmax,
+                                      _PV_VIEWS[view_key], size=(300, 250))
+            panels.append(_PILImg.fromarray(panel))
+
+        # Composite 4 panels into one image
+        pw, ph = panels[0].size
+        canvas = _PILImg.new('RGB', (pw * 4, ph), color=(248, 248, 248))
+        for i, panel in enumerate(panels):
+            canvas.paste(panel, (i * pw, 0))
+        pil = canvas
+
+    if title:
+        # Draw title using matplotlib (lightweight, just text)
+        fig, ax = plt.subplots(figsize=(pil.width / 100, 0.3), facecolor='#f8f8f8')
+        ax.text(0.5, 0.5, title, ha='center', va='center', fontsize=10,
+                transform=ax.transAxes)
+        ax.set_axis_off()
+        title_buf = io.BytesIO()
+        fig.savefig(title_buf, format='png', dpi=100, bbox_inches='tight',
+                    facecolor='#f8f8f8')
+        plt.close(fig)
+        title_buf.seek(0)
+        title_img = _PILImg.open(title_buf)
+        final = _PILImg.new('RGB', (pil.width, pil.height + title_img.height),
+                           color=(248, 248, 248))
+        final.paste(title_img, ((pil.width - title_img.width) // 2, 0))
+        final.paste(pil, (0, title_img.height))
+        pil = final
+
+    buf = io.BytesIO()
+    pil.save(buf, format='PNG')
+    return base64.b64encode(buf.getvalue()).decode('utf-8')
+
+
+# ---------- Nilearn fallback renderer ----------
+
+def _render_brain_nilearn_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 (TribeV2-style).
-
-    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.
-    """
+    """Render with nilearn's plot_surf_stat_map (slower matplotlib 3D fallback)."""
     if not _NILEARN_AVAILABLE or _voxel_coords is None:
         return None
+    try:
+        _plot_surf_stat_map
+    except NameError:
+        return None
+
     fs    = _get_fsavg5()
     tex_l, tex_r = _voxels_to_surface(values, _voxel_coords)
-    # Percentile-based vmax stretches the colormap over actual signal range
     vmax  = float(np.nanpercentile(np.abs(values), 98)) or 1e-6
-    # bg_on_data=True keeps sulcal texture visible everywhere — looks like a brain
     kwargs = dict(cmap='RdBu_r', colorbar=False, vmin=-vmax, vmax=vmax,
                   bg_on_data=True)
 
-    # Views chosen to maximise visual cortex visibility (NSD = occipital / posterior)
-    # (elev, azim) follow TribeV2 VIEW_DICT convention
     _VIEWS_FULL = [
-        (tex_l, 'infl_left',  'sulc_left',  'left',  (0, -135)),  # L posterior-lateral
-        (tex_l, 'infl_left',  'sulc_left',  'left',  (0, 0)),      # L medial
-        (tex_r, 'infl_right', 'sulc_right', 'right', (0, 180)),    # R medial
-        (tex_r, 'infl_right', 'sulc_right', 'right', (0, -45)),    # R posterior-lateral
+        (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')
-        # Posterior view of left hemisphere shows occipital cortex directly
         _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)
@@ -148,6 +292,25 @@ def _render_brain_surface_b64(values: np.ndarray, title: str = '',
     return base64.b64encode(buf.getvalue()).decode('utf-8')
 
 
+# ---------- Unified dispatch ----------
+
+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.
+
+    Tries pyvista (fast VTK offscreen) first, falls back to nilearn (matplotlib 3D).
+    Returns base64 PNG or None.
+    """
+    b64 = _render_brain_pyvista_b64(values, title=title, compact=compact,
+                                     cbar_label=cbar_label)
+    if b64 is not None:
+        return b64
+    return _render_brain_nilearn_b64(values, title=title, compact=compact,
+                                      cbar_label=cbar_label, figsize=figsize,
+                                      dpi=dpi)
+
+
 # ---------- Phi (brain alignment) ----------
 
 _phi_cv          = None   # (C, V) concept-by-voxel matrix, memory-mapped
@@ -335,6 +498,10 @@ def dynadiff_request(sample_idx: int, steerings: list, seed: int) -> dict:
 
 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

scripts/explorer/datasets.py

@@ -180,6 +180,16 @@ def _load_dataset(path: str, label: str, *,
         entry['heatmap_patch_grid'] = d.get('patch_grid', 16)
         has_hm = 'no'
 
+    # Brain render sidecar (precomputed compact phi map PNGs)
+    brain_render_sidecar = stem + '_brain_renders.pt'
+    if os.path.exists(brain_render_sidecar):
+        print(f"  Loading brain render sidecar {os.path.basename(brain_render_sidecar)} ...")
+        br = torch.load(brain_render_sidecar, map_location='cpu', weights_only=False)
+        entry['phi_map_cache'] = {int(k): v for k, v in br.items()}
+        print(f"  Cached {len(entry['phi_map_cache'])} phi map renders")
+    else:
+        entry['phi_map_cache'] = {}
+
     has_clip = 'yes' if entry['clip_embeds'] is not None else 'no'
     print(f"  d={d_model}, n={entry['n_images']}, backbone={entry['backbone']}, "
           f"clip={has_clip}, heatmaps={has_hm}")

scripts/explorer/panels/steering.py

@@ -307,18 +307,31 @@ if HAS_DYNADIFF:
                 list(_dd_source.data['lam']),
                 list(_dd_source.data['threshold']))
 
+    _steer_render_token = [0]  # mutable counter to discard stale renders
+
     def _update_steer_brain():
         feats, lams, thrs = _steerings_from_source()
         if not feats:
             steer_brain_div.text = ''
             return
-        combined = compute_steering_direction(feats, lams, thrs)
-        b64 = render_fmri_brain_compact_b64(
-            combined, 'Steering Direction (φ sum)')
-        steer_brain_div.text = (
-            f'<img src="data:image/png;base64,{b64}" '
-            f'style="max-width:100%"/>'
-            if b64 else '')
+        _steer_render_token[0] += 1
+        my_token = _steer_render_token[0]
+        steer_brain_div.text = ''
+        doc = curdoc()
+
+        def _bg():
+            combined = compute_steering_direction(feats, lams, thrs)
+            b64 = render_fmri_brain_compact_b64(
+                combined, 'Steering Direction (φ sum)')
+            def _apply():
+                if _steer_render_token[0] == my_token:
+                    steer_brain_div.text = (
+                        f'<img src="data:image/png;base64,{b64}" '
+                        f'style="max-width:100%"/>'
+                        if b64 else '')
+            doc.add_next_tick_callback(_apply)
+
+        threading.Thread(target=_bg, daemon=True).start()
 
     def _update_steered_brain():
         fmri = _Session.gt_fmri
@@ -329,13 +342,21 @@ if HAS_DYNADIFF:
         if not feats:
             steered_brain_div.text = ''
             return
-        steered = compute_steered_fmri(fmri, feats, lams, thrs)
-        b64 = render_fmri_brain_compact_b64(
-            steered, 'Expected Steered Brain')
-        steered_brain_div.text = (
-            f'<img src="data:image/png;base64,{b64}" '
-            f'style="max-width:100%"/>'
-            if b64 else '')
+        steered_brain_div.text = ''
+        doc = curdoc()
+
+        def _bg():
+            steered = compute_steered_fmri(fmri, feats, lams, thrs)
+            b64 = render_fmri_brain_compact_b64(
+                steered, 'Expected Steered Brain')
+            def _apply():
+                steered_brain_div.text = (
+                    f'<img src="data:image/png;base64,{b64}" '
+                    f'style="max-width:100%"/>'
+                    if b64 else '')
+            doc.add_next_tick_callback(_apply)
+
+        threading.Thread(target=_bg, daemon=True).start()
 
     def _update_active_tiles():
         feats = list(_dd_source.data['feat'])