Integrate ESMFold inference backend into VizFold (structure export + metadata)

[jayvenn21]

Integrate ESMFold inference backend into VizFold (Fixes #43)

Context

This PR introduces the full ESMFold backend for VizFold — inference, hook-based trace extraction, an interactive visualization dashboard, and HPC deployment tooling. It uses the HuggingFace EsmForProteinFolding model and exports outputs in the VizFold archive format.

This addresses Issue #43 end-to-end: inference, trace extraction across all model components, structure module tracing, a FastAPI bridge, and a React frontend for interactive exploration.

---

What's Included

HuggingFace-based ESMFold integration

• Uses facebook/esmfold_v1 via transformers (pinned >=4.36.0)
• Avoids OpenFold CUDA build dependency for improved portability
• ESMFoldRunner implements the BackendBase interface (load_model, run_inference, supports_attention, supports_activations)

GPU-enabled inference on ICE (SLURM-tested)

• Successfully validated with --device cuda on PACE ICE H100 nodes
• Handles large checkpoint download via configurable HF_HOME
• Fail-fast sanity check in SLURM script if torch/transformers not importable

Structure export

• structure/predicted.pdb (full atom37 PDB via HF openfold_utils)
• structure/predicted.pt (coordinate tensor)
• Fallback minimal CA-only PDB with valid 3-letter residue codes

Trace extraction via forward hooks

• ESM-2 encoder attention: hooks on model.esm.encoder.layer[i].attention.self, monkey-patching forward to force output_attentions=True so real attention probabilities [B, H, N, N] are captured
• ESM-2 activations: per-layer hidden states [B, N, D] with <cls> and <eos> special tokens stripped to align with attention maps
• Attention slice validation against expected_seq_len with warnings on mismatch
• Head filtering performed in-hook to reduce GPU memory usage
• Layer and head subsetting via --layers and --heads

Evoformer trunk intermediates

• Per-block sequence state [L, C_s] and pair state [L, L, C_z] via hooks on trunk.blocks[i]
• Final trunk representations s_s and s_z captured from recycling iterations
• Per-recycle s_s and s_z tensors archived as recycle_{i}_s_s / recycle_{i}_s_z

Structure module tracing

• IPA (Invariant Point Attention) attention matrices [H, N, N] captured via CapturingSoftmax patching
• Per-recycle backbone positions and single representations
• Handles both dict and dataclass outputs from HuggingFace via _extract_from_output helper

Interactive visualization dashboard

• React + Vite frontend with 3Dmol.js WebGL structure viewer and Plotly.js heatmaps
• Trunk Evolution mode (s_z averaged across 128 hidden channels) and ESM-2 Attention mode (averaged across heads)
• Bidirectional crosshair sync: click a residue in 3D → crosshair on heatmap; click a heatmap cell → highlight residue in 3D
• Camera-preserving style switching (Cartoon/Sticks/Spheres) without viewer recreation
• Responsive heatmap layout that fills available panel space
• Loading and error states per panel (backend-down detection, missing trace data)
• Configurable API URL via VITE_API_URL environment variable

FastAPI bridge server

• Serves meta.json, PDB structure, and trace tensors over HTTP
• Server-side averaging for s_z (across hidden channels) and attention (across heads) to minimize payload size
• Path traversal protection via safe_join() + explicit ..///\\ checks
• Configurable output directory via --dir CLI argument

Metadata + archive format

• meta.json consistent with VizFold archive schema
• trace/index.json mapping layers to files, shapes, dtypes
• trace/summary.json with per-layer attention entropy and activation norms
• trace_formats field indicating which output formats were produced (pt, txt)
• Shapes recorded for attention, activations, trunk, and structure module outputs

CLI support

• run_pretrained_esmf.py with: --fasta, --out, --device, --dtype, --trace_mode, --layers, --heads, --save_fp16, --seed, --deterministic, --structure_traces

Tests, docs, and HPC support

• tests/test_esmf_smoke.py — CLI help, missing FASTA, schema, import, full smoke-run, and end-to-end validation tests
• docs/esmfold.md — setup, usage examples (attention, activations, structure module traces), output layout, and meta.json schema
• docs/hpc_ice.md — PACE ICE deployment guide
• scripts/hpc/ice/run_esmf_ice.slurm — SLURM batch script with environment validation

---

Design Decision

This implementation uses the HuggingFace ESMFold model rather than directly building OpenFold, which avoids CUDA toolchain mismatches on HPC systems and simplifies reproducibility across environments.

---

Team Contributions (Integrated)

The following team PRs have been reviewed, merged, and integrated into this branch:

• @rohan5986 (PR #3, PR #6, PR #8): Captured s_s and s_z at every recycling iteration via trunk hooks, enforced use_safetensors=True for HPC stability (CVE-2025-32434), built the complete interactive dashboard (3Dmol.js + Plotly heatmaps + FastAPI bridge), added path traversal security, env-based API URL config
• @JeevanandanRamasamy (PR #1, PR #4, PR #7): Extracted Evoformer trunk per-block intermediates, added VizFold-compatible .txt attention export with --top_k support, fixed output_attentions positional arg bug, implemented BackendBase interface, added attention slice validation and in-hook head filtering, stripped <cls>/<eos> from activations, added structure module dict/dataclass handling
• @Mose-Kim02 (PR #5): Validation test suite for ESMFold trace extraction pipeline, reproducibility documentation
• @jayvenn21 (PR #9): Refactored frontend into StructureViewer/TraceExplorer/TimelineControls components, added camera persistence, responsive heatmap, bidirectional crosshair sync, colorbar labels, loading/error states, CSS cleanup

---

Validation Evidence

Validated on PACE ICE with GPU and locally on CPU. ESMFold inference is deterministic; repeated runs produce identical archives.

Example run (attention + activations):

HF_HOME=/tmp/$USER/hf_cache \
python run_pretrained_esmf.py \
  --fasta examples/monomer/fasta_dir_6KWC/6KWC.fasta \
  --out outputs/test_trace \
  --device cuda \
  --trace_mode attention+activations

Example run (structure module traces):

python run_pretrained_esmf.py \
  --fasta examples/monomer/fasta_dir_6KWC/6KWC.fasta \
  --out outputs/esmf_6KWC \
  --trace_mode attention+activations \
  --structure_traces \
  --save_fp16

---

Output Archive

outputs/demo_traces/
├── meta.json
├── logs.txt
├── structure/
│   ├── predicted.pdb
│   └── predicted.pt
├── trace/
│   ├── index.json
│   ├── summary.json
│   ├── attention/              (36 layers)
│   ├── activations/            (36 layers + recycle_*_s_s/s_z)
│   ├── trunk/                  (per-block seq/pair + final s_s/s_z)
│   │   ├── block_000_seq.pt
│   │   ├── block_000_pair.pt
│   │   ├── s_s.pt
│   │   └── s_z.pt
│   └── structure_module/       (with --structure_traces)
│       ├── ipa_attention/
│       │   └── recycle_00_block_00.pt
│       └── backbone/
│           ├── recycle_00_positions.pt
│           └── recycle_00_states.pt
├── attention_files/            (.txt format for VizFold viz tools)
│   └── msa_row_attn_layer0.txt
└── frontend/                   (React dashboard)
    ├── src/
    │   ├── App.jsx
    │   └── components/
    │       ├── StructureViewer.jsx
    │       ├── TraceExplorer.jsx
    │       └── TimelineControls.jsx
    └── server.py               (FastAPI bridge)

---

Architecture Diagram

---

Demo Video

esmfolddemo.mp4

---

Screenshots

---

What's Next

Most planned extensions from the original PR are complete. Remaining tasks:

• Visualization validation — confirm existing VizFold notebooks and PyMOL scripts work with ESMFold traces
• Cross-model comparison — tooling to compare OpenFold and ESMFold archives side-by-side
• Boltz integration — merge Boltz-2 backend work into the main VizFold repo (coordinating with Boltz team)

[JeevanandanRamasamy]

This is a solid addition and the architecture is clean. Documentation is very thorough. I highlighted some potential minor issues in my review.