Issue #8 Visualize Intermediate Representations

[SreeDan]

Attempts to close #8

This PR adds a complete visualization layer (viz/) to the OpenFold fork, along with the extraction infrastructure needed to capture intermediate Evoformer representations and attention weights during a model run. All new code is additive. We didn't change the existing internals of any OpenFold modules except two small hook registration lines in openfold/__init__.py and openfold/utils/import_weights.py.

Architecture is in 3 layers:

model run
   │  EvoformerRecorder hooks (openfold/utils/evoformer_instrumentation.py)
   │  save_attention_topk text files + reps.pt
   ▼
representation_tensor_utils.py
   │  validate · select_channel · aggregate · normalize
   │  prepare_heatmap_data · prepare_lineplot_data
   ▼
viz/integrations.py
   │  *_from_representation  ·  *_from_artifact
   ▼
viz/plots/  →  matplotlib Figure  →  notebook / Flask UI

Layer 1:

openfold/utils/evoformer_instrumentation.py registers PyTorch forward hooks on the Evoformer stack via EvoformerRecorder. During a model run it captures MSA and pair representation tensors at configurable layers (clone --> detach --> CPU --> optional dtype cast) and writes them to reps.pt via torch.save. Attention weights are written as sparse top-K text files by save_attention_topk. This approach adds zero overhead when enabled=False and leaves the model weights untouched.

openfold/utils/evoformer_run_artifact.py wraps the on-disk outputs. EvoformerRunArtifact provides load_reps() and get_attention_matrix(kind, layer, head, ...) which reconstructs a dense N×N matrix from the sparse text files (mean over heads, optional single-residue slice). The artifact is the handoff point that the viz layer reads from.

run_evoformer_hook_pretrained_openfold.py is the end-to-end inference script: loads pretrained weights, attaches hooks, runs the model, and writes all artifacts to outputs/.

Layer 2:

Sits between extraction and rendering. Designed so the same call works for a raw tensor from a model run or a synthetic array from _fakes.py

Layer 3 - Plot functions

Core plots (viz/plots/)

Function	Input shape	Purpose
plot_heatmap	(R, C)	Residue-by-residue map with colorbar and optional red highlight lines
plot_heatmap_grid	(K, R, C)	K panels in a grid, optional shared colorbar
plot_line	(N,)	Single residue-indexed signal
plot_lines	(K, N)	K overlaid signals with legend
plot_layer_trajectory	(L, N)	One channel across L layers, one line per residue
plot_histogram	any shape	Value distribution (flattened, NaNs dropped)

Integration bridge (viz/integrations.py)

• heatmap_from_representation / line_from_representation / lines_from_representation / pair_channel_grid — take a raw pair/msa/single tensor, route it through prepare_*, and call the appropriate plot function. Axis labels and titles are set automatically per kind.
• attention_heatmap_from_artifact / representation_heatmap_from_artifact / representation_line_from_artifact / representation_tensor_from_artifact / figure_from_artifact — read directly from an EvoformerRunArtifact. figure_from_artifact is a single-call dispatcher: pass attn_kind= for attention or rep_kind= for representations.

Flask UI (viz/web_interface.py, viz/render_for_ui.py)

render_for_ui.py is a CLI that pre-generates 96 PNGs (48 layers × 2 attention types).

web_interface.py serves the PNGs with layer/type dropdowns. Configurable via environment variables (VIZFOLD_IMAGE_DIR, VIZFOLD_PROT, VIZFOLD_TRI_IDX, VIZFOLD_NUM_LAYERS) so it can point at any run without editing source.

Synthetic data helpers (viz/_fakes.py)

Shape-matching synthetic tensors (fake_attention_heads, fake_pair_channel, fake_single_channels, fake_layer_trajectory, fake_distribution) so the entire pipeline can be exercised without a GPU or model weights.

Each layer is independently testable and the plot functions accept plain NumPy arrays, so any layer can be bypassed.

Testing:

# Install viz dependencies
pip install -r viz/requirements.txt

# --- Unit / integration tests ---
python -m unittest tests.test_representation_tensor_utils -v
python -m unittest tests.test_viz_integration -v
python -m unittest tests.test_viz_artifact_bridge -v

# --- Flask UI (synthetic mode, no model run needed) ---
python -m viz.render_for_ui --mode demo --n-res 96
# Expected: "[demo] wrote 96 PNGs into outputs/attention_images_6KWC_demo_tri_18"

python viz/web_interface.py
# Open http://localhost:5001
# Verify: layer dropdown 0-47, both attention types load images, protein name in heading

# --- Security check: path traversal must be blocked ---
curl -s -o /dev/null -w "%{http_code}" "http://localhost:5001/image?path=/etc/passwd"
# Expected: 403

# --- Bad query params must not 500 ---
curl -s -o /dev/null -w "%{http_code}" "http://localhost:5001/?attn_type=EVIL&layer=abc"
# Expected: 200 (falls back to msa_row, layer 0)

# --- Env-var config ---
VIZFOLD_NUM_LAYERS=10 python viz/web_interface.py
# Layer dropdown should show only 0-9

# --- Real model run (if a GPU + weights are available) ---
python run_evoformer_hook_pretrained_openfold.py \
    --config_preset model_1 \
    --output_dir outputs/run_6KWC \
    --demo_attn
python -m viz.render_for_ui --mode real \
    --attn-dir outputs/attention_files_6KWC_demo_tri_18
python viz/web_interface.py
# Verify real attention heatmaps load in the UI

Example of end-to-end flow:

demo.mp4