scripts/explorer/README.md

# SAE Feature Explorer

Interactive Bokeh app for exploring SAE (Sparse Autoencoder) features.
Launch with `run_explorer.sh` or directly:

```bash
bokeh serve scripts/explorer --port 5006 --args --data path/to/explorer_data.pt --image-dir /path/to/images
```

The app is then accessible at `http://localhost:5006/explorer`.

---

## File Structure

```
explorer/
├── main.py              # Entry point: loads data, wires callbacks, builds layout
├── args.py              # All CLI argument definitions (argparse)
├── state.py             # Shared mutable state and dataset registry
├── datasets.py          # Dataset loading from .pt files
├── inference.py         # Lazy GPU backbone+SAE and CLIP model loaders
├── brain.py             # Brain-alignment data (Phi) and DynaDiff setup
├── rendering.py         # Pure image/HTML rendering functions
├── widgets.py           # Shared Bokeh widgets used across multiple panels
└── panels/
    ├── umap.py          # UMAP scatter plot and controls
    ├── feature.py       # Feature detail view (heatmaps, stats, cortical profile)
    ├── feature_list.py  # Sortable feature table, name search, Gemini auto-interp
    ├── patch.py         # Patch explorer (click image regions → active features)
    ├── clip_search.py   # Free-text CLIP feature search
    └── dynadiff.py      # DynaDiff brain-steering panel
```

---

## Module Responsibilities

### `main.py`
The Bokeh server entry point. Responsible for three things only:
1. Calling `load_all_datasets()` before any panels are imported
2. Wiring the cross-panel callbacks that need to touch multiple panels (UMAP tap, dataset switch, Go/Random buttons)
3. Assembling the three-column page layout and registering it with `curdoc()`

### `args.py`
Defines and parses all CLI arguments. Exposes a single `args` object imported by any module that needs a CLI value. Parsed once at startup.

### `state.py`
Two things live here:
- `_S` — a plain class used as a mutable namespace for session scalars (active dataset index, render token, search filter, etc.). No `global` statements needed anywhere else.
- `_all_datasets` — the list of loaded dataset dicts. `active_ds()` returns `_all_datasets[_S.active]` and is the standard way for any panel to read the current dataset.
- `display_name(feat)` — returns the best label for a feature (manual → auto-interp → empty).

### `datasets.py`
Contains a single `_load_dataset()` function that handles both regular `explorer_data.pt` files and `brain_meis.pt` files (previously two separate, mostly-duplicate functions). Every loaded dataset is a plain dict with a fixed schema. Derived arrays (`freq`, `log_freq`, `live_mask`, `umap_backup`, etc.) are computed once at load time and stored in the dict, so nothing is recomputed when switching datasets.

### `inference.py`
Lazy loaders for the two optional GPU resources:
- `get_clip()` — loads the CLIP model on first free-text search query
- `get_gpu_runner()` — loads the backbone + SAE on first patch-explorer use
- `run_gpu_inference(pil_img)` — runs an image through the loaded backbone+SAE

Neither is loaded unless actually needed.

### `brain.py`
Loaded at import time from `--phi-dir` and `--dynadiff-dir`. Sets the module-level flags `HAS_PHI` and `HAS_DYNADIFF` that panels check to show/hide features. Provides helper functions for reading per-feature phi data (`phi_voxel_row`, `phi_c_for_feat`, `phi_c_vals`) and rendering brain scatter plots (`render_cortical_profile`, `render_steering_preview`).

### `rendering.py`
Pure functions — no Bokeh widget dependencies. Safe to call from worker threads. Covers:
- Image loading (`load_image`, `resolve_img_path`, `parse_img_label`)
- Heatmap blending (`render_heatmap_overlay`, `render_zoomed_overlay`)
- PIL/Bokeh conversion (`pil_to_data_url`, `pil_to_bokeh_rgba`)
- HTML builders (`make_image_grid_html`, `make_compare_aggregations_html`, `make_steering_html`, `status_html`)
- Layout helper (`make_collapsible`)

`render_zoomed_overlay` takes `zoom_patches` and `alpha` as explicit parameters rather than reading from widgets, keeping it a pure function.

### `widgets.py`
Bokeh widget instances that are read or written by more than one panel. Creating them here avoids duplication and makes the sharing explicit. Includes `feature_input`, `go_button`, `random_btn`, `zoom_slider`, `heatmap_alpha_slider`, `nsd_subset_toggle`, and `view_select`.

### `panels/umap.py`
Owns the UMAP `ColumnDataSource`, the Bokeh figure, the color mapper, and the type/color `Select` controls. Exposes `rebuild_umap_source(ds)` called by `main.py` on dataset switch.

### `panels/feature.py`
Owns the middle column: `stats_div`, `status_div`, `top_heatmap_div`, `mean_heatmap_div`, `compare_agg_div`, `brain_div`. The core function `update_feature_display(feat)` schedules rendering via `add_next_tick_callback` to keep the UI responsive. Uses a render token to discard stale renders if the user clicks quickly. `select_and_display(feat)` additionally syncs the UMAP highlight.

### `panels/feature_list.py`
The sortable feature DataTable on the left, plus:
- Name search (filters the table)
- Manual name editing (auto-saved to JSON, optionally pushed to HuggingFace)
- Gemini auto-interp button (calls Gemini API in a worker thread)

### `panels/patch.py`
The patch explorer: load an image, click or drag-paint patches, see the top SAE features for that region. Activations come from `inference.run_gpu_inference()` with a per-dataset LRU cache.

### `panels/clip_search.py`
Free-text CLIP search against precomputed per-feature image embeddings. Builds a stub Div if no CLIP embeddings are present in the active dataset.

### `panels/dynadiff.py`
Brain-steering panel. Users build a feature list with per-feature λ and brain-voxel threshold values, choose a sample index, and trigger DynaDiff reconstruction. Runs in a worker thread with result pushed back via `add_next_tick_callback`. Builds stub widgets if `HAS_DYNADIFF` is False.

---

## Data Flow

```
CLI args (args.py)
    │
    ▼
load_all_datasets() ──► _all_datasets (state.py)
                              │
                              ▼
                         active_ds() ◄── _S.active (state.py)
                              │
                    ┌─────────┼──────────┐
                    ▼         ▼          ▼
               umap.py   feature.py  feature_list.py  ...
```

All panels read data through `active_ds()['key']`. On a dataset switch, `main.py` updates `_S.active` and calls each panel's `rebuild_*` function — no global variable rebinding required.

---

## Cross-Panel Calls

Panels occasionally need to call functions defined in other panels (e.g., clicking a row in the CLIP results should trigger the feature detail view). These calls are always made from inside callback functions using lazy imports, which avoids circular imports at module load time:

```python
# Inside a callback in clip_search.py — imported lazily, not at module level
def _on_result_select(attr, old, new):
    from .feature import select_and_display
    select_and_display(feat)
```

The one exception is `main.py`, which imports from all panels and wires the few truly cross-cutting callbacks (UMAP tap → feature display, Go/Random buttons).