AlloGen public release: Q_theta scorer + PXDesign guidance + Colab demo

Single commit, no prior history.

Contents:
- Q_theta scorer (graph transformer, SE(3)-invariant + ESM-2 conditioning,
~898K params, v4-S2 target-swap checkpoint via Git LFS).
- PXDesign guidance scripts (Langevin / SMC / TDS / classifier) under
code/scripts/pxdesign_guidance/.
- CaM inference sample (96 binder-CaM graphs + matching ESM-2 features).
- Colab demo at notebooks/AlloGen_CaM_demo.ipynb (one-click for biology users:
load scorer, score 96 designs, view ROC/best-of-K, guidance recipe).
- README with method figure, inference quickstart, full Python scoring API.

MIT license.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

.gitattributes

@@ -0,0 +1,14 @@
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.tar.gz filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text
+*.pdf filter=lfs diff=lfs merge=lfs -text
+*.svg -text

.gitignore

@@ -0,0 +1,41 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+.pytest_cache/
+.coverage
+*.egg-info/
+
+# Env
+.env
+.venv/
+venv/
+env/
+
+# OS
+.DS_Store
+Thumbs.db
+
+# IDE
+.idea/
+.vscode/
+*.swp
+
+# Logs / runs / caches
+*.log
+logs/
+outputs/
+wandb/
+.ipynb_checkpoints/
+
+# Local scoring outputs
+results/
+/tmp_*
+
+# Misc
+*.tmp
+*.bak
+.allogen_test
+.agent*_done.txt

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Hanqun Cao
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,133 @@
+---
+license: mit
+tags:
+  - protein-design
+  - allosteric
+  - state-selectivity
+  - guided-generation
+  - rfdiffusion
+  - pxdesign
+  - proteina
+library_name: pytorch
+---
+
+# AlloGen
+
+<p align="center">
+  <img src="figures/allogen_main.png" alt="AlloGen method overview" width="100%"/>
+</p>
+
+State-selectivity scoring + guided generation for allosteric binder design.
+
+🧪 **One-click demo for biology users:**
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/#fileId=https%3A//huggingface.co/ChatterjeeLab/AlloGen/raw/main/notebooks/AlloGen_CaM_demo.ipynb) — score CaM binders and run Q_θ-guided PXDesign sampling in 5 minutes. Notebook lives at [`notebooks/AlloGen_CaM_demo.ipynb`](notebooks/AlloGen_CaM_demo.ipynb).
+
+AlloGen trains a scorer Q_θ(X, Y) ∈ (0,1) that ranks how well a binder Y discriminates a target's **holo** (active) state X¹ from its **apo** (inactive) state X⁰. The selectivity score is:
+
+    S(Y) = Q_θ(X¹, Y) − Q_θ(X⁰, Y)
+
+Q_θ serves as both a re-ranker (best-of-K) and a gradient signal for guided generation on top of frozen priors (RFdiffusion, PXDesign, Proteina-ComplexA) via Langevin, SMC, TDS, or classifier guidance.
+
+This repository accompanies the paper *AlloGen: State-Selective Scoring for Allosteric Binder Design* (NeurIPS 2026).
+
+## Installation
+
+```bash
+conda env create -f environment.yml
+conda activate allogen
+```
+
+Or pip-only:
+
+```bash
+python -m venv .venv && source .venv/bin/activate
+pip install -r requirements.txt
+```
+
+Python 3.10 + PyTorch 2.x are required. A CUDA GPU is recommended for guidance, but CPU works for scoring single designs.
+
+## Inference quickstart
+
+```bash
+# Score the bundled CaM inference sample against the v4-S2 (target-swap) checkpoint
+python code/scripts/evaluate.py \
+    --target cam \
+    --checkpoint checkpoints/Q_theta_phase2.pt \
+    --data_dir data/sample/ \
+    --outdir /tmp/cam_inference \
+    --no_wandb
+```
+
+See [`inference.md`](inference.md) for the scoring API + guidance command lines.
+
+## Repo layout
+
+```
+code/
+  data/           dataset / graph construction, PDB I/O, target YAMLs
+  models/         Q_θ scorer (graph transformer) + differentiable wrapper
+  trainers/       two-phase training loop (DockQ regression + selectivity)
+  utils/          PDB I/O, backbone frames, SAM optimizer
+  scripts/        evaluate, rescore, PXDesign guidance (see scripts/README.md)
+checkpoints/      Q_θ paper weights (v4-S2 target-swap split, via Git LFS)
+data/sample/      tiny CaM inference sample (test split only)
+```
+
+## Checkpoints
+
+Paper weights for the **v4-S2 target-swap** split are bundled via **Git LFS**:
+
+```bash
+git lfs install
+git lfs pull
+```
+
+| File | Use |
+|---|---|
+| `checkpoints/Q_theta_phase1.pt` | Phase 1 (DockQ regression) intermediate checkpoint |
+| `checkpoints/Q_theta_phase2.pt` | Phase 2 (selectivity) — main paper result |
+| `checkpoints/Q_theta_train_curve.csv` | Training curve metadata |
+
+## Scoring a single design
+
+```python
+import sys; sys.path.insert(0, 'code')
+from models.differentiable_features import DifferentiableQTheta
+
+scorer = DifferentiableQTheta(
+    checkpoint='checkpoints/Q_theta_phase2.pt',
+    device='cuda:0',
+)
+scorer.load_receptor(
+    holo_path='your_holo.pdb', rec_chain='A',
+    apo_path='your_apo.pdb',   apo_chain='A',
+)
+q_holo = scorer.score('design.pdb', binder_chain='B', state='holo')
+q_apo  = scorer.score('design.pdb', binder_chain='B', state='apo')
+print(f'S = {q_holo - q_apo:.3f}')
+```
+
+## Guidance methods
+
+The shipped guidance code wraps **PXDesign** as the prior and uses Q_θ as the gradient / classifier signal. All four method variants (Langevin, SMC, TDS, classifier guidance) live in `code/scripts/pxdesign_guidance/`.
+
+See [`inference.md`](inference.md) §3 for command lines.
+
+To deploy Q_θ with **RFdiffusion**, **Proteina-ComplexA**, or any other backbone prior, see [`code/scripts/README.md`](code/scripts/README.md) — Q_θ exposes `DifferentiableQTheta` for `∇_x S(x)`, and the PXDesign code is a worked template to mirror.
+
+## Citation
+
+```bibtex
+@inproceedings{cao2026allogen,
+  title     = {AlloGen: State-Selective Scoring for Allosteric Binder Design},
+  author    = {Cao, Hanqun and others},
+  booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
+  year      = {2026}
+}
+```
+
+(BibTeX key will be finalized at camera-ready.)
+
+## License
+
+MIT — see [`LICENSE`](LICENSE).

checkpoints/Q_theta_phase1.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1684955f481c1406b12cc0e1ec3509a2a2e2def8b0a9071ec0c96be00d330e7c
+size 3617774

checkpoints/Q_theta_phase2.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:716e4716c0014a46cfd4a2e26c238486c8ad1e3a03809320247cfb734538f8c6
+size 3618158

checkpoints/Q_theta_train_curve.csv

@@ -0,0 +1,16 @@
+epoch,loss,test_rho,cam_rho,rho_cam,rho_bcl2,rho_era,rho_mdm2,rho_ran,rho_a2a,rho_pai1,rho_integrin
+1,5.527279376983643,0.44785173068985473,0.48303455690607044,0.48303455690607044,0.5733270876635986,0.6666325520384175,0.655984079559599,0.4957646288057514,0.28072319745279745,0.32040567098387046,0.10694207210873281
+2,5.452569961547852,0.45406694660326774,0.4863130719981931,0.4863130719981931,0.5755749254666933,0.6683595793753047,0.6593234065151792,0.49658672598782,0.29627094377326,0.33010796670396186,0.11999895300572927
+3,5.449374318122864,0.44742845709394635,0.4966343232141348,0.4966343232141348,0.5758960451528496,0.6612355916106455,0.6531535587702088,0.4753625218665822,0.23926254059823043,0.33456826378855364,0.14331481175036578
+4,5.44292688369751,0.4461415166784438,0.49274867569754505,0.49274867569754505,0.5728775201029797,0.6614514700277563,0.6451799702468982,0.46946643826626333,0.242717595336111,0.34013213953325055,0.1445583242167464
+5,5.441892623901367,0.43839115930587513,0.49080585193925014,0.49080585193925014,0.5653633194469204,0.6586450506053149,0.6361543905961949,0.4544470120828291,0.21334963006412605,0.33976427997988223,0.1485997397324834
+6,5.439353764057159,0.4284122153367078,0.46931336411311275,0.46931336411311275,0.5655559912586142,0.6597244426908693,0.634838828707037,0.4423570277236172,0.19089177426790224,0.3362236317787114,0.1283926621537984
+7,5.439265787601471,0.42722475631179213,0.45923496586695794,0.45923496586695794,0.5661982306309269,0.6614514700277563,0.635386597507575,0.4425217082123247,0.20816704795730515,0.3344762989002116,0.1103617313912795
+8,5.439032256603241,0.423813755390132,0.45146367083377825,0.45146367083377825,0.5662624545681583,0.6633943757817542,0.6325423252001796,0.4388258604749022,0.2029844658504843,0.33617764933454036,0.09885924107725882
+9,5.435689151287079,0.4221364421714867,0.45122081786399143,0.45122081786399143,0.5656202151958456,0.660156199525091,0.6294276063720167,0.4378717142280125,0.19952941111260372,0.33502808823026414,0.09823748484406851
+10,5.43562650680542,0.419534099250895,0.4498851265301637,0.4498851265301637,0.5641430646395261,0.662099105279089,0.6279539020257997,0.4329763310874754,0.1822541374232008,0.33530398289529045,0.10165714412661521
+11,5.435124695301056,0.41942099211035866,0.451949376773352,0.451949376773352,0.5633723773927508,0.6612355916106455,0.6249514872613453,0.43204285159336153,0.1822541374232008,0.3360397020020272,0.10352241282618613
+12,5.440709412097931,0.41971570222256677,0.4518279502884586,0.4518279502884586,0.5642072885767574,0.6612355916106455,0.6261524531671271,0.43317581854869713,0.1822541374232008,0.33534996533946143,0.10352241282618613
+13,5.430392742156982,0.4198042517465396,0.4523136562280323,0.4523136562280323,0.5629228098321319,0.6627467405304217,0.6259874349510655,0.4334288217301594,0.1822541374232008,0.33525800045111936,0.10352241282618613
+14,5.436960756778717,0.4199312463250485,0.4515850973186718,0.4515850973186718,0.5639503928278323,0.6627467405304217,0.6264251916075623,0.4329262960644863,0.1822541374232008,0.3360397020020272,0.10352241282618613
+15,5.433559775352478,0.419866834088955,0.4515850973186718,0.4515850973186718,0.5635650492044447,0.6627467405304217,0.6264251916075623,0.432842324243296,0.1822541374232008,0.3359937195578562,0.10352241282618613

code/data/dataset.py

@@ -0,0 +1,832 @@
+"""
+PyTorch Dataset for two-state complex scoring.
+
+Loads preprocessed graph data and provides batched tensors
+with padding for variable-sized interface graphs.
+"""
+
+import os
+import json
+import pickle
+import numpy as np
+import torch
+from torch.utils.data import Dataset, DataLoader
+
+# Global ESM embedding cache: {file_path: tensor}
+_ESM_CACHE = {}
+
+
+def preload_esm_cache(esm_dir, targets):
+    """Preload all ESM .pt files into global cache before DataLoader workers fork.
+
+    This ensures forked workers inherit the populated cache via copy-on-write,
+    avoiding redundant I/O across workers.
+    """
+    import glob as glob_mod
+    n = 0
+    for target in targets:
+        target_dir = os.path.join(esm_dir, target)
+        if not os.path.isdir(target_dir):
+            continue
+        for pt_file in glob_mod.glob(os.path.join(target_dir, '*.pt')):
+            if pt_file not in _ESM_CACHE:
+                _ESM_CACHE[pt_file] = torch.load(pt_file, map_location='cpu', weights_only=True)
+                n += 1
+    return n
+
+
+def load_esm_for_sample(sample, esm_dir, target_name, max_nodes=128):
+    """Load and index ESM-2 embeddings for a sample's interface residues.
+
+    Returns: esm_feats [max_nodes, 1280] or None if unavailable.
+    """
+    graph = sample['graph']
+    rec_idx = graph.get('rec_iface_idx')
+    binder_idx = graph.get('binder_iface_idx')
+    if rec_idx is None or binder_idx is None:
+        return None
+
+    # Get PDB ID (strip chain suffix like "2G1T_AE" -> "2G1T")
+    pdb_id = sample.get('pdb', '')
+    base_pdb = pdb_id.split('_')[0] if '_' in pdb_id else pdb_id
+    rec_chain = sample.get('rec_chain_id', 'A')
+    binder_chain = sample.get('binder_chain_id', 'B')
+
+    # Load ESM embeddings (cached)
+    rec_path = os.path.join(esm_dir, target_name, f'{base_pdb}_{rec_chain}.pt')
+    binder_path = os.path.join(esm_dir, target_name, f'{base_pdb}_{binder_chain}.pt')
+
+    def _load_cached(path):
+        if path not in _ESM_CACHE:
+            if not os.path.exists(path):
+                return None
+            _ESM_CACHE[path] = torch.load(path, map_location='cpu', weights_only=True)
+        return _ESM_CACHE[path]
+
+    rec_esm = _load_cached(rec_path)
+    binder_esm = _load_cached(binder_path)
+    if rec_esm is None or binder_esm is None:
+        return None
+
+    esm_dim = rec_esm.shape[-1]  # 1280
+    n_rec = len(rec_idx)
+    n_binder = len(binder_idx)
+
+    # Index ESM embeddings by interface residue indices (clamp to valid range)
+    rec_idx_safe = np.clip(rec_idx, 0, len(rec_esm) - 1)
+    binder_idx_safe = np.clip(binder_idx, 0, len(binder_esm) - 1)
+
+    esm_feats = np.zeros((max_nodes, esm_dim), dtype=np.float32)
+    esm_feats[:n_rec] = rec_esm[rec_idx_safe].numpy()
+    esm_feats[n_rec:n_rec + n_binder] = binder_esm[binder_idx_safe].numpy()
+
+    return esm_feats
+
+
+def load_rosetta_labels(rosetta_dir, target):
+    """Load Rosetta dG labels for a target and normalize to [0,1]."""
+    path = os.path.join(rosetta_dir, f'{target}_rosetta.json')
+    if not os.path.exists(path):
+        return None
+    with open(path) as f:
+        raw = json.load(f)
+    if not raw:
+        return None
+    # Filter outliers: dG values outside [-500, 500] are failed Rosetta runs
+    dG_MIN, dG_MAX = -500.0, 500.0
+    # Normalize: sigmoid(-dG / tau) maps dG to [0,1]
+    # More negative dG = better binding = higher score
+    tau = 15.0  # temperature; dG=-30 -> 0.88, dG=-15 -> 0.73, dG=0 -> 0.5
+    labels = {}
+    for pdb_id, metrics in raw.items():
+        dG = metrics.get('dG_separated', 0.0)
+        if not np.isfinite(dG) or dG < dG_MIN or dG > dG_MAX:
+            continue  # skip failed Rosetta runs
+        labels[pdb_id] = 1.0 / (1.0 + np.exp(dG / tau))
+        labels[pdb_id.upper()] = labels[pdb_id]
+        labels[pdb_id.lower()] = labels[pdb_id]
+    return labels
+
+
+def apply_rosetta_labels(samples, rosetta_labels, label_source='rosetta', alpha=0.5):
+    """Replace or combine sample labels with Rosetta-derived labels."""
+    if rosetta_labels is None:
+        return
+    n_replaced = 0
+    for s in samples:
+        pdb_id = s.get('pdb', '')
+        # Strip chain suffixes: "2G1T_AE" -> "2G1T"
+        base_pdb = pdb_id.split('_')[0] if '_' in pdb_id else pdb_id
+        rosetta_val = rosetta_labels.get(base_pdb) or rosetta_labels.get(base_pdb.upper())
+        if rosetta_val is None:
+            continue
+        if s['type'] == 'positive':
+            new_label = rosetta_val
+        elif s['type'].startswith('negative'):
+            new_label = 0.0  # apo mismatch stays 0
+            continue
+        elif s['type'].startswith('decoy'):
+            # Scale Rosetta label by DockQ-proxy quality
+            new_label = s['label'] * rosetta_val
+        else:
+            continue
+        if label_source == 'rosetta':
+            s['label'] = float(new_label)
+        elif label_source == 'combined':
+            s['label'] = float(alpha * s['label'] + (1 - alpha) * new_label)
+        n_replaced += 1
+    return n_replaced
+
+
+class TwoStateComplexDataset(Dataset):
+    """
+    Dataset of protein complex interface graphs with two-state labels.
+
+    Each sample contains:
+        node_feats: [N, node_dim] interface residue features
+        edge_feats: [N, N, edge_dim] pairwise SE(3)-invariant features
+        node_mask: [N] bool
+        label: scalar float in [0, 1] (DockQ proxy / selectivity label)
+        type: str (positive / negative_apo / decoy_*)
+        pdb: str
+    """
+
+    def __init__(self, data_path: str, max_nodes: int = 128, augment: bool = False,
+                 rosetta_labels: dict = None, label_source: str = 'dockq',
+                 esm_dir: str = None, target_name: str = None,
+                 binder_dropout: float = 0.0):
+        with open(data_path, 'rb') as f:
+            self.samples = pickle.load(f)
+        self.max_nodes = max_nodes
+        self.augment = augment
+        self.esm_dir = esm_dir
+        self.target_name = target_name
+        self.binder_dropout = binder_dropout
+        if label_source != 'dockq' and rosetta_labels:
+            apply_rosetta_labels(self.samples, rosetta_labels, label_source)
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        sample = self.samples[idx]
+        graph = sample['graph']
+
+        node_feats = graph['node_feats']  # [N, node_dim]
+        edge_feats = graph['edge_feats']  # [N, N, edge_dim]
+        node_mask = graph['node_mask']    # [N]
+
+        N = len(node_feats)
+        assert N <= self.max_nodes, f"Too many nodes: {N} > {self.max_nodes}"
+
+        # Pad to max_nodes
+        node_dim = node_feats.shape[-1]
+        edge_dim = edge_feats.shape[-1]
+
+        node_feats_pad = np.zeros((self.max_nodes, node_dim), dtype=np.float32)
+        edge_feats_pad = np.zeros((self.max_nodes, self.max_nodes, edge_dim), dtype=np.float32)
+        node_mask_pad = np.zeros(self.max_nodes, dtype=bool)
+
+        node_feats_pad[:N] = node_feats
+        edge_feats_pad[:N, :N] = edge_feats
+        node_mask_pad[:N] = node_mask
+
+        # Optional: random coordinate noise augmentation
+        if self.augment:
+            noise = np.random.randn(*node_feats_pad.shape) * 0.01
+            node_feats_pad = node_feats_pad + noise.astype(np.float32)
+
+        # Binder-dropout: simulate backbone-only designs by masking binder
+        # sequence features (AA one-hot → UNK, chi angles → 0)
+        apply_binder_drop = (self.binder_dropout > 0
+                             and np.random.rand() < self.binder_dropout)
+        if apply_binder_drop:
+            n_rec = graph.get('n_rec', N // 2)
+            # Zero out binder AA one-hot (dims 0-20), set UNK (dim 20 = 1)
+            node_feats_pad[n_rec:N, :21] = 0.0
+            node_feats_pad[n_rec:N, 20] = 1.0  # UNK
+            # Zero out binder chi angles (dims 27-30)
+            node_feats_pad[n_rec:N, 27:31] = 0.0
+            # Keep backbone torsions (dims 21-26) and chain indicator (dim 31)
+
+        result = {
+            'node_feats': torch.from_numpy(node_feats_pad),   # [max_nodes, node_dim]
+            'edge_feats': torch.from_numpy(edge_feats_pad),   # [max_nodes, max_nodes, edge_dim]
+            'node_mask': torch.from_numpy(node_mask_pad),     # [max_nodes]
+            'label': torch.tensor(sample['label'], dtype=torch.float32),
+            'type': sample['type'],
+            'pdb': sample['pdb'],
+        }
+
+        # ESM-2 features (lazy load; zero-fill if unavailable)
+        if self.esm_dir:
+            esm = load_esm_for_sample(sample, self.esm_dir,
+                                      self.target_name or '', self.max_nodes)
+            if esm is not None:
+                esm_feats = esm
+            else:
+                esm_feats = np.zeros((self.max_nodes, 1280), dtype=np.float32)
+            # Zero binder ESM if binder-dropout active
+            if apply_binder_drop:
+                n_rec = graph.get('n_rec', N // 2)
+                n_binder = graph.get('n_binder', N - n_rec)
+                esm_feats[n_rec:n_rec + n_binder] = 0.0
+            result['esm_feats'] = torch.from_numpy(esm_feats)
+
+        return result
+
+
+def collate_fn(batch):
+    """Collate a list of samples into batched tensors."""
+    node_feats = torch.stack([s['node_feats'] for s in batch])
+    edge_feats = torch.stack([s['edge_feats'] for s in batch])
+    node_mask = torch.stack([s['node_mask'] for s in batch])
+    labels = torch.stack([s['label'] for s in batch])
+    types = [s['type'] for s in batch]
+    pdbs = [s['pdb'] for s in batch]
+
+    result = {
+        'node_feats': node_feats,  # [B, N, node_dim]
+        'edge_feats': edge_feats,  # [B, N, N, edge_dim]
+        'node_mask': node_mask,    # [B, N]
+        'label': labels,           # [B]
+        'type': types,
+        'pdb': pdbs,
+    }
+
+    # Stack ESM features if present (handle mixed availability with zero-fill)
+    has_esm = any('esm_feats' in s for s in batch)
+    if has_esm:
+        esm_list = []
+        for s in batch:
+            if 'esm_feats' in s:
+                esm_list.append(s['esm_feats'])
+            else:
+                # Get shape from a sample that has ESM
+                ref = next(x['esm_feats'] for x in batch if 'esm_feats' in x)
+                esm_list.append(torch.zeros_like(ref))
+        result['esm_feats'] = torch.stack(esm_list)
+
+    return result
+
+
+class TwoStateDatasetPaired(Dataset):
+    """
+    Paired dataset: returns (positive, negative) pairs for selectivity training.
+    Groups samples by PDB ID and pairs positive (holo) with negative (apo) examples.
+    """
+
+    def __init__(self, data_path: str, max_nodes: int = 128, augment: bool = False,
+                 esm_dir: str = None, target_name: str = None,
+                 binder_dropout: float = 0.0):
+        with open(data_path, 'rb') as f:
+            samples = pickle.load(f)
+        self.max_nodes = max_nodes
+        self.augment = augment
+        self.esm_dir = esm_dir
+        self.target_name = target_name
+        self.binder_dropout = binder_dropout
+
+        # Group by PDB
+        from collections import defaultdict
+        by_pdb = defaultdict(lambda: {'positive': [], 'negative': [], 'decoy': []})
+        for s in samples:
+            pdb = s['pdb']
+            t = s['type']
+            if t == 'positive':
+                by_pdb[pdb]['positive'].append(s)
+            elif t.startswith('negative'):
+                by_pdb[pdb]['negative'].append(s)
+            elif t.startswith('decoy'):
+                by_pdb[pdb]['decoy'].append(s)
+
+        # Build pairs: (positive, negative) per PDB
+        self.pairs = []
+        for pdb, groups in by_pdb.items():
+            if len(groups['positive']) > 0 and len(groups['negative']) > 0:
+                for pos in groups['positive']:
+                    for neg in groups['negative']:
+                        self.pairs.append((pos, neg))
+            # Also add (positive, decoy_large_rmsd) pairs
+            if len(groups['positive']) > 0 and len(groups['decoy']) > 0:
+                large_decoys = [s for s in groups['decoy'] if 'rmsd' in s['type'] and
+                                float(s['type'].replace('decoy_rmsd', '')) > 4.0]
+                for pos in groups['positive']:
+                    for neg in large_decoys[:3]:  # limit to 3 hard decoys per positive
+                        self.pairs.append((pos, neg))
+
+    def __len__(self):
+        return len(self.pairs)
+
+    def _prepare(self, sample, apply_binder_drop=False):
+        graph = sample['graph']
+        node_feats = graph['node_feats']
+        edge_feats = graph['edge_feats']
+        node_mask = graph['node_mask']
+        N = len(node_feats)
+        node_dim = node_feats.shape[-1]
+        edge_dim = edge_feats.shape[-1]
+
+        node_feats_pad = np.zeros((self.max_nodes, node_dim), dtype=np.float32)
+        edge_feats_pad = np.zeros((self.max_nodes, self.max_nodes, edge_dim), dtype=np.float32)
+        node_mask_pad = np.zeros(self.max_nodes, dtype=bool)
+
+        n = min(N, self.max_nodes)
+        node_feats_pad[:n] = node_feats[:n]
+        edge_feats_pad[:n, :n] = edge_feats[:n, :n]
+        node_mask_pad[:n] = node_mask[:n]
+
+        # Binder-dropout: simulate backbone-only designs
+        if apply_binder_drop:
+            n_rec = graph.get('n_rec', n // 2)
+            node_feats_pad[n_rec:n, :21] = 0.0
+            node_feats_pad[n_rec:n, 20] = 1.0  # UNK
+            node_feats_pad[n_rec:n, 27:31] = 0.0
+
+        result = {
+            'node_feats': torch.from_numpy(node_feats_pad),
+            'edge_feats': torch.from_numpy(edge_feats_pad),
+            'node_mask': torch.from_numpy(node_mask_pad),
+            'label': torch.tensor(sample['label'], dtype=torch.float32),
+            'contact_energy': torch.tensor(
+                sample.get('contact_energy', 0.5), dtype=torch.float32
+            ),
+        }
+
+        # ESM-2 features (zero-fill if unavailable)
+        if self.esm_dir:
+            esm = load_esm_for_sample(sample, self.esm_dir,
+                                      self.target_name or '', self.max_nodes)
+            if esm is not None:
+                esm_feats = esm
+            else:
+                esm_feats = np.zeros((self.max_nodes, 1280), dtype=np.float32)
+            if apply_binder_drop:
+                n_rec = graph.get('n_rec', n // 2)
+                n_binder = graph.get('n_binder', n - n_rec)
+                esm_feats[n_rec:n_rec + n_binder] = 0.0
+            result['esm_feats'] = torch.from_numpy(esm_feats)
+
+        return result
+
+    def __getitem__(self, idx):
+        pos_sample, neg_sample = self.pairs[idx]
+        # Same dropout decision for both pos and neg in a pair
+        drop = (self.binder_dropout > 0
+                and np.random.rand() < self.binder_dropout)
+        return {
+            'pos': self._prepare(pos_sample, apply_binder_drop=drop),
+            'neg': self._prepare(neg_sample, apply_binder_drop=drop),
+        }
+
+
+def collate_paired_fn(batch):
+    """Collate paired (positive, negative) samples."""
+    pos_batch = {
+        'node_feats': torch.stack([s['pos']['node_feats'] for s in batch]),
+        'edge_feats': torch.stack([s['pos']['edge_feats'] for s in batch]),
+        'node_mask': torch.stack([s['pos']['node_mask'] for s in batch]),
+        'label': torch.stack([s['pos']['label'] for s in batch]),
+        'contact_energy': torch.stack([s['pos']['contact_energy'] for s in batch]),
+    }
+    neg_batch = {
+        'node_feats': torch.stack([s['neg']['node_feats'] for s in batch]),
+        'edge_feats': torch.stack([s['neg']['edge_feats'] for s in batch]),
+        'node_mask': torch.stack([s['neg']['node_mask'] for s in batch]),
+        'label': torch.stack([s['neg']['label'] for s in batch]),
+        'contact_energy': torch.stack([s['neg']['contact_energy'] for s in batch]),
+    }
+    # ESM features (handle mixed availability)
+    has_pos_esm = any('esm_feats' in s['pos'] for s in batch)
+    if has_pos_esm:
+        def _stack_esm(batch_list, key):
+            esm_list = []
+            ref = next((x[key]['esm_feats'] for x in batch_list if 'esm_feats' in x[key]), None)
+            for s in batch_list:
+                if 'esm_feats' in s[key]:
+                    esm_list.append(s[key]['esm_feats'])
+                else:
+                    esm_list.append(torch.zeros_like(ref))
+            return torch.stack(esm_list)
+        pos_batch['esm_feats'] = _stack_esm(batch, 'pos')
+        neg_batch['esm_feats'] = _stack_esm(batch, 'neg')
+    return {'pos': pos_batch, 'neg': neg_batch}
+
+
+class PathAwareDatasetPaired(Dataset):
+    """
+    Paired dataset with transition-path frames for path-aware Phase 2 training.
+
+    Extends TwoStateDatasetPaired: each sample returns (positive, negative, path_frames)
+    where path_frames is a list of prepared graph dicts for intermediate conformations
+    stored in the positive sample's 'path_graphs' field.
+    """
+
+    def __init__(self, data_path: str, max_nodes: int = 128, augment: bool = False):
+        with open(data_path, 'rb') as f:
+            samples = pickle.load(f)
+        self.max_nodes = max_nodes
+        self.augment = augment
+
+        from collections import defaultdict
+        by_pdb = defaultdict(lambda: {'positive': [], 'negative': [], 'decoy': []})
+        for s in samples:
+            pdb = s['pdb']
+            t = s['type']
+            if t == 'positive':
+                by_pdb[pdb]['positive'].append(s)
+            elif t.startswith('negative'):
+                by_pdb[pdb]['negative'].append(s)
+            elif t.startswith('decoy'):
+                by_pdb[pdb]['decoy'].append(s)
+
+        self.pairs = []
+        for pdb, groups in by_pdb.items():
+            if len(groups['positive']) > 0 and len(groups['negative']) > 0:
+                for pos in groups['positive']:
+                    for neg in groups['negative']:
+                        self.pairs.append((pos, neg))
+            if len(groups['positive']) > 0 and len(groups['decoy']) > 0:
+                large_decoys = [s for s in groups['decoy'] if 'rmsd' in s['type'] and
+                                float(s['type'].replace('decoy_rmsd', '')) > 4.0]
+                for pos in groups['positive']:
+                    for neg in large_decoys[:3]:
+                        self.pairs.append((pos, neg))
+
+    def _prepare(self, sample):
+        graph = sample['graph']
+        node_feats = graph['node_feats']
+        edge_feats = graph['edge_feats']
+        node_mask = graph['node_mask']
+        N = len(node_feats)
+        node_dim = node_feats.shape[-1]
+        edge_dim = edge_feats.shape[-1]
+
+        node_feats_pad = np.zeros((self.max_nodes, node_dim), dtype=np.float32)
+        edge_feats_pad = np.zeros((self.max_nodes, self.max_nodes, edge_dim), dtype=np.float32)
+        node_mask_pad = np.zeros(self.max_nodes, dtype=bool)
+
+        n = min(N, self.max_nodes)
+        node_feats_pad[:n] = node_feats[:n]
+        edge_feats_pad[:n, :n] = edge_feats[:n, :n]
+        node_mask_pad[:n] = node_mask[:n]
+
+        return {
+            'node_feats': torch.from_numpy(node_feats_pad),
+            'edge_feats': torch.from_numpy(edge_feats_pad),
+            'node_mask': torch.from_numpy(node_mask_pad),
+            'label': torch.tensor(sample.get('label', 0.0), dtype=torch.float32),
+            'contact_energy': torch.tensor(
+                sample.get('contact_energy', 0.5), dtype=torch.float32
+            ),
+        }
+
+    def _prepare_graph_only(self, path_entry):
+        """Prepare a path frame graph (no label/contact_energy needed)."""
+        graph = path_entry['graph']
+        node_feats = graph['node_feats']
+        edge_feats = graph['edge_feats']
+        node_mask = graph['node_mask']
+        N = len(node_feats)
+        node_dim = node_feats.shape[-1]
+        edge_dim = edge_feats.shape[-1]
+
+        node_feats_pad = np.zeros((self.max_nodes, node_dim), dtype=np.float32)
+        edge_feats_pad = np.zeros((self.max_nodes, self.max_nodes, edge_dim), dtype=np.float32)
+        node_mask_pad = np.zeros(self.max_nodes, dtype=bool)
+
+        n = min(N, self.max_nodes)
+        node_feats_pad[:n] = node_feats[:n]
+        edge_feats_pad[:n, :n] = edge_feats[:n, :n]
+        node_mask_pad[:n] = node_mask[:n]
+
+        return {
+            'node_feats': torch.from_numpy(node_feats_pad),
+            'edge_feats': torch.from_numpy(edge_feats_pad),
+            'node_mask': torch.from_numpy(node_mask_pad),
+        }
+
+    def __len__(self):
+        return len(self.pairs)
+
+    def __getitem__(self, idx):
+        pos_sample, neg_sample = self.pairs[idx]
+        result = {
+            'pos': self._prepare(pos_sample),
+            'neg': self._prepare(neg_sample),
+        }
+
+        # Prepare path frames if available
+        path_graphs = pos_sample.get('path_graphs', [])
+        prepared_paths = []
+        path_taus = []
+        for pg in path_graphs:
+            prepared_paths.append(self._prepare_graph_only(pg))
+            path_taus.append(pg['tau'])
+
+        result['path'] = prepared_paths
+        result['path_taus'] = path_taus
+
+        return result
+
+
+def collate_path_paired_fn(batch):
+    """Collate paired samples with variable-length path frames."""
+    pos_batch = {
+        'node_feats': torch.stack([s['pos']['node_feats'] for s in batch]),
+        'edge_feats': torch.stack([s['pos']['edge_feats'] for s in batch]),
+        'node_mask': torch.stack([s['pos']['node_mask'] for s in batch]),
+        'label': torch.stack([s['pos']['label'] for s in batch]),
+        'contact_energy': torch.stack([s['pos']['contact_energy'] for s in batch]),
+    }
+    neg_batch = {
+        'node_feats': torch.stack([s['neg']['node_feats'] for s in batch]),
+        'edge_feats': torch.stack([s['neg']['edge_feats'] for s in batch]),
+        'node_mask': torch.stack([s['neg']['node_mask'] for s in batch]),
+        'label': torch.stack([s['neg']['label'] for s in batch]),
+        'contact_energy': torch.stack([s['neg']['contact_energy'] for s in batch]),
+    }
+
+    # Collate path frames: find max K across batch, pad shorter ones
+    max_k = max((len(s['path']) for s in batch), default=0)
+    path_batches = []
+    path_taus = []
+
+    if max_k > 0:
+        # Build a zero-filled placeholder for padding (graph-only keys)
+        ref = batch[0]['path'][0] if batch[0]['path'] else batch[0]['pos']
+        zero_placeholder = {
+            'node_feats': torch.zeros_like(ref['node_feats']),
+            'edge_feats': torch.zeros_like(ref['edge_feats']),
+            'node_mask': torch.zeros_like(ref['node_mask']),
+        }
+
+        for k_idx in range(max_k):
+            frames_at_k = []
+            taus_at_k = []
+            for s in batch:
+                if k_idx < len(s['path']):
+                    frames_at_k.append(s['path'][k_idx])
+                    taus_at_k.append(s['path_taus'][k_idx])
+                else:
+                    frames_at_k.append(zero_placeholder)
+                    taus_at_k.append(1.0)
+
+            path_batches.append({
+                'node_feats': torch.stack([f['node_feats'] for f in frames_at_k]),
+                'edge_feats': torch.stack([f['edge_feats'] for f in frames_at_k]),
+                'node_mask': torch.stack([f['node_mask'] for f in frames_at_k]),
+            })
+            path_taus.append(taus_at_k[0])
+
+    result = {'pos': pos_batch, 'neg': neg_batch}
+    if path_batches:
+        result['path'] = path_batches
+        result['path_taus'] = path_taus
+    return result
+
+
+class MultiTargetDataset(Dataset):
+    """
+    Pooled dataset combining samples from multiple targets.
+    Supports balanced sampling across targets.
+    """
+
+    def __init__(self, data_paths: list, max_nodes: int = 128, augment: bool = False,
+                 balance: bool = True, rosetta_dir: str = None, label_source: str = 'dockq',
+                 esm_dir: str = None, binder_dropout: float = 0.0):
+        """
+        Args:
+            data_paths: list of (target_name, pkl_path) tuples
+            max_nodes: max interface graph size
+            augment: apply noise augmentation
+            balance: if True, oversample smaller targets to balance
+            rosetta_dir: directory containing Rosetta label JSONs
+            label_source: 'dockq', 'rosetta', or 'combined'
+        """
+        self.max_nodes = max_nodes
+        self.augment = augment
+        self.esm_dir = esm_dir
+        self.binder_dropout = binder_dropout
+
+        # Load all samples with target labels
+        self.samples = []
+        self.target_indices = {}  # target_name -> list of indices
+
+        for target_name, path in data_paths:
+            if not os.path.exists(path):
+                continue
+            with open(path, 'rb') as f:
+                target_samples = pickle.load(f)
+
+            # Apply Rosetta labels if requested
+            if label_source != 'dockq' and rosetta_dir:
+                rl = load_rosetta_labels(rosetta_dir, target_name)
+                if rl:
+                    apply_rosetta_labels(target_samples, rl, label_source)
+
+            start_idx = len(self.samples)
+            for s in target_samples:
+                s['_target'] = target_name
+                self.samples.append(s)
+            end_idx = len(self.samples)
+            self.target_indices[target_name] = list(range(start_idx, end_idx))
+
+        # Build balanced sampling weights
+        if balance and len(self.target_indices) > 1:
+            non_empty = {k: v for k, v in self.target_indices.items() if len(v) > 0}
+            max_count = max(len(idxs) for idxs in non_empty.values()) if non_empty else 1
+            self.weights = np.zeros(len(self.samples))
+            for target_name, idxs in self.target_indices.items():
+                if len(idxs) == 0:
+                    continue
+                weight = max_count / len(idxs)
+                for i in idxs:
+                    self.weights[i] = weight
+            self.weights /= self.weights.sum()
+        else:
+            self.weights = None
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        sample = self.samples[idx]
+        graph = sample['graph']
+        node_feats = graph['node_feats']
+        edge_feats = graph['edge_feats']
+        node_mask = graph['node_mask']
+        N = len(node_feats)
+        node_dim = node_feats.shape[-1]
+        edge_dim = edge_feats.shape[-1]
+
+        node_feats_pad = np.zeros((self.max_nodes, node_dim), dtype=np.float32)
+        edge_feats_pad = np.zeros((self.max_nodes, self.max_nodes, edge_dim), dtype=np.float32)
+        node_mask_pad = np.zeros(self.max_nodes, dtype=bool)
+
+        n = min(N, self.max_nodes)
+        node_feats_pad[:n] = node_feats[:n]
+        edge_feats_pad[:n, :n] = edge_feats[:n, :n]
+        node_mask_pad[:n] = node_mask[:n]
+
+        if self.augment:
+            noise = np.random.randn(*node_feats_pad.shape) * 0.01
+            node_feats_pad = node_feats_pad + noise.astype(np.float32)
+
+        # Binder-dropout: simulate backbone-only designs
+        apply_binder_drop = (self.binder_dropout > 0
+                             and np.random.rand() < self.binder_dropout)
+        if apply_binder_drop:
+            n_rec = graph.get('n_rec', N // 2)
+            node_feats_pad[n_rec:N, :21] = 0.0
+            node_feats_pad[n_rec:N, 20] = 1.0  # UNK
+            node_feats_pad[n_rec:N, 27:31] = 0.0
+
+        result = {
+            'node_feats': torch.from_numpy(node_feats_pad),
+            'edge_feats': torch.from_numpy(edge_feats_pad),
+            'node_mask': torch.from_numpy(node_mask_pad),
+            'label': torch.tensor(sample['label'], dtype=torch.float32),
+            'type': sample['type'],
+            'pdb': sample['pdb'],
+            'target': sample.get('_target', 'unknown'),
+        }
+
+        # ESM-2 features (zero-fill if unavailable)
+        if self.esm_dir:
+            target_name = sample.get('_target', 'unknown')
+            esm = load_esm_for_sample(sample, self.esm_dir, target_name, self.max_nodes)
+            if esm is not None:
+                esm_feats = esm
+            else:
+                esm_feats = np.zeros((self.max_nodes, 1280), dtype=np.float32)
+            if apply_binder_drop:
+                n_rec = graph.get('n_rec', N // 2)
+                n_binder = graph.get('n_binder', N - n_rec)
+                esm_feats[n_rec:n_rec + n_binder] = 0.0
+            result['esm_feats'] = torch.from_numpy(esm_feats)
+
+        return result
+
+    @staticmethod
+    def get_pooled_dataloaders(data_dir, targets, batch_size=16, max_nodes=128,
+                                num_workers=4, paired=False,
+                                rosetta_dir=None, label_source='dockq',
+                                esm_dir=None, binder_dropout=0.0):
+        """Build pooled dataloaders from multiple targets.
+
+        Args:
+            data_dir: root data directory
+            targets: list of target names
+            batch_size: batch size
+            max_nodes: max interface nodes
+            num_workers: dataloader workers
+            paired: if True, build paired dataloaders for Phase 2
+            rosetta_dir: directory with Rosetta label JSONs
+            label_source: 'dockq', 'rosetta', or 'combined'
+        """
+        from torch.utils.data import WeightedRandomSampler
+
+        # Preload ESM embeddings into global cache before creating datasets/workers
+        if esm_dir:
+            n_loaded = preload_esm_cache(esm_dir, targets)
+
+        loaders = {}
+        for split in ['train', 'val', 'test']:
+            data_paths = []
+            for target in targets:
+                path = os.path.join(data_dir, target, f"{split}.pkl")
+                if os.path.exists(path):
+                    data_paths.append((target, path))
+
+            if not data_paths:
+                continue
+
+            augment = (split == 'train')
+            bd = binder_dropout if split == 'train' else 0.0
+
+            if paired:
+                # For paired mode, concatenate paired datasets
+                all_pairs = []
+                for target, path in data_paths:
+                    ds = TwoStateDatasetPaired(path, max_nodes=max_nodes, augment=augment,
+                                               esm_dir=esm_dir, target_name=target,
+                                               binder_dropout=bd)
+                    all_pairs.append(ds)
+
+                if not all_pairs:
+                    continue
+
+                # Use ConcatDataset
+                from torch.utils.data import ConcatDataset
+                concat_ds = ConcatDataset(all_pairs)
+                p_batch = min(batch_size, max(1, len(concat_ds) // 2))
+                loaders[split] = DataLoader(
+                    concat_ds, batch_size=p_batch,
+                    shuffle=(split == 'train'),
+                    num_workers=num_workers,
+                    collate_fn=collate_paired_fn,
+                    pin_memory=True,
+                )
+            else:
+                dataset = MultiTargetDataset(data_paths, max_nodes=max_nodes,
+                                             augment=augment, balance=(split == 'train'),
+                                             rosetta_dir=rosetta_dir, label_source=label_source,
+                                             esm_dir=esm_dir, binder_dropout=bd)
+
+                sampler = None
+                shuffle = (split == 'train')
+                if split == 'train' and dataset.weights is not None:
+                    sampler = WeightedRandomSampler(
+                        weights=dataset.weights,
+                        num_samples=len(dataset),
+                        replacement=True
+                    )
+                    shuffle = False
+
+                loaders[split] = DataLoader(
+                    dataset, batch_size=batch_size,
+                    shuffle=shuffle, sampler=sampler,
+                    num_workers=num_workers,
+                    collate_fn=collate_fn,
+                    pin_memory=True,
+                    drop_last=(split == 'train' and len(dataset) > batch_size),
+                )
+
+        return loaders
+
+
+def get_dataloaders(data_dir: str, target: str, batch_size: int = 16,
+                    max_nodes: int = 128, num_workers: int = 4,
+                    paired: bool = False, esm_dir: str = None,
+                    binder_dropout: float = 0.0):
+    """Build train/val/test dataloaders for a given target."""
+    loaders = {}
+    for split in ['train', 'val', 'test']:
+        path = os.path.join(data_dir, target, f"{split}.pkl")
+        if not os.path.exists(path):
+            continue
+
+        augment = (split == 'train')
+        bd = binder_dropout if split == 'train' else 0.0
+        if paired and split == 'train':
+            dataset = TwoStateDatasetPaired(path, max_nodes=max_nodes, augment=augment,
+                                            esm_dir=esm_dir, target_name=target,
+                                            binder_dropout=bd)
+            collate = collate_paired_fn
+        else:
+            dataset = TwoStateComplexDataset(path, max_nodes=max_nodes, augment=augment,
+                                             esm_dir=esm_dir, target_name=target,
+                                             binder_dropout=bd)
+            collate = collate_fn
+
+        loaders[split] = DataLoader(
+            dataset,
+            batch_size=batch_size,
+            shuffle=(split == 'train'),
+            num_workers=num_workers,
+            collate_fn=collate,
+            pin_memory=True,
+            drop_last=(split == 'train' and len(dataset) > batch_size),
+        )
+    return loaders

code/models/differentiable_features.py

@@ -0,0 +1,622 @@
+"""
+Differentiable feature extraction for Q_theta guidance.
+
+This module re-implements the key feature extraction functions from features.py
+and pdb_utils.py using PyTorch operations, enabling gradient computation through
+Q_theta with respect to backbone coordinates.
+
+The differentiable path:
+    coords (N,4,3) → backbone frames → torsions, distances, directions, rotations
+    → node_feats, edge_feats → Q_theta → score → backward() → ∇coords
+
+Non-differentiable features (AA one-hot, chain_id, seq_sep, same_chain) are
+treated as constants.
+"""
+
+import os
+import torch
+import torch.nn.functional as F
+import numpy as np
+
+
+# ── Differentiable backbone frame computation ────────────────────────────────
+
+def compute_backbone_frames_torch(coords, mask):
+    """
+    Compute SE(3)-equivariant backbone frames from N, CA, C atoms.
+    Differentiable w.r.t. coords.
+
+    Args:
+        coords: [N, 4, 3] backbone coords (N, CA, C, O) — requires_grad=True for binder
+        mask: [N] bool tensor
+
+    Returns:
+        origins: [N, 3] = CA positions
+        rotations: [N, 3, 3] = rotation matrices (columns are x, y, z axes)
+    """
+    N_res = coords.shape[0]
+    device = coords.device
+
+    origins = coords[:, 1, :]  # CA positions [N, 3]
+    rotations = torch.eye(3, device=device, dtype=coords.dtype).unsqueeze(0).expand(N_res, -1, -1).clone()
+
+    ca = coords[:, 1, :]  # [N, 3]
+    n_atom = coords[:, 0, :]   # [N, 3]
+    c_atom = coords[:, 2, :]   # [N, 3]
+
+    # z-axis: CA -> C
+    z = c_atom - ca  # [N, 3]
+    z_norm = torch.norm(z, dim=-1, keepdim=True).clamp(min=1e-6)  # [N, 1]
+    z = z / z_norm  # [N, 3]
+
+    # y-axis: CA -> N, orthogonalized against z
+    y = n_atom - ca  # [N, 3]
+    y_proj = (y * z).sum(dim=-1, keepdim=True)  # [N, 1]
+    y = y - y_proj * z  # [N, 3]
+    y_norm = torch.norm(y, dim=-1, keepdim=True).clamp(min=1e-6)  # [N, 1]
+    y = y / y_norm  # [N, 3]
+
+    # x-axis: y cross z
+    x = torch.cross(y, z, dim=-1)  # [N, 3]
+
+    # Stack columns: [N, 3, 3] where columns are x, y, z
+    rot = torch.stack([x, y, z], dim=-1)  # [N, 3, 3]
+
+    # Apply mask: identity for masked residues
+    mask_f = mask.float().unsqueeze(-1).unsqueeze(-1)  # [N, 1, 1]
+    eye = torch.eye(3, device=device, dtype=coords.dtype).unsqueeze(0)  # [1, 3, 3]
+    rotations = rot * mask_f + eye * (1 - mask_f)
+
+    return origins, rotations
+
+
+# ── Differentiable torsion angle computation ─────────────────────────────────
+
+def _dihedral_torch(p0, p1, p2, p3):
+    """
+    Compute dihedral angle for batches of 4 points. Returns sin, cos.
+    Differentiable w.r.t. all inputs.
+
+    Args:
+        p0, p1, p2, p3: [N, 3] tensors
+
+    Returns:
+        sin_angle: [N]
+        cos_angle: [N]
+    """
+    b1 = p1 - p0  # [N, 3]
+    b2 = p2 - p1
+    b3 = p3 - p2
+
+    n1 = torch.cross(b1, b2, dim=-1)  # [N, 3]
+    n2 = torch.cross(b2, b3, dim=-1)
+
+    n1_norm = torch.norm(n1, dim=-1, keepdim=True).clamp(min=1e-8)
+    n2_norm = torch.norm(n2, dim=-1, keepdim=True).clamp(min=1e-8)
+    n1 = n1 / n1_norm
+    n2 = n2 / n2_norm
+
+    b2_norm = torch.norm(b2, dim=-1, keepdim=True).clamp(min=1e-8)
+    m1 = torch.cross(n1, b2 / b2_norm, dim=-1)  # [N, 3]
+
+    cos_angle = (n1 * n2).sum(dim=-1)   # [N]
+    sin_angle = (m1 * n2).sum(dim=-1)   # [N]
+
+    return sin_angle, cos_angle
+
+
+def compute_torsion_angles_torch(coords, mask):
+    """
+    Compute backbone torsion angles (phi, psi, omega) as sin/cos pairs.
+    Differentiable w.r.t. coords.
+
+    Args:
+        coords: [N, 4, 3] backbone coords (N, CA, C, O)
+        mask: [N] bool tensor
+
+    Returns:
+        torsions: [N, 6] (sin_phi, cos_phi, sin_psi, cos_psi, sin_omega, cos_omega)
+    """
+    N = coords.shape[0]
+    device = coords.device
+    torsions = torch.zeros(N, 6, device=device, dtype=coords.dtype)
+
+    if N < 2:
+        return torsions
+
+    n_atoms = coords[:, 0, :]   # N atoms [N, 3]
+    ca_atoms = coords[:, 1, :]  # CA atoms
+    c_atoms = coords[:, 2, :]   # C atoms
+
+    # Phi: C_{i-1} - N_i - CA_i - C_i  (for i >= 1)
+    if N > 1:
+        phi_mask = mask[1:] & mask[:-1]  # [N-1]
+        sin_phi, cos_phi = _dihedral_torch(
+            c_atoms[:-1],   # C_{i-1}
+            n_atoms[1:],    # N_i
+            ca_atoms[1:],   # CA_i
+            c_atoms[1:]     # C_i
+        )
+        torsions[1:, 0] = sin_phi * phi_mask.float()
+        torsions[1:, 1] = cos_phi * phi_mask.float()
+
+    # Psi: N_i - CA_i - C_i - N_{i+1}  (for i < N-1)
+    if N > 1:
+        psi_mask = mask[:-1] & mask[1:]  # [N-1]
+        sin_psi, cos_psi = _dihedral_torch(
+            n_atoms[:-1],   # N_i
+            ca_atoms[:-1],  # CA_i
+            c_atoms[:-1],   # C_i
+            n_atoms[1:]     # N_{i+1}
+        )
+        torsions[:-1, 2] = sin_psi * psi_mask.float()
+        torsions[:-1, 3] = cos_psi * psi_mask.float()
+
+    # Omega: CA_{i-1} - C_{i-1} - N_i - CA_i  (for i >= 1)
+    if N > 1:
+        omega_mask = mask[1:] & mask[:-1]  # [N-1]
+        sin_omega, cos_omega = _dihedral_torch(
+            ca_atoms[:-1],  # CA_{i-1}
+            c_atoms[:-1],   # C_{i-1}
+            n_atoms[1:],    # N_i
+            ca_atoms[1:]    # CA_i
+        )
+        torsions[1:, 4] = sin_omega * omega_mask.float()
+        torsions[1:, 5] = cos_omega * omega_mask.float()
+
+    return torsions
+
+
+# ── Differentiable RBF distance encoding ─────────────────────────────────────
+
+def rbf_encode_torch(distances, d_min=0.0, d_max=20.0, n_bins=16):
+    """
+    RBF encoding of distances using Gaussian basis functions.
+    Differentiable w.r.t. distances.
+
+    Args:
+        distances: [...] tensor
+    Returns:
+        encoded: [..., n_bins] tensor
+    """
+    centers = torch.linspace(d_min, d_max, n_bins, device=distances.device, dtype=distances.dtype)
+    sigma = (d_max - d_min) / (n_bins - 1)
+    return torch.exp(-((distances.unsqueeze(-1) - centers) ** 2) / (2 * sigma ** 2))
+
+
+# ── Differentiable edge feature computation ──────────────────────────────────
+
+def compute_edge_features_torch(origins, rotations, seq_idx, chain_ids, mask,
+                                 n_bins_rbf=16, n_bins_sep=8, max_sep=32):
+    """
+    Compute SE(3)-invariant edge features between all residue pairs.
+    Differentiable w.r.t. origins and rotations (which derive from coords).
+
+    Args:
+        origins: [N, 3] CA positions
+        rotations: [N, 3, 3] backbone frame rotations
+        seq_idx: [N] int tensor — sequence indices (non-differentiable)
+        chain_ids: [N] int tensor — chain labels (non-differentiable)
+        mask: [N] bool tensor
+
+    Returns:
+        edge_feats: [N, N, 37]
+    """
+    N = origins.shape[0]
+    device = origins.device
+    dtype = origins.dtype
+
+    # --- Distance features (differentiable) ---
+    diff = origins.unsqueeze(1) - origins.unsqueeze(0)  # [N, N, 3]
+    dist = torch.norm(diff, dim=-1).clamp(min=1e-8)  # [N, N]
+    dist_rbf = rbf_encode_torch(dist, d_min=0., d_max=20., n_bins=n_bins_rbf)  # [N, N, 16]
+
+    # --- Direction in local frame (differentiable) ---
+    unit_diff = diff / dist.unsqueeze(-1)  # [N, N, 3]
+    # local_dir[i,j] = R_i^T @ (ca_j - ca_i) / dist
+    # rotations: [N, 3, 3], unit_diff: [N, N, 3]
+    local_dir = torch.einsum('ikl,ijl->ijk', rotations, unit_diff)  # [N, N, 3]
+
+    # --- Relative rotation (differentiable) ---
+    # rel_rot[i,j] = R_i^T @ R_j -> [N, N, 3, 3] -> flatten to [N, N, 9]
+    rel_rot = torch.einsum('ikl,jlm->ijkm', rotations, rotations)  # [N, N, 3, 3]
+    rel_rot_flat = rel_rot.reshape(N, N, 9)  # [N, N, 9]
+
+    # --- Sequence separation (non-differentiable, constant) ---
+    sep = seq_idx.unsqueeze(1) - seq_idx.unsqueeze(0)  # [N, N]
+    bins = torch.linspace(-max_sep, max_sep, n_bins_sep + 1, device=device)
+    sep_clipped = sep.float().clamp(-max_sep, max_sep)
+    # Bin encoding via soft assignment (but really we just use hard binning)
+    sep_enc = torch.zeros(N, N, n_bins_sep, device=device, dtype=dtype)
+    bin_idx = torch.bucketize(sep_clipped, bins) - 1
+    bin_idx = bin_idx.clamp(0, n_bins_sep - 1)
+    # Scatter one-hot
+    sep_enc.scatter_(2, bin_idx.unsqueeze(-1).long(), 1.0)
+
+    # Cross-chain pairs get sep=0
+    same_chain = (chain_ids.unsqueeze(1) == chain_ids.unsqueeze(0))  # [N, N]
+    cross_chain = ~same_chain
+    sep_enc[cross_chain] = 0.0
+
+    # --- Same chain indicator (non-differentiable, constant) ---
+    same_chain_feat = same_chain.float().unsqueeze(-1)  # [N, N, 1]
+
+    # --- Concatenate ---
+    edge_feats = torch.cat([
+        dist_rbf,       # [N, N, 16]
+        local_dir,      # [N, N, 3]
+        rel_rot_flat,   # [N, N, 9]
+        sep_enc,        # [N, N, 8]
+        same_chain_feat # [N, N, 1]
+    ], dim=-1)          # [N, N, 37]
+
+    # Zero out edges involving masked residues
+    mask_2d = mask.unsqueeze(1) & mask.unsqueeze(0)  # [N, N]
+    edge_feats = edge_feats * mask_2d.unsqueeze(-1).float()
+
+    return edge_feats
+
+
+# ── Full differentiable interface graph builder ──────────────────────────────
+
+def build_differentiable_interface_graph(
+    rec_coords, rec_mask, rec_aa_idx, rec_chi,
+    binder_coords, binder_mask, binder_aa_idx, binder_chi,
+    cutoff=8.0, max_nodes=128
+):
+    """
+    Build interface graph with differentiable features w.r.t. binder_coords.
+    Receptor coords are treated as constants (detached).
+
+    Args:
+        rec_coords: [N_rec, 4, 3] — receptor backbone coords (constant, no grad)
+        rec_mask: [N_rec] bool
+        rec_aa_idx: [N_rec] int — amino acid indices (constant)
+        rec_chi: [N_rec, 4] — chi1/chi2 sin/cos (constant)
+        binder_coords: [N_binder, 4, 3] — binder backbone coords (requires_grad)
+        binder_mask: [N_binder] bool
+        binder_aa_idx: [N_binder] int — amino acid indices (constant, UNK for designed)
+        binder_chi: [N_binder, 4] — chi1/chi2 sin/cos (zeros for backbone-only)
+        cutoff: interface distance cutoff (Å)
+        max_nodes: maximum nodes per chain in the graph
+
+    Returns:
+        node_feats: [1, N_total, 32] tensor
+        edge_feats: [1, N_total, N_total, 37] tensor
+        node_mask: [1, N_total] bool tensor
+        n_rec: int
+        n_binder: int
+        or None if no interface
+    """
+    device = binder_coords.device
+    dtype = binder_coords.dtype
+    NUM_AA = 21
+
+    # ── Find interface residues (differentiable distances but hard threshold) ──
+    rec_ca = rec_coords[:, 1, :]      # [N_rec, 3]
+    binder_ca = binder_coords[:, 1, :]  # [N_binder, 3]
+
+    # Pairwise CA distances
+    dist_mat = torch.cdist(rec_ca.unsqueeze(0), binder_ca.unsqueeze(0)).squeeze(0)  # [N_rec, N_binder]
+    # Mask invalid residues
+    dist_mat = dist_mat.clone()
+    dist_mat[~rec_mask, :] = float('inf')
+    dist_mat[:, ~binder_mask] = float('inf')
+
+    rec_iface = (dist_mat < cutoff).any(dim=1)      # [N_rec]
+    binder_iface = (dist_mat < cutoff).any(dim=0)    # [N_binder]
+
+    rec_iface_idx = torch.where(rec_iface)[0]
+    binder_iface_idx = torch.where(binder_iface)[0]
+
+    # Truncate if too many
+    if len(rec_iface_idx) > max_nodes // 2:
+        rec_iface_idx = rec_iface_idx[:max_nodes // 2]
+    if len(binder_iface_idx) > max_nodes // 2:
+        binder_iface_idx = binder_iface_idx[:max_nodes // 2]
+
+    n_rec = len(rec_iface_idx)
+    n_binder = len(binder_iface_idx)
+    n_total = n_rec + n_binder
+
+    if n_total == 0:
+        return None
+
+    # ── Extract interface subsets ──
+    rec_iface_coords = rec_coords[rec_iface_idx]          # [n_rec, 4, 3]
+    binder_iface_coords = binder_coords[binder_iface_idx]  # [n_binder, 4, 3]
+    rec_iface_mask = rec_mask[rec_iface_idx]
+    binder_iface_mask = binder_mask[binder_iface_idx]
+
+    # ── Compute backbone frames (differentiable) ──
+    rec_origins, rec_rotations = compute_backbone_frames_torch(rec_iface_coords, rec_iface_mask)
+    binder_origins, binder_rotations = compute_backbone_frames_torch(binder_iface_coords, binder_iface_mask)
+
+    # ── Compute torsion angles (differentiable) ──
+    rec_torsion = compute_torsion_angles_torch(rec_iface_coords, rec_iface_mask)     # [n_rec, 6]
+    binder_torsion = compute_torsion_angles_torch(binder_iface_coords, binder_iface_mask)  # [n_binder, 6]
+
+    # ── Node features ──
+    # AA one-hot (non-differentiable constant)
+    rec_aa_onehot = F.one_hot(rec_aa_idx[rec_iface_idx].long(), NUM_AA).float()       # [n_rec, 21]
+    binder_aa_onehot = F.one_hot(binder_aa_idx[binder_iface_idx].long(), NUM_AA).float()  # [n_binder, 21]
+
+    # Chi angles (constant for receptor, zeros for backbone-only binder)
+    rec_chi_iface = rec_chi[rec_iface_idx]              # [n_rec, 4]
+    binder_chi_iface = binder_chi[binder_iface_idx]     # [n_binder, 4]
+
+    # Chain indicator
+    rec_chain_feat = torch.zeros(n_rec, 1, device=device, dtype=dtype)
+    binder_chain_feat = torch.ones(n_binder, 1, device=device, dtype=dtype)
+
+    # Concatenate node features: [AA(21) + torsions(6) + chi(4) + chain(1)] = 32
+    rec_node = torch.cat([rec_aa_onehot, rec_torsion, rec_chi_iface, rec_chain_feat], dim=-1)
+    binder_node = torch.cat([binder_aa_onehot, binder_torsion, binder_chi_iface, binder_chain_feat], dim=-1)
+    node_feats = torch.cat([rec_node, binder_node], dim=0)  # [N_total, 32]
+    node_mask_flat = torch.cat([rec_iface_mask, binder_iface_mask], dim=0)  # [N_total]
+
+    # ── Edge features (differentiable) ──
+    all_origins = torch.cat([rec_origins, binder_origins], dim=0)      # [N_total, 3]
+    all_rotations = torch.cat([rec_rotations, binder_rotations], dim=0)  # [N_total, 3, 3]
+
+    # Sequence indices
+    rec_seq_idx = rec_iface_idx
+    binder_seq_idx = binder_iface_idx + rec_coords.shape[0]
+    all_seq_idx = torch.cat([rec_seq_idx, binder_seq_idx], dim=0)
+
+    # Chain IDs
+    all_chain_ids = torch.cat([
+        torch.zeros(n_rec, device=device, dtype=torch.long),
+        torch.ones(n_binder, device=device, dtype=torch.long)
+    ], dim=0)
+
+    edge_feats = compute_edge_features_torch(
+        all_origins, all_rotations, all_seq_idx, all_chain_ids, node_mask_flat
+    )  # [N_total, N_total, 37]
+
+    # Add batch dimension
+    return {
+        'node_feats': node_feats.unsqueeze(0),      # [1, N, 32]
+        'edge_feats': edge_feats.unsqueeze(0),      # [1, N, N, 37]
+        'node_mask': node_mask_flat.unsqueeze(0),    # [1, N]
+        'n_rec': n_rec,
+        'n_binder': n_binder,
+    }
+
+
+# ── Differentiable Q_theta scoring function ──────────────────────────────────
+
+class DifferentiableQTheta:
+    """
+    Wraps the Q_theta scorer for differentiable scoring w.r.t. binder backbone
+    coordinates. Receptor structures are pre-loaded and cached.
+
+    Usage:
+        dq = DifferentiableQTheta(checkpoint_path, device)
+        dq.load_receptor(holo_pdb, chain='A', label='holo')
+        dq.load_receptor(apo_pdb, chain='A', label='apo')
+
+        binder_coords = torch.tensor(...)  # [N_binder, 4, 3], requires_grad=True
+        score_holo = dq.score(binder_coords, binder_mask, binder_aa_idx, 'holo')
+        score_apo = dq.score(binder_coords, binder_mask, binder_aa_idx, 'apo')
+        selectivity = score_holo - score_apo
+        selectivity.backward()
+        # binder_coords.grad now contains ∂S/∂coords
+    """
+
+    def __init__(self, checkpoint_path, device='cuda:0', esm_dir=None):
+        import sys, os
+        _code_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
+        if _code_dir not in sys.path:
+            sys.path.insert(0, _code_dir)
+        from models.scorer import build_model
+
+        self.device = torch.device(device)
+        ckpt = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
+        self.config = ckpt['config']
+        self.model = build_model(self.config)
+        self.model.load_state_dict(ckpt['model_state'])
+        self.model = self.model.to(self.device)
+        self.model.eval()
+
+        # ESM feature support
+        self.use_esm = self.config.get('esm_dim', 0) > 0
+        self.esm_dim = self.config.get('esm_dim', 0)
+        self.esm_dir = esm_dir or os.path.join(os.environ.get('ALLOGEN_ROOT', '.'), 'data/esm2_embeddings')
+
+        # Cache receptor data
+        self.receptors = {}  # label -> {coords, mask, aa_idx, chi, esm_emb?}
+
+    def load_receptor(self, pdb_path, chain='A', label='holo',
+                       esm_target=None, esm_key=None):
+        """Pre-load and cache receptor structure, optionally with ESM embeddings.
+
+        Args:
+            pdb_path: path to receptor PDB
+            chain: chain ID
+            label: cache key
+            esm_target: target name for ESM dir (e.g., 'abl' for data/esm2_embeddings/abl/)
+            esm_key: ESM embedding file key (e.g., '6XR7_A'). If None, auto-derived.
+        """
+        import sys, os
+        _code_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
+        if _code_dir not in sys.path:
+            sys.path.insert(0, _code_dir)
+        from utils.pdb_utils import (
+            load_structure, get_residues, get_backbone_coords,
+            get_aa_indices, compute_chi_angles
+        )
+
+        model = load_structure(pdb_path)
+        chain_obj = model[chain]
+        residues = get_residues(chain_obj)
+        coords, mask = get_backbone_coords(residues)
+        aa_idx = get_aa_indices(residues)
+        chi = compute_chi_angles(residues, mask)
+
+        rec_data = {
+            'coords': torch.from_numpy(coords).float().to(self.device),
+            'mask': torch.from_numpy(mask).bool().to(self.device),
+            'aa_idx': torch.from_numpy(aa_idx).long().to(self.device),
+            'chi': torch.from_numpy(chi).float().to(self.device),
+            'residues': residues,
+        }
+
+        # Load ESM embeddings if model uses ESM
+        if self.use_esm and esm_target:
+            pdb_id = os.path.basename(pdb_path).replace('.pdb', '')
+            if esm_key is None:
+                esm_key = f'{pdb_id}_{chain}'
+            esm_path = os.path.join(self.esm_dir, esm_target, f'{esm_key}.pt')
+            if os.path.exists(esm_path):
+                esm_emb = torch.load(esm_path, map_location=self.device, weights_only=True)
+                # Truncate/pad to match residue count
+                n_res = len(residues)
+                if esm_emb.shape[0] > n_res:
+                    esm_emb = esm_emb[:n_res]
+                elif esm_emb.shape[0] < n_res:
+                    pad = torch.zeros(n_res - esm_emb.shape[0], esm_emb.shape[1],
+                                     device=self.device)
+                    esm_emb = torch.cat([esm_emb, pad], dim=0)
+                rec_data['esm_emb'] = esm_emb.float()
+            else:
+                rec_data['esm_emb'] = torch.zeros(len(residues), self.esm_dim,
+                                                    device=self.device)
+
+        self.receptors[label] = rec_data
+
+    def load_receptor_from_coords(self, coords, mask, aa_idx=None, chi=None,
+                                   label='path'):
+        """
+        Load a receptor from raw backbone coords (not from PDB file).
+
+        Used for interpolated path frames that don't have PDB files.
+        If aa_idx is None, uses all-ALA (index 0). If chi is None, uses zeros.
+
+        Args:
+            coords: [N, 4, 3] numpy or torch backbone coords (N, CA, C, O)
+            mask: [N] numpy or torch bool
+            aa_idx: [N] numpy or torch int (default: all-ALA = 0)
+            chi: [N, 4] numpy or torch float (default: zeros)
+            label: str key for caching
+        """
+        import numpy as np
+
+        # Convert numpy to torch if needed
+        if isinstance(coords, np.ndarray):
+            coords = torch.from_numpy(coords).float()
+        if isinstance(mask, np.ndarray):
+            mask = torch.from_numpy(mask).bool()
+
+        N = coords.shape[0]
+
+        if aa_idx is None:
+            aa_idx = torch.zeros(N, dtype=torch.long)  # all-ALA
+        elif isinstance(aa_idx, np.ndarray):
+            aa_idx = torch.from_numpy(aa_idx).long()
+
+        if chi is None:
+            chi = torch.zeros(N, 4, dtype=coords.dtype)
+        elif isinstance(chi, np.ndarray):
+            chi = torch.from_numpy(chi).float()
+
+        self.receptors[label] = {
+            'coords': coords.to(self.device),
+            'mask': mask.to(self.device),
+            'aa_idx': aa_idx.to(self.device),
+            'chi': chi.to(self.device),
+        }
+
+    def score(self, binder_coords, binder_mask, binder_aa_idx=None,
+              binder_chi=None, receptor_label='holo', cutoff=8.0):
+        """
+        Score binder against a cached receptor. Differentiable w.r.t. binder_coords.
+
+        Args:
+            binder_coords: [N_binder, 4, 3] tensor (can have requires_grad=True)
+            binder_mask: [N_binder] bool tensor
+            binder_aa_idx: [N_binder] int tensor (default: all UNK)
+            binder_chi: [N_binder, 4] tensor (default: zeros)
+            receptor_label: key into cached receptors
+            cutoff: interface distance cutoff
+
+        Returns:
+            score: scalar tensor in (0, 1), differentiable w.r.t. binder_coords
+        """
+        rec = self.receptors[receptor_label]
+        N_binder = binder_coords.shape[0]
+
+        if binder_aa_idx is None:
+            binder_aa_idx = torch.full((N_binder,), 20, device=self.device, dtype=torch.long)  # UNK
+        if binder_chi is None:
+            binder_chi = torch.zeros(N_binder, 4, device=self.device, dtype=binder_coords.dtype)
+
+        graph = build_differentiable_interface_graph(
+            rec_coords=rec['coords'],
+            rec_mask=rec['mask'],
+            rec_aa_idx=rec['aa_idx'],
+            rec_chi=rec['chi'],
+            binder_coords=binder_coords,
+            binder_mask=binder_mask,
+            binder_aa_idx=binder_aa_idx,
+            binder_chi=binder_chi,
+            cutoff=cutoff,
+        )
+
+        if graph is None:
+            # No interface — return zero score with gradient
+            return torch.zeros(1, device=self.device, dtype=binder_coords.dtype, requires_grad=True).squeeze()
+
+        # Build ESM features if model uses ESM
+        esm_feats = None
+        if self.use_esm:
+            n_rec = graph['n_rec']
+            n_binder = graph['n_binder']
+            n_total = n_rec + n_binder
+            # Receptor ESM: use cached if available, else zeros
+            if 'esm_emb' in rec:
+                rec_esm = rec['esm_emb']
+                # Need to select interface residues (same indices as structural features)
+                # The graph was built with rec_iface_idx — we need those indices
+                # For simplicity, use zeros for now and rely on the projection layer
+                # to handle the zero binder ESM gracefully
+                rec_esm_full = rec_esm  # [N_rec_total, 1280]
+            else:
+                rec_esm_full = torch.zeros(rec['coords'].shape[0], self.esm_dim,
+                                          device=self.device)
+            # Binder ESM: zeros (designed backbone, no sequence)
+            binder_esm = torch.zeros(binder_coords.shape[0], self.esm_dim,
+                                     device=self.device)
+            # We need interface indices to select — rebuild them
+            rec_ca = rec['coords'][:, 1, :]
+            binder_ca = binder_coords[:, 1, :]
+            dist_mat = torch.cdist(rec_ca.unsqueeze(0), binder_ca.unsqueeze(0)).squeeze(0)
+            dist_mat_c = dist_mat.clone()
+            dist_mat_c[~rec['mask'], :] = float('inf')
+            dist_mat_c[:, ~binder_mask] = float('inf')
+            rec_iface = (dist_mat_c < cutoff).any(dim=1)
+            binder_iface = (dist_mat_c < cutoff).any(dim=0)
+            rec_iface_idx = torch.where(rec_iface)[0][:n_rec]
+            binder_iface_idx = torch.where(binder_iface)[0][:n_binder]
+
+            rec_esm_iface = rec_esm_full[rec_iface_idx]  # [n_rec, 1280]
+            binder_esm_iface = binder_esm[binder_iface_idx]  # [n_binder, 1280]
+            esm_combined = torch.cat([rec_esm_iface, binder_esm_iface], dim=0)  # [n_total, 1280]
+            esm_feats = esm_combined.unsqueeze(0)  # [1, n_total, 1280]
+
+        score = self.model(graph['node_feats'], graph['edge_feats'], graph['node_mask'],
+                          esm_feats=esm_feats)
+        return score.squeeze()  # scalar
+
+    def selectivity_margin(self, binder_coords, binder_mask,
+                           binder_aa_idx=None, binder_chi=None,
+                           holo_label='holo', apo_label='apo', cutoff=8.0):
+        """
+        Compute selectivity margin S = Q(holo, Y) - Q(apo, Y).
+        Differentiable w.r.t. binder_coords.
+        """
+        q_holo = self.score(binder_coords, binder_mask, binder_aa_idx, binder_chi,
+                            holo_label, cutoff)
+        q_apo = self.score(binder_coords, binder_mask, binder_aa_idx, binder_chi,
+                           apo_label, cutoff)
+        return q_holo - q_apo, q_holo, q_apo

code/models/features.py

@@ -0,0 +1,250 @@
+"""
+SE(3)-invariant feature extraction for interface graphs.
+Node and edge features used by the Q_theta scorer.
+"""
+
+import os
+import sys
+import numpy as np
+
+# Ensure utils is importable (for both direct and package imports)
+_CODE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+if _CODE_DIR not in sys.path:
+    sys.path.insert(0, _CODE_DIR)
+
+from utils.pdb_utils import (
+    rbf_encode, compute_backbone_frames, compute_torsion_angles,
+    get_aa_indices, compute_chi_angles, get_cb_positions, NUM_AA
+)
+
+# Feature dimensions
+# one-hot AA (21) + backbone torsions (6) + chi1 sin/cos (2) + chi2 sin/cos (2) + chain indicator (1) = 32
+NODE_DIM = NUM_AA + 6 + 4 + 1  # = 32
+EDGE_DIM = 16 + 3 + 9 + 8 + 1  # RBF dist (16) + direction (3) + rel rotation (9) + seq sep (8) + same chain (1) = 37
+
+MAX_SEQ_SEP = 32  # bins for sequence separation
+
+
+def seq_sep_encode(sep, n_bins=8, max_sep=MAX_SEQ_SEP):
+    """Bin-encode sequence separation."""
+    bins = np.linspace(-max_sep, max_sep, n_bins + 1)
+    sep_clipped = np.clip(sep, -max_sep, max_sep)
+    encoded = np.zeros(n_bins, dtype=np.float32)
+    bin_idx = np.digitize(sep_clipped, bins) - 1
+    bin_idx = np.clip(bin_idx, 0, n_bins - 1)
+    encoded[bin_idx] = 1.0
+    return encoded
+
+
+def extract_node_features(residues, coords, mask, torsion_angles, chi_angles, chain_id):
+    """
+    Compute per-residue node features.
+
+    Args:
+        residues: list of Bio.PDB residues
+        coords: [N, 4, 3] backbone coords
+        mask: [N] bool
+        torsion_angles: [N, 6] sin/cos of phi, psi, omega
+        chi_angles: [N, 4] sin/cos of chi1, chi2
+        chain_id: 0 = receptor, 1 = binder
+
+    Returns:
+        node_feats: [N, NODE_DIM]  (NODE_DIM = 32)
+    """
+    N = len(residues)
+    aa_idx = get_aa_indices(residues)
+
+    # One-hot amino acid
+    aa_onehot = np.zeros((N, NUM_AA), dtype=np.float32)
+    for i in range(N):
+        if mask[i]:
+            aa_onehot[i, aa_idx[i]] = 1.0
+
+    # Chain indicator
+    chain_feat = np.full((N, 1), chain_id, dtype=np.float32)
+
+    # Concatenate
+    node_feats = np.concatenate([
+        aa_onehot,          # [N, 21]
+        torsion_angles,     # [N, 6]
+        chi_angles,         # [N, 4]
+        chain_feat,         # [N, 1]
+    ], axis=-1)
+
+    return node_feats  # [N, 32]
+
+
+def extract_edge_features(coords_i, frames_i, coords_j, frames_j,
+                          seq_idx_i, seq_idx_j, chain_i, chain_j, mask_i, mask_j):
+    """
+    Compute SE(3)-invariant edge features between residue sets i and j.
+    Vectorized over all pairs.
+
+    Args:
+        coords_i: [N_i, 4, 3] backbone coords of set i (full interface)
+        frames_i: (origins_i [N_i, 3], rotations_i [N_i, 3, 3])
+        coords_j: [N_j, 4, 3]
+        frames_j: (origins_j [N_j, 3], rotations_j [N_j, 3, 3])
+        seq_idx_i: [N_i] integer sequence indices (for sequence separation)
+        seq_idx_j: [N_j] integer sequence indices
+        chain_i: int (0 or 1)
+        chain_j: int (0 or 1)
+        mask_i: [N_i] bool
+        mask_j: [N_j] bool
+
+    Returns:
+        edge_feats: [N_i, N_j, EDGE_DIM]
+    """
+    N_i, N_j = len(coords_i), len(coords_j)
+    origins_i, rotations_i = frames_i
+    origins_j, rotations_j = frames_j
+
+    ca_i = origins_i  # [N_i, 3]
+    ca_j = origins_j  # [N_j, 3]
+
+    # --- Distance features ---
+    diff = ca_j[None, :, :] - ca_i[:, None, :]  # [N_i, N_j, 3]
+    dist = np.sqrt((diff ** 2).sum(axis=-1))     # [N_i, N_j]
+    dist_rbf = rbf_encode(dist, d_min=0., d_max=20., n_bins=16)  # [N_i, N_j, 16]
+
+    # --- Direction in local frame of i ---
+    # unit vector from i to j in global frame
+    unit_diff = diff / (dist[..., None] + 1e-8)  # [N_i, N_j, 3]
+    # rotate by R_i^T to get local direction
+    # rotations_i: [N_i, 3, 3], unit_diff: [N_i, N_j, 3]
+    # local_dir[i,j] = R_i^T @ (ca_j - ca_i) / dist
+    local_dir = np.einsum('ikl,ijl->ijk', rotations_i, unit_diff)  # [N_i, N_j, 3]
+
+    # --- Relative rotation: R_i^T R_j ---
+    # rotations_i: [N_i, 3, 3], rotations_j: [N_j, 3, 3]
+    # rel_rot[i,j] = R_i^T @ R_j -> [N_i, N_j, 3, 3] -> flatten to [N_i, N_j, 9]
+    rel_rot = np.einsum('ikl,jlm->ijkm', rotations_i, rotations_j)  # [N_i, N_j, 3, 3]
+    rel_rot_flat = rel_rot.reshape(N_i, N_j, 9)  # [N_i, N_j, 9]
+
+    # --- Sequence separation ---
+    sep = seq_idx_j[None, :] - seq_idx_i[:, None]  # [N_i, N_j]
+    # Encode each pair (loop over all; use vectorized bin assignment)
+    sep_flat = sep.reshape(-1)
+    sep_enc = np.array([seq_sep_encode(s) for s in sep_flat])  # [N_i*N_j, 8]
+    sep_enc = sep_enc.reshape(N_i, N_j, 8)
+
+    # Cross-chain pairs get sep=0 by convention if different chains
+    if chain_i != chain_j:
+        sep_enc[:] = 0.0
+
+    # --- Same chain indicator ---
+    same_chain = float(chain_i == chain_j)
+    same_chain_feat = np.full((N_i, N_j, 1), same_chain, dtype=np.float32)
+
+    # --- Concatenate ---
+    edge_feats = np.concatenate([
+        dist_rbf,       # [N_i, N_j, 16]
+        local_dir,      # [N_i, N_j, 3]
+        rel_rot_flat,   # [N_i, N_j, 9]
+        sep_enc,        # [N_i, N_j, 8]
+        same_chain_feat # [N_i, N_j, 1]
+    ], axis=-1)         # [N_i, N_j, 37]
+
+    # Zero out edges involving masked residues
+    edge_feats[~mask_i, :, :] = 0.0
+    edge_feats[:, ~mask_j, :] = 0.0
+
+    return edge_feats.astype(np.float32)
+
+
+def build_interface_graph(rec_residues, rec_coords, rec_mask,
+                          binder_residues, binder_coords, binder_mask,
+                          rec_interface_mask, binder_interface_mask,
+                          max_nodes: int = 128):
+    """
+    Build a joint interface graph combining receptor and binder interface residues.
+
+    Returns a dict with:
+        node_feats: [N_total, NODE_DIM]
+        edge_feats: [N_total, N_total, EDGE_DIM]
+        node_mask: [N_total] bool
+        n_rec: int (number of receptor interface nodes)
+        n_binder: int (number of binder interface nodes)
+    """
+    # Select interface residues
+    rec_iface_idx = np.where(rec_interface_mask)[0]
+    binder_iface_idx = np.where(binder_interface_mask)[0]
+
+    # Truncate if too many
+    if len(rec_iface_idx) > max_nodes // 2:
+        rec_iface_idx = rec_iface_idx[:max_nodes // 2]
+    if len(binder_iface_idx) > max_nodes // 2:
+        binder_iface_idx = binder_iface_idx[:max_nodes // 2]
+
+    n_rec = len(rec_iface_idx)
+    n_binder = len(binder_iface_idx)
+    n_total = n_rec + n_binder
+
+    if n_total == 0:
+        return None
+
+    # Extract coords for interface residues
+    rec_iface_coords = rec_coords[rec_iface_idx]    # [n_rec, 4, 3]
+    binder_iface_coords = binder_coords[binder_iface_idx]  # [n_binder, 4, 3]
+    rec_iface_mask = rec_mask[rec_iface_idx]
+    binder_iface_mask = binder_mask[binder_iface_idx]
+
+    # Compute backbone frames
+    rec_origins, rec_rotations = compute_backbone_frames(rec_iface_coords, rec_iface_mask)
+    binder_origins, binder_rotations = compute_backbone_frames(binder_iface_coords, binder_iface_mask)
+
+    # Compute torsion angles
+    # We need full-chain coords for proper phi/psi computation, but use local approximation here
+    rec_torsion = compute_torsion_angles(rec_iface_coords, rec_iface_mask)
+    binder_torsion = compute_torsion_angles(binder_iface_coords, binder_iface_mask)
+
+    # Extract residues
+    rec_iface_residues = [rec_residues[i] for i in rec_iface_idx]
+    binder_iface_residues = [binder_residues[i] for i in binder_iface_idx]
+
+    # Compute sidechain chi1/chi2 angles
+    rec_chi = compute_chi_angles(rec_iface_residues, rec_iface_mask)
+    binder_chi = compute_chi_angles(binder_iface_residues, binder_iface_mask)
+
+    # Node features
+    rec_node_feats = extract_node_features(
+        rec_iface_residues, rec_iface_coords, rec_iface_mask, rec_torsion, rec_chi, chain_id=0
+    )  # [n_rec, NODE_DIM]
+    binder_node_feats = extract_node_features(
+        binder_iface_residues, binder_iface_coords, binder_iface_mask, binder_torsion, binder_chi, chain_id=1
+    )  # [n_binder, NODE_DIM]
+
+    node_feats = np.concatenate([rec_node_feats, binder_node_feats], axis=0)  # [N, NODE_DIM]
+    node_mask = np.concatenate([rec_iface_mask, binder_iface_mask], axis=0)
+
+    # Edge features (4 blocks: RR, RB, BR, BB)
+    all_coords = np.concatenate([rec_iface_coords, binder_iface_coords], axis=0)
+    all_mask = node_mask
+    all_origins = np.concatenate([rec_origins, binder_origins], axis=0)
+    all_rotations = np.concatenate([rec_rotations, binder_rotations], axis=0)
+    all_seq_idx = np.concatenate([rec_iface_idx, binder_iface_idx + len(rec_residues)], axis=0)
+    all_chain = np.array([0] * n_rec + [1] * n_binder, dtype=np.int32)
+
+    # Compute full NxN edge features
+    frames_all = (all_origins, all_rotations)
+    edge_feats = extract_edge_features(
+        all_coords, frames_all,
+        all_coords, frames_all,
+        all_seq_idx, all_seq_idx,
+        -1, -1,  # chain handled via all_chain array below
+        all_mask, all_mask
+    )  # [N, N, EDGE_DIM]
+
+    # Patch same_chain feature (last dim) using actual chain IDs
+    same_chain_feat = (all_chain[:, None] == all_chain[None, :]).astype(np.float32)
+    edge_feats[:, :, -1] = same_chain_feat
+
+    return {
+        'node_feats': node_feats.astype(np.float32),    # [N, NODE_DIM]
+        'edge_feats': edge_feats.astype(np.float32),    # [N, N, EDGE_DIM]
+        'node_mask': node_mask,                          # [N]
+        'n_rec': n_rec,
+        'n_binder': n_binder,
+        'rec_iface_idx': rec_iface_idx,                  # [n_rec] original residue indices
+        'binder_iface_idx': binder_iface_idx,            # [n_binder] original residue indices
+    }

code/models/scorer.py

@@ -0,0 +1,585 @@
+"""
+Q_theta: State-selectivity scorer for Allo-Designer.
+
+Architecture: Dense Edge-Biased Graph Transformer
+  - Input: padded interface graph (node feats + pairwise edge feats)
+  - SE(3)-invariant features (all features from distances/angles in backbone frames)
+  - Output: Q_theta(X, Y) in (0,1) = probability-like compatibility/selectivity score
+
+No torch_geometric dependency: uses dense attention with edge biases.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+
+class RBFLayer(nn.Module):
+    """Learnable RBF embedding for edge distances."""
+    def __init__(self, n_bins: int = 16, d_min: float = 0., d_max: float = 20.):
+        super().__init__()
+        centers = torch.linspace(d_min, d_max, n_bins)
+        self.register_buffer('centers', centers)
+        self.log_sigma = nn.Parameter(torch.zeros(1))
+
+    def forward(self, dist):
+        # dist: [...] -> [..., n_bins]
+        sigma = torch.exp(self.log_sigma)
+        return torch.exp(-((dist.unsqueeze(-1) - self.centers) ** 2) / (2 * sigma ** 2))
+
+
+class EdgeBiasedMHA(nn.Module):
+    """
+    Multi-Head Self-Attention with additive edge biases.
+    Implements the core equation:
+        A_ij = (Q_i K_j^T / sqrt(d)) + b_ij
+    where b_ij is computed from edge features.
+    """
+    def __init__(self, d_model: int, n_heads: int, d_edge: int, dropout: float = 0.1):
+        super().__init__()
+        assert d_model % n_heads == 0
+        self.n_heads = n_heads
+        self.d_head = d_model // n_heads
+        self.scale = math.sqrt(self.d_head)
+
+        self.qkv_proj = nn.Linear(d_model, 3 * d_model, bias=False)
+        self.out_proj = nn.Linear(d_model, d_model)
+        self.edge_proj = nn.Linear(d_edge, n_heads)  # edge features -> per-head bias
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, edge_feats, mask=None):
+        """
+        x: [B, N, d_model]
+        edge_feats: [B, N, N, d_edge]
+        mask: [B, N] bool (True = valid residue)
+        """
+        B, N, D = x.shape
+        H = self.n_heads
+
+        # QKV projection
+        qkv = self.qkv_proj(x).reshape(B, N, 3, H, self.d_head).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)  # each [B, H, N, d_head]
+
+        # Scaled dot-product attention logits
+        attn_logits = (q @ k.transpose(-2, -1)) / self.scale  # [B, H, N, N]
+
+        # Edge bias: [B, N, N, H] -> [B, H, N, N]
+        edge_bias = self.edge_proj(edge_feats).permute(0, 3, 1, 2)  # [B, H, N, N]
+        attn_logits = attn_logits + edge_bias
+
+        # Padding mask: mask out padded positions
+        if mask is not None:
+            # mask: [B, N] True=valid; padding=False
+            padding = ~mask  # [B, N] True=padding
+            attn_logits = attn_logits.masked_fill(
+                padding[:, None, None, :],  # [B, 1, 1, N]
+                float('-inf')
+            )
+
+        attn_weights = self.dropout(F.softmax(attn_logits, dim=-1))
+
+        # Handle all-padding rows (NaN -> 0)
+        attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
+
+        out = (attn_weights @ v)  # [B, H, N, d_head]
+        out = out.transpose(1, 2).reshape(B, N, D)  # [B, N, D]
+        return self.out_proj(out)
+
+
+class InterfaceTransformerLayer(nn.Module):
+    """Single layer of edge-biased transformer with pre-norm."""
+    def __init__(self, d_model: int, n_heads: int, d_edge: int, ff_mult: int = 4, dropout: float = 0.1):
+        super().__init__()
+        self.attn = EdgeBiasedMHA(d_model, n_heads, d_edge, dropout)
+        self.ff = nn.Sequential(
+            nn.Linear(d_model, d_model * ff_mult),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(d_model * ff_mult, d_model),
+        )
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.drop = nn.Dropout(dropout)
+
+    def forward(self, x, edge_feats, mask=None):
+        x = x + self.drop(self.attn(self.norm1(x), edge_feats, mask))
+        x = x + self.drop(self.ff(self.norm2(x)))
+        return x
+
+
+class GATLayer(nn.Module):
+    """Multi-head GAT layer with pre-norm. No edge features in attention."""
+    def __init__(self, d_model: int, n_heads: int, ff_mult: int = 4, dropout: float = 0.1):
+        super().__init__()
+        assert d_model % n_heads == 0
+        self.n_heads = n_heads
+        self.d_head = d_model // n_heads
+
+        self.W = nn.Linear(d_model, d_model, bias=False)
+        self.a_l = nn.Parameter(torch.randn(n_heads, self.d_head))
+        self.a_r = nn.Parameter(torch.randn(n_heads, self.d_head))
+        nn.init.xavier_uniform_(self.a_l.unsqueeze(0))
+        nn.init.xavier_uniform_(self.a_r.unsqueeze(0))
+        self.out_proj = nn.Linear(d_model, d_model)
+        self.leaky_relu = nn.LeakyReLU(0.2)
+        self.attn_drop = nn.Dropout(dropout)
+
+        self.ff = nn.Sequential(
+            nn.Linear(d_model, d_model * ff_mult), nn.GELU(),
+            nn.Dropout(dropout), nn.Linear(d_model * ff_mult, d_model),
+        )
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.drop = nn.Dropout(dropout)
+
+    def forward(self, x, edge_feats, mask=None):
+        B, N, D = x.shape
+        H = self.n_heads
+
+        h = self.norm1(x)
+        Wh = self.W(h).view(B, N, H, self.d_head)           # [B, N, H, d_head]
+        e_l = (Wh * self.a_l).sum(-1)                        # [B, N, H]
+        e_r = (Wh * self.a_r).sum(-1)                        # [B, N, H]
+        attn = self.leaky_relu(e_l.unsqueeze(2) + e_r.unsqueeze(1))  # [B, N, N, H]
+        attn = attn.permute(0, 3, 1, 2)                      # [B, H, N, N]
+
+        if mask is not None:
+            attn = attn.masked_fill(~mask[:, None, None, :], float('-inf'))
+
+        attn = self.attn_drop(F.softmax(attn, dim=-1))
+        attn = torch.nan_to_num(attn, nan=0.0)
+
+        out = torch.einsum('bhnm,bmhd->bnhd', attn, Wh)
+        out = out.reshape(B, N, D)
+        x = x + self.drop(self.out_proj(out))
+        x = x + self.drop(self.ff(self.norm2(x)))
+        return x
+
+
+class GCNLayer(nn.Module):
+    """GCN layer with edge-weighted message passing and pre-norm."""
+    def __init__(self, d_model: int, d_edge: int, ff_mult: int = 4, dropout: float = 0.1):
+        super().__init__()
+        self.msg_proj = nn.Linear(d_model, d_model, bias=False)
+        self.edge_weight = nn.Linear(d_edge, 1)
+
+        self.ff = nn.Sequential(
+            nn.Linear(d_model, d_model * ff_mult), nn.GELU(),
+            nn.Dropout(dropout), nn.Linear(d_model * ff_mult, d_model),
+        )
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.drop = nn.Dropout(dropout)
+
+    def forward(self, x, edge_feats, mask=None):
+        B, N, D = x.shape
+        h = self.norm1(x)
+        msg = self.msg_proj(h)                                # [B, N, D]
+
+        w = self.edge_weight(edge_feats).squeeze(-1)          # [B, N, N]
+        if mask is not None:
+            w = w.masked_fill(~mask[:, None, :], float('-inf'))
+        w = F.softmax(w, dim=-1)
+        w = torch.nan_to_num(w, nan=0.0)
+
+        agg = torch.bmm(w, msg)                               # [B, N, D]
+        x = x + self.drop(agg)
+        x = x + self.drop(self.ff(self.norm2(x)))
+        return x
+
+
+class CrossChainTransformerLayer(nn.Module):
+    """Cross-chain attention: each node attends only to nodes from the other chain."""
+    def __init__(self, d_model: int, n_heads: int, d_edge: int, ff_mult: int = 4, dropout: float = 0.1):
+        super().__init__()
+        assert d_model % n_heads == 0
+        self.n_heads = n_heads
+        self.d_head = d_model // n_heads
+        self.scale = math.sqrt(self.d_head)
+
+        self.qkv_proj = nn.Linear(d_model, 3 * d_model, bias=False)
+        self.out_proj = nn.Linear(d_model, d_model)
+        self.edge_proj = nn.Linear(d_edge, n_heads)
+        self.attn_drop = nn.Dropout(dropout)
+
+        self.ff = nn.Sequential(
+            nn.Linear(d_model, d_model * ff_mult), nn.GELU(),
+            nn.Dropout(dropout), nn.Linear(d_model * ff_mult, d_model),
+        )
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.drop = nn.Dropout(dropout)
+
+    def forward(self, x, edge_feats, mask=None, chain_mask=None):
+        """
+        x: [B, N, d_model]
+        edge_feats: [B, N, N, d_edge]
+        mask: [B, N] bool (True = valid)
+        chain_mask: [B, N] float (0=receptor, 1=binder)
+        """
+        B, N, D = x.shape
+        H = self.n_heads
+
+        h = self.norm1(x)
+        qkv = self.qkv_proj(h).reshape(B, N, 3, H, self.d_head).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)  # each [B, H, N, d_head]
+
+        attn_logits = (q @ k.transpose(-2, -1)) / self.scale  # [B, H, N, N]
+        edge_bias = self.edge_proj(edge_feats).permute(0, 3, 1, 2)  # [B, H, N, N]
+        attn_logits = attn_logits + edge_bias
+
+        # Mask padding
+        if mask is not None:
+            attn_logits = attn_logits.masked_fill(~mask[:, None, None, :], float('-inf'))
+
+        # Cross-chain mask: block same-chain attention
+        if chain_mask is not None:
+            same_chain = (chain_mask.unsqueeze(1) == chain_mask.unsqueeze(2))  # [B, N, N]
+            attn_logits = attn_logits.masked_fill(same_chain[:, None, :, :], float('-inf'))
+
+        attn_weights = self.attn_drop(F.softmax(attn_logits, dim=-1))
+        attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
+
+        out = (attn_weights @ v).transpose(1, 2).reshape(B, N, D)
+        x = x + self.drop(self.out_proj(out))
+        x = x + self.drop(self.ff(self.norm2(x)))
+        return x
+
+
+class EdgeUpdateLayer(nn.Module):
+    """Updates edge features using node representations each layer.
+    Memory-efficient: projects nodes to low-dim before outer product."""
+    def __init__(self, d_model: int, d_edge: int, dropout: float = 0.1):
+        super().__init__()
+        d_proj = min(32, d_model // 4)  # Low-dim projection to save memory
+        self.proj_i = nn.Linear(d_model, d_proj, bias=False)
+        self.proj_j = nn.Linear(d_model, d_proj, bias=False)
+        self.edge_mlp = nn.Sequential(
+            nn.Linear(2 * d_proj + d_edge, d_edge),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(d_edge, d_edge),
+        )
+        self.norm = nn.LayerNorm(d_edge)
+
+    def forward(self, h, e, mask=None):
+        B, N, D = h.shape
+        hi = self.proj_i(h).unsqueeze(2).expand(-1, -1, N, -1)  # [B, N, N, d_proj]
+        hj = self.proj_j(h).unsqueeze(1).expand(-1, N, -1, -1)  # [B, N, N, d_proj]
+        inp = torch.cat([hi, hj, self.norm(e)], dim=-1)
+        e = e + self.edge_mlp(inp)
+        return e
+
+
+class InterfaceGNN(nn.Module):
+    """
+    Q_theta scorer: SE(3)-invariant dense graph transformer for interface scoring.
+
+    Input:
+        node_feats: [B, N, node_dim] per-residue features
+        edge_feats: [B, N, N, edge_dim] pairwise edge features
+        mask: [B, N] bool (True = valid residue, False = padding)
+
+    Output:
+        scores: [B] in (0, 1) = Q_theta(X, Y)
+    """
+    def __init__(
+        self,
+        node_dim: int = 28,
+        edge_dim: int = 37,
+        hidden_dim: int = 128,
+        n_layers: int = 4,
+        n_heads: int = 8,
+        ff_mult: int = 4,
+        dropout: float = 0.1,
+        backbone: str = 'transformer',
+        pooling: str = 'meanmax',   # 'meanmax' or 'attention'
+        edge_update: bool = False,
+        esm_dim: int = 0,          # 0 = no ESM; >0 = ESM embedding dim to project
+        esm_proj_dim: int = 128,   # projection dim for ESM features
+        esm_dropout: float = 0.0,  # dropout on ESM projection
+    ):
+        super().__init__()
+        actual_node_dim = node_dim + (esm_proj_dim if esm_dim > 0 else 0)
+        self.esm_dim = esm_dim
+        if esm_dim > 0:
+            layers = [
+                nn.Linear(esm_dim, esm_proj_dim),
+                nn.LayerNorm(esm_proj_dim),
+                nn.GELU(),
+            ]
+            if esm_dropout > 0:
+                layers.append(nn.Dropout(esm_dropout))
+            self.esm_proj = nn.Sequential(*layers)
+        self.node_embed = nn.Sequential(
+            nn.Linear(actual_node_dim, hidden_dim),
+            nn.LayerNorm(hidden_dim),
+            nn.GELU(),
+        )
+        self.edge_embed = nn.Sequential(
+            nn.Linear(edge_dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, hidden_dim // 2),
+        )
+        d_edge_hidden = hidden_dim // 2
+
+        if backbone == 'transformer':
+            self.layers = nn.ModuleList([
+                InterfaceTransformerLayer(hidden_dim, n_heads, d_edge_hidden, ff_mult, dropout)
+                for _ in range(n_layers)
+            ])
+        elif backbone == 'gat':
+            self.layers = nn.ModuleList([
+                GATLayer(hidden_dim, n_heads, ff_mult, dropout)
+                for _ in range(n_layers)
+            ])
+        elif backbone == 'gcn':
+            self.layers = nn.ModuleList([
+                GCNLayer(hidden_dim, d_edge_hidden, ff_mult, dropout)
+                for _ in range(n_layers)
+            ])
+        elif backbone == 'crosschain':
+            # Interleave self-attention and cross-chain attention
+            layers = []
+            for i in range(n_layers):
+                if i % 2 == 0:
+                    layers.append(InterfaceTransformerLayer(hidden_dim, n_heads, d_edge_hidden, ff_mult, dropout))
+                else:
+                    layers.append(CrossChainTransformerLayer(hidden_dim, n_heads, d_edge_hidden, ff_mult, dropout))
+            self.layers = nn.ModuleList(layers)
+        else:
+            raise ValueError(f"Unknown backbone: {backbone}")
+
+        self.norm_out = nn.LayerNorm(hidden_dim)
+
+        # Edge update layers (optional)
+        self.edge_update = edge_update
+        if edge_update:
+            self.edge_update_layers = nn.ModuleList([
+                EdgeUpdateLayer(hidden_dim, d_edge_hidden, dropout)
+                for _ in range(n_layers)
+            ])
+
+        # Pooling
+        self.pooling = pooling
+        if pooling == 'attention':
+            self.attn_pool = nn.Sequential(
+                nn.Linear(hidden_dim, hidden_dim // 2),
+                nn.Tanh(),
+                nn.Linear(hidden_dim // 2, 1),
+            )
+            pool_dim = hidden_dim
+        else:
+            pool_dim = 2 * hidden_dim
+
+        # Scoring head
+        self.head = nn.Sequential(
+            nn.Linear(pool_dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, hidden_dim // 2),
+            nn.GELU(),
+            nn.Linear(hidden_dim // 2, 1),
+        )
+
+    def forward(self, node_feats, edge_feats, mask, esm_feats=None):
+        """
+        node_feats: [B, N, node_dim]
+        edge_feats: [B, N, N, edge_dim]
+        mask: [B, N] bool
+        esm_feats: [B, N, esm_dim] optional ESM-2 embeddings
+        Returns: scores [B] in (0, 1)
+        """
+        B, N, _ = node_feats.shape
+
+        # Extract chain mask for cross-chain attention (last dim = chain indicator)
+        chain_mask = node_feats[:, :, -1]  # [B, N] float: 0=receptor, 1=binder
+
+        # Optionally concatenate projected ESM features
+        if self.esm_dim > 0 and esm_feats is not None:
+            esm_proj = self.esm_proj(esm_feats)  # [B, N, 128]
+            node_feats = torch.cat([node_feats, esm_proj], dim=-1)
+
+        # Embed nodes and edges
+        h = self.node_embed(node_feats)          # [B, N, hidden_dim]
+        e = self.edge_embed(edge_feats)          # [B, N, N, hidden_dim//2]
+
+        # Graph transformer layers (with optional edge updates)
+        for i, layer in enumerate(self.layers):
+            if isinstance(layer, CrossChainTransformerLayer):
+                h = layer(h, e, mask, chain_mask=chain_mask)
+            else:
+                h = layer(h, e, mask)
+            if self.edge_update:
+                e = self.edge_update_layers[i](h, e, mask)
+
+        h = self.norm_out(h)  # [B, N, hidden_dim]
+
+        # Pooling
+        mask_f = mask.float().unsqueeze(-1)  # [B, N, 1]
+
+        if self.pooling == 'attention':
+            # Learned attention pooling
+            attn_logits = self.attn_pool(h).squeeze(-1)  # [B, N]
+            attn_logits = attn_logits.masked_fill(~mask, float('-inf'))
+            attn_weights = F.softmax(attn_logits, dim=-1).unsqueeze(-1)  # [B, N, 1]
+            attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
+            h_pool = (h * attn_weights).sum(dim=1)  # [B, hidden_dim]
+        else:
+            # Mean + max pooling
+            h_masked = h * mask_f
+            h_mean = h_masked.sum(dim=1) / (mask_f.sum(dim=1) + 1e-8)
+            h_max_input = h_masked + (1 - mask_f) * (-1e9)
+            h_max = h_max_input.max(dim=1).values
+            h_pool = torch.cat([h_mean, h_max], dim=-1)  # [B, 2*hidden_dim]
+
+        # Score
+        logits = self.head(h_pool).squeeze(-1)  # [B]
+        scores = torch.sigmoid(logits)          # [B] in (0, 1)
+        return scores
+
+
+class AlloDesignerScorer(nn.Module):
+    """
+    Full Q_theta model wrapper with loss computation.
+
+    Implements the two-stage training objective:
+        Phase 1: DockQ regression (MSE loss)
+        Phase 2: Selectivity margin ranking (contrastive loss)
+
+    The selectivity margin from the paper (Eq. 3):
+        S_theta(Y; X+, N) = logit(Q(X+, Y)) - log sum_X- exp(logit(Q(X-, Y)))
+    """
+    def __init__(self, node_dim=28, edge_dim=37, hidden_dim=128,
+                 n_layers=4, n_heads=8, dropout=0.1, backbone='transformer',
+                 pooling='meanmax', edge_update=False, esm_dim=0,
+                 esm_proj_dim=128, esm_dropout=0.0):
+        super().__init__()
+        self.gnn = InterfaceGNN(node_dim, edge_dim, hidden_dim, n_layers, n_heads,
+                                dropout=dropout, backbone=backbone,
+                                pooling=pooling, edge_update=edge_update,
+                                esm_dim=esm_dim, esm_proj_dim=esm_proj_dim,
+                                esm_dropout=esm_dropout)
+
+    def forward(self, node_feats, edge_feats, mask, esm_feats=None):
+        return self.gnn(node_feats, edge_feats, mask, esm_feats=esm_feats)
+
+    def compute_dockq_loss(self, scores, dockq_labels):
+        """Phase 1: MSE regression loss against DockQ labels."""
+        return F.mse_loss(scores, dockq_labels.float())
+
+    def compute_selectivity_loss(self, pos_scores, neg_scores_list, margin: float = 0.2):
+        """
+        Phase 2: Selectivity margin loss.
+
+        For each binder Y:
+            pos_score = Q(X+, Y)
+            neg_scores = [Q(X-, Y) for X- in N]
+
+        Loss = -mean(S_theta) where
+            S_theta = logit(pos_score) - log sum exp(logit(neg_scores))
+
+        Also computes a soft margin loss:
+            L_margin = mean(max(0, margin - (pos_score - neg_score)))
+        """
+        # logit = log(p / (1-p))
+        eps = 1e-6
+        pos_logit = torch.log(pos_scores.clamp(eps, 1 - eps) / (1 - pos_scores).clamp(eps))
+
+        # neg_scores_list: list of [B] tensors
+        neg_logits = torch.stack([
+            torch.log(s.clamp(eps, 1 - eps) / (1 - s).clamp(eps))
+            for s in neg_scores_list
+        ], dim=-1)  # [B, n_neg]
+
+        # InfoNCE-style selectivity margin
+        log_denom = torch.logsumexp(neg_logits, dim=-1)  # [B]
+        selectivity = pos_logit - log_denom              # [B]
+        selectivity_loss = -selectivity.mean()
+
+        # Soft margin loss (averaged over all negatives)
+        margin_losses = []
+        for neg_scores in neg_scores_list:
+            margin_losses.append(F.relu(margin - (pos_scores - neg_scores)))
+        margin_loss = torch.stack(margin_losses, dim=-1).mean()
+
+        return selectivity_loss + margin_loss
+
+    def compute_path_selectivity_loss(self, pos_scores, neg_scores_list,
+                                       path_scores_list, path_taus,
+                                       margin=0.2, path_lambda=0.5):
+        """
+        Extended selectivity loss with path monotonicity regularization.
+
+        Args:
+            pos_scores: [B] Q(X1, Y) -- goal state scores
+            neg_scores_list: list of [B] -- Q(X0, Y), Q(X_cryptic, Y), etc.
+            path_scores_list: list of [B] -- Q(X_tau, Y) for each path frame
+            path_taus: list of float -- tau values for each path frame (sorted)
+            margin: margin for ranking loss
+            path_lambda: weight for path monotonicity loss
+
+        Returns:
+            total_loss: selectivity loss + path_lambda * monotonicity loss
+            loss_dict: breakdown of loss components
+        """
+        # Standard selectivity loss (unchanged)
+        select_loss = self.compute_selectivity_loss(pos_scores, neg_scores_list, margin)
+
+        # Path monotonicity loss: ensure Q increases with tau
+        loss_monotone = torch.tensor(0.0, device=pos_scores.device)
+        if path_scores_list and path_lambda > 0:
+            small_margin = 0.05
+            # Consecutive path frames should be monotonically increasing
+            for i in range(len(path_scores_list) - 1):
+                loss_monotone = loss_monotone + F.relu(
+                    path_scores_list[i] - path_scores_list[i + 1] + small_margin
+                ).mean()
+            # Last path frame should be less than positive (holo) score
+            loss_monotone = loss_monotone + F.relu(
+                path_scores_list[-1] - pos_scores + margin
+            ).mean()
+            # First path frame should be greater than negative (apo) score
+            if neg_scores_list:
+                loss_monotone = loss_monotone + F.relu(
+                    neg_scores_list[0] - path_scores_list[0] + small_margin
+                ).mean()
+
+        total = select_loss + path_lambda * loss_monotone
+        return total, {
+            'loss_selectivity': select_loss.item(),
+            'loss_path_monotone': loss_monotone.item(),
+        }
+
+    def compute_combined_loss(self, pos_scores, neg_scores_list, dockq_labels,
+                              lambda_rank: float = 1.0):
+        """Combined Phase 1 + Phase 2 loss."""
+        # Regression loss on all scores (pos + neg get appropriate labels)
+        dockq_loss = self.compute_dockq_loss(pos_scores, dockq_labels)
+
+        # Selectivity loss
+        select_loss = self.compute_selectivity_loss(pos_scores, neg_scores_list)
+
+        return dockq_loss + lambda_rank * select_loss, {
+            'loss_dockq': dockq_loss.item(),
+            'loss_selectivity': select_loss.item(),
+        }
+
+
+def build_model(config: dict) -> AlloDesignerScorer:
+    """Build the Q_theta scorer from a config dict."""
+    return AlloDesignerScorer(
+        node_dim=config.get('node_dim', 32),
+        edge_dim=config.get('edge_dim', 37),
+        hidden_dim=config.get('hidden_dim', 128),
+        n_layers=config.get('n_layers', 4),
+        n_heads=config.get('n_heads', 8),
+        dropout=config.get('dropout', 0.1),
+        backbone=config.get('backbone', 'transformer'),
+        pooling=config.get('pooling', 'meanmax'),
+        edge_update=config.get('edge_update', False),
+        esm_dim=config.get('esm_dim', 0),
+        esm_proj_dim=config.get('esm_proj_dim', 128),
+        esm_dropout=config.get('esm_dropout', 0.0),
+    )

code/requirements.txt

@@ -0,0 +1,22 @@
+# Core
+torch>=2.0.0
+numpy>=1.24.0
+
+# Protein structure
+biopython>=1.80
+
+# ML utilities
+scipy>=1.10.0
+scikit-learn>=1.3.0
+
+# Experiment tracking
+wandb>=0.12.0
+
+# Config
+pyyaml>=6.0
+
+# Visualization
+matplotlib>=3.7.0
+
+# Optional accelerations
+einops>=0.6.0

code/scripts/README.md

@@ -0,0 +1,55 @@
+# `code/scripts/` — entry points
+
+This public release ships only the inference and sampling code for Q_θ.
+
+| File / dir | Purpose |
+|---|---|
+| `evaluate.py` | Score binders in a pre-built `*.pkl` test set with a Q_θ checkpoint; reports Spearman ρ, AUC, selectivity gap. |
+| `rescore.py` | Re-score raw PDB designs (binder + holo + apo) with Q_θ. |
+| `pxdesign_guidance/` | PXDesign-prior guidance with Q_θ (Langevin / SMC / TDS / classifier). |
+
+Training, baseline scoring (ProteinMPNN / ESM-IF / Rosetta / DFIRE / energy panel), guidance for RFdiffusion / Proteina-ComplexA, and paper-figure aggregation are **not** shipped; the inference path above is the only supported surface for the public release.
+
+---
+
+## Deploying Q_θ with other base models
+
+Q_θ provides two interfaces:
+
+1. **Re-ranker (best-of-K).** Given K candidate binders from any prior, score each with `S(Y) = Q_θ(X¹, Y) − Q_θ(X⁰, Y)` and pick the top. No gradient signal needed; the prior is unmodified.
+2. **Gradient signal for guidance.** Compute `∇_Y S(Y)` via `DifferentiableQTheta` (in `code/models/differentiable_features.py`) and inject into the prior's sampler (Langevin step, SMC weight, TDS twist, classifier guidance score).
+
+The `pxdesign_guidance/` subdir is a worked example of interface (2) wrapping PXDesign. To plug Q_θ into another prior, mirror that pattern:
+
+### RFdiffusion
+
+1. Clone RFdiffusion: <https://github.com/RosettaCommons/RFdiffusion>.
+2. Follow its install + checkpoint download.
+3. In RFdiffusion's diffusion loop, after each denoising step, materialize the predicted backbone, build the holo/apo graph inputs expected by `DifferentiableQTheta`, and either:
+   - Apply a Langevin nudge: `x ← x + η · ∇_x S(x)`.
+   - Add a classifier-guidance term to the denoiser's `xt-1` mean: `μ' = μ + s · σ² · ∇_x log p(y|x)`, where `log p(y|x) ≈ S(x)` (Q_θ is treated as the log-likelihood of "is good binder").
+4. Reference template: `pxdesign_guidance/guided_pxdesign.py`.
+
+### Proteina-ComplexA
+
+1. Clone Proteina: <https://github.com/proteinabio/proteina-complexa> (or the released artifact).
+2. Use its ComplexA mode that emits binder coords conditioned on a receptor.
+3. Same plug pattern as RFdiffusion — wrap the sampler with `DifferentiableQTheta` for guidance, or run unguided and re-rank with `evaluate.py` / `rescore.py`.
+
+### Any backbone prior
+
+The only contract Q_θ enforces:
+
+- Receptor input is a PDB with holo and apo coordinates.
+- Binder input is a PDB (or coords) with chain id distinct from receptor's.
+- For guidance, expose differentiable Cα + backbone coordinates so `∇_x S(x)` flows.
+
+See `code/models/differentiable_features.py:DifferentiableQTheta` for the exact interface (`load_receptor(holo_path, apo_path, …)`, `score(design_path, binder_chain, state)`, `.differentiable_score(coords, …)`).
+
+---
+
+## Why other guidance scripts aren't shipped
+
+The RFdiffusion / Proteina guidance variants in our internal tree depend on those projects' un-released CIF formats and patched samplers; we don't want to ship modified third-party code. The PXDesign variants we do ship use only PXDesign's public API and are self-contained.
+
+For citation / reproduction context, see the paper §4 (guidance methods).

code/scripts/evaluate.py

@@ -0,0 +1,332 @@
+"""
+Evaluation script for the trained Q_theta scorer.
+
+Computes:
+  1. Selectivity metrics (gap, ranking accuracy, AUC)
+  2. DockQ correlation (Spearman/Pearson)
+  3. Score distributions (violin plots)
+  4. Best-of-K analysis (as function of K)
+  5. Per-target breakdown
+
+Usage:
+    python code/scripts/evaluate.py \
+        --target cam \
+        --checkpoint checkpoints/Q_theta_phase2.pt \
+        --data_dir data/processed \
+        --gpu 7
+"""
+
+import os
+import sys
+import argparse
+import logging
+import json
+import numpy as np
+import torch
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+from scipy.stats import spearmanr, pearsonr
+from sklearn.metrics import roc_auc_score, roc_curve
+
+_CODE_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
+if _CODE_DIR not in sys.path:
+    sys.path.insert(0, _CODE_DIR)
+
+from models.scorer import build_model
+from data.dataset import TwoStateComplexDataset, collate_fn
+from torch.utils.data import DataLoader
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
+logger = logging.getLogger(__name__)
+
+
+def compute_best_of_k(pos_scores, K_values=None, threshold=0.7):
+    """
+    Simulate best-of-K selection: what fraction of draws contain at least one good binder?
+    Assumes pos_scores are from a distribution of candidate binders for goal state X+.
+    """
+    if K_values is None:
+        K_values = [1, 2, 5, 10, 20, 50, 100]
+    results = {}
+    n = len(pos_scores)
+    n_trials = 1000
+
+    for K in K_values:
+        successes = 0
+        for _ in range(n_trials):
+            idxs = np.random.choice(n, size=min(K, n), replace=False)
+            best_score = pos_scores[idxs].max()
+            if best_score >= threshold:
+                successes += 1
+        results[K] = successes / n_trials
+
+    return results
+
+
+def compute_selectivity_margin(pos_scores, neg_scores):
+    """Compute per-sample selectivity margin S_theta."""
+    eps = 1e-6
+    pos_logit = np.log(pos_scores.clip(eps, 1-eps) / (1-pos_scores).clip(eps))
+    neg_logit = np.log(neg_scores.clip(eps, 1-eps) / (1-neg_scores).clip(eps))
+    selectivity = pos_logit - np.log(np.exp(neg_logit) + 1e-8)
+    return selectivity
+
+
+def plot_score_distributions(pos_scores, neg_scores, decoy_scores=None,
+                              title='Score Distributions', outpath=None):
+    """Violin plot of score distributions for different complex types."""
+    fig, ax = plt.subplots(figsize=(8, 6))
+
+    data = [pos_scores, neg_scores]
+    labels = ['Positive\n(X+, Y)', 'Negative\n(X0, Y)']
+    colors = ['#2196F3', '#F44336']
+
+    if decoy_scores is not None and len(decoy_scores) > 0:
+        data.append(decoy_scores)
+        labels.append('Decoys\n(X+, Y~)')
+        colors.append('#FF9800')
+
+    parts = ax.violinplot(data, positions=range(len(data)), showmedians=True)
+    for i, (pc, c) in enumerate(zip(parts['bodies'], colors)):
+        pc.set_facecolor(c)
+        pc.set_alpha(0.7)
+
+    ax.set_xticks(range(len(data)))
+    ax.set_xticklabels(labels)
+    ax.set_ylabel('Q_theta Score', fontsize=12)
+    ax.set_title(title, fontsize=14)
+    ax.set_ylim(0, 1)
+    ax.axhline(0.5, color='gray', linestyle='--', alpha=0.5, label='Decision boundary')
+    ax.legend()
+
+    # Add mean + std annotations
+    for i, (d, c) in enumerate(zip(data, colors)):
+        ax.text(i, 0.02, f'μ={d.mean():.2f}\nσ={d.std():.2f}',
+                ha='center', fontsize=9, color=c)
+
+    plt.tight_layout()
+    if outpath:
+        plt.savefig(outpath, dpi=150, bbox_inches='tight')
+        logger.info(f"Saved plot to {outpath}")
+    plt.close()
+
+
+def plot_roc_curve(labels, scores, title='ROC Curve', outpath=None):
+    """Plot ROC curve for positive vs negative classification."""
+    fpr, tpr, _ = roc_curve(labels, scores)
+    auc = roc_auc_score(labels, scores)
+
+    fig, ax = plt.subplots(figsize=(6, 6))
+    ax.plot(fpr, tpr, 'b-', lw=2, label=f'AUC = {auc:.3f}')
+    ax.plot([0, 1], [0, 1], 'k--', lw=1)
+    ax.set_xlabel('False Positive Rate')
+    ax.set_ylabel('True Positive Rate')
+    ax.set_title(title)
+    ax.legend()
+    plt.tight_layout()
+    if outpath:
+        plt.savefig(outpath, dpi=150, bbox_inches='tight')
+    plt.close()
+    return auc
+
+
+def plot_best_of_k(results, outpath=None):
+    """Plot best-of-K success rate as a function of K."""
+    Ks = sorted(results.keys())
+    success_rates = [results[K] for K in Ks]
+
+    fig, ax = plt.subplots(figsize=(8, 5))
+    ax.semilogx(Ks, success_rates, 'b-o', lw=2, markersize=8)
+    ax.set_xlabel('K (number of candidates)', fontsize=12)
+    ax.set_ylabel('Success rate (best score > 0.7)', fontsize=12)
+    ax.set_title('Best-of-K Analysis', fontsize=14)
+    ax.set_ylim(0, 1.05)
+    ax.grid(True, alpha=0.3)
+    ax.axhline(0.8, color='red', linestyle='--', alpha=0.5, label='80% success')
+    ax.legend()
+    plt.tight_layout()
+    if outpath:
+        plt.savefig(outpath, dpi=150, bbox_inches='tight')
+    plt.close()
+
+
+@torch.no_grad()
+def evaluate(model, loader, device):
+    """Run model on a dataset and collect all predictions."""
+    model.eval()
+    all_scores, all_labels, all_types, all_pdbs = [], [], [], []
+
+    for batch in loader:
+        esm_feats = batch['esm_feats'].to(device) if 'esm_feats' in batch else None
+        scores = model(
+            batch['node_feats'].to(device),
+            batch['edge_feats'].to(device),
+            batch['node_mask'].to(device),
+            esm_feats=esm_feats,
+        )
+        all_scores.extend(scores.cpu().numpy().tolist())
+        all_labels.extend(batch['label'].numpy().tolist())
+        all_types.extend(batch['type'])
+        all_pdbs.extend(batch['pdb'])
+
+    return (np.array(all_scores), np.array(all_labels),
+            np.array(all_types), np.array(all_pdbs))
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Evaluate Allo-Designer Q_theta scorer')
+    parser.add_argument('--target', default='cam',
+                        help='Target name (cam, abl, era, or any custom target with data in data/processed/)')
+    parser.add_argument('--all_targets', action='store_true',
+                        help='Evaluate on all available targets and produce aggregated results')
+    parser.add_argument('--checkpoint', required=True, help='Path to model checkpoint')
+    parser.add_argument('--data_dir', default='data/processed')
+    parser.add_argument('--split', choices=['val', 'test'], default='test')
+    parser.add_argument('--batch_size', type=int, default=32)
+    parser.add_argument('--gpu', type=int, default=7)
+    parser.add_argument('--outdir', default='results')
+    parser.add_argument('--bok_threshold', type=float, default=0.7,
+                        help='Score threshold for best-of-K (default 0.7; use per-target value for calibrated results)')
+    parser.add_argument('--esm_dir', default=None,
+                        help='Path to ESM-2 embedding cache (auto-detected at <data_dir>/esm2_embeddings if omitted)')
+    parser.add_argument('--no_wandb', action='store_true', help='(ignored; here for CLI compatibility)')
+    args = parser.parse_args()
+
+    # Auto-detect ESM dir under data_dir
+    if args.esm_dir is None:
+        cand = os.path.join(args.data_dir, 'esm2_embeddings')
+        if os.path.isdir(cand):
+            args.esm_dir = cand
+
+    device = torch.device(f'cuda:{args.gpu}' if torch.cuda.is_available() else 'cpu')
+    os.makedirs(args.outdir, exist_ok=True)
+    os.makedirs(f'{args.outdir}/figures', exist_ok=True)
+    os.makedirs(f'{args.outdir}/tables', exist_ok=True)
+
+    # Load model
+    state = torch.load(args.checkpoint, map_location=device)
+    config = state.get('config', {})
+    model = build_model(config).to(device)
+    model.load_state_dict(state['model_state'])
+    logger.info(f"Loaded model from {args.checkpoint}")
+
+    # Load dataset
+    data_path = os.path.join(args.data_dir, args.target, f'{args.split}.pkl')
+    if not os.path.exists(data_path):
+        logger.error(f"Data not found: {data_path}")
+        sys.exit(1)
+
+    dataset = TwoStateComplexDataset(data_path, max_nodes=128,
+                                     esm_dir=args.esm_dir, target_name=args.target)
+    loader = DataLoader(
+        dataset, batch_size=args.batch_size, shuffle=False,
+        num_workers=2, collate_fn=collate_fn
+    )
+
+    # Run evaluation
+    logger.info(f"Evaluating on {len(dataset)} samples...")
+    scores, labels, types, pdbs = evaluate(model, loader, device)
+
+    # Separate by type
+    pos_mask = (types == 'positive')
+    neg_apo_mask = (types == 'negative_apo')
+    decoy_mask = np.array(['decoy' in t for t in types])
+
+    pos_scores = scores[pos_mask]
+    neg_scores = scores[neg_apo_mask]
+    decoy_scores = scores[decoy_mask]
+
+    logger.info(f"\n{'='*50}")
+    logger.info(f"Results for {args.target} ({args.split})")
+    logger.info(f"{'='*50}")
+    logger.info(f"Positive samples: {pos_mask.sum()}")
+    logger.info(f"Negative (apo) samples: {neg_apo_mask.sum()}")
+    logger.info(f"Decoy samples: {decoy_mask.sum()}")
+
+    # --- Core metrics ---
+    metrics = {}
+
+    # 1. Spearman correlation with DockQ labels
+    sp, p_val = spearmanr(scores, labels)
+    metrics['spearman_all'] = float(sp)
+    metrics['spearman_pval'] = float(p_val)
+    logger.info(f"\nSpearman(Q_theta, DockQ): {sp:.3f} (p={p_val:.3e})")
+
+    # 2. Selectivity gap (positive vs negative_apo)
+    if pos_mask.sum() > 0 and neg_apo_mask.sum() > 0:
+        gap = float(pos_scores.mean() - neg_scores.mean())
+        ranking_acc = float((pos_scores.mean() > neg_scores).mean() if len(neg_scores) > 0 else 0.5)
+        metrics['selectivity_gap'] = gap
+        metrics['pos_score_mean'] = float(pos_scores.mean())
+        metrics['neg_score_mean'] = float(neg_scores.mean())
+        metrics['pos_score_std'] = float(pos_scores.std())
+        metrics['neg_score_std'] = float(neg_scores.std())
+        logger.info(f"Selectivity gap (pos - neg): {gap:.3f}")
+        logger.info(f"  Pos: {pos_scores.mean():.3f} ± {pos_scores.std():.3f}")
+        logger.info(f"  Neg: {neg_scores.mean():.3f} ± {neg_scores.std():.3f}")
+
+    # 3. AUC for positive vs negative
+    if pos_mask.sum() > 0 and neg_apo_mask.sum() > 0:
+        pn_scores = np.concatenate([pos_scores, neg_scores])
+        pn_labels = np.concatenate([np.ones(len(pos_scores)), np.zeros(len(neg_scores))])
+        auc = roc_auc_score(pn_labels, pn_scores)
+        metrics['auc_pos_vs_neg'] = float(auc)
+        logger.info(f"AUC (pos vs neg_apo): {auc:.3f}")
+
+        # ROC curve
+        plot_roc_curve(
+            pn_labels, pn_scores,
+            title=f'ROC: Positive vs Negative Apo ({args.target.upper()})',
+            outpath=f'{args.outdir}/figures/roc_{args.target}_{args.split}.png'
+        )
+
+    # 4. AUC for quality classification (DockQ > 0.5)
+    binary = (labels > 0.5).astype(int)
+    if binary.sum() > 0 and binary.sum() < len(binary):
+        auc_quality = roc_auc_score(binary, scores)
+        metrics['auc_quality'] = float(auc_quality)
+        logger.info(f"AUC (quality>0.5): {auc_quality:.3f}")
+
+    # 5. Best-of-K analysis
+    if len(pos_scores) > 0:
+        bok_results = compute_best_of_k(pos_scores, K_values=[1, 2, 5, 10, 20, 50],
+                                         threshold=args.bok_threshold)
+        metrics['best_of_k'] = {str(K): float(v) for K, v in bok_results.items()}
+        logger.info(f"\nBest-of-K success rates:")
+        for K, rate in bok_results.items():
+            logger.info(f"  K={K:3d}: {rate:.3f}")
+        plot_best_of_k(
+            bok_results,
+            outpath=f'{args.outdir}/figures/best_of_k_{args.target}_{args.split}.png'
+        )
+
+    # 6. Score distributions plot
+    plot_score_distributions(
+        pos_scores if len(pos_scores) > 0 else np.array([]),
+        neg_scores if len(neg_scores) > 0 else np.array([]),
+        decoy_scores if len(decoy_scores) > 0 else None,
+        title=f'Q_theta Score Distributions ({args.target.upper()})',
+        outpath=f'{args.outdir}/figures/score_dist_{args.target}_{args.split}.png'
+    )
+
+    # Save metrics
+    out_json = f'{args.outdir}/tables/eval_{args.target}_{args.split}.json'
+    with open(out_json, 'w') as f:
+        json.dump(metrics, f, indent=2)
+    logger.info(f"\nSaved metrics to {out_json}")
+
+    # Print summary table
+    logger.info(f"\n{'='*50}")
+    logger.info("SUMMARY TABLE")
+    logger.info(f"{'='*50}")
+    logger.info(f"{'Metric':<30} {'Value':>10}")
+    logger.info(f"{'-'*42}")
+    for k, v in metrics.items():
+        if isinstance(v, float):
+            logger.info(f"{k:<30} {v:>10.4f}")
+    logger.info(f"{'='*50}")
+
+
+if __name__ == '__main__':
+    main()

code/scripts/pxdesign_guidance/__init__.py

@@ -0,0 +1 @@
+# PXDesign + Q_theta guidance integration

code/scripts/pxdesign_guidance/convert_cif_to_pdb.py

@@ -0,0 +1,132 @@
+"""
+Convert PXDesign CIF outputs to PDB format for evaluation pipeline.
+
+PXDesign outputs .cif files with:
+- Chain IDs like A0/B0 (multi-char, not PDB-compatible)
+- Non-standard residue name 'xpb' for designed binder residues
+
+This script converts them to PDB format with:
+- Single-char chain IDs (A, B)
+- Preserved residue names (xpb is kept; eval tools handle it)
+
+Usage:
+    python code/scripts/pxdesign_guidance/convert_cif_to_pdb.py
+"""
+import os
+import sys
+from glob import glob
+
+from Bio.PDB import MMCIFParser, PDBIO, Select
+
+_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+_PROJECT_DIR = os.path.abspath(os.path.join(_SCRIPT_DIR, '../../..'))
+
+
+class ChainRenamer(Select):
+    """Rename multi-char chain IDs to single-char for PDB format."""
+    def __init__(self, chain_map):
+        self.chain_map = chain_map
+
+    def accept_chain(self, chain):
+        return 1
+
+    def accept_residue(self, residue):
+        return 1
+
+    def accept_atom(self, atom):
+        return 1
+
+
+def convert_cif_to_pdb(cif_path, pdb_path):
+    """Convert a single CIF file to PDB format."""
+    parser = MMCIFParser(QUIET=True)
+    structure = parser.get_structure('s', cif_path)
+    model = structure[0]
+
+    # Build chain ID mapping (A0->A, B0->B, etc.)
+    chain_map = {}
+    used_ids = set()
+    for chain in model.get_chains():
+        old_id = chain.id
+        # Use first character
+        new_id = old_id[0] if old_id else 'A'
+        # Avoid duplicates
+        while new_id in used_ids:
+            new_id = chr(ord(new_id) + 1)
+        used_ids.add(new_id)
+        chain_map[old_id] = new_id
+
+    # Rename chains and fix non-standard residue names
+    chains_to_rename = list(model.get_chains())
+    for chain in chains_to_rename:
+        old_id = chain.id
+        new_id = chain_map.get(old_id, old_id)
+        if old_id != new_id:
+            chain.id = new_id
+        # Rename 'xpb' residues to 'GLY' (backbone-only binder residues)
+        for residue in chain.get_residues():
+            if residue.resname.strip().lower() == 'xpb':
+                residue.resname = 'GLY'
+
+    # Write PDB
+    io = PDBIO()
+    io.set_structure(structure)
+    io.save(pdb_path)
+    return True
+
+
+def convert_directory(src_dir, method_name):
+    """Convert all CIF files in a directory tree to PDB."""
+    cif_files = sorted(glob(os.path.join(src_dir, '**/*.cif'), recursive=True))
+    cif_files = [f for f in cif_files if 'sample' in os.path.basename(f).lower()]
+
+    if not cif_files:
+        print(f"  No CIF files found in {src_dir}")
+        return 0
+
+    # Create converted_pdbs directory
+    converted_dir = os.path.join(src_dir, 'converted_pdbs')
+    os.makedirs(converted_dir, exist_ok=True)
+
+    n_converted = 0
+    for cif_path in cif_files:
+        basename = os.path.basename(cif_path).replace('.cif', '.pdb')
+        # For TDS/SMC with round subdirs, include round info
+        rel_path = os.path.relpath(cif_path, src_dir)
+        parts = rel_path.split(os.sep)
+        if any(p.startswith('round_') for p in parts):
+            round_part = [p for p in parts if p.startswith('round_')][0]
+            basename = f"{round_part}_{basename}"
+
+        pdb_path = os.path.join(converted_dir, basename)
+        try:
+            convert_cif_to_pdb(cif_path, pdb_path)
+            n_converted += 1
+        except Exception as e:
+            print(f"  Failed {cif_path}: {e}")
+
+    print(f"  Converted {n_converted}/{len(cif_files)} CIF -> PDB in {converted_dir}")
+    return n_converted
+
+
+def main():
+    methods = {
+        'pxdesign_guided': os.path.join(_PROJECT_DIR, 'results/pxdesign_guided'),
+        'pxdesign_tds': os.path.join(_PROJECT_DIR, 'results/pxdesign_tds'),
+        'pxdesign_smc': os.path.join(_PROJECT_DIR, 'results/pxdesign_smc'),
+    }
+    # Langevin outputs are already PDB (post-hoc refinement)
+
+    total = 0
+    for name, src_dir in methods.items():
+        print(f"\n{name}:")
+        if os.path.exists(src_dir):
+            total += convert_directory(src_dir, name)
+        else:
+            print(f"  Directory not found: {src_dir}")
+
+    print(f"\nTotal converted: {total}")
+
+
+if __name__ == '__main__':
+    main()

code/scripts/pxdesign_guidance/guided_pxdesign.py

@@ -0,0 +1,408 @@
+"""
+PXDesign + Q_theta Classifier Guidance.
+
+Monkey-patches PXDesign's diffusion sampling loop to inject Q_theta selectivity
+gradient after each denoising step. This steers the diffusion trajectory toward
+binder backbones that are conformationally selective.
+
+The patched diffusion loop:
+    x_denoised = denoise_net(x_noisy, t_hat, ...)
+    grad = ∇_{x_denoised}[Q(holo,Y) - Q(apo,Y)]   # <-- INJECTED
+    x_denoised = x_denoised + scale(t) * grad        # <-- INJECTED
+    delta = (x_noisy - x_denoised) / t_hat
+    x_l = x_noisy + eta * dt * delta
+
+Usage:
+    python code/scripts/pxdesign_guidance/guided_pxdesign.py \
+        --input experiments/pxdesign_cam/output/cam_binder.json \
+        --qtheta_checkpoint results/checkpoints_cam_v3/best_phase2.pt \
+        --ref_holo data/pdbs/cam_holo/3CLN.pdb \
+        --ref_apo data/pdbs/cam_apo/1CFD.pdb \
+        --guidance_scale 1.0 \
+        --N_sample 50 --N_step 400 \
+        --gpu 0
+"""
+
+import os
+import sys
+import argparse
+import json
+import logging
+import time
+import shutil
+from typing import Callable, Optional, Union
+from functools import partial
+
+import numpy as np
+import torch
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
+logger = logging.getLogger(__name__)
+
+# ── Paths ────────────────────────────────────────────────────────────────────
+_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+_ALLO_CODE_DIR = os.path.abspath(os.path.join(_SCRIPT_DIR, '..', '..'))
+_ALLO_ROOT = os.path.abspath(os.path.join(_ALLO_CODE_DIR, '..'))
+_PXDESIGN_DIR = os.environ.get('PXDESIGN_DIR', '')
+
+if _ALLO_CODE_DIR not in sys.path:
+    sys.path.insert(0, _ALLO_CODE_DIR)
+if _PXDESIGN_DIR not in sys.path:
+    sys.path.insert(0, _PXDESIGN_DIR)
+
+
+def guided_sample_diffusion(
+    denoise_net: Callable,
+    input_feature_dict: dict,
+    s_inputs: torch.Tensor,
+    s_trunk: torch.Tensor,
+    z_trunk: torch.Tensor,
+    noise_schedule: torch.Tensor,
+    N_sample: int = 1,
+    gamma0: float = 0.8,
+    gamma_min: float = 1.0,
+    noise_scale_lambda: float = 1.003,
+    step_scale_eta: Union[float, dict] = {"type": "const", "min": 1.5, "max": 1.5},
+    diffusion_chunk_size: Optional[int] = None,
+    inplace_safe: bool = False,
+    attn_chunk_size: Optional[int] = None,
+    # Guidance parameters (injected via partial)
+    guidance_module=None,
+    guidance_scale: float = 1.0,
+    guidance_start: float = 0.8,
+    guidance_end: float = 0.1,
+) -> torch.Tensor:
+    """
+    Modified PXDesign sample_diffusion with Q_theta classifier guidance.
+
+    Same as original generator.sample_diffusion but with gradient injection
+    after each denoising step. The gradient is scaled by a schedule that
+    applies stronger guidance at high noise levels (early steps).
+    """
+    from protenix.model.utils import centre_random_augmentation
+
+    N_atom = input_feature_dict["atom_to_token_idx"].size(-1)
+    batch_shape = s_inputs.shape[:-2]
+    device = s_inputs.device
+    dtype = s_inputs.dtype
+
+    logger.info(f"Guided sampling: scale={guidance_scale}, "
+                f"window=[{guidance_end:.1f}, {guidance_start:.1f}]")
+
+    def _chunk_sample_diffusion_guided(chunk_n_sample, inplace_safe):
+        x_l = noise_schedule[0] * torch.randn(
+            size=(*batch_shape, chunk_n_sample, N_atom, 3),
+            device=device, dtype=dtype
+        )
+        T = len(noise_schedule)
+
+        for step_t, (c_tau_last, c_tau) in enumerate(
+            zip(noise_schedule[:-1], noise_schedule[1:])
+        ):
+            # Centre random augmentation
+            x_l = (
+                centre_random_augmentation(x_input_coords=x_l, N_sample=1)
+                .squeeze(dim=-3)
+                .to(dtype)
+            )
+
+            # Predictor step: add noise
+            gamma = float(gamma0) if c_tau > gamma_min else 0
+            t_hat = c_tau_last * (gamma + 1)
+            delta_noise_level = torch.sqrt(t_hat**2 - c_tau_last**2)
+            x_noisy = x_l + noise_scale_lambda * delta_noise_level * torch.randn(
+                size=x_l.shape, device=device, dtype=dtype
+            )
+
+            # Reshape t_hat for network
+            t_hat_tensor = (
+                t_hat.reshape((1,) * (len(batch_shape) + 1))
+                .expand(*batch_shape, chunk_n_sample)
+                .to(dtype)
+            )
+
+            # Denoise
+            x_denoised = denoise_net(
+                x_noisy=x_noisy,
+                t_hat_noise_level=t_hat_tensor,
+                input_feature_dict=input_feature_dict,
+                s_inputs=s_inputs,
+                s_trunk=s_trunk,
+                z_trunk=z_trunk,
+                chunk_size=attn_chunk_size,
+                inplace_safe=inplace_safe,
+            )
+
+            # ── Q_theta guidance injection ──────────────────────────────
+            if guidance_module is not None:
+                # Compute progress fraction (0=start/high noise, 1=end/low noise)
+                progress = step_t / (T - 1) if T > 1 else 1.0
+
+                # Apply guidance only within the specified window
+                if guidance_end <= (1.0 - progress) <= guidance_start:
+                    # Handle batch dimensions
+                    x_for_grad = x_denoised
+                    if x_for_grad.dim() > 3:
+                        x_for_grad = x_for_grad.squeeze(0)
+
+                    # Scale: stronger at high noise, weaker near convergence
+                    noise_fraction = 1.0 - progress
+                    scale = guidance_scale * noise_fraction
+
+                    try:
+                        # Compute gradient for first sample (or all if small batch)
+                        n_guide = min(chunk_n_sample, 4)
+                        grad_accum = torch.zeros_like(x_for_grad)
+
+                        for si in range(n_guide):
+                            grad, margin = guidance_module.compute_guidance_gradient(
+                                x_for_grad, input_feature_dict,
+                                t_hat=t_hat, sample_idx=si
+                            )
+                            grad_accum[si] = grad[si] if grad.shape[0] > si else grad[0]
+
+                        # Broadcast gradient to remaining samples
+                        if n_guide < chunk_n_sample and n_guide > 0:
+                            avg_grad = grad_accum[:n_guide].mean(dim=0, keepdim=True)
+                            grad_accum[n_guide:] = avg_grad.expand(
+                                chunk_n_sample - n_guide, -1, -1)
+
+                        # Normalize gradient to prevent explosion
+                        grad_norm = grad_accum.norm(dim=-1, keepdim=True).clamp(min=1e-8)
+                        grad_normalized = grad_accum / grad_norm
+                        avg_norm = grad_norm.mean().item()
+
+                        # Apply guidance
+                        if avg_norm > 1e-6:
+                            # Scale by average gradient magnitude to keep step size reasonable
+                            x_denoised = x_denoised + scale * avg_norm * grad_normalized
+
+                            if step_t % 50 == 0:
+                                logger.info(
+                                    f"  Step {step_t}/{T}: margin={margin:.3f}, "
+                                    f"grad_norm={avg_norm:.4f}, scale={scale:.3f}")
+                    except Exception as e:
+                        if step_t % 100 == 0:
+                            logger.debug(f"  Step {step_t}: guidance failed: {e}")
+            # ── End guidance ────────────────────────────────────────────
+
+            # Euler step
+            delta = (x_noisy - x_denoised) / t_hat_tensor[..., None, None]
+            dt = c_tau - t_hat_tensor
+            if isinstance(step_scale_eta, float):
+                eta = step_scale_eta
+            elif step_scale_eta["type"] == "const":
+                assert step_scale_eta["min"] == step_scale_eta["max"]
+                eta = step_scale_eta["min"]
+            else:
+                eta_min, eta_max = step_scale_eta["min"], step_scale_eta["max"]
+                if step_scale_eta["type"] == "linear":
+                    eta = eta_min + (eta_max - eta_min) * (step_t / T)
+                elif step_scale_eta["type"] == "poly":
+                    eta = eta_min + (eta_max - eta_min) * (step_t / T) ** 2
+                elif step_scale_eta["type"] == "cos":
+                    eta = eta_min + 0.5 * (eta_max - eta_min) * (
+                        1 - np.cos(np.pi * step_t / T))
+                elif step_scale_eta["type"] == "piecewise":
+                    eta = eta_min if step_t / T < 0.5 else eta_max
+                elif step_scale_eta["type"] == "piecewise_65":
+                    eta = eta_min if step_t / T < 0.65 else eta_max
+                elif step_scale_eta["type"] == "piecewise_70":
+                    eta = eta_min if step_t / T < 0.70 else eta_max
+                else:
+                    raise ValueError("Unsupported eta schedule!")
+            x_l = x_noisy + eta * dt[..., None, None] * delta
+
+        return x_l
+
+    # Chunked sampling
+    if diffusion_chunk_size is None:
+        x_l = _chunk_sample_diffusion_guided(N_sample, inplace_safe=inplace_safe)
+    else:
+        x_l = []
+        no_chunks = N_sample // diffusion_chunk_size + (
+            N_sample % diffusion_chunk_size != 0)
+        for i in range(no_chunks):
+            chunk_n_sample = (
+                diffusion_chunk_size
+                if i < no_chunks - 1
+                else N_sample - i * diffusion_chunk_size
+            )
+            chunk_x_l = _chunk_sample_diffusion_guided(
+                chunk_n_sample, inplace_safe=inplace_safe)
+            x_l.append(chunk_x_l)
+        x_l = torch.cat(x_l, -3)
+
+    return x_l
+
+
+def run_guided_pxdesign(args):
+    """Run PXDesign with Q_theta classifier guidance."""
+    if 'CUDA_VISIBLE_DEVICES' not in os.environ:
+        os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)
+
+    # Import PXDesign components
+    from pxdesign.runner.inference import InferenceRunner, main as pxdesign_main
+    from pxdesign.utils.infer import (
+        get_configs, convert_to_bioassembly_dict, download_inference_cache, derive_seed
+    )
+    from pxdesign.utils.inputs import process_input_file
+    from protenix.config import save_config
+    from protenix.utils.seed import seed_everything
+    from protenix.utils.torch_utils import autocasting_disable_decorator
+
+    from qtheta_pxdesign import QThetaPXDesignGuidance
+
+    # Set up output directory
+    outdir = args.outdir if os.path.isabs(args.outdir) else os.path.join(_ALLO_ROOT, args.outdir)
+    os.makedirs(outdir, exist_ok=True)
+
+    # Build PXDesign CLI arguments
+    pxdesign_argv = [
+        '--dump_dir', outdir,
+        '--input', args.input,
+        '--dtype', 'bf16',
+        '--N_sample', str(args.N_sample),
+        '--N_step', str(args.N_step),
+    ]
+
+    configs = get_configs(pxdesign_argv)
+    configs.input_json_path = process_input_file(
+        configs.input_json_path, out_dir=outdir)
+    download_inference_cache(configs)
+
+    # Convert inputs
+    save_config(configs, os.path.join(outdir, "config.yaml"))
+    with open(configs.input_json_path, "r") as f:
+        orig_inputs = json.load(f)
+    for x in orig_inputs:
+        convert_to_bioassembly_dict(x, outdir)
+    configs.input_json_path = os.path.join(outdir, "input_tasks.json")
+    with open(configs.input_json_path, "w") as f:
+        json.dump(orig_inputs, f, indent=4)
+
+    # Create runner
+    runner = InferenceRunner(configs)
+
+    # Initialize Q_theta guidance
+    guidance = QThetaPXDesignGuidance(
+        checkpoint=args.qtheta_checkpoint if os.path.isabs(args.qtheta_checkpoint) else os.path.join(_ALLO_ROOT, args.qtheta_checkpoint),
+        ref_holo=args.ref_holo if os.path.isabs(args.ref_holo) else os.path.join(_ALLO_ROOT, args.ref_holo),
+        ref_apo=args.ref_apo if os.path.isabs(args.ref_apo) else os.path.join(_ALLO_ROOT, args.ref_apo),
+        ref_chain=args.ref_chain,
+        device='cuda:0',  # After CUDA_VISIBLE_DEVICES remapping
+        esm_target=args.esm_target,
+    )
+
+    # Monkey-patch the sample_diffusion function
+    from pxdesign.model import generator as pxdesign_generator
+    import pxdesign.model.pxdesign as pxdesign_model
+
+    # Create guided version with guidance params bound
+    guided_fn = partial(
+        guided_sample_diffusion,
+        guidance_module=guidance,
+        guidance_scale=args.guidance_scale,
+        guidance_start=args.guidance_start,
+        guidance_end=args.guidance_end,
+    )
+
+    # Patch the module-level function in generator.py
+    pxdesign_generator.sample_diffusion = guided_fn
+
+    # CRITICAL: pxdesign.py does `from pxdesign.model.generator import sample_diffusion`
+    # which creates a local binding in pxdesign.model.pxdesign namespace.
+    # We must patch that local binding too, otherwise the ProtenixDesign.sample_diffusion()
+    # method will still call the original unpatched function.
+    pxdesign_model.sample_diffusion = guided_fn
+
+    logger.info("PXDesign diffusion loop patched with Q_theta guidance")
+
+    # Run inference
+    seeds = [derive_seed(time.time_ns())] if not configs.seeds else configs.seeds
+    for seed in seeds:
+        logger.info(f"Running guided inference with seed {seed}")
+        seed_everything(seed=seed, deterministic=False)
+        runner._inference(seed)
+
+    # Score all generated designs
+    logger.info("Scoring generated designs...")
+    from glob import glob
+
+    pdb_dir = outdir
+    pdbs = []
+    for ext in ('*.pdb', '*.cif'):
+        pdbs.extend(glob(os.path.join(pdb_dir, '**/' + ext), recursive=True))
+    pdbs = sorted([p for p in pdbs if 'sample' in os.path.basename(p).lower()])
+
+    results = []
+    for i, pdb_path in enumerate(pdbs):
+        design_id = os.path.basename(pdb_path).replace('.pdb', '').replace('.cif', '')
+        result = guidance.score_design(pdb_path)
+        if result is not None:
+            result['design_id'] = design_id
+            result['pdb_path'] = pdb_path
+            results.append(result)
+            logger.info(
+                f"[{i+1}/{len(pdbs)}] {design_id}: "
+                f"Q+={result['q_holo']:.3f} Q-={result['q_apo']:.3f} "
+                f"S={result['margin']:+.3f}")
+
+    # Save results
+    if results:
+        results.sort(key=lambda x: x['margin'], reverse=True)
+        margins = np.array([r['margin'] for r in results])
+
+        summary = {
+            'method': 'PXDesign + Classifier Guidance',
+            'n_designs': len(results),
+            'guidance_scale': args.guidance_scale,
+            'guidance_window': [args.guidance_end, args.guidance_start],
+            'margin_mean': float(margins.mean()),
+            'margin_std': float(margins.std()),
+            'frac_positive': float((margins > 0).mean()),
+            'q_holo_mean': float(np.mean([r['q_holo'] for r in results])),
+            'q_apo_mean': float(np.mean([r['q_apo'] for r in results])),
+        }
+
+        with open(os.path.join(outdir, 'guided_scores.json'), 'w') as f:
+            json.dump(results, f, indent=2)
+        with open(os.path.join(outdir, 'guided_summary.json'), 'w') as f:
+            json.dump(summary, f, indent=2)
+
+        logger.info(f"\n{'='*60}")
+        logger.info(f"PXDesign + Classifier Guidance Results ({len(results)} designs)")
+        logger.info(f"  Margin: {margins.mean():.3f} ± {margins.std():.3f}")
+        logger.info(f"  Fraction S > 0: {(margins > 0).mean():.1%}")
+        logger.info(f"  Q(holo) mean: {summary['q_holo_mean']:.3f}")
+        logger.info(f"{'='*60}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description='PXDesign + Q_theta Classifier Guidance')
+    parser.add_argument('--input', default='experiments/pxdesign_cam/output/cam_binder.json',
+                        help='PXDesign input JSON')
+    parser.add_argument('--qtheta_checkpoint',
+                        default='results/checkpoints_cam_v3/best_phase2.pt')
+    parser.add_argument('--ref_holo', default='data/pdbs/cam_holo/3CLN.pdb')
+    parser.add_argument('--ref_apo', default='data/pdbs/cam_apo/1CFD.pdb')
+    parser.add_argument('--ref_chain', default='A')
+    parser.add_argument('--guidance_scale', type=float, default=1.0,
+                        help='Guidance gradient scale')
+    parser.add_argument('--guidance_start', type=float, default=0.8,
+                        help='Start guidance at this noise fraction (high noise)')
+    parser.add_argument('--guidance_end', type=float, default=0.1,
+                        help='Stop guidance at this noise fraction (low noise)')
+    parser.add_argument('--N_sample', type=int, default=50)
+    parser.add_argument('--N_step', type=int, default=400)
+    parser.add_argument('--gpu', type=int, default=0)
+    parser.add_argument('--outdir', default='results/pxdesign_guided')
+    parser.add_argument('--esm_target', default='cam',
+                        help='Subdir under data/esm2_embeddings (e.g., adk, cam)')
+    args = parser.parse_args()
+
+    run_guided_pxdesign(args)
+
+
+if __name__ == '__main__':
+    main()

code/scripts/pxdesign_guidance/iterative_refinement.py

@@ -0,0 +1,338 @@
+"""
+Iterative Refinement via Langevin Noise-Refine Cycles.
+
+Inspired by ProDifEvo (Uehara et al., ICML 2025): repeatedly perturb and
+refine structures through Q_theta gradient ascent. Each cycle adds noise
+for diversity, then refines with Langevin dynamics toward higher selectivity.
+
+This allows designs to escape local optima and explore better selectivity
+regions that single-shot generation cannot reach.
+
+Pipeline:
+  1. Start from existing PXDesign outputs (seed structures)
+  2. Align binder to reference receptor frames
+  3. Run Langevin refinement with Q_theta gradient
+  4. Score the refined output
+  5. Repeat for K iterations, keeping best designs
+
+Usage:
+    python code/scripts/pxdesign_guidance/iterative_refinement.py \
+        --input_dir results/pxdesign_guided/converted_pdbs \
+        --qtheta_checkpoint results/checkpoints_cam_v3/best_phase2.pt \
+        --ref_holo data/pdbs/cam_holo/3CLN.pdb \
+        --ref_apo data/pdbs/cam_apo/1CFD.pdb \
+        --n_iterations 3 --n_designs 10 \
+        --gpu 6
+"""
+import os
+import sys
+import json
+import logging
+import numpy as np
+import torch
+from glob import glob
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
+logger = logging.getLogger(__name__)
+
+_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+_ALLO_CODE_DIR = os.path.abspath(os.path.join(_SCRIPT_DIR, '..', '..'))
+_ALLO_ROOT = os.path.abspath(os.path.join(_ALLO_CODE_DIR, '..'))
+
+if _ALLO_CODE_DIR not in sys.path:
+    sys.path.insert(0, _ALLO_CODE_DIR)
+
+
+def score_designs(pdb_paths, guidance):
+    """Score a list of PDB paths with Q_theta."""
+    results = []
+    for pdb_path in pdb_paths:
+        result = guidance.score_design(pdb_path)
+        if result is not None:
+            result['pdb_path'] = pdb_path
+            result['design_id'] = os.path.basename(pdb_path).replace('.pdb', '').replace('.cif', '')
+            results.append(result)
+    return results
+
+
+def run_langevin_cycle(pdb_paths, guidance, n_steps=50, step_size=0.005,
+                       iteration=0, outdir='results/iterative_refinement'):
+    """Run Langevin refinement cycle on binder backbone coords using Q_theta.
+
+    Uses guidance.dq (DifferentiableQTheta) for differentiable scoring.
+    Aligns binder to holo/apo reference frames for dual-state scoring.
+    """
+    from utils.pdb_utils import (load_structure, get_residues, get_backbone_coords,
+                                  get_aa_indices, align_structures)
+
+    refined_results = []
+    os.makedirs(outdir, exist_ok=True)
+
+    for pdb_path in pdb_paths:
+        try:
+            model = load_structure(pdb_path)
+            chains = {c.id: c for c in model.get_chains()}
+
+            binder_chain = None
+            for cid in sorted(chains.keys()):
+                if cid != 'A':
+                    binder_chain = cid
+                    break
+            if binder_chain is None:
+                continue
+
+            rec_res = get_residues(chains['A'])
+            if not rec_res:
+                rec_res = get_residues(chains['A'], only_standard=False)
+            binder_res = get_residues(chains[binder_chain])
+            if not binder_res:
+                binder_res = get_residues(chains[binder_chain], only_standard=False)
+            if len(binder_res) < 5:
+                continue
+
+            binder_coords, binder_mask = get_backbone_coords(binder_res)
+            rec_coords, _ = get_backbone_coords(rec_res)
+
+            try:
+                aa_idx = get_aa_indices(binder_res)
+            except Exception:
+                aa_idx = np.zeros(len(binder_res), dtype=np.int64)
+
+            # Compute alignment transforms
+            rec_ca = rec_coords[:, 1, :]
+            ref_holo_ca = guidance.ref_holo_ca.cpu().numpy()
+            ref_apo_ca = guidance.ref_apo_ca.cpu().numpy()
+            n_h = min(len(rec_ca), len(ref_holo_ca))
+            n_a = min(len(rec_ca), len(ref_apo_ca))
+            if n_h < 5 or n_a < 5:
+                continue
+
+            _, R_h = align_structures(rec_ca[:n_h], ref_holo_ca[:n_h])
+            center_h = rec_ca[:n_h].mean(0)
+            ref_center_h = ref_holo_ca[:n_h].mean(0)
+            aligned_holo = (binder_coords.reshape(-1, 3) - center_h) @ R_h.T + ref_center_h
+            aligned_holo = aligned_holo.reshape(-1, 4, 3)
+
+            _, R_a = align_structures(rec_ca[:n_a], ref_apo_ca[:n_a])
+            center_a = rec_ca[:n_a].mean(0)
+            ref_center_a = ref_apo_ca[:n_a].mean(0)
+
+            device = guidance.device
+            dq = guidance.dq
+
+            # Precompute alignment tensors (detached constants)
+            R_h_t = torch.from_numpy(R_h).float().to(device)
+            R_a_t = torch.from_numpy(R_a).float().to(device)
+            center_h_t = torch.from_numpy(center_h).float().to(device)
+            ref_center_h_t = torch.from_numpy(ref_center_h).float().to(device)
+            center_a_t = torch.from_numpy(center_a).float().to(device)
+            ref_center_a_t = torch.from_numpy(ref_center_a).float().to(device)
+
+            # Work in holo-aligned frame
+            coords_t = torch.from_numpy(aligned_holo.copy()).float().to(device)
+            mask_t = torch.from_numpy(binder_mask).bool().to(device)
+            aa_t = torch.from_numpy(aa_idx).long().to(device)
+
+            # Add noise for diversity (constant, small)
+            noise = torch.randn_like(coords_t) * 0.05
+            coords_t = coords_t + noise
+
+            best_margin = -float('inf')
+            best_coords = coords_t.clone()
+
+            def project_bond_lengths(coords, target_dist=3.8, n_iters=5):
+                """Project CA-CA distances to target_dist via SHAKE-like iteration."""
+                with torch.no_grad():
+                    for _ in range(n_iters):
+                        ca = coords[:, 1, :].clone()
+                        for i in range(len(ca) - 1):
+                            delta = ca[i+1] - ca[i]
+                            d = delta.norm()
+                            if d < 1e-6:
+                                continue
+                            correction = 0.5 * (d - target_dist) / d * delta
+                            coords[i, :, :] += correction.unsqueeze(0)
+                            coords[i+1, :, :] -= correction.unsqueeze(0)
+                return coords
+
+            for step in range(n_steps):
+                coords_t = coords_t.detach().requires_grad_(True)
+
+                with torch.enable_grad():
+                    q_holo = dq.score(coords_t, mask_t, binder_aa_idx=aa_t,
+                                       receptor_label='holo')
+
+                    # Transform holo-frame → original → apo-frame
+                    flat_t = coords_t.reshape(-1, 3)
+                    original = (flat_t - ref_center_h_t) @ R_h_t + center_h_t
+                    apo_aligned = (original - center_a_t) @ R_a_t.T + ref_center_a_t
+                    coords_apo = apo_aligned.reshape(-1, 4, 3)
+
+                    q_apo = dq.score(coords_apo, mask_t, binder_aa_idx=aa_t,
+                                      receptor_label='apo')
+                    margin = q_holo - q_apo
+                    margin.backward()
+
+                grad = coords_t.grad
+                if grad is None or torch.isnan(grad).any():
+                    continue
+
+                grad_norm = grad.norm().clamp(min=1e-8)
+
+                if margin.item() > best_margin:
+                    best_margin = margin.item()
+                    best_coords = coords_t.detach().clone()
+
+                if step % 10 == 0:
+                    logger.info(f"  [{os.path.basename(pdb_path)}] Step {step}: "
+                                f"Q+={q_holo.item():.3f} Q-={q_apo.item():.3f} "
+                                f"S={margin.item():.3f} |g|={grad_norm.item():.4f}")
+
+                with torch.no_grad():
+                    coords_t = coords_t + step_size * grad / grad_norm
+                    # Annealed Langevin noise (small)
+                    noise_scale = step_size * 0.05 * (1 - step / n_steps)
+                    coords_t = coords_t + noise_scale * torch.randn_like(coords_t)
+                    # Hard projection: enforce CA-CA = 3.8A
+                    coords_t = project_bond_lengths(coords_t)
+
+            # Write refined backbone PDB
+            final_coords = best_coords.detach().cpu().numpy()
+            basename = os.path.basename(pdb_path).replace('.pdb', '')
+            out_path = os.path.join(outdir, f'{basename}_iter{iteration}.pdb')
+
+            atom_names = [' N  ', ' CA ', ' C  ', ' O  ']
+            elements = ['N', 'C', 'C', 'O']
+            with open(out_path, 'w') as f:
+                atom_num = 1
+                for i in range(len(final_coords)):
+                    if not binder_mask[i]:
+                        continue
+                    for j, (aname, elem) in enumerate(zip(atom_names, elements)):
+                        x, y, z = final_coords[i, j]
+                        f.write(f"ATOM  {atom_num:5d} {aname} ALA B{i+1:4d}    "
+                                f"{x:8.3f}{y:8.3f}{z:8.3f}  1.00  0.00           {elem}\n")
+                        atom_num += 1
+                f.write("END\n")
+
+            # Score refined design
+            result = guidance.score_design(out_path)
+            if result is not None:
+                result['pdb_path'] = out_path
+                result['iteration'] = iteration
+                result['best_margin_during_opt'] = best_margin
+                refined_results.append(result)
+                logger.info(f"  -> Refined: S={result['margin']:.3f} "
+                            f"(best during opt: {best_margin:.3f})")
+
+        except Exception as e:
+            logger.warning(f"Failed to refine {pdb_path}: {e}")
+            import traceback
+            traceback.print_exc()
+
+    return refined_results
+
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--input_dir',
+                        default='results/pxdesign_guided/converted_pdbs')
+    parser.add_argument('--qtheta_checkpoint',
+                        default='results/checkpoints_cam_v3/best_phase2.pt')
+    parser.add_argument('--ref_holo', default='data/pdbs/cam_holo/3CLN.pdb')
+    parser.add_argument('--ref_apo', default='data/pdbs/cam_apo/1CFD.pdb')
+    parser.add_argument('--ref_chain', default='A')
+    parser.add_argument('--n_iterations', type=int, default=4,
+                        help='Number of refine cycles')
+    parser.add_argument('--n_designs', type=int, default=20,
+                        help='Number of designs to refine')
+    parser.add_argument('--n_steps', type=int, default=50,
+                        help='Langevin steps per iteration')
+    parser.add_argument('--step_size', type=float, default=0.005)
+    parser.add_argument('--gpu', type=int, default=6)
+    parser.add_argument('--outdir', default='results/iterative_refinement')
+    args = parser.parse_args()
+
+    os.chdir(_ALLO_ROOT)
+
+    from scripts.pxdesign_guidance.qtheta_pxdesign import QThetaPXDesignGuidance
+
+    outdir = args.outdir
+    os.makedirs(outdir, exist_ok=True)
+
+    # Initialize scorer
+    guidance = QThetaPXDesignGuidance(
+        checkpoint=args.qtheta_checkpoint,
+        ref_holo=args.ref_holo,
+        ref_apo=args.ref_apo,
+        ref_chain=args.ref_chain,
+        device=f'cuda:{args.gpu}',
+    )
+    guidance._lazy_init()
+
+    # Collect input designs
+    input_pdbs = sorted(glob(os.path.join(args.input_dir, '*.pdb')))[:args.n_designs]
+    logger.info(f"Selected {len(input_pdbs)} designs for iterative refinement")
+
+    # Score initial designs
+    logger.info("Scoring initial designs...")
+    initial_results = score_designs(input_pdbs, guidance)
+    initial_margins = [r['margin'] for r in initial_results]
+    logger.info(f"Initial: S={np.mean(initial_margins):.3f}\u00b1{np.std(initial_margins):.3f}")
+
+    all_iteration_results = {'initial': initial_results}
+
+    # Iterative refinement
+    current_pdbs = input_pdbs
+    for iteration in range(args.n_iterations):
+        logger.info(f"\n{'='*50}")
+        logger.info(f"Iteration {iteration + 1}/{args.n_iterations}")
+        logger.info(f"{'='*50}")
+
+        iter_results = run_langevin_cycle(
+            current_pdbs, guidance,
+            n_steps=args.n_steps,
+            step_size=args.step_size,
+            iteration=iteration,
+            outdir=outdir,
+        )
+
+        if iter_results:
+            margins = [r['margin'] for r in iter_results]
+            logger.info(f"Iteration {iteration}: S={np.mean(margins):.3f}\u00b1{np.std(margins):.3f}")
+            all_iteration_results[f'iteration_{iteration}'] = iter_results
+
+            # Use refined designs as input for next iteration
+            current_pdbs = [r['pdb_path'] for r in iter_results]
+
+    # Summary
+    logger.info(f"\n{'='*60}")
+    logger.info("Iterative Refinement Summary")
+    logger.info(f"{'='*60}")
+    for key, results in all_iteration_results.items():
+        if results:
+            margins = [r['margin'] for r in results]
+            logger.info(f"{key:15s}: S={np.mean(margins):.3f}\u00b1{np.std(margins):.3f}, "
+                        f"N={len(results)}, S>0={100*np.mean([m>0 for m in margins]):.0f}%")
+
+    # Save results
+    out_path = os.path.join(outdir, 'iterative_refinement_summary.json')
+    summary = {}
+    for key, results in all_iteration_results.items():
+        if results:
+            margins = [r['margin'] for r in results]
+            summary[key] = {
+                'n': len(results),
+                'margin_mean': float(np.mean(margins)),
+                'margin_std': float(np.std(margins)),
+                'margin_max': float(np.max(margins)),
+                'frac_positive': float(np.mean([m > 0 for m in margins])),
+            }
+    with open(out_path, 'w') as f:
+        json.dump(summary, f, indent=2)
+    logger.info(f"\nSaved to {out_path}")
+
+
+if __name__ == '__main__':
+    main()

code/scripts/pxdesign_guidance/langevin_pxdesign.py

@@ -0,0 +1,374 @@
+"""
+PXDesign + Langevin Refinement.
+
+Post-hoc gradient ascent on existing PXDesign binder backbones using Q_theta
+selectivity gradient:
+    x_{t+1} = x_t + η · ∇_x[Q(holo,Y) - Q(apo,Y)] + √(2η) · ε
+
+Takes PXDesign outputs (which have full sidechains), extracts backbone coords,
+refines them via Langevin dynamics, and outputs refined backbone-only PDBs.
+
+Usage:
+    python code/scripts/pxdesign_guidance/langevin_pxdesign.py \
+        --designs_dir experiments/pxdesign_cam/output/ \
+        --qtheta_checkpoint results/checkpoints_cam_v3/best_phase2.pt \
+        --ref_holo data/pdbs/cam_holo/3CLN.pdb \
+        --ref_apo data/pdbs/cam_apo/1CFD.pdb \
+        --n_steps 100 --step_size 0.01 \
+        --gpu 0
+"""
+
+import os
+import sys
+import argparse
+import json
+import logging
+import numpy as np
+import torch
+from glob import glob
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
+logger = logging.getLogger(__name__)
+
+_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+_ALLO_CODE_DIR = os.path.abspath(os.path.join(_SCRIPT_DIR, '..', '..'))
+_ALLO_ROOT = os.path.abspath(os.path.join(_ALLO_CODE_DIR, '..'))
+
+if _ALLO_CODE_DIR not in sys.path:
+    sys.path.insert(0, _ALLO_CODE_DIR)
+
+from utils.pdb_utils import (
+    load_structure, get_residues, get_backbone_coords,
+    get_aa_indices, align_structures
+)
+
+
+def write_backbone_pdb(coords, mask, out_path, chain='B'):
+    """Write backbone PDB (N, CA, C, O) from [N, 4, 3] numpy coords."""
+    atom_names = [' N  ', ' CA ', ' C  ', ' O  ']
+    elements = ['N', 'C', 'C', 'O']
+    with open(out_path, 'w') as f:
+        atom_idx = 1
+        for i in range(len(coords)):
+            if not mask[i]:
+                continue
+            for j, (aname, elem) in enumerate(zip(atom_names, elements)):
+                x, y, z = coords[i, j, :]
+                f.write(
+                    f"ATOM  {atom_idx:5d} {aname:4s} ALA {chain}{i+1:4d}    "
+                    f"{x:8.3f}{y:8.3f}{z:8.3f}  1.00  0.00           {elem}\n"
+                )
+                atom_idx += 1
+        f.write("END\n")
+
+
+def find_pxdesign_pdbs(designs_dir):
+    """Find all PXDesign output PDB files."""
+    pdbs = sorted(glob(os.path.join(designs_dir, '**/*.pdb'), recursive=True))
+    pdbs = [p for p in pdbs if 'sample' in os.path.basename(p).lower()
+            or 'design' in os.path.basename(p).lower()
+            or 'rank' in os.path.basename(p).lower()]
+    if not pdbs:
+        pdbs = sorted(glob(os.path.join(designs_dir, '**/*.pdb'), recursive=True))
+    return pdbs
+
+
+def langevin_refine(dq, binder_coords_init, binder_mask, binder_aa_idx,
+                     rec_coords, rec_mask, ref_holo_ca, ref_apo_ca,
+                     n_steps=100, step_size=0.01, noise_scale=0.0,
+                     device='cuda:0'):
+    """
+    Langevin refinement of binder backbone coordinates.
+
+    Args:
+        dq: DifferentiableQTheta scorer
+        binder_coords_init: [N_binder, 4, 3] numpy — initial binder backbone
+        binder_mask: [N_binder] numpy bool
+        binder_aa_idx: [N_binder] numpy int
+        rec_coords: [N_rec, 4, 3] numpy — receptor backbone
+        rec_mask: [N_rec] numpy bool
+        ref_holo_ca: [N_ref, 3] torch — holo reference CA
+        ref_apo_ca: [N_ref, 3] torch — apo reference CA
+        n_steps: int
+        step_size: float (η)
+        noise_scale: float (for stochastic Langevin, 0 = gradient ascent)
+        device: str
+
+    Returns:
+        best_coords: [N_binder, 4, 3] numpy — refined coords
+        trajectory: list of dicts with step info
+    """
+    device = torch.device(device)
+
+    # Convert to tensors
+    x = torch.from_numpy(binder_coords_init.copy()).float().to(device)
+    mask_t = torch.from_numpy(binder_mask).bool().to(device)
+    aa_t = torch.from_numpy(binder_aa_idx).long().to(device)
+    rec_ca = torch.from_numpy(rec_coords[:, 1, :]).float().to(device)
+
+    best_margin = -float('inf')
+    best_coords = binder_coords_init.copy()
+    best_q_holo = 0.0
+    best_q_apo = 0.0
+    trajectory = []
+
+    for step in range(n_steps):
+        x_grad = x.clone().requires_grad_(True)
+
+        try:
+            with torch.enable_grad():
+                # Align to holo reference
+                n_align_h = min(len(rec_ca), len(ref_holo_ca))
+                if n_align_h < 5:
+                    break
+                from qtheta_pxdesign import differentiable_kabsch
+                R_h, t_h = differentiable_kabsch(rec_ca[:n_align_h].detach(),
+                                                   ref_holo_ca[:n_align_h].detach())
+                R_h, t_h = R_h.detach(), t_h.detach()
+                aligned_holo = x_grad.reshape(-1, 3) @ R_h.T + t_h
+                aligned_holo = aligned_holo.reshape(-1, 4, 3)
+
+                q_holo = dq.score(aligned_holo, mask_t, binder_aa_idx=aa_t,
+                                   receptor_label='holo')
+
+                # Align to apo reference
+                n_align_a = min(len(rec_ca), len(ref_apo_ca))
+                R_a, t_a = differentiable_kabsch(rec_ca[:n_align_a].detach(),
+                                                   ref_apo_ca[:n_align_a].detach())
+                R_a, t_a = R_a.detach(), t_a.detach()
+                aligned_apo = x_grad.reshape(-1, 3) @ R_a.T + t_a
+                aligned_apo = aligned_apo.reshape(-1, 4, 3)
+
+                q_apo = dq.score(aligned_apo, mask_t, binder_aa_idx=aa_t,
+                                  receptor_label='apo')
+
+                margin = q_holo - q_apo
+                margin.backward()
+
+            grad = x_grad.grad
+            if grad is None or torch.isnan(grad).any():
+                continue
+
+            # Gradient ascent step
+            x = x + step_size * grad
+
+            # Optional noise for stochastic Langevin
+            if noise_scale > 0:
+                x = x + noise_scale * np.sqrt(2 * step_size) * torch.randn_like(x)
+
+            current_margin = margin.item()
+            step_info = {
+                'step': step,
+                'q_holo': q_holo.item(),
+                'q_apo': q_apo.item(),
+                'margin': current_margin,
+                'grad_norm': grad.norm().item(),
+            }
+            trajectory.append(step_info)
+
+            if current_margin > best_margin:
+                best_margin = current_margin
+                best_coords = x.detach().cpu().numpy()
+                best_q_holo = q_holo.item()
+                best_q_apo = q_apo.item()
+
+            if step % 20 == 0:
+                logger.info(
+                    f"  Step {step:3d}: Q+={q_holo.item():.3f} Q-={q_apo.item():.3f} "
+                    f"S={current_margin:+.3f} |∇|={grad.norm().item():.4f}")
+
+        except Exception as e:
+            logger.debug(f"  Step {step}: {e}")
+            continue
+
+    return best_coords, trajectory, best_margin, best_q_holo, best_q_apo
+
+
+def main():
+    parser = argparse.ArgumentParser(description='PXDesign + Langevin Refinement')
+    parser.add_argument('--designs_dir', default='experiments/pxdesign_cam/output/')
+    parser.add_argument('--qtheta_checkpoint',
+                        default='results/checkpoints_cam_v3/best_phase2.pt')
+    parser.add_argument('--ref_holo', default='data/pdbs/cam_holo/3CLN.pdb')
+    parser.add_argument('--ref_apo', default='data/pdbs/cam_apo/1CFD.pdb')
+    parser.add_argument('--ref_chain', default='A')
+    parser.add_argument('--n_steps', type=int, default=100)
+    parser.add_argument('--step_size', type=float, default=0.01)
+    parser.add_argument('--noise_scale', type=float, default=0.0,
+                        help='Noise scale for stochastic Langevin (0=gradient ascent)')
+    parser.add_argument('--gpu', type=int, default=0)
+    parser.add_argument('--outdir', default='results/pxdesign_langevin')
+    args = parser.parse_args()
+
+    os.chdir(_ALLO_ROOT)
+
+    device = f'cuda:{args.gpu}'
+
+    from models.differentiable_features import DifferentiableQTheta
+
+    # Load scorer
+    dq = DifferentiableQTheta(args.qtheta_checkpoint, device=device)
+    dq.load_receptor(args.ref_holo, chain=args.ref_chain, label='holo')
+    dq.load_receptor(args.ref_apo, chain=args.ref_chain, label='apo')
+
+    # Load reference CA coords
+    holo_model = load_structure(args.ref_holo)
+    holo_res = get_residues(holo_model[args.ref_chain])
+    holo_coords, _ = get_backbone_coords(holo_res)
+    ref_holo_ca = torch.from_numpy(holo_coords[:, 1, :]).float().to(device)
+
+    apo_model = load_structure(args.ref_apo)
+    apo_res = get_residues(apo_model[args.ref_chain])
+    apo_coords, _ = get_backbone_coords(apo_res)
+    ref_apo_ca = torch.from_numpy(apo_coords[:, 1, :]).float().to(device)
+
+    # Find designs
+    pdbs = find_pxdesign_pdbs(args.designs_dir)
+    logger.info(f"Found {len(pdbs)} PXDesign outputs to refine")
+
+    outdir = args.outdir
+    os.makedirs(outdir, exist_ok=True)
+
+    all_results = []
+    for i, pdb_path in enumerate(pdbs):
+        design_id = os.path.basename(pdb_path).replace('.pdb', '').replace('.cif', '')
+        logger.info(f"\n[{i+1}/{len(pdbs)}] Refining {design_id}...")
+
+        try:
+            model = load_structure(pdb_path)
+            chains = {c.get_id(): c for c in model.get_chains()}
+            chain_ids = sorted(chains.keys())
+
+            # Identify chains
+            ref_len = len(holo_res)
+            rec_chain_id, binder_chain_id = None, None
+            for cid in chain_ids:
+                cres = get_residues(chains[cid])
+                if abs(len(cres) - ref_len) < ref_len * 0.3:
+                    rec_chain_id = cid
+                else:
+                    binder_chain_id = cid
+
+            if rec_chain_id is None or binder_chain_id is None:
+                if len(chain_ids) >= 2:
+                    rec_chain_id, binder_chain_id = chain_ids[0], chain_ids[1]
+                else:
+                    logger.warning(f"Skipping {design_id}: cannot identify chains")
+                    continue
+
+            rec_res = get_residues(chains[rec_chain_id])
+            binder_res = get_residues(chains[binder_chain_id])
+
+            rec_coords_np, rec_mask = get_backbone_coords(rec_res)
+            binder_coords_np, binder_mask = get_backbone_coords(binder_res)
+            aa_idx = get_aa_indices(binder_res)
+
+            # Score before refinement
+            rec_ca = rec_coords_np[:, 1, :]
+            n_align = min(len(rec_ca), len(holo_coords[:, 1, :]))
+            _, R_h = align_structures(rec_ca[:n_align], holo_coords[:n_align, 1, :])
+            center_h = rec_ca[:n_align].mean(0)
+            ref_center_h = holo_coords[:n_align, 1, :].mean(0)
+
+            aligned_init = (binder_coords_np.reshape(-1, 3) - center_h) @ R_h.T + ref_center_h
+            aligned_init = aligned_init.reshape(-1, 4, 3)
+            with torch.no_grad():
+                q_h_init = dq.score(
+                    torch.from_numpy(aligned_init).float().to(device),
+                    torch.from_numpy(binder_mask).bool().to(device),
+                    binder_aa_idx=torch.from_numpy(aa_idx).long().to(device),
+                    receptor_label='holo').item()
+
+            n_align_a = min(len(rec_ca), len(apo_coords[:, 1, :]))
+            _, R_a = align_structures(rec_ca[:n_align_a], apo_coords[:n_align_a, 1, :])
+            center_a = rec_ca[:n_align_a].mean(0)
+            ref_center_a = apo_coords[:n_align_a, 1, :].mean(0)
+            aligned_init_a = (binder_coords_np.reshape(-1, 3) - center_a) @ R_a.T + ref_center_a
+            aligned_init_a = aligned_init_a.reshape(-1, 4, 3)
+            with torch.no_grad():
+                q_a_init = dq.score(
+                    torch.from_numpy(aligned_init_a).float().to(device),
+                    torch.from_numpy(binder_mask).bool().to(device),
+                    binder_aa_idx=torch.from_numpy(aa_idx).long().to(device),
+                    receptor_label='apo').item()
+
+            margin_init = q_h_init - q_a_init
+
+            # Run Langevin refinement
+            refined_coords, trajectory, best_margin, best_qh, best_qa = langevin_refine(
+                dq, binder_coords_np, binder_mask, aa_idx,
+                rec_coords_np, rec_mask, ref_holo_ca, ref_apo_ca,
+                n_steps=args.n_steps, step_size=args.step_size,
+                noise_scale=args.noise_scale, device=device,
+            )
+
+            # Use best-margin values (matching the saved best_coords PDB)
+            margin_final = best_margin if trajectory else margin_init
+
+            # Save refined PDB
+            out_pdb = os.path.join(outdir, f'{design_id}_refined.pdb')
+            write_backbone_pdb(refined_coords, binder_mask, out_pdb)
+
+            result = {
+                'design_id': design_id,
+                'pdb_path': pdb_path,
+                'refined_pdb': out_pdb,
+                'q_holo_init': q_h_init,
+                'q_apo_init': q_a_init,
+                'margin_init': margin_init,
+                'q_holo_final': best_qh if trajectory else q_h_init,
+                'q_apo_final': best_qa if trajectory else q_a_init,
+                'margin_final': margin_final,
+                'margin_delta': margin_final - margin_init,
+                'n_steps_converged': len(trajectory),
+                'n_res': len(binder_res),
+            }
+            all_results.append(result)
+
+            logger.info(
+                f"  {design_id}: S_init={margin_init:+.3f} -> S_final={margin_final:+.3f} "
+                f"(Δ={margin_final - margin_init:+.3f})")
+
+        except Exception as e:
+            logger.warning(f"Failed to refine {design_id}: {e}")
+            continue
+
+    # Summary
+    if all_results:
+        all_results.sort(key=lambda x: x['margin_final'], reverse=True)
+        margins_init = np.array([r['margin_init'] for r in all_results])
+        margins_final = np.array([r['margin_final'] for r in all_results])
+        deltas = margins_final - margins_init
+
+        summary = {
+            'method': 'PXDesign + Langevin',
+            'n_designs': len(all_results),
+            'n_steps': args.n_steps,
+            'step_size': args.step_size,
+            'margin_init_mean': float(margins_init.mean()),
+            'margin_final_mean': float(margins_final.mean()),
+            'margin_delta_mean': float(deltas.mean()),
+            'frac_improved': float((deltas > 0).mean()),
+            'frac_positive_init': float((margins_init > 0).mean()),
+            'frac_positive_final': float((margins_final > 0).mean()),
+            'q_holo_final_mean': float(np.mean([r['q_holo_final'] for r in all_results])),
+        }
+
+        with open(os.path.join(outdir, 'langevin_scores.json'), 'w') as f:
+            json.dump(all_results, f, indent=2)
+        with open(os.path.join(outdir, 'langevin_summary.json'), 'w') as f:
+            json.dump(summary, f, indent=2)
+
+        logger.info(f"\n{'='*60}")
+        logger.info(f"PXDesign + Langevin Results ({len(all_results)} designs)")
+        logger.info(f"  Margin init:  {margins_init.mean():.3f} ± {margins_init.std():.3f}")
+        logger.info(f"  Margin final: {margins_final.mean():.3f} ± {margins_final.std():.3f}")
+        logger.info(f"  Δ margin:     {deltas.mean():+.3f} ± {deltas.std():.3f}")
+        logger.info(f"  % improved:   {(deltas > 0).mean():.1%}")
+        logger.info(f"  S>0 init/final: {(margins_init > 0).mean():.1%} / "
+                     f"{(margins_final > 0).mean():.1%}")
+        logger.info(f"{'='*60}")
+
+
+if __name__ == '__main__':
+    main()

code/scripts/pxdesign_guidance/qtheta_pxdesign.py

@@ -0,0 +1,477 @@
+"""
+Core Q_theta guidance module for PXDesign integration.
+
+Provides differentiable Q_theta scoring for PXDesign's atom coordinate format.
+Key responsibilities:
+  - Extract binder backbone (N, CA, C, O) from PXDesign's flat atom array
+  - Align binder to reference receptor frames via differentiable Kabsch
+  - Compute selectivity gradient ∇[Q(holo,Y) - Q(apo,Y)] w.r.t. atom coords
+  - Works in pxdesign env (PyTorch 2.3.1) using pure-PyTorch scorer (no e3nn)
+
+Usage:
+    guidance = QThetaPXDesignGuidance(
+        checkpoint='results/checkpoints_cam_v3/best_phase2.pt',
+        ref_holo='data/pdbs/cam_holo/3CLN.pdb',
+        ref_apo='data/pdbs/cam_apo/1CFD.pdb',
+        ref_chain='A',
+        device='cuda:0',
+    )
+    # Inside PXDesign diffusion loop:
+    grad = guidance.compute_guidance_gradient(x_denoised, input_feature_dict, t_hat)
+    x_denoised = x_denoised + scale * grad
+"""
+
+import os
+import sys
+import logging
+import numpy as np
+import torch
+
+logger = logging.getLogger(__name__)
+
+# Add Allo-Designer code directory to path
+_ALLO_CODE_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))
+if _ALLO_CODE_DIR not in sys.path:
+    sys.path.insert(0, _ALLO_CODE_DIR)
+
+
+def differentiable_kabsch(mobile, target):
+    """
+    Differentiable Kabsch alignment using SVD.
+
+    Args:
+        mobile: [N, 3] tensor (points to align FROM)
+        target: [N, 3] tensor (points to align TO)
+
+    Returns:
+        R: [3, 3] rotation matrix
+        t: [3] translation vector
+        Such that aligned = (mobile - mobile_center) @ R.T + target_center
+    """
+    mobile_center = mobile.mean(dim=0)
+    target_center = target.mean(dim=0)
+
+    mobile_centered = mobile - mobile_center
+    target_centered = target - target_center
+
+    H = mobile_centered.T @ target_centered  # [3, 3]
+    U, S, Vh = torch.linalg.svd(H)
+    V = Vh.T
+
+    # Ensure proper rotation (det > 0)
+    d = torch.det(V @ U.T)
+    sign_matrix = torch.diag(torch.tensor([1.0, 1.0, torch.sign(d)],
+                                           device=mobile.device, dtype=mobile.dtype))
+    R = V @ sign_matrix @ U.T  # [3, 3]
+    t = target_center - mobile_center @ R.T  # [3]
+
+    return R, t
+
+
+class QThetaPXDesignGuidance:
+    """
+    Q_theta guidance for PXDesign diffusion process.
+
+    Lazily initializes the scorer and reference structures on first use.
+    Handles extraction of binder backbone from PXDesign's flat atom array
+    and alignment to reference receptor frames.
+    """
+
+    def __init__(self, checkpoint, ref_holo, ref_apo, ref_chain='A',
+                 device='cuda:0', cutoff=8.0, esm_target='cam'):
+        self.checkpoint = checkpoint
+        self.ref_holo = ref_holo
+        self.ref_apo = ref_apo
+        self.ref_chain = ref_chain
+        self.device = torch.device(device)
+        self.cutoff = cutoff
+        self.esm_target = esm_target
+
+        self._initialized = False
+        self.dq = None
+        self.ref_holo_ca = None
+        self.ref_apo_ca = None
+
+    def _lazy_init(self):
+        """Initialize Q_theta scorer and load reference structures."""
+        if self._initialized:
+            return
+
+        from models.differentiable_features import DifferentiableQTheta
+        from utils.pdb_utils import load_structure, get_residues, get_backbone_coords
+
+        logger.info(f"Loading Q_theta checkpoint: {self.checkpoint}")
+        self.dq = DifferentiableQTheta(self.checkpoint, device=str(self.device))
+        self.dq.load_receptor(self.ref_holo, chain=self.ref_chain, label='holo',
+                             esm_target=self.esm_target)
+        self.dq.load_receptor(self.ref_apo, chain=self.ref_chain, label='apo',
+                             esm_target=self.esm_target)
+
+        # Cache reference CA coords for alignment
+        holo_model = load_structure(self.ref_holo)
+        holo_res = get_residues(holo_model[self.ref_chain])
+        holo_coords, _ = get_backbone_coords(holo_res)
+        self.ref_holo_ca = torch.from_numpy(holo_coords[:, 1, :]).float().to(self.device)
+
+        apo_model = load_structure(self.ref_apo)
+        apo_res = get_residues(apo_model[self.ref_chain])
+        apo_coords, _ = get_backbone_coords(apo_res)
+        self.ref_apo_ca = torch.from_numpy(apo_coords[:, 1, :]).float().to(self.device)
+
+        self._initialized = True
+        logger.info(f"Q_theta guidance initialized: holo={len(holo_res)} res, apo={len(apo_res)} res")
+
+    def extract_binder_backbone(self, x_coords, input_feature_dict):
+        """
+        Extract binder backbone atoms (N, CA, C, O) from PXDesign's flat atom array.
+
+        PXDesign stores all atoms in a flat [N_atom, 3] array. Entity annotations
+        identify which atoms belong to the designed binder (entity_id=2 typically,
+        or the last entity). We extract backbone atoms for each binder residue.
+
+        Args:
+            x_coords: [N_sample, N_atom, 3] — current coordinates from diffusion
+            input_feature_dict: dict with atom_to_token_idx, entity_id, etc.
+
+        Returns:
+            binder_bb: [N_sample, N_binder_res, 4, 3] — backbone coords (N, CA, C, O)
+            binder_mask: [N_binder_res] — validity mask
+            rec_bb: [N_rec_res, 4, 3] — receptor backbone coords (from condition)
+            rec_mask: [N_rec_res] — receptor validity mask
+            binder_atom_indices: [N_binder_bb_atoms] — indices into flat atom array
+        """
+        atom_to_token = input_feature_dict['atom_to_token_idx']  # [N_atom]
+        if atom_to_token.dim() > 1:
+            atom_to_token = atom_to_token.squeeze(0)
+
+        # Identify binder vs receptor tokens
+        # In PXDesign: design_token_mask=True for binder tokens
+        design_token_mask = input_feature_dict.get('design_token_mask', None)
+        if design_token_mask is not None:
+            if design_token_mask.dim() > 1:
+                design_token_mask = design_token_mask.squeeze(0)
+            binder_tokens = torch.where(design_token_mask)[0]
+            rec_tokens = torch.where(~design_token_mask)[0]
+        else:
+            # Fallback: use entity_id (binder is typically entity_id=2, the last entity)
+            entity_id = input_feature_dict['entity_id']
+            if entity_id.dim() > 1:
+                entity_id = entity_id.squeeze(0)
+            max_entity = entity_id.max()
+            binder_tokens = torch.where(entity_id == max_entity)[0]
+            rec_tokens = torch.where(entity_id != max_entity)[0]
+
+        # Map tokens to atoms
+        # For standard amino acids, atom order within each token is:
+        # N(0), CA(1), C(2), O(3), CB(4), ...
+        # We need atoms 0-3 (N, CA, C, O) per token
+
+        # Get atom indices for each binder token
+        n_binder_res = len(binder_tokens)
+        if n_binder_res == 0:
+            return None
+
+        # Find atoms belonging to each binder residue
+        binder_bb_list = []
+        binder_atom_idx_list = []
+        for tok_idx in binder_tokens:
+            atom_indices = torch.where(atom_to_token == tok_idx.item())[0]
+            if len(atom_indices) >= 4:
+                # First 4 atoms are N, CA, C, O for standard amino acids
+                bb_atoms = atom_indices[:4]
+                binder_bb_list.append(bb_atoms)
+                binder_atom_idx_list.append(bb_atoms)
+
+        if not binder_bb_list:
+            return None
+
+        n_binder_res = len(binder_bb_list)
+        binder_bb_indices = torch.stack(binder_bb_list)  # [N_binder, 4]
+        all_binder_atom_indices = torch.cat(binder_atom_idx_list)  # [N_binder * 4]
+
+        # Extract binder backbone coords for all samples
+        # x_coords: [N_sample, N_atom, 3]
+        binder_bb = x_coords[:, binder_bb_indices, :]  # [N_sample, N_binder, 4, 3]
+        binder_mask = torch.ones(n_binder_res, dtype=torch.bool, device=x_coords.device)
+
+        # Extract receptor backbone from x_coords or condition_coordinate.
+        # PXDesign stores condition_coordinate in label_dict (not input_feature_dict),
+        # so we extract receptor backbone from x_coords directly. In the diffusion
+        # process, receptor atoms are conditioned at their reference positions.
+        # Try condition_coordinate first (if available), then fall back to x_coords.
+        cond_coords = input_feature_dict.get('condition_coordinate', None)
+        if cond_coords is None:
+            # Also try label_dict nesting
+            label_dict = input_feature_dict.get('label_dict', None)
+            if label_dict is not None:
+                cond_coords = label_dict.get('condition_coordinate', None)
+
+        rec_bb = None
+        rec_mask = None
+
+        # Get receptor backbone atoms
+        rec_bb_list = []
+        for tok_idx in rec_tokens:
+            atom_indices = torch.where(atom_to_token == tok_idx.item())[0]
+            if len(atom_indices) >= 4:
+                rec_bb_list.append(atom_indices[:4])
+
+        if rec_bb_list:
+            rec_bb_indices = torch.stack(rec_bb_list)  # [N_rec, 4]
+
+            if cond_coords is not None:
+                if cond_coords.dim() > 2:
+                    cond_coords = cond_coords.squeeze(0)
+                rec_bb = cond_coords[rec_bb_indices, :]  # [N_rec, 4, 3]
+            else:
+                # Fallback: extract receptor coords from x_coords (sample 0)
+                # Receptor atoms are conditioned and constant across samples
+                rec_bb = x_coords[0, rec_bb_indices, :].detach()  # [N_rec, 4, 3]
+
+            rec_mask = torch.ones(len(rec_bb_list), dtype=torch.bool,
+                                  device=x_coords.device)
+
+        return {
+            'binder_bb': binder_bb,           # [N_sample, N_binder, 4, 3]
+            'binder_mask': binder_mask,        # [N_binder]
+            'rec_bb': rec_bb,                  # [N_rec, 4, 3] or None
+            'rec_mask': rec_mask,              # [N_rec] or None
+            'binder_atom_indices': binder_bb_indices,  # [N_binder, 4]
+            'all_binder_atom_indices': all_binder_atom_indices,  # [N_binder * 4]
+        }
+
+    def align_and_score(self, binder_bb, rec_bb, rec_mask, receptor_label):
+        """
+        Align binder to a reference receptor frame and score with Q_theta.
+
+        Uses the receptor chain from the design to compute Kabsch alignment
+        to the reference receptor, then transforms the binder accordingly.
+
+        Args:
+            binder_bb: [N_binder, 4, 3] — binder backbone coords (requires_grad)
+            rec_bb: [N_rec, 4, 3] — receptor backbone coords
+            rec_mask: [N_rec] bool
+            receptor_label: 'holo' or 'apo'
+
+        Returns:
+            score: scalar tensor, differentiable w.r.t. binder_bb
+        """
+        if receptor_label == 'holo':
+            ref_ca = self.ref_holo_ca
+        else:
+            ref_ca = self.ref_apo_ca
+
+        # Get CA atoms from receptor
+        rec_ca = rec_bb[:, 1, :]  # [N_rec, 3]
+
+        # Use overlapping residues for alignment (take min length)
+        n_align = min(len(rec_ca), len(ref_ca))
+        if n_align < 5:
+            return torch.zeros(1, device=binder_bb.device, dtype=binder_bb.dtype,
+                               requires_grad=True).squeeze()
+
+        mobile_ca = rec_ca[:n_align].detach()
+        target_ca = ref_ca[:n_align].detach()
+
+        # Compute Kabsch alignment (detached — no gradient through rotation)
+        R, t = differentiable_kabsch(mobile_ca, target_ca)
+        R = R.detach()
+        t = t.detach()
+
+        # Apply transform to binder (gradient flows through binder_bb)
+        binder_flat = binder_bb.reshape(-1, 3)  # [N_binder*4, 3]
+        aligned = binder_flat @ R.T + t  # [N_binder*4, 3]
+        aligned_bb = aligned.reshape(-1, 4, 3)  # [N_binder, 4, 3]
+
+        # Score with Q_theta
+        binder_mask = torch.ones(aligned_bb.shape[0], dtype=torch.bool,
+                                  device=binder_bb.device)
+        score = self.dq.score(aligned_bb, binder_mask, receptor_label=receptor_label,
+                              cutoff=self.cutoff)
+        return score
+
+    def compute_guidance_gradient(self, x_denoised, input_feature_dict, t_hat=None,
+                                   sample_idx=0):
+        """
+        Compute Q_theta selectivity gradient for guidance.
+
+        Args:
+            x_denoised: [N_sample, N_atom, 3] — denoised coordinates from diffusion net
+            input_feature_dict: PXDesign input features dict
+            t_hat: current noise level (for logging/scaling)
+            sample_idx: which sample to compute gradient for (or -1 for all)
+
+        Returns:
+            gradient: [N_sample, N_atom, 3] — gradient to add to x_denoised
+                      (non-zero only at binder backbone atom positions)
+            margin: float — current selectivity margin
+        """
+        self._lazy_init()
+
+        extraction = self.extract_binder_backbone(x_denoised.detach(), input_feature_dict)
+        if extraction is None:
+            return torch.zeros_like(x_denoised), 0.0
+
+        binder_bb = extraction['binder_bb']      # [N_sample, N_binder, 4, 3]
+        binder_mask = extraction['binder_mask']   # [N_binder]
+        rec_bb = extraction['rec_bb']             # [N_rec, 4, 3]
+        rec_mask = extraction['rec_mask']          # [N_rec]
+        binder_atom_indices = extraction['binder_atom_indices']  # [N_binder, 4]
+
+        if rec_bb is None:
+            return torch.zeros_like(x_denoised), 0.0
+
+        N_sample = x_denoised.shape[0]
+        gradient = torch.zeros_like(x_denoised)
+        margins = []
+
+        # Ensure receptor is float32 for Q_theta scoring
+        if rec_bb is not None:
+            rec_bb = rec_bb.float()
+
+        # Process each sample
+        indices = range(N_sample) if sample_idx == -1 else [sample_idx]
+        for si in indices:
+            # Make binder coords differentiable, cast to float32 for Q_theta
+            binder_si = binder_bb[si].clone().float().requires_grad_(True)  # [N_binder, 4, 3]
+
+            try:
+                with torch.enable_grad():
+                    q_holo = self.align_and_score(binder_si, rec_bb, rec_mask, 'holo')
+                    q_apo = self.align_and_score(binder_si, rec_bb, rec_mask, 'apo')
+                    margin = q_holo - q_apo
+                    margin.backward()
+
+                if binder_si.grad is not None and not torch.isnan(binder_si.grad).any():
+                    # Map gradient back to full atom array
+                    grad_bb = binder_si.grad  # [N_binder, 4, 3]
+                    for ri in range(len(binder_atom_indices)):
+                        for ai in range(4):
+                            atom_idx = binder_atom_indices[ri, ai]
+                            gradient[si, atom_idx] = grad_bb[ri, ai]
+                    margins.append(margin.item())
+                else:
+                    margins.append(0.0)
+            except Exception as e:
+                logger.debug(f"Gradient computation failed for sample {si}: {e}")
+                margins.append(0.0)
+
+        avg_margin = np.mean(margins) if margins else 0.0
+        return gradient, avg_margin
+
+    def score_design(self, pdb_path, rec_chain='A', binder_chain='B'):
+        """
+        Score a single PXDesign output PDB/CIF (post-hoc, no gradient).
+
+        Handles PXDesign CIF files which use chain IDs like 'A0'/'B0' and
+        non-standard residue name 'xpb' for designed binder residues.
+
+        Returns:
+            dict with q_holo, q_apo, margin, or None on failure
+        """
+        self._lazy_init()
+
+        from utils.pdb_utils import (
+            load_structure, get_residues, get_backbone_coords,
+            get_aa_indices, align_structures
+        )
+
+        try:
+            model = load_structure(pdb_path)
+            chains = {c.get_id(): c for c in model.get_chains()}
+
+            if len(chains) < 2:
+                return None
+
+            chain_ids = sorted(chains.keys())
+
+            # Identify receptor and binder
+            # PXDesign CIF uses chain IDs like 'A0', 'B0' instead of 'A', 'B'
+            rc, bc = None, None
+            if rec_chain in chains and binder_chain in chains:
+                rc, bc = rec_chain, binder_chain
+            else:
+                # Match by residue count: receptor matches reference length,
+                # binder is the other chain
+                ref_model = load_structure(self.ref_holo)
+                ref_res = get_residues(ref_model[self.ref_chain])
+                ref_len = len(ref_res)
+                for cid in chain_ids:
+                    # Try standard residues first, then all residues
+                    cres = get_residues(chains[cid])
+                    if not cres:
+                        cres = get_residues(chains[cid], only_standard=False)
+                    n_res = len(cres)
+                    if n_res > 0 and abs(n_res - ref_len) < ref_len * 0.3:
+                        rc = cid
+                    elif n_res > 0:
+                        bc = cid
+                if rc is None or bc is None:
+                    rc, bc = chain_ids[0], chain_ids[1]
+
+            rec_res = get_residues(chains[rc])
+            if not rec_res:
+                rec_res = get_residues(chains[rc], only_standard=False)
+
+            # For binder: PXDesign uses 'xpb' residue names (non-standard)
+            binder_res = get_residues(chains[bc])
+            if not binder_res:
+                binder_res = get_residues(chains[bc], only_standard=False)
+
+            if not rec_res or not binder_res:
+                return None
+
+            rec_coords, rec_mask = get_backbone_coords(rec_res)
+            binder_coords, binder_mask = get_backbone_coords(binder_res)
+
+            # Handle amino acid indices: use get_aa_indices for standard AAs,
+            # default to GLY (7) for non-standard (PXDesign 'xpb')
+            try:
+                aa_idx = get_aa_indices(binder_res)
+            except Exception:
+                aa_idx = np.zeros(len(binder_res), dtype=np.int64)  # default to ALA
+
+            device = self.device
+
+            # Align to holo
+            rec_ca = rec_coords[:, 1, :]
+            ref_holo_ca_np = self.ref_holo_ca.cpu().numpy()
+            n_align = min(len(rec_ca), len(ref_holo_ca_np))
+            if n_align < 5:
+                return None
+            _, R_h = align_structures(rec_ca[:n_align], ref_holo_ca_np[:n_align])
+            center_h = rec_ca[:n_align].mean(0)
+            ref_center_h = ref_holo_ca_np[:n_align].mean(0)
+            aligned_holo = (binder_coords.reshape(-1, 3) - center_h) @ R_h.T + ref_center_h
+            aligned_holo = aligned_holo.reshape(-1, 4, 3)
+
+            # Align to apo
+            ref_apo_ca_np = self.ref_apo_ca.cpu().numpy()
+            n_align_a = min(len(rec_ca), len(ref_apo_ca_np))
+            _, R_a = align_structures(rec_ca[:n_align_a], ref_apo_ca_np[:n_align_a])
+            center_a = rec_ca[:n_align_a].mean(0)
+            ref_center_a = ref_apo_ca_np[:n_align_a].mean(0)
+            aligned_apo = (binder_coords.reshape(-1, 3) - center_a) @ R_a.T + ref_center_a
+            aligned_apo = aligned_apo.reshape(-1, 4, 3)
+
+            with torch.no_grad():
+                coords_h = torch.from_numpy(aligned_holo).float().to(device)
+                coords_a = torch.from_numpy(aligned_apo).float().to(device)
+                mask_t = torch.from_numpy(binder_mask).bool().to(device)
+                aa_t = torch.from_numpy(aa_idx).long().to(device)
+
+                q_holo = self.dq.score(coords_h, mask_t, binder_aa_idx=aa_t,
+                                        receptor_label='holo').item()
+                q_apo = self.dq.score(coords_a, mask_t, binder_aa_idx=aa_t,
+                                       receptor_label='apo').item()
+
+            return {
+                'q_holo': q_holo,
+                'q_apo': q_apo,
+                'margin': q_holo - q_apo,
+                'n_res': len(binder_res),
+            }
+
+        except Exception as e:
+            logger.warning(f"Error scoring {pdb_path}: {e}")
+            return None

code/scripts/pxdesign_guidance/smc_pxdesign.py

@@ -0,0 +1,262 @@
+"""
+PXDesign + SMC Reranking.
+
+Post-hoc Sequential Monte Carlo: generate multiple batches of vanilla PXDesign
+binders, score with Q_theta, and rank by selectivity margin. No modification
+to the PXDesign diffusion process — pure generate-score-rank pipeline.
+
+This is the simplest Q_theta integration strategy: generate a large pool of
+candidates and select the best ones by selectivity score.
+
+Usage:
+    python code/scripts/pxdesign_guidance/smc_pxdesign.py \
+        --input experiments/pxdesign_cam/output/cam_binder.json \
+        --qtheta_checkpoint results/checkpoints_cam_v3/best_phase2.pt \
+        --ref_holo data/pdbs/cam_holo/3CLN.pdb \
+        --ref_apo data/pdbs/cam_apo/1CFD.pdb \
+        --n_particles 16 --n_rounds 4 \
+        --gpu 0
+"""
+
+import os
+import sys
+import argparse
+import json
+import logging
+import shutil
+import subprocess
+from glob import glob
+
+import numpy as np
+import torch
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
+logger = logging.getLogger(__name__)
+
+_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+_ALLO_CODE_DIR = os.path.abspath(os.path.join(_SCRIPT_DIR, '..', '..'))
+_ALLO_ROOT = os.path.abspath(os.path.join(_ALLO_CODE_DIR, '..'))
+
+if _ALLO_CODE_DIR not in sys.path:
+    sys.path.insert(0, _ALLO_CODE_DIR)
+
+
+def run_pxdesign_batch(input_json, outdir, n_sample, n_step, gpu):
+    """Run vanilla PXDesign via CLI subprocess."""
+    pxdesign_python = 'python'
+
+    # Use pxdesign CLI
+    cmd = [
+        pxdesign_python, '-m', 'pxdesign.runner.cli', 'infer',
+        '-o', outdir,
+        '-i', input_json,
+        '--dtype', 'bf16',
+        '--N_sample', str(n_sample),
+        '--N_step', str(n_step),
+    ]
+
+    env = os.environ.copy()
+    # Inherit CUDA_VISIBLE_DEVICES from parent
+
+    logger.info(f"Running PXDesign: {n_sample} samples -> {outdir}")
+    result = subprocess.run(cmd, capture_output=True, text=True, env=env,
+                            timeout=7200)
+
+    if result.returncode != 0:
+        # Try alternative invocation via module
+        cmd_alt = [
+            pxdesign_python, '-m', 'pxdesign.runner.inference',
+            '--dump_dir', outdir,
+            '--input', input_json,
+            '--dtype', 'bf16',
+            '--N_sample', str(n_sample),
+            '--N_step', str(n_step),
+        ]
+        result = subprocess.run(cmd_alt, capture_output=True, text=True, env=env,
+                                timeout=7200)
+        if result.returncode != 0:
+            logger.error(f"PXDesign failed:\nstdout: {result.stdout[-1000:]}\nstderr: {result.stderr[-1000:]}")
+            return False
+    return True
+
+
+def collect_pdbs(outdir):
+    """Collect PDB/CIF files from PXDesign output."""
+    pdbs = []
+    for ext in ('*.pdb', '*.cif'):
+        pdbs.extend(glob(os.path.join(outdir, '**/' + ext), recursive=True))
+    pdbs = sorted(pdbs)
+    filtered = [p for p in pdbs if 'sample' in os.path.basename(p).lower()
+                or 'design' in os.path.basename(p).lower()
+                or 'rank' in os.path.basename(p).lower()]
+    return filtered if filtered else pdbs
+
+
+def smc_particle_filter(args):
+    """Run SMC reranking with PXDesign."""
+    os.chdir(_ALLO_ROOT)
+
+    from qtheta_pxdesign import QThetaPXDesignGuidance
+
+    outdir = args.outdir
+    os.makedirs(outdir, exist_ok=True)
+
+    # Initialize scorer
+    guidance = QThetaPXDesignGuidance(
+        checkpoint=args.qtheta_checkpoint,
+        ref_holo=args.ref_holo,
+        ref_apo=args.ref_apo,
+        ref_chain=args.ref_chain,
+        device=f'cuda:{args.gpu}',
+    )
+    guidance._lazy_init()
+
+    all_designs = []
+    round_summaries = []
+
+    for round_idx in range(args.n_rounds):
+        round_dir = os.path.join(outdir, f'round_{round_idx}')
+        os.makedirs(round_dir, exist_ok=True)
+
+        logger.info(f"\n{'='*60}")
+        logger.info(f"SMC Round {round_idx + 1}/{args.n_rounds}")
+        logger.info(f"{'='*60}")
+
+        # Generate particles via vanilla PXDesign
+        gen_dir = os.path.join(round_dir, 'generated')
+        success = run_pxdesign_batch(
+            input_json=args.input,
+            outdir=gen_dir,
+            n_sample=args.n_particles,
+            n_step=args.N_step,
+            gpu=args.gpu,
+        )
+
+        if not success:
+            # If subprocess fails, try using existing PXDesign outputs
+            logger.warning(f"Round {round_idx} generation failed. "
+                           f"Checking for existing outputs...")
+            pdbs = collect_pdbs(args.designs_dir) if hasattr(args, 'designs_dir') else []
+            if not pdbs:
+                continue
+        else:
+            pdbs = collect_pdbs(gen_dir)
+
+        if not pdbs:
+            logger.warning(f"No PDBs found in round {round_idx}")
+            continue
+
+        # Score all particles
+        logger.info(f"Scoring {len(pdbs)} particles...")
+        round_results = []
+        for pdb_path in pdbs:
+            result = guidance.score_design(pdb_path)
+            if result is not None:
+                result['pdb_path'] = pdb_path
+                result['design_id'] = os.path.basename(pdb_path).replace('.pdb', '').replace('.cif', '')
+                result['round'] = round_idx
+                round_results.append(result)
+
+        if not round_results:
+            continue
+
+        margins = np.array([r['margin'] for r in round_results])
+
+        round_summary = {
+            'round': round_idx,
+            'n_particles': len(round_results),
+            'margin_mean': float(margins.mean()),
+            'margin_std': float(margins.std()),
+            'margin_max': float(margins.max()),
+            'frac_positive': float((margins > 0).mean()),
+        }
+        round_summaries.append(round_summary)
+
+        logger.info(f"Round {round_idx}: margin={margins.mean():.3f}±{margins.std():.3f}, "
+                     f"max={margins.max():.3f}, S>0={round_summary['frac_positive']:.1%}")
+
+        all_designs.extend(round_results)
+
+    # Final ranking and summary
+    if all_designs:
+        all_designs.sort(key=lambda x: x['margin'], reverse=True)
+        all_margins = np.array([d['margin'] for d in all_designs])
+        holo_scores = np.array([d['q_holo'] for d in all_designs])
+
+        # Best-of-K
+        bok = {}
+        for K in [1, 2, 5, 10]:
+            n_trials = 2000
+            n_avail = len(all_margins)
+            successes = sum(
+                1 for _ in range(n_trials)
+                if all_margins[np.random.choice(n_avail, min(K, n_avail), replace=False)].max() > 0
+            )
+            bok[K] = successes / n_trials
+
+        summary = {
+            'method': 'PXDesign + SMC',
+            'n_rounds': args.n_rounds,
+            'n_particles_per_round': args.n_particles,
+            'total_designs': len(all_designs),
+            'margin_mean': float(all_margins.mean()),
+            'margin_std': float(all_margins.std()),
+            'margin_max': float(all_margins.max()),
+            'frac_positive': float((all_margins > 0).mean()),
+            'q_holo_mean': float(holo_scores.mean()),
+            'q_apo_mean': float(np.mean([d['q_apo'] for d in all_designs])),
+            'best_of_k': {str(k): v for k, v in bok.items()},
+            'round_summaries': round_summaries,
+            'top5': all_designs[:5],
+        }
+
+        with open(os.path.join(outdir, 'smc_scores.json'), 'w') as f:
+            json.dump(all_designs, f, indent=2)
+        with open(os.path.join(outdir, 'smc_summary.json'), 'w') as f:
+            json.dump(summary, f, indent=2)
+
+        # Copy best designs
+        best_dir = os.path.join(outdir, 'best_designs')
+        os.makedirs(best_dir, exist_ok=True)
+        for i, d in enumerate(all_designs[:20]):
+            if os.path.exists(d['pdb_path']):
+                dest = os.path.join(best_dir, f'rank_{i:02d}_{d["design_id"]}.pdb')
+                shutil.copy2(d['pdb_path'], dest)
+
+        logger.info(f"\n{'='*60}")
+        logger.info(f"PXDesign + SMC Results ({len(all_designs)} total designs)")
+        logger.info(f"  Margin: {all_margins.mean():.3f} ± {all_margins.std():.3f}")
+        logger.info(f"  Max margin: {all_margins.max():.3f}")
+        logger.info(f"  Fraction S > 0: {(all_margins > 0).mean():.1%}")
+        logger.info(f"  Q(holo) mean: {holo_scores.mean():.3f}")
+        logger.info(f"  Best-of-K:")
+        for k, v in sorted(bok.items()):
+            logger.info(f"    K={k:3d}: {v:.3f}")
+        logger.info(f"{'='*60}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description='PXDesign + SMC Reranking')
+    parser.add_argument('--input', default='experiments/pxdesign_cam/output/cam_binder.json',
+                        help='PXDesign input JSON')
+    parser.add_argument('--designs_dir', default='experiments/pxdesign_cam/output/',
+                        help='Existing PXDesign outputs (fallback if generation fails)')
+    parser.add_argument('--qtheta_checkpoint',
+                        default='results/checkpoints_cam_v3/best_phase2.pt')
+    parser.add_argument('--ref_holo', default='data/pdbs/cam_holo/3CLN.pdb')
+    parser.add_argument('--ref_apo', default='data/pdbs/cam_apo/1CFD.pdb')
+    parser.add_argument('--ref_chain', default='A')
+    parser.add_argument('--n_particles', type=int, default=16,
+                        help='Particles per round')
+    parser.add_argument('--n_rounds', type=int, default=4,
+                        help='Number of SMC rounds')
+    parser.add_argument('--N_step', type=int, default=400)
+    parser.add_argument('--gpu', type=int, default=0)
+    parser.add_argument('--outdir', default='results/pxdesign_smc')
+    args = parser.parse_args()
+
+    smc_particle_filter(args)
+
+
+if __name__ == '__main__':
+    main()

code/scripts/pxdesign_guidance/tds_pxdesign.py

@@ -0,0 +1,323 @@
+"""
+PXDesign + Twisted Diffusion Sampling (TDS).
+
+Multi-round particle filtering with guided PXDesign:
+  Round r:
+    1. Generate N particles via PXDesign with Q_theta classifier guidance
+    2. Score each particle with Q_theta selectivity margin
+    3. Compute importance weights w_i ~ exp(margin_i / temperature)
+    4. Resample particles (keep best, discard worst)
+    5. Add perturbation noise for diversity
+
+This combines in-process guidance (the "twisted proposal") with post-hoc
+importance-weighted resampling for highest-quality designs.
+
+Usage:
+    python code/scripts/pxdesign_guidance/tds_pxdesign.py \
+        --input experiments/pxdesign_cam/output/cam_binder.json \
+        --qtheta_checkpoint results/checkpoints_cam_v3/best_phase2.pt \
+        --ref_holo data/pdbs/cam_holo/3CLN.pdb \
+        --ref_apo data/pdbs/cam_apo/1CFD.pdb \
+        --n_particles 16 --n_rounds 4 \
+        --guidance_scale 0.5 \
+        --gpu 0
+"""
+
+import os
+import sys
+import argparse
+import json
+import logging
+import shutil
+import subprocess
+from glob import glob
+
+import numpy as np
+import torch
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
+logger = logging.getLogger(__name__)
+
+_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+_ALLO_CODE_DIR = os.path.abspath(os.path.join(_SCRIPT_DIR, '..', '..'))
+_ALLO_ROOT = os.path.abspath(os.path.join(_ALLO_CODE_DIR, '..'))
+
+if _ALLO_CODE_DIR not in sys.path:
+    sys.path.insert(0, _ALLO_CODE_DIR)
+
+
+def compute_ess(log_weights):
+    """Compute effective sample size from log-weights."""
+    log_weights = log_weights - log_weights.max()
+    weights = np.exp(log_weights)
+    weights = weights / weights.sum()
+    return 1.0 / (weights ** 2).sum()
+
+
+def run_guided_pxdesign_batch(input_json, outdir, n_sample, n_step,
+                               gpu, guidance_args):
+    """Run guided PXDesign as a subprocess."""
+    pxdesign_python = 'python'
+    cmd = [
+        pxdesign_python,
+        os.path.join(_SCRIPT_DIR, 'guided_pxdesign.py'),
+        '--input', input_json,
+        '--qtheta_checkpoint', guidance_args['checkpoint'],
+        '--ref_holo', guidance_args['ref_holo'],
+        '--ref_apo', guidance_args['ref_apo'],
+        '--ref_chain', guidance_args['ref_chain'],
+        '--guidance_scale', str(guidance_args['guidance_scale']),
+        '--guidance_start', str(guidance_args.get('guidance_start', 0.8)),
+        '--guidance_end', str(guidance_args.get('guidance_end', 0.1)),
+        '--N_sample', str(n_sample),
+        '--N_step', str(n_step),
+        '--gpu', str(gpu),
+        '--outdir', outdir,
+    ]
+
+    env = os.environ.copy()
+    # Inherit CUDA_VISIBLE_DEVICES from parent
+
+    logger.info(f"Running guided PXDesign: {n_sample} samples -> {outdir}")
+    result = subprocess.run(cmd, capture_output=True, text=True, env=env,
+                            timeout=7200)
+
+    if result.returncode != 0:
+        logger.error(f"PXDesign failed:\n{result.stderr[-2000:]}")
+        return False
+    return True
+
+
+def run_vanilla_pxdesign_batch(input_json, outdir, n_sample, n_step, gpu):
+    """Run vanilla PXDesign (no guidance) as a subprocess."""
+    pxdesign_env = 'python'
+
+    cmd = [
+        pxdesign_env, '-m', 'pxdesign.runner.inference',
+        '--dump_dir', outdir,
+        '--input', input_json,
+        '--dtype', 'bf16',
+        '--N_sample', str(n_sample),
+        '--N_step', str(n_step),
+    ]
+
+    env = os.environ.copy()
+    # Inherit CUDA_VISIBLE_DEVICES from parent
+
+    logger.info(f"Running vanilla PXDesign: {n_sample} samples -> {outdir}")
+    result = subprocess.run(cmd, capture_output=True, text=True, env=env,
+                            timeout=7200)
+
+    if result.returncode != 0:
+        logger.error(f"PXDesign failed:\n{result.stderr[-2000:]}")
+        return False
+    return True
+
+
+def collect_pdbs(outdir):
+    """Collect PDB/CIF paths from PXDesign output directory."""
+    pdbs = []
+    for ext in ('*.pdb', '*.cif'):
+        pdbs.extend(glob(os.path.join(outdir, '**/' + ext), recursive=True))
+    pdbs = sorted(pdbs)
+    filtered = [p for p in pdbs if 'sample' in os.path.basename(p).lower()
+                or 'design' in os.path.basename(p).lower()
+                or 'rank' in os.path.basename(p).lower()]
+    return filtered if filtered else pdbs
+
+
+def tds_particle_filter(args):
+    """Run TDS particle filtering with PXDesign."""
+    from qtheta_pxdesign import QThetaPXDesignGuidance
+
+    outdir = os.path.join(_ALLO_ROOT, args.outdir)
+    os.makedirs(outdir, exist_ok=True)
+
+    # Initialize scorer
+    guidance = QThetaPXDesignGuidance(
+        checkpoint=os.path.join(_ALLO_ROOT, args.qtheta_checkpoint),
+        ref_holo=os.path.join(_ALLO_ROOT, args.ref_holo),
+        ref_apo=os.path.join(_ALLO_ROOT, args.ref_apo),
+        ref_chain=args.ref_chain,
+        device=f'cuda:{args.gpu}',
+    )
+    guidance._lazy_init()
+
+    guidance_args = {
+        'checkpoint': args.qtheta_checkpoint,
+        'ref_holo': args.ref_holo,
+        'ref_apo': args.ref_apo,
+        'ref_chain': args.ref_chain,
+        'guidance_scale': args.guidance_scale,
+        'guidance_start': args.guidance_start,
+        'guidance_end': args.guidance_end,
+    }
+
+    all_designs = []
+    round_summaries = []
+
+    for round_idx in range(args.n_rounds):
+        round_dir = os.path.join(outdir, f'round_{round_idx}')
+        os.makedirs(round_dir, exist_ok=True)
+
+        logger.info(f"\n{'='*60}")
+        logger.info(f"TDS Round {round_idx + 1}/{args.n_rounds}")
+        logger.info(f"{'='*60}")
+
+        # Generate particles via guided PXDesign
+        gen_dir = os.path.join(round_dir, 'generated')
+        success = run_guided_pxdesign_batch(
+            input_json=os.path.join(_ALLO_ROOT, args.input),
+            outdir=gen_dir,
+            n_sample=args.n_particles,
+            n_step=args.N_step,
+            gpu=args.gpu,
+            guidance_args=guidance_args,
+        )
+
+        if not success:
+            logger.warning(f"Round {round_idx} generation failed, skipping")
+            continue
+
+        # Collect and score particles
+        pdbs = collect_pdbs(gen_dir)
+        if not pdbs:
+            logger.warning(f"No PDBs found in round {round_idx}")
+            continue
+
+        logger.info(f"Scoring {len(pdbs)} particles...")
+        round_results = []
+        for pdb_path in pdbs:
+            result = guidance.score_design(pdb_path)
+            if result is not None:
+                result['pdb_path'] = pdb_path
+                result['design_id'] = os.path.basename(pdb_path).replace('.pdb', '').replace('.cif', '')
+                result['round'] = round_idx
+                round_results.append(result)
+
+        if not round_results:
+            logger.warning(f"No scorable designs in round {round_idx}")
+            continue
+
+        margins = np.array([r['margin'] for r in round_results])
+
+        # Compute importance weights
+        log_weights = margins / args.temperature
+        ess = compute_ess(log_weights)
+
+        round_summary = {
+            'round': round_idx,
+            'n_particles': len(round_results),
+            'margin_mean': float(margins.mean()),
+            'margin_std': float(margins.std()),
+            'margin_max': float(margins.max()),
+            'frac_positive': float((margins > 0).mean()),
+            'ess': float(ess),
+        }
+        round_summaries.append(round_summary)
+
+        logger.info(f"Round {round_idx}: margin={margins.mean():.3f}±{margins.std():.3f}, "
+                     f"max={margins.max():.3f}, S>0={round_summary['frac_positive']:.1%}, "
+                     f"ESS={ess:.1f}/{len(round_results)}")
+
+        # Add to design pool
+        all_designs.extend(round_results)
+
+        # Resample for next round (top-K selection for PXDesign since
+        # we can't easily perturb and re-denoise)
+        if round_idx < args.n_rounds - 1:
+            # Copy best designs to inform next round
+            # For PXDesign, each round generates fresh samples with guidance
+            # Resampling influence is through the guidance strength
+            # Increase guidance scale for later rounds
+            guidance_args['guidance_scale'] = args.guidance_scale * (1.0 + 0.2 * (round_idx + 1))
+            logger.info(f"Increasing guidance scale to {guidance_args['guidance_scale']:.2f} "
+                        f"for next round")
+
+    # Final summary
+    if all_designs:
+        all_designs.sort(key=lambda x: x['margin'], reverse=True)
+        all_margins = np.array([d['margin'] for d in all_designs])
+        holo_scores = np.array([d['q_holo'] for d in all_designs])
+
+        # Best-of-K
+        bok = {}
+        for K in [1, 2, 5, 10]:
+            n_trials = 2000
+            n_avail = len(all_margins)
+            successes = sum(
+                1 for _ in range(n_trials)
+                if all_margins[np.random.choice(n_avail, min(K, n_avail), replace=False)].max() > 0
+            )
+            bok[K] = successes / n_trials
+
+        summary = {
+            'method': 'PXDesign + TDS',
+            'n_rounds': args.n_rounds,
+            'n_particles_per_round': args.n_particles,
+            'total_designs': len(all_designs),
+            'guidance_scale': args.guidance_scale,
+            'temperature': args.temperature,
+            'margin_mean': float(all_margins.mean()),
+            'margin_std': float(all_margins.std()),
+            'margin_max': float(all_margins.max()),
+            'frac_positive': float((all_margins > 0).mean()),
+            'q_holo_mean': float(holo_scores.mean()),
+            'best_of_k': {str(k): v for k, v in bok.items()},
+            'round_summaries': round_summaries,
+            'top5': all_designs[:5],
+        }
+
+        with open(os.path.join(outdir, 'tds_scores.json'), 'w') as f:
+            json.dump(all_designs, f, indent=2)
+        with open(os.path.join(outdir, 'tds_summary.json'), 'w') as f:
+            json.dump(summary, f, indent=2)
+
+        # Copy best designs to top-level
+        best_dir = os.path.join(outdir, 'best_designs')
+        os.makedirs(best_dir, exist_ok=True)
+        for i, d in enumerate(all_designs[:20]):
+            if os.path.exists(d['pdb_path']):
+                dest = os.path.join(best_dir, f'rank_{i:02d}_{d["design_id"]}.pdb')
+                shutil.copy2(d['pdb_path'], dest)
+
+        logger.info(f"\n{'='*60}")
+        logger.info(f"PXDesign + TDS Results ({len(all_designs)} total designs)")
+        logger.info(f"  Margin: {all_margins.mean():.3f} ± {all_margins.std():.3f}")
+        logger.info(f"  Max margin: {all_margins.max():.3f}")
+        logger.info(f"  Fraction S > 0: {(all_margins > 0).mean():.1%}")
+        logger.info(f"  Q(holo) mean: {holo_scores.mean():.3f}")
+        logger.info(f"  Best-of-K:")
+        for k, v in sorted(bok.items()):
+            logger.info(f"    K={k:3d}: {v:.3f}")
+        logger.info(f"{'='*60}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description='PXDesign + TDS')
+    parser.add_argument('--input', default='experiments/pxdesign_cam/output/cam_binder.json')
+    parser.add_argument('--qtheta_checkpoint',
+                        default='results/checkpoints_cam_v3/best_phase2.pt')
+    parser.add_argument('--ref_holo', default='data/pdbs/cam_holo/3CLN.pdb')
+    parser.add_argument('--ref_apo', default='data/pdbs/cam_apo/1CFD.pdb')
+    parser.add_argument('--ref_chain', default='A')
+    parser.add_argument('--n_particles', type=int, default=16,
+                        help='Particles per round')
+    parser.add_argument('--n_rounds', type=int, default=4,
+                        help='Number of TDS rounds')
+    parser.add_argument('--guidance_scale', type=float, default=0.5,
+                        help='Initial guidance scale')
+    parser.add_argument('--guidance_start', type=float, default=0.8)
+    parser.add_argument('--guidance_end', type=float, default=0.1)
+    parser.add_argument('--temperature', type=float, default=0.5,
+                        help='Temperature for importance weights')
+    parser.add_argument('--N_step', type=int, default=400)
+    parser.add_argument('--gpu', type=int, default=0)
+    parser.add_argument('--outdir', default='results/pxdesign_tds')
+    args = parser.parse_args()
+
+    tds_particle_filter(args)
+
+
+if __name__ == '__main__':
+    main()

code/scripts/rescore.py

@@ -0,0 +1,178 @@
+#!/usr/bin/env python3
+"""
+Re-score binder PDB designs with a Q_theta checkpoint.
+
+Walks a directory of designs (binder PDB + sibling holo / apo receptor PDBs),
+runs each through DifferentiableQTheta, and writes per-design
+S = Q_theta(holo) - Q_theta(apo) plus the raw holo/apo scores to JSON.
+
+Usage:
+    python code/scripts/rescore.py \\
+        --checkpoint checkpoints/Q_theta_phase2.pt \\
+        --gpu 0
+"""
+import os, sys, json, argparse, glob, logging
+import numpy as np
+import torch
+from pathlib import Path
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
+logger = logging.getLogger(__name__)
+
+BASE = str(Path(__file__).resolve().parent.parent.parent)
+sys.path.insert(0, os.path.join(BASE, 'code'))
+sys.path.insert(0, BASE)
+
+from models.differentiable_features import DifferentiableQTheta
+from utils.pdb_utils import load_structure, get_residues, get_backbone_coords, get_aa_indices, align_structures
+
+HOLO_PDB = os.path.join(BASE, 'data/pdbs/cam_holo/3CLN.pdb')
+APO_PDB = os.path.join(BASE, 'data/pdbs/cam_apo/1CFD.pdb')
+
+
+def score_pdb_list(dq, pdb_list, ref_resnums, ref_coords, device):
+    """Score a list of design PDB files."""
+    results = []
+    for pdb_path in pdb_list:
+        name = os.path.basename(pdb_path).replace(".pdb", "")
+        try:
+            design_model = load_structure(pdb_path)
+            chains = [c.id for c in design_model.get_chains()]
+            rec_chain = 'A' if 'A' in chains else chains[0]
+            binder_chain = 'B' if 'B' in chains else [c for c in chains if c != rec_chain][0]
+
+            rec_res = get_residues(design_model[rec_chain])
+            binder_res = get_residues(design_model[binder_chain])
+            rec_coords_d, _ = get_backbone_coords(rec_res)
+            binder_coords, binder_mask = get_backbone_coords(binder_res)
+            binder_aa_idx = get_aa_indices(binder_res)
+
+            design_resnums = {r.get_id()[1]: i for i, r in enumerate(rec_res)}
+            common = sorted(set(design_resnums.keys()) & set(ref_resnums.keys()))
+            if len(common) < 10:
+                logger.warning(f"  Skip {name}: <10 common residues")
+                continue
+
+            d_ca = rec_coords_d[[design_resnums[r] for r in common], 1]
+            r_ca = ref_coords[[ref_resnums[r] for r in common], 1]
+            mobile_center = d_ca.mean(0)
+            ref_center = r_ca.mean(0)
+            _, R = align_structures(d_ca, r_ca)
+
+            flat = binder_coords.reshape(-1, 3) - mobile_center
+            aligned_binder = (flat @ R.T + ref_center).reshape(-1, 4, 3)
+
+            coords_t = torch.from_numpy(aligned_binder).float().to(device)
+            mask_t = torch.from_numpy(binder_mask).bool().to(device)
+            aa_t = torch.from_numpy(binder_aa_idx).long().to(device)
+
+            with torch.no_grad():
+                q_holo = dq.score(coords_t, mask_t, binder_aa_idx=aa_t,
+                                  receptor_label='holo').item()
+                q_apo = dq.score(coords_t, mask_t, binder_aa_idx=aa_t,
+                                 receptor_label='apo').item()
+            S = q_holo - q_apo
+            results.append({"design": name, "Q_holo": q_holo, "Q_apo": q_apo, "S": S})
+        except Exception as e:
+            logger.warning(f"  Skip {name}: {e}")
+    return results
+
+
+def summarize(results, label):
+    if not results:
+        return {}
+    S = [r["S"] for r in results]
+    return {
+        "method": label, "n": len(S),
+        "S_mean": float(np.mean(S)), "S_std": float(np.std(S)),
+        "S_pos_pct": float(np.mean([s > 0 for s in S]) * 100),
+        "Q_holo_mean": float(np.mean([r["Q_holo"] for r in results])),
+        "Q_apo_mean": float(np.mean([r["Q_apo"] for r in results])),
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--gpu", type=int, default=7)
+    parser.add_argument("--checkpoint", default="checkpoints/Q_theta_phase2.pt")
+    args = parser.parse_args()
+
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
+    device = "cuda:0"
+
+    logger.info(f"Loading Q_theta from {args.checkpoint}")
+    dq = DifferentiableQTheta(checkpoint_path=args.checkpoint, device=device,
+                               esm_dir=os.path.join(BASE, "data/esm2_embeddings"))
+    dq.load_receptor(HOLO_PDB, chain='A', label='holo', esm_target='cam')
+    dq.load_receptor(APO_PDB, chain='A', label='apo', esm_target='cam')
+
+    ref_model = load_structure(HOLO_PDB)
+    ref_res = get_residues(ref_model['A'])
+    ref_coords, _ = get_backbone_coords(ref_res)
+    ref_resnums = {r.get_id()[1]: i for i, r in enumerate(ref_res)}
+
+    output_dir = os.path.join(BASE, "results/v2_strict_holdout/scoring")
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Define design directories
+    design_sets = {
+        "vanilla": os.path.join(BASE, "results/independent_validation/vanilla/holo_pdbs"),
+        "langevin": os.path.join(BASE, "results/langevin_refinement/refined_pdbs"),
+        "classifier": os.path.join(BASE, "results/guided_diffusion/guided"),
+        "smc_r3": os.path.join(BASE, "results/smc_guidance/cam/round_3"),
+    }
+
+    # Also check for TDS and PXDesign
+    tds_dirs = glob.glob(os.path.join(BASE, "results/tds_guidance/cam/designs"))
+    if tds_dirs:
+        design_sets["tds"] = tds_dirs[0]
+
+    # PXDesign directories
+    for px_method in ["pxdesign_scoring", "pxdesign_classifier", "pxdesign_tds",
+                       "pxdesign_smc", "pxdesign_langevin"]:
+        px_dir = os.path.join(BASE, f"results_familysplit/design_bd30/{px_method}")
+        if not os.path.exists(px_dir):
+            px_dir = os.path.join(BASE, f"results/{px_method}")
+        if os.path.exists(px_dir):
+            pdbs = glob.glob(os.path.join(px_dir, "*.pdb"))
+            if pdbs:
+                design_sets[px_method] = px_dir
+
+    all_results = {}
+    summaries = []
+
+    for method, pdb_dir in design_sets.items():
+        if not os.path.exists(pdb_dir):
+            logger.warning(f"  {method}: directory not found ({pdb_dir})")
+            continue
+        pdbs = sorted(glob.glob(os.path.join(pdb_dir, "*.pdb")))
+        if not pdbs:
+            logger.warning(f"  {method}: no PDB files")
+            continue
+
+        logger.info(f"\n=== {method} ({len(pdbs)} designs) ===")
+        results = score_pdb_list(dq, pdbs, ref_resnums, ref_coords, device)
+        s = summarize(results, method)
+        if s:
+            summaries.append(s)
+            logger.info(f"  {method}: n={s['n']}, S̄={s['S_mean']:.3f}±{s['S_std']:.3f}, "
+                       f"S>0={s['S_pos_pct']:.0f}%, Q+={s['Q_holo_mean']:.3f}, Q-={s['Q_apo_mean']:.3f}")
+        all_results[method] = {"results": results, "summary": s}
+
+    # Save
+    with open(os.path.join(output_dir, "rescore_v2_all.json"), "w") as f:
+        json.dump(all_results, f, indent=2)
+
+    # Print summary table
+    print("\n" + "=" * 70)
+    print("V2 RESCORING SUMMARY (strict holdout, CaM OOD)")
+    print("=" * 70)
+    print(f"{'Method':20s} {'n':>4s} {'S̄':>8s} {'±σ':>6s} {'S>0%':>6s} {'Q+':>6s} {'Q-':>6s}")
+    print("-" * 70)
+    for s in sorted(summaries, key=lambda x: x['S_mean'], reverse=True):
+        print(f"{s['method']:20s} {s['n']:4d} {s['S_mean']:8.3f} {s['S_std']:6.3f} "
+              f"{s['S_pos_pct']:5.1f}% {s['Q_holo_mean']:6.3f} {s['Q_apo_mean']:6.3f}")
+
+
+if __name__ == "__main__":
+    main()

code/trainers/trainer.py

@@ -0,0 +1,674 @@
+"""
+Trainer for the Q_theta state-selectivity scorer.
+
+Implements two-phase training:
+  Phase 1: DockQ regression (learn complex quality from all data)
+  Phase 2: Selectivity fine-tuning (learn to rank X+ > X- for the same binder)
+
+Integrates with Weights & Biases for experiment tracking.
+"""
+
+import os
+import time
+import logging
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.optim import AdamW
+from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR
+from scipy.stats import spearmanr
+from sklearn.metrics import roc_auc_score
+
+import wandb
+
+logger = logging.getLogger(__name__)
+
+
+class AverageMeter:
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0.0
+        self.avg = 0.0
+        self.sum = 0.0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+
+class AlloDesignerTrainer:
+    """
+    Two-phase trainer for Q_theta.
+
+    Phase 1 (DockQ regression):
+      - Minimizes MSE(Q_theta(X, Y), DockQ_label) on all complex types
+      - Learns general complex quality
+
+    Phase 2 (Selectivity fine-tuning):
+      - Minimizes selectivity margin loss on paired (pos, neg) data
+      - Learns to rank Q(X+, Y) > Q(X-, Y)
+      - Combined: L = L_regression + lambda_rank * L_selectivity
+    """
+
+    def __init__(self, model, config, device='cuda'):
+        self.model = model.to(device)
+        self.config = config
+        self.device = device
+        self.use_sam = config.get('optimizer', 'adamw') == 'sam'
+
+        # Optimizer
+        if self.use_sam:
+            from utils.sam import SAM
+            self.optimizer = SAM(
+                model.parameters(),
+                base_optimizer=AdamW,
+                rho=0.05,
+                lr=config.get('lr', 1e-4),
+                weight_decay=config.get('weight_decay', 1e-4),
+                betas=(0.9, 0.999),
+            )
+            # SAM wraps AdamW; scheduler goes on base_optimizer
+            sched_optimizer = self.optimizer.base_optimizer
+        else:
+            self.optimizer = AdamW(
+                model.parameters(),
+                lr=config.get('lr', 1e-4),
+                weight_decay=config.get('weight_decay', 1e-4),
+                betas=(0.9, 0.999),
+            )
+            sched_optimizer = self.optimizer
+
+        # Learning rate scheduler (warmup + cosine)
+        n_warmup = config.get('warmup_steps', 100)
+        n_total = config.get('max_steps', 5000)
+
+        warmup_sched = LinearLR(sched_optimizer, start_factor=0.01, end_factor=1.0, total_iters=n_warmup)
+        cosine_sched = CosineAnnealingLR(sched_optimizer, T_max=n_total - n_warmup, eta_min=1e-6)
+        self.scheduler = SequentialLR(sched_optimizer, [warmup_sched, cosine_sched], milestones=[n_warmup])
+
+        self.global_step = 0
+        self.best_val_metric = -float('inf')
+        self.checkpoint_dir = config.get('checkpoint_dir', 'results/checkpoints')
+        os.makedirs(self.checkpoint_dir, exist_ok=True)
+
+    # ------------------------------------------------------------------ #
+    # Phase 1: DockQ regression
+    # ------------------------------------------------------------------ #
+
+    def train_step_phase1(self, batch):
+        """Single training step for Phase 1 (DockQ regression)."""
+        self.model.train()
+        node_feats = batch['node_feats'].to(self.device)    # [B, N, node_dim]
+        edge_feats = batch['edge_feats'].to(self.device)    # [B, N, N, edge_dim]
+        node_mask = batch['node_mask'].to(self.device)      # [B, N]
+        labels = batch['label'].to(self.device)             # [B]
+        esm_feats = batch['esm_feats'].to(self.device) if 'esm_feats' in batch else None
+
+        self.optimizer.zero_grad()
+
+        scores = self.model(node_feats, edge_feats, node_mask, esm_feats=esm_feats)  # [B]
+        loss = nn.functional.mse_loss(scores, labels)
+
+        loss.backward()
+        nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
+
+        if self.use_sam:
+            self.optimizer.first_step()
+            # Second forward-backward pass
+            scores2 = self.model(node_feats, edge_feats, node_mask, esm_feats=esm_feats)
+            loss2 = nn.functional.mse_loss(scores2, labels)
+            self.optimizer.zero_grad()
+            loss2.backward()
+            nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
+            self.optimizer.second_step()
+        else:
+            self.optimizer.step()
+
+        self.scheduler.step()
+        self.global_step += 1
+
+        return {'loss': loss.item(), 'scores': scores.detach(), 'labels': labels}
+
+    def run_phase1(self, train_loader, val_loader, n_epochs: int = 30, run_name: str = 'phase1'):
+        """Phase 1 training loop."""
+        logger.info(f"Starting Phase 1 (DockQ regression) for {n_epochs} epochs")
+        wandb.define_metric('phase1/step')
+        wandb.define_metric('phase1/*', step_metric='phase1/step')
+
+        for epoch in range(n_epochs):
+            # Train
+            train_meter = AverageMeter()
+            all_scores, all_labels = [], []
+
+            for batch in train_loader:
+                result = self.train_step_phase1(batch)
+                train_meter.update(result['loss'], n=len(result['scores']))
+                all_scores.append(result['scores'].cpu().numpy())
+                all_labels.append(result['labels'].cpu().numpy())
+
+                if self.global_step % 50 == 0:
+                    wandb.log({
+                        'phase1/train_loss': result['loss'],
+                        'phase1/lr': self.optimizer.param_groups[0]['lr'],
+                        'phase1/step': self.global_step,
+                    })
+
+            # Compute Spearman corr on training data
+            all_scores = np.concatenate(all_scores)
+            all_labels = np.concatenate(all_labels)
+            train_spearman = spearmanr(all_scores, all_labels).correlation
+
+            # Validate
+            val_metrics = self.evaluate_phase1(val_loader)
+
+            logger.info(
+                f"Phase1 Epoch {epoch+1}/{n_epochs} | "
+                f"Train Loss: {train_meter.avg:.4f} | "
+                f"Train Spearman: {train_spearman:.3f} | "
+                f"Val Loss: {val_metrics['val_loss']:.4f} | "
+                f"Val Spearman: {val_metrics['val_spearman']:.3f} | "
+                f"Val AUC: {val_metrics.get('val_auc', 0):.3f}"
+            )
+
+            wandb.log({
+                'phase1/epoch': epoch + 1,
+                'phase1/train_loss_epoch': train_meter.avg,
+                'phase1/train_spearman': train_spearman,
+                **{f'phase1/{k}': v for k, v in val_metrics.items()},
+            })
+
+            # Checkpoint best model
+            if val_metrics['val_spearman'] > self.best_val_metric:
+                self.best_val_metric = val_metrics['val_spearman']
+                self.save_checkpoint('best_phase1.pt', extra={'epoch': epoch, 'phase': 1})
+                logger.info(f"  -> New best Phase 1 model (val_spearman={self.best_val_metric:.3f})")
+
+        logger.info("Phase 1 training complete.")
+
+    @torch.no_grad()
+    def evaluate_phase1(self, loader):
+        """Evaluate Phase 1 model on val/test set."""
+        self.model.eval()
+        all_scores, all_labels = [], []
+        total_loss = 0.0
+        n_batches = 0
+
+        for batch in loader:
+            node_feats = batch['node_feats'].to(self.device)
+            edge_feats = batch['edge_feats'].to(self.device)
+            node_mask = batch['node_mask'].to(self.device)
+            labels = batch['label'].to(self.device)
+            esm_feats = batch['esm_feats'].to(self.device) if 'esm_feats' in batch else None
+
+            scores = self.model(node_feats, edge_feats, node_mask, esm_feats=esm_feats)
+            loss = nn.functional.mse_loss(scores, labels)
+
+            total_loss += loss.item()
+            n_batches += 1
+            all_scores.append(scores.cpu().numpy())
+            all_labels.append(labels.cpu().numpy())
+
+        all_scores = np.concatenate(all_scores)
+        all_labels = np.concatenate(all_labels)
+
+        spearman = spearmanr(all_scores, all_labels).correlation
+        if np.isnan(spearman):
+            spearman = 0.0
+
+        metrics = {
+            'val_loss': total_loss / max(n_batches, 1),
+            'val_spearman': float(spearman),
+        }
+
+        # AUC for binary quality (label > 0.5 = positive)
+        binary_labels = (all_labels > 0.5).astype(int)
+        if binary_labels.sum() > 0 and binary_labels.sum() < len(binary_labels):
+            try:
+                metrics['val_auc'] = roc_auc_score(binary_labels, all_scores)
+            except Exception:
+                pass
+
+        return metrics
+
+    # ------------------------------------------------------------------ #
+    # Phase 2: Selectivity fine-tuning
+    # ------------------------------------------------------------------ #
+
+    def train_step_phase2(self, batch, lambda_rank: float = 1.0, margin: float = 0.2,
+                          lambda_ddg: float = 0.1):
+        """Single training step for Phase 2 (selectivity margin + ddG auxiliary)."""
+        self.model.train()
+
+        pos = batch['pos']
+        neg = batch['neg']
+
+        pos_node = pos['node_feats'].to(self.device)
+        pos_edge = pos['edge_feats'].to(self.device)
+        pos_mask = pos['node_mask'].to(self.device)
+        pos_label = pos['label'].to(self.device)
+        pos_ce = pos.get('contact_energy', None)
+        if pos_ce is not None:
+            pos_ce = pos_ce.to(self.device)
+
+        neg_node = neg['node_feats'].to(self.device)
+        neg_edge = neg['edge_feats'].to(self.device)
+        neg_mask = neg['node_mask'].to(self.device)
+        pos_esm = pos['esm_feats'].to(self.device) if 'esm_feats' in pos else None
+        neg_esm = neg['esm_feats'].to(self.device) if 'esm_feats' in neg else None
+
+        self.optimizer.zero_grad()
+
+        pos_scores = self.model(pos_node, pos_edge, pos_mask, esm_feats=pos_esm)   # [B]
+        neg_scores = self.model(neg_node, neg_edge, neg_mask, esm_feats=neg_esm)   # [B]
+
+        # Regression loss on positive examples
+        loss_reg = nn.functional.mse_loss(pos_scores, pos_label)
+
+        # Selectivity margin loss: pos_score - neg_score > margin
+        loss_margin = nn.functional.relu(margin - (pos_scores - neg_scores)).mean()
+
+        # InfoNCE-style selectivity loss
+        eps = 1e-6
+        pos_logit = torch.log(pos_scores.clamp(eps, 1 - eps) / (1 - pos_scores).clamp(eps))
+        neg_logit = torch.log(neg_scores.clamp(eps, 1 - eps) / (1 - neg_scores).clamp(eps))
+        log_denom = torch.stack([pos_logit, neg_logit], dim=-1).logsumexp(dim=-1)
+        infonce_loss = -(pos_logit - log_denom).mean()
+
+        # ddG auxiliary loss: MSE against contact-energy proxy (physics-informed soft label)
+        loss_ddg = torch.tensor(0.0, device=self.device)
+        if pos_ce is not None and pos_ce.shape[0] > 0:
+            # pos_ce is a contact-energy-based ddG proxy in [0, 1]
+            # Align positive score toward the contact energy signal
+            loss_ddg = nn.functional.mse_loss(pos_scores, pos_ce)
+
+        loss = loss_reg + lambda_rank * (loss_margin + infonce_loss) + lambda_ddg * loss_ddg
+
+        loss.backward()
+        nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
+
+        if self.use_sam:
+            self.optimizer.first_step()
+            # Second forward-backward for SAM
+            pos_scores2 = self.model(pos_node, pos_edge, pos_mask, esm_feats=pos_esm)
+            neg_scores2 = self.model(neg_node, neg_edge, neg_mask, esm_feats=neg_esm)
+            loss_reg2 = nn.functional.mse_loss(pos_scores2, pos_label)
+            loss_margin2 = nn.functional.relu(margin - (pos_scores2 - neg_scores2)).mean()
+            eps2 = 1e-6
+            pl2 = torch.log(pos_scores2.clamp(eps2, 1-eps2) / (1-pos_scores2).clamp(eps2))
+            nl2 = torch.log(neg_scores2.clamp(eps2, 1-eps2) / (1-neg_scores2).clamp(eps2))
+            ld2 = torch.stack([pl2, nl2], dim=-1).logsumexp(dim=-1)
+            infonce2 = -(pl2 - ld2).mean()
+            loss2 = loss_reg2 + lambda_rank * (loss_margin2 + infonce2)
+            self.optimizer.zero_grad()
+            loss2.backward()
+            nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
+            self.optimizer.second_step()
+        else:
+            self.optimizer.step()
+
+        self.scheduler.step()
+        self.global_step += 1
+
+        selectivity_gap = (pos_scores - neg_scores).mean().item()
+
+        return {
+            'loss': loss.item(),
+            'loss_reg': loss_reg.item(),
+            'loss_margin': loss_margin.item(),
+            'loss_infonce': infonce_loss.item(),
+            'loss_ddg': loss_ddg.item(),
+            'selectivity_gap': selectivity_gap,
+            'pos_scores': pos_scores.detach(),
+            'neg_scores': neg_scores.detach(),
+        }
+
+    def train_step_phase2_v2(self, batch, lambda_rank: float = 1.0, margin: float = 0.2,
+                             lambda_ddg: float = 0.0, lambda_path: float = 0.5):
+        """Phase 2 training step with multi-negative + path monotonicity."""
+        self.model.train()
+
+        pos = batch['pos']
+        neg = batch['neg']
+
+        pos_node = pos['node_feats'].to(self.device)
+        pos_edge = pos['edge_feats'].to(self.device)
+        pos_mask = pos['node_mask'].to(self.device)
+        pos_label = pos['label'].to(self.device)
+        pos_ce = pos.get('contact_energy', None)
+        if pos_ce is not None:
+            pos_ce = pos_ce.to(self.device)
+
+        neg_node = neg['node_feats'].to(self.device)
+        neg_edge = neg['edge_feats'].to(self.device)
+        neg_mask = neg['node_mask'].to(self.device)
+        pos_esm = pos['esm_feats'].to(self.device) if 'esm_feats' in pos else None
+        neg_esm = neg['esm_feats'].to(self.device) if 'esm_feats' in neg else None
+
+        self.optimizer.zero_grad()
+
+        pos_scores = self.model(pos_node, pos_edge, pos_mask, esm_feats=pos_esm)
+        neg_scores = self.model(neg_node, neg_edge, neg_mask, esm_feats=neg_esm)
+
+        # Score path frames if present
+        path_scores = []
+        path_taus = batch.get('path_taus', [])
+        for path_frame in batch.get('path', []):
+            p_node = path_frame['node_feats'].to(self.device)
+            p_edge = path_frame['edge_feats'].to(self.device)
+            p_mask = path_frame['node_mask'].to(self.device)
+            p_score = self.model(p_node, p_edge, p_mask)
+            path_scores.append(p_score)
+
+        # Regression loss on positive examples
+        loss_reg = nn.functional.mse_loss(pos_scores, pos_label)
+
+        # Selectivity margin loss
+        loss_margin = nn.functional.relu(margin - (pos_scores - neg_scores)).mean()
+
+        # InfoNCE-style selectivity loss
+        eps = 1e-6
+        pos_logit = torch.log(pos_scores.clamp(eps, 1 - eps) / (1 - pos_scores).clamp(eps))
+        neg_logit = torch.log(neg_scores.clamp(eps, 1 - eps) / (1 - neg_scores).clamp(eps))
+        log_denom = torch.stack([pos_logit, neg_logit], dim=-1).logsumexp(dim=-1)
+        infonce_loss = -(pos_logit - log_denom).mean()
+
+        # ddG auxiliary loss
+        loss_ddg = torch.tensor(0.0, device=self.device)
+        if pos_ce is not None and pos_ce.shape[0] > 0 and lambda_ddg > 0:
+            loss_ddg = nn.functional.mse_loss(pos_scores, pos_ce)
+
+        # Path monotonicity loss
+        loss_path = torch.tensor(0.0, device=self.device)
+        if path_scores and lambda_path > 0:
+            small_margin = 0.05
+            for i in range(len(path_scores) - 1):
+                loss_path = loss_path + nn.functional.relu(
+                    path_scores[i] - path_scores[i + 1] + small_margin
+                ).mean()
+            # Last path frame < positive score
+            loss_path = loss_path + nn.functional.relu(
+                path_scores[-1] - pos_scores + margin
+            ).mean()
+            # First path frame > negative score
+            loss_path = loss_path + nn.functional.relu(
+                neg_scores - path_scores[0] + small_margin
+            ).mean()
+
+        loss = (loss_reg + lambda_rank * (loss_margin + infonce_loss)
+                + lambda_ddg * loss_ddg + lambda_path * loss_path)
+
+        loss.backward()
+        nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
+        self.optimizer.step()
+        self.scheduler.step()
+        self.global_step += 1
+
+        selectivity_gap = (pos_scores - neg_scores).mean().item()
+
+        return {
+            'loss': loss.item(),
+            'loss_reg': loss_reg.item(),
+            'loss_margin': loss_margin.item(),
+            'loss_infonce': infonce_loss.item(),
+            'loss_ddg': loss_ddg.item(),
+            'loss_path': loss_path.item(),
+            'selectivity_gap': selectivity_gap,
+            'pos_scores': pos_scores.detach(),
+            'neg_scores': neg_scores.detach(),
+        }
+
+    def run_phase2_path(self, train_loader, val_loader, n_epochs: int = 20,
+                        lambda_rank: float = 1.0, margin: float = 0.2,
+                        lambda_ddg: float = 0.0, lambda_path: float = 0.5):
+        """Phase 2 with path-aware training loop."""
+        logger.info(f"Starting Phase 2 (path-aware) for {n_epochs} epochs "
+                    f"[lambda_rank={lambda_rank}, lambda_path={lambda_path}]")
+        self.best_val_metric = -float('inf')
+
+        for epoch in range(n_epochs):
+            loss_meter = AverageMeter()
+            gap_meter = AverageMeter()
+            path_meter = AverageMeter()
+
+            for batch in train_loader:
+                result = self.train_step_phase2_v2(
+                    batch, lambda_rank, margin, lambda_ddg, lambda_path)
+                B = len(result['pos_scores'])
+                loss_meter.update(result['loss'], B)
+                gap_meter.update(result['selectivity_gap'], B)
+                path_meter.update(result['loss_path'], B)
+
+                if self.global_step % 50 == 0:
+                    wandb.log({
+                        'phase2/train_loss': result['loss'],
+                        'phase2/loss_margin': result['loss_margin'],
+                        'phase2/loss_infonce': result['loss_infonce'],
+                        'phase2/loss_path': result['loss_path'],
+                        'phase2/selectivity_gap': result['selectivity_gap'],
+                        'phase2/lr': self.optimizer.param_groups[0]['lr'],
+                        'phase2/step': self.global_step,
+                    })
+
+            val_metrics = self.evaluate_phase2(val_loader)
+
+            logger.info(
+                f"Phase2-Path Epoch {epoch+1}/{n_epochs} | "
+                f"Loss: {loss_meter.avg:.4f} | "
+                f"Gap: {gap_meter.avg:.3f} | "
+                f"Path: {path_meter.avg:.4f} | "
+                f"Val Gap: {val_metrics['val_selectivity_gap']:.3f} | "
+                f"Val Acc: {val_metrics['val_ranking_acc']:.3f}"
+            )
+
+            wandb.log({
+                'phase2/epoch': epoch + 1,
+                'phase2/train_loss_epoch': loss_meter.avg,
+                'phase2/train_gap_epoch': gap_meter.avg,
+                'phase2/train_path_loss_epoch': path_meter.avg,
+                **{f'phase2/{k}': v for k, v in val_metrics.items()},
+            })
+
+            if val_metrics['val_selectivity_gap'] > self.best_val_metric:
+                self.best_val_metric = val_metrics['val_selectivity_gap']
+                self.save_checkpoint('best_phase2.pt', extra={'epoch': epoch, 'phase': 2})
+                logger.info(f"  -> New best Phase 2 model (val_gap={self.best_val_metric:.3f})")
+
+        logger.info("Phase 2 (path-aware) training complete.")
+
+    def run_phase2(self, train_loader, val_loader, n_epochs: int = 20,
+                   lambda_rank: float = 1.0, margin: float = 0.2,
+                   lambda_ddg: float = 0.1):
+        """Phase 2 training loop (selectivity fine-tuning + ddG auxiliary)."""
+        logger.info(f"Starting Phase 2 (selectivity fine-tuning) for {n_epochs} epochs "
+                    f"[lambda_rank={lambda_rank}, lambda_ddg={lambda_ddg}]")
+        self.best_val_metric = -float('inf')
+
+        for epoch in range(n_epochs):
+            loss_meter = AverageMeter()
+            gap_meter = AverageMeter()
+
+            for batch in train_loader:
+                result = self.train_step_phase2(batch, lambda_rank, margin, lambda_ddg)
+                B = len(result['pos_scores'])
+                loss_meter.update(result['loss'], B)
+                gap_meter.update(result['selectivity_gap'], B)
+
+                if self.global_step % 50 == 0:
+                    wandb.log({
+                        'phase2/train_loss': result['loss'],
+                        'phase2/loss_margin': result['loss_margin'],
+                        'phase2/loss_infonce': result['loss_infonce'],
+                        'phase2/loss_ddg': result['loss_ddg'],
+                        'phase2/selectivity_gap': result['selectivity_gap'],
+                        'phase2/lr': self.optimizer.param_groups[0]['lr'],
+                        'phase2/step': self.global_step,
+                    })
+
+            # Validate
+            val_metrics = self.evaluate_phase2(val_loader)
+
+            logger.info(
+                f"Phase2 Epoch {epoch+1}/{n_epochs} | "
+                f"Loss: {loss_meter.avg:.4f} | "
+                f"Gap: {gap_meter.avg:.3f} | "
+                f"Val Gap: {val_metrics['val_selectivity_gap']:.3f} | "
+                f"Val Acc: {val_metrics['val_ranking_acc']:.3f}"
+            )
+
+            wandb.log({
+                'phase2/epoch': epoch + 1,
+                'phase2/train_loss_epoch': loss_meter.avg,
+                'phase2/train_gap_epoch': gap_meter.avg,
+                **{f'phase2/{k}': v for k, v in val_metrics.items()},
+            })
+
+            # Checkpoint
+            if val_metrics['val_selectivity_gap'] > self.best_val_metric:
+                self.best_val_metric = val_metrics['val_selectivity_gap']
+                self.save_checkpoint('best_phase2.pt', extra={'epoch': epoch, 'phase': 2})
+                logger.info(f"  -> New best Phase 2 model (val_gap={self.best_val_metric:.3f})")
+
+        logger.info("Phase 2 training complete.")
+
+    @torch.no_grad()
+    def evaluate_phase2(self, loader):
+        """Evaluate selectivity on paired (pos, neg) val set."""
+        self.model.eval()
+        all_pos_scores, all_neg_scores = [], []
+
+        for batch in loader:
+            if 'pos' not in batch:
+                continue
+            pos = batch['pos']
+            neg = batch['neg']
+
+            pos_esm = pos['esm_feats'].to(self.device) if 'esm_feats' in pos else None
+            neg_esm = neg['esm_feats'].to(self.device) if 'esm_feats' in neg else None
+            pos_scores = self.model(
+                pos['node_feats'].to(self.device),
+                pos['edge_feats'].to(self.device),
+                pos['node_mask'].to(self.device),
+                esm_feats=pos_esm
+            )
+            neg_scores = self.model(
+                neg['node_feats'].to(self.device),
+                neg['edge_feats'].to(self.device),
+                neg['node_mask'].to(self.device),
+                esm_feats=neg_esm
+            )
+            all_pos_scores.append(pos_scores.cpu().numpy())
+            all_neg_scores.append(neg_scores.cpu().numpy())
+
+        if not all_pos_scores:
+            return {'val_selectivity_gap': 0.0, 'val_ranking_acc': 0.5}
+
+        all_pos = np.concatenate(all_pos_scores)
+        all_neg = np.concatenate(all_neg_scores)
+
+        gap = float((all_pos - all_neg).mean())
+        acc = float((all_pos > all_neg).mean())
+
+        return {
+            'val_selectivity_gap': gap,
+            'val_ranking_acc': acc,
+            'val_pos_score_mean': float(all_pos.mean()),
+            'val_neg_score_mean': float(all_neg.mean()),
+        }
+
+    # ------------------------------------------------------------------ #
+    # Checkpointing
+    # ------------------------------------------------------------------ #
+
+    def save_checkpoint(self, filename: str, extra: dict = None):
+        path = os.path.join(self.checkpoint_dir, filename)
+        state = {
+            'model_state': self.model.state_dict(),
+            'optimizer_state': self.optimizer.state_dict(),
+            'global_step': self.global_step,
+            'config': self.config,
+        }
+        if extra:
+            state.update(extra)
+        torch.save(state, path)
+        logger.debug(f"Saved checkpoint: {path}")
+
+    def load_checkpoint(self, filename: str):
+        path = os.path.join(self.checkpoint_dir, filename)
+        if not os.path.exists(path):
+            logger.warning(f"Checkpoint not found: {path}")
+            return False
+        state = torch.load(path, map_location=self.device)
+        self.model.load_state_dict(state['model_state'])
+        self.optimizer.load_state_dict(state['optimizer_state'])
+        self.global_step = state.get('global_step', 0)
+        logger.info(f"Loaded checkpoint from {path} (step {self.global_step})")
+        return True
+
+    # ------------------------------------------------------------------ #
+    # Full evaluation (test set)
+    # ------------------------------------------------------------------ #
+
+    @torch.no_grad()
+    def evaluate_test(self, test_loader, phase: int = 2):
+        """Full evaluation on test set with all metrics."""
+        self.model.eval()
+        all_scores, all_labels, all_types = [], [], []
+
+        for batch in test_loader:
+            if 'pos' in batch:
+                # Paired batch
+                for key in ['pos', 'neg']:
+                    d = batch[key]
+                    d_esm = d['esm_feats'].to(self.device) if 'esm_feats' in d else None
+                    scores = self.model(
+                        d['node_feats'].to(self.device),
+                        d['edge_feats'].to(self.device),
+                        d['node_mask'].to(self.device),
+                        esm_feats=d_esm
+                    )
+                    all_scores.extend(scores.cpu().numpy().tolist())
+                    all_labels.extend(d['label'].numpy().tolist())
+                    all_types.extend(['pos' if key == 'pos' else 'neg'] * len(scores))
+            else:
+                esm_feats = batch['esm_feats'].to(self.device) if 'esm_feats' in batch else None
+                scores = self.model(
+                    batch['node_feats'].to(self.device),
+                    batch['edge_feats'].to(self.device),
+                    batch['node_mask'].to(self.device),
+                    esm_feats=esm_feats
+                )
+                all_scores.extend(scores.cpu().numpy().tolist())
+                all_labels.extend(batch['label'].numpy().tolist())
+                all_types.extend(batch['type'])
+
+        all_scores = np.array(all_scores)
+        all_labels = np.array(all_labels)
+
+        metrics = {}
+
+        # Spearman correlation (all samples)
+        metrics['test_spearman'] = float(spearmanr(all_scores, all_labels).correlation or 0)
+
+        # AUC (binary: label > 0.5 = positive quality)
+        binary = (all_labels > 0.5).astype(int)
+        if binary.sum() > 0 and binary.sum() < len(binary):
+            try:
+                metrics['test_auc'] = float(roc_auc_score(binary, all_scores))
+            except Exception:
+                pass
+
+        # Selectivity gap (pos vs neg_apo pairs)
+        pos_mask = np.array([t == 'pos' or t == 'positive' for t in all_types])
+        neg_mask = np.array([t == 'neg' or t == 'negative_apo' for t in all_types])
+        if pos_mask.sum() > 0 and neg_mask.sum() > 0:
+            metrics['test_selectivity_gap'] = float(all_scores[pos_mask].mean() - all_scores[neg_mask].mean())
+
+        logger.info(f"Test evaluation: {metrics}")
+        wandb.log({f'test/{k}': v for k, v in metrics.items()})
+
+        return metrics, all_scores, all_labels, all_types

code/utils/anm.py

@@ -0,0 +1,208 @@
+"""
+Anisotropic Network Model (ANM) for conformational path interpolation.
+
+From-scratch implementation using scipy eigendecomposition.
+Projects the apo→holo displacement onto low-frequency normal modes
+to create physically motivated interpolation paths.
+"""
+
+import numpy as np
+from scipy.linalg import eigh
+
+
+def compute_anm_modes(ca_coords, cutoff=15.0, n_modes=10):
+    """
+    Build elastic network Hessian and compute normal modes via eigendecomposition.
+
+    Args:
+        ca_coords: [N, 3] CA atom coordinates
+        cutoff: distance cutoff for spring connections (Angstroms)
+        n_modes: number of non-trivial modes to return
+
+    Returns:
+        eigenvalues: [n_modes] array of eigenvalues (force constants)
+        eigenvectors: [n_modes, N, 3] mode displacement vectors
+    """
+    N = len(ca_coords)
+    if N < 4:
+        return np.zeros(n_modes), np.zeros((n_modes, N, 3))
+
+    # Build 3N x 3N Hessian with uniform spring constant (gamma=1)
+    H = np.zeros((3 * N, 3 * N), dtype=np.float64)
+
+    for i in range(N):
+        for j in range(i + 1, N):
+            diff = ca_coords[j] - ca_coords[i]
+            dist = np.linalg.norm(diff)
+            if dist > cutoff or dist < 1e-6:
+                continue
+
+            # Outer product of unit displacement vector
+            unit = diff / dist
+            block = np.outer(unit, unit)  # [3, 3]
+
+            # Off-diagonal: H[i,j] = -gamma * (r_ij ⊗ r_ij) / |r_ij|^2
+            # With uniform gamma=1 and unit vectors, this simplifies to:
+            ii, jj = 3 * i, 3 * j
+            H[ii:ii+3, jj:jj+3] = -block
+            H[jj:jj+3, ii:ii+3] = -block
+
+            # Diagonal: accumulate
+            H[ii:ii+3, ii:ii+3] += block
+            H[jj:jj+3, jj:jj+3] += block
+
+    # Eigendecompose — first 6 modes are trivial (3 translation + 3 rotation)
+    n_total = min(6 + n_modes, 3 * N)
+    eigenvalues, eigvecs = eigh(H, subset_by_index=[0, n_total - 1])
+
+    # Skip the 6 trivial zero-frequency modes
+    start = min(6, len(eigenvalues) - 1)
+    n_available = len(eigenvalues) - start
+    n_return = min(n_modes, n_available)
+
+    evals = eigenvalues[start:start + n_return]
+    evecs = eigvecs[:, start:start + n_return]  # [3N, n_return]
+
+    # Reshape eigenvectors to [n_modes, N, 3]
+    mode_vectors = np.zeros((n_return, N, 3))
+    for k in range(n_return):
+        mode_vectors[k] = evecs[:, k].reshape(N, 3)
+
+    # Pad if fewer modes available than requested
+    if n_return < n_modes:
+        pad_evals = np.zeros(n_modes)
+        pad_evals[:n_return] = evals
+        pad_modes = np.zeros((n_modes, N, 3))
+        pad_modes[:n_return] = mode_vectors
+        return pad_evals, pad_modes
+
+    return evals, mode_vectors
+
+
+def _kabsch_align(mobile_ca, ref_ca):
+    """Kabsch alignment of mobile onto ref (CA atoms only)."""
+    t_mobile = mobile_ca.mean(axis=0)
+    t_ref = ref_ca.mean(axis=0)
+
+    m = mobile_ca - t_mobile
+    r = ref_ca - t_ref
+
+    H = m.T @ r
+    U, S, Vt = np.linalg.svd(H)
+    d = np.linalg.det(Vt.T @ U.T)
+    sign = np.array([1.0, 1.0, np.sign(d)])
+    R = Vt.T @ np.diag(sign) @ U.T
+
+    return R, t_mobile, t_ref
+
+
+def _reconstruct_oxygen(coords):
+    """Reconstruct O atom from N, CA, C with ideal C=O geometry."""
+    C_pos = coords[:, 2, :]
+    CA_pos = coords[:, 1, :]
+    C_CA = C_pos - CA_pos
+    C_CA_norm = np.linalg.norm(C_CA, axis=-1, keepdims=True)
+    C_CA_norm = np.maximum(C_CA_norm, 1e-8)
+    O_pos = C_pos + (C_CA / C_CA_norm) * 1.24
+    coords[:, 3, :] = O_pos
+    return coords
+
+
+def anm_backbone_path(coords_x0, coords_x1, mask_x0, mask_x1,
+                       n_frames=5, n_modes=10, cutoff=15.0):
+    """
+    Interpolate backbone along dominant ANM modes from X0 toward X1.
+
+    Low-frequency modes capture global domain motions (e.g., CaM hinge bending),
+    creating physically informed paths where large-scale motions precede local
+    adjustments.
+
+    Args:
+        coords_x0: [N0, 4, 3] backbone coords (N, CA, C, O) for apo state
+        coords_x1: [N1, 4, 3] backbone coords for holo state
+        mask_x0: [N0] bool
+        mask_x1: [N1] bool
+        n_frames: number of intermediate frames (excluding endpoints)
+        n_modes: number of ANM modes to use for projection
+        cutoff: ANM spring cutoff in Angstroms
+
+    Returns:
+        path_frames: list of (coords_tau, mask_tau, tau) tuples
+            Same interface as interpolate_backbone_path
+    """
+    n_common = min(len(coords_x0), len(coords_x1))
+    c0 = coords_x0[:n_common].copy()
+    c1 = coords_x1[:n_common].copy()
+    m0 = mask_x0[:n_common]
+    m1 = mask_x1[:n_common]
+
+    common_mask = m0 & m1
+    if common_mask.sum() < 5:
+        return []
+
+    # Kabsch-align X0 onto X1 using valid CA atoms
+    ca0 = c0[common_mask, 1, :]
+    ca1 = c1[common_mask, 1, :]
+    R, t_mobile, t_ref = _kabsch_align(ca0, ca1)
+
+    # Apply alignment to all X0 backbone atoms
+    flat0 = c0.reshape(-1, 3)
+    aligned0 = (flat0 - t_mobile) @ R.T + t_ref
+    c0_aligned = aligned0.reshape(n_common, 4, 3)
+
+    # Compute apo→holo displacement (CA atoms, valid residues only)
+    ca0_aligned = c0_aligned[common_mask, 1, :]  # [N_valid, 3]
+    ca1_valid = c1[common_mask, 1, :]
+
+    displacement = ca1_valid - ca0_aligned  # [N_valid, 3]
+
+    # Compute ANM modes of the aligned apo structure
+    eigenvalues, mode_vectors = compute_anm_modes(
+        ca0_aligned, cutoff=cutoff, n_modes=n_modes
+    )  # mode_vectors: [n_modes, N_valid, 3]
+
+    # Project displacement onto each mode
+    # d_k = sum_i mode_k[i] . displacement[i]
+    projections = np.zeros(n_modes)
+    for k in range(n_modes):
+        projections[k] = np.sum(mode_vectors[k] * displacement)
+
+    # Reconstruct mode-projected displacement: d_mode = sum_k d_k * mode_k
+    mode_displacement = np.zeros_like(displacement)  # [N_valid, 3]
+    for k in range(n_modes):
+        mode_displacement += projections[k] * mode_vectors[k]
+
+    # Residual displacement not captured by modes
+    residual = displacement - mode_displacement
+
+    # Generate intermediate frames
+    taus = np.linspace(0, 1, n_frames + 2)[1:-1]
+    path_frames = []
+
+    for tau in taus:
+        # Apply mode-projected + residual displacement at each tau
+        # Mode component applies smoothly; residual is linear
+        ca_interp = ca0_aligned + tau * mode_displacement + tau * residual
+
+        # Build full backbone by interpolating all 4 atom types
+        coords_tau = (1.0 - tau) * c0_aligned + tau * c1
+        # Override CA positions with ANM-interpolated values
+        coords_tau[common_mask, 1, :] = ca_interp
+
+        # Adjust N, C positions relative to CA shift
+        # The N/CA/C triangle is preserved by blending the ANM CA shift
+        # with the linear interpolation of N and C
+        ca_shift = ca_interp - ((1.0 - tau) * ca0_aligned + tau * ca1_valid)
+        coords_tau[common_mask, 0, :] += ca_shift  # N atoms
+        coords_tau[common_mask, 2, :] += ca_shift  # C atoms
+
+        # Reconstruct O from N, CA, C
+        coords_tau = _reconstruct_oxygen(coords_tau)
+
+        path_frames.append((
+            coords_tau.astype(np.float32),
+            common_mask.copy(),
+            float(tau),
+        ))
+
+    return path_frames

code/utils/path_utils.py

@@ -0,0 +1,448 @@
+"""
+Transition-path interpolation utilities for conformational induction.
+
+Provides:
+  - Kabsch-aligned backbone interpolation between two conformational states
+  - Gaussian Schrödinger Bridge (DSB) stochastic interpolation
+  - Precomputed frame loading (for AlphaFlow / AFsample2)
+  - Unified dispatcher: generate_path_frames()
+  - Per-residue displacement computation (for allosteric hinge weighting)
+  - Monotonically increasing path weight generation
+
+Used by the path-aware training, guidance, and refinement modules.
+"""
+
+import os
+import logging
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+def _kabsch_align(mobile_ca, ref_ca):
+    """
+    Kabsch alignment of mobile onto ref (CA atoms only).
+
+    Args:
+        mobile_ca: [N, 3] array
+        ref_ca: [N, 3] array
+
+    Returns:
+        R: [3, 3] rotation matrix
+        t_mobile: [3] mobile centroid
+        t_ref: [3] ref centroid
+        Such that: aligned = (mobile - t_mobile) @ R.T + t_ref
+    """
+    t_mobile = mobile_ca.mean(axis=0)
+    t_ref = ref_ca.mean(axis=0)
+
+    m = mobile_ca - t_mobile
+    r = ref_ca - t_ref
+
+    H = m.T @ r
+    U, S, Vt = np.linalg.svd(H)
+    d = np.linalg.det(Vt.T @ U.T)
+    sign = np.array([1.0, 1.0, np.sign(d)])
+    R = Vt.T @ np.diag(sign) @ U.T
+
+    return R, t_mobile, t_ref
+
+
+def interpolate_backbone_path(coords_x0, coords_x1, mask_x0, mask_x1, n_frames=5):
+    """
+    Generate intermediate backbone conformations along the X0 -> X1 path.
+
+    1. Find common valid residues between X0 and X1
+    2. Kabsch-align X0 onto X1 using CA atoms
+    3. Linearly interpolate backbone coords at n_frames equally-spaced tau values
+    4. Reconstruct O from N/CA/C with ideal geometry
+
+    Args:
+        coords_x0: [N0, 4, 3] backbone coords (N, CA, C, O) for state 0
+        coords_x1: [N1, 4, 3] backbone coords for state 1
+        mask_x0: [N0] bool
+        mask_x1: [N1] bool
+        n_frames: number of intermediate frames (excluding endpoints)
+
+    Returns:
+        path_frames: list of (coords_tau, mask_tau, tau) tuples
+            coords_tau: [N_common, 4, 3] interpolated backbone coords
+            mask_tau: [N_common] bool
+            tau: float in (0, 1) exclusive
+    """
+    # Use common length
+    n_common = min(len(coords_x0), len(coords_x1))
+    c0 = coords_x0[:n_common].copy()
+    c1 = coords_x1[:n_common].copy()
+    m0 = mask_x0[:n_common]
+    m1 = mask_x1[:n_common]
+
+    # Valid in both states
+    common_mask = m0 & m1
+    if common_mask.sum() < 5:
+        return []
+
+    # Kabsch-align X0 onto X1 using valid CA atoms
+    ca0 = c0[common_mask, 1, :]  # CA atoms
+    ca1 = c1[common_mask, 1, :]
+
+    R, t_mobile, t_ref = _kabsch_align(ca0, ca1)
+
+    # Apply alignment to all X0 backbone atoms
+    n_res = n_common
+    flat0 = c0.reshape(-1, 3)
+    aligned0 = (flat0 - t_mobile) @ R.T + t_ref
+    c0_aligned = aligned0.reshape(n_res, 4, 3)
+
+    # Generate intermediate frames
+    taus = np.linspace(0, 1, n_frames + 2)[1:-1]  # exclude endpoints
+    path_frames = []
+
+    for tau in taus:
+        # Linear interpolation: X_tau = (1 - tau) * X0_aligned + tau * X1
+        coords_tau = (1.0 - tau) * c0_aligned + tau * c1
+
+        # Reconstruct O from N, CA, C with ideal C=O bond geometry
+        C_pos = coords_tau[:, 2, :]   # C atoms
+        CA_pos = coords_tau[:, 1, :]  # CA atoms
+        C_CA = C_pos - CA_pos
+        C_CA_norm = np.linalg.norm(C_CA, axis=-1, keepdims=True)
+        C_CA_norm = np.maximum(C_CA_norm, 1e-8)
+        O_pos = C_pos + (C_CA / C_CA_norm) * 1.24  # ideal C=O bond length
+        coords_tau[:, 3, :] = O_pos
+
+        path_frames.append((
+            coords_tau.astype(np.float32),
+            common_mask.copy(),
+            float(tau),
+        ))
+
+    return path_frames
+
+
+def compute_residue_displacements(coords_x0, coords_x1, mask_x0, mask_x1):
+    """
+    Per-residue CA displacement between X0 and X1 after Kabsch alignment.
+
+    Args:
+        coords_x0: [N0, 4, 3] backbone coords for state 0
+        coords_x1: [N1, 4, 3] backbone coords for state 1
+        mask_x0: [N0] bool
+        mask_x1: [N1] bool
+
+    Returns:
+        displacements: [N_common] array of per-residue CA RMSD
+        common_mask: [N_common] bool — which residues are valid
+    """
+    n_common = min(len(coords_x0), len(coords_x1))
+    c0 = coords_x0[:n_common]
+    c1 = coords_x1[:n_common]
+    m0 = mask_x0[:n_common]
+    m1 = mask_x1[:n_common]
+    common_mask = m0 & m1
+
+    if common_mask.sum() < 5:
+        return np.zeros(n_common), common_mask
+
+    ca0 = c0[common_mask, 1, :]
+    ca1 = c1[common_mask, 1, :]
+
+    R, t_mobile, t_ref = _kabsch_align(ca0, ca1)
+
+    # Align all CA of X0
+    all_ca0 = c0[:, 1, :]
+    aligned_ca0 = (all_ca0 - t_mobile) @ R.T + t_ref
+
+    # Per-residue displacement
+    all_ca1 = c1[:, 1, :]
+    displacements = np.linalg.norm(aligned_ca0 - all_ca1, axis=-1)
+
+    # Zero out invalid residues
+    displacements[~common_mask] = 0.0
+
+    return displacements.astype(np.float32), common_mask
+
+
+def generate_path_weights(n_frames, mode='linear'):
+    """
+    Generate monotonically increasing weights for path frames.
+
+    The weights increase toward tau=1 (the goal state), so that
+    intermediate conformations closer to X1 are weighted more heavily.
+
+    Args:
+        n_frames: number of intermediate frames
+        mode: weight schedule
+            'linear': w_tau = tau
+            'quadratic': w_tau = tau^2
+            'exponential': w_tau = (exp(tau) - 1) / (e - 1)
+            'uniform': w_tau = 1/n_frames (equal weighting)
+
+    Returns:
+        weights: [n_frames] numpy array, normalized to sum to 1
+    """
+    if n_frames == 0:
+        return np.array([], dtype=np.float32)
+
+    taus = np.linspace(0, 1, n_frames + 2)[1:-1]  # same as interpolation
+
+    if mode == 'linear':
+        weights = taus.copy()
+    elif mode == 'quadratic':
+        weights = taus ** 2
+    elif mode == 'exponential':
+        weights = (np.exp(taus) - 1.0) / (np.e - 1.0)
+    elif mode == 'uniform':
+        weights = np.ones(n_frames, dtype=np.float32)
+    else:
+        raise ValueError(f"Unknown weight mode: {mode}")
+
+    # Normalize to sum to 1
+    total = weights.sum()
+    if total > 0:
+        weights = weights / total
+
+    return weights.astype(np.float32)
+
+
+# ---------------------------------------------------------------------------
+# Gaussian Schrödinger Bridge (AlignDSB) interpolation
+# ---------------------------------------------------------------------------
+
+def dsb_backbone_path(coords_x0, coords_x1, mask_x0, mask_x1,
+                       n_frames=5, sigma=0.5, n_samples=20, seed=42):
+    """
+    Gaussian Schrödinger Bridge with t*(1-t) variance schedule.
+
+    Analytic formula (no neural network):
+        X_t = (1-t) * X0_aligned + t * X1 + sqrt(t * (1-t)) * sigma * Z
+
+    Variance peaks at t=0.5 (maximum uncertainty mid-transition) and vanishes
+    at endpoints. sigma controls noise amplitude in Angstroms.
+
+    For each tau, samples n_samples noisy interpolations and selects the
+    median (by RMSD to the mean) for robustness.
+
+    Args:
+        coords_x0: [N0, 4, 3] backbone coords for state 0
+        coords_x1: [N1, 4, 3] backbone coords for state 1
+        mask_x0: [N0] bool
+        mask_x1: [N1] bool
+        n_frames: number of intermediate frames
+        sigma: noise amplitude (Angstroms)
+        n_samples: number of samples per frame for median selection
+        seed: random seed
+
+    Returns:
+        path_frames: list of (coords_tau, mask_tau, tau) tuples
+    """
+    rng = np.random.RandomState(seed)
+
+    n_common = min(len(coords_x0), len(coords_x1))
+    c0 = coords_x0[:n_common].copy()
+    c1 = coords_x1[:n_common].copy()
+    m0 = mask_x0[:n_common]
+    m1 = mask_x1[:n_common]
+
+    common_mask = m0 & m1
+    if common_mask.sum() < 5:
+        return []
+
+    # Kabsch-align X0 onto X1
+    ca0 = c0[common_mask, 1, :]
+    ca1 = c1[common_mask, 1, :]
+    R, t_mobile, t_ref = _kabsch_align(ca0, ca1)
+
+    flat0 = c0.reshape(-1, 3)
+    aligned0 = (flat0 - t_mobile) @ R.T + t_ref
+    c0_aligned = aligned0.reshape(n_common, 4, 3)
+
+    taus = np.linspace(0, 1, n_frames + 2)[1:-1]
+    path_frames = []
+
+    for tau in taus:
+        noise_scale = np.sqrt(tau * (1.0 - tau)) * sigma
+
+        # Generate n_samples noisy interpolations
+        samples = []
+        for _ in range(n_samples):
+            Z = rng.randn(n_common, 4, 3).astype(np.float64)
+            X_t = (1.0 - tau) * c0_aligned + tau * c1 + noise_scale * Z
+            samples.append(X_t)
+
+        samples = np.array(samples)  # [n_samples, N, 4, 3]
+        mean_sample = samples.mean(axis=0)  # [N, 4, 3]
+
+        # Select median sample by RMSD to mean (CA atoms)
+        rmsds = []
+        for s in samples:
+            diff = s[common_mask, 1, :] - mean_sample[common_mask, 1, :]
+            rmsd = np.sqrt((diff ** 2).sum() / common_mask.sum())
+            rmsds.append(rmsd)
+        median_idx = np.argsort(rmsds)[len(rmsds) // 2]
+        coords_tau = samples[median_idx]
+
+        # Reconstruct O from N, CA, C
+        C_pos = coords_tau[:, 2, :]
+        CA_pos = coords_tau[:, 1, :]
+        C_CA = C_pos - CA_pos
+        C_CA_norm = np.linalg.norm(C_CA, axis=-1, keepdims=True)
+        C_CA_norm = np.maximum(C_CA_norm, 1e-8)
+        coords_tau[:, 3, :] = C_pos + (C_CA / C_CA_norm) * 1.24
+
+        path_frames.append((
+            coords_tau.astype(np.float32),
+            common_mask.copy(),
+            float(tau),
+        ))
+
+    return path_frames
+
+
+# ---------------------------------------------------------------------------
+# Precomputed frame loading (for AlphaFlow / AFsample2)
+# ---------------------------------------------------------------------------
+
+def load_precomputed_frames(target, method, precomputed_dir,
+                             coords_x0, coords_x1, mask_x0, mask_x1,
+                             n_frames=5):
+    """
+    Load pre-generated frames from .npz and Kabsch-align to this complex's
+    receptor coordinate frame.
+
+    Expected file: {precomputed_dir}/{target}/{method}/frames.npz
+    with keys: 'frames' [n_frames, N_ref, 4, 3], 'taus' [n_frames],
+               'mask' [N_ref] bool
+
+    Args:
+        target: target name (e.g. 'cam')
+        method: method name ('alphaflow' or 'afsample2')
+        precomputed_dir: root directory for precomputed frames
+        coords_x0, coords_x1: apo/holo backbone coords for alignment
+        mask_x0, mask_x1: residue masks
+        n_frames: number of frames to return
+
+    Returns:
+        path_frames: list of (coords_tau, mask_tau, tau) tuples
+    """
+    npz_path = os.path.join(precomputed_dir, target, method, 'frames.npz')
+    if not os.path.exists(npz_path):
+        logger.warning(f"Precomputed frames not found: {npz_path}, "
+                       f"falling back to linear interpolation")
+        return interpolate_backbone_path(coords_x0, coords_x1,
+                                          mask_x0, mask_x1, n_frames)
+
+    data = np.load(npz_path)
+    pre_frames = data['frames']   # [K, N_ref, 4, 3]
+    pre_taus = data['taus']       # [K]
+    pre_mask = data['mask']       # [N_ref]
+
+    n_common = min(len(coords_x0), len(coords_x1), len(pre_mask))
+    m0 = mask_x0[:n_common]
+    m1 = mask_x1[:n_common]
+    pm = pre_mask[:n_common]
+    common_mask = m0 & m1 & pm
+
+    if common_mask.sum() < 5:
+        logger.warning(f"Too few common residues for {target}/{method}, "
+                       f"falling back to linear")
+        return interpolate_backbone_path(coords_x0, coords_x1,
+                                          mask_x0, mask_x1, n_frames)
+
+    # Align precomputed frames to the holo receptor (X1) coordinate frame
+    # The precomputed frames were generated from the reference apo sequence
+    # and may be in a different coordinate frame
+    ref_ca = coords_x1[:n_common][common_mask, 1, :]  # holo CA as reference
+
+    path_frames = []
+    K = min(len(pre_frames), n_frames)
+
+    # Select n_frames evenly spaced from available frames
+    if len(pre_frames) > n_frames:
+        indices = np.linspace(0, len(pre_frames) - 1, n_frames).astype(int)
+    else:
+        indices = np.arange(K)
+
+    for idx in indices:
+        frame = pre_frames[idx, :n_common].copy()  # [N_common, 4, 3]
+        tau = float(pre_taus[idx])
+
+        # Kabsch-align frame CA to holo CA
+        frame_ca = frame[common_mask, 1, :]
+        R, t_frame, t_ref = _kabsch_align(frame_ca, ref_ca)
+
+        flat_frame = frame.reshape(-1, 3)
+        aligned = (flat_frame - t_frame) @ R.T + t_ref
+        frame_aligned = aligned.reshape(n_common, 4, 3)
+
+        # Reconstruct O
+        C_pos = frame_aligned[:, 2, :]
+        CA_pos = frame_aligned[:, 1, :]
+        C_CA = C_pos - CA_pos
+        C_CA_norm = np.linalg.norm(C_CA, axis=-1, keepdims=True)
+        C_CA_norm = np.maximum(C_CA_norm, 1e-8)
+        frame_aligned[:, 3, :] = C_pos + (C_CA / C_CA_norm) * 1.24
+
+        path_frames.append((
+            frame_aligned.astype(np.float32),
+            common_mask.copy(),
+            tau,
+        ))
+
+    return path_frames
+
+
+# ---------------------------------------------------------------------------
+# Unified dispatcher
+# ---------------------------------------------------------------------------
+
+def generate_path_frames(coords_x0, coords_x1, mask_x0, mask_x1,
+                          method='linear', n_frames=5,
+                          precomputed_dir=None, target=None, **kwargs):
+    """
+    Dispatch to method-specific frame generation.
+
+    Args:
+        coords_x0, coords_x1: [N, 4, 3] backbone coords for apo/holo
+        mask_x0, mask_x1: [N] bool masks
+        method: one of 'linear', 'alphaflow', 'afsample2', 'dsb', 'anm'
+        n_frames: number of intermediate frames
+        precomputed_dir: directory for precomputed frames (alphaflow/afsample2)
+        target: target name (needed for precomputed methods)
+        **kwargs: method-specific parameters (sigma, n_modes, etc.)
+
+    Returns:
+        path_frames: list of (coords_tau, mask_tau, tau) tuples
+    """
+    if method == 'linear':
+        return interpolate_backbone_path(
+            coords_x0, coords_x1, mask_x0, mask_x1, n_frames)
+
+    elif method in ('alphaflow', 'afsample2'):
+        if precomputed_dir is None:
+            raise ValueError(f"precomputed_dir required for method '{method}'")
+        if target is None:
+            raise ValueError(f"target name required for method '{method}'")
+        return load_precomputed_frames(
+            target, method, precomputed_dir,
+            coords_x0, coords_x1, mask_x0, mask_x1, n_frames)
+
+    elif method == 'dsb':
+        return dsb_backbone_path(
+            coords_x0, coords_x1, mask_x0, mask_x1,
+            n_frames=n_frames,
+            sigma=kwargs.get('sigma', 0.5),
+            n_samples=kwargs.get('n_samples', 20),
+            seed=kwargs.get('seed', 42))
+
+    elif method == 'anm':
+        from utils.anm import anm_backbone_path
+        return anm_backbone_path(
+            coords_x0, coords_x1, mask_x0, mask_x1,
+            n_frames=n_frames,
+            n_modes=kwargs.get('n_modes', 10),
+            cutoff=kwargs.get('cutoff', 15.0))
+
+    else:
+        raise ValueError(f"Unknown path method: '{method}'. "
+                         f"Choose from: linear, alphaflow, afsample2, dsb, anm")

code/utils/pdb_utils.py

@@ -0,0 +1,472 @@
+"""
+PDB parsing utilities for Allo-Designer.
+Extracts backbone geometry, computes local frames, and identifies interface residues.
+"""
+
+import numpy as np
+from Bio import PDB
+from Bio.PDB import PDBParser, MMCIFParser, PDBIO
+from Bio.PDB.Polypeptide import is_aa
+import warnings
+warnings.filterwarnings("ignore", category=PDB.PDBExceptions.PDBConstructionWarning)
+
+AA3_TO_IDX = {
+    'ALA': 0, 'ARG': 1, 'ASN': 2, 'ASP': 3, 'CYS': 4,
+    'GLN': 5, 'GLU': 6, 'GLY': 7, 'HIS': 8, 'ILE': 9,
+    'LEU': 10, 'LYS': 11, 'MET': 12, 'PHE': 13, 'PRO': 14,
+    'SER': 15, 'THR': 16, 'TRP': 17, 'TYR': 18, 'VAL': 19,
+    'UNK': 20,
+}
+NUM_AA = 21  # 20 standard + UNK
+
+
+def load_structure(pdb_path: str, model_id: int = 0):
+    """Load a PDB/CIF file and return the first model."""
+    if pdb_path.endswith('.cif') or pdb_path.endswith('.mmcif'):
+        parser = MMCIFParser(QUIET=True)
+    else:
+        parser = PDBParser(QUIET=True)
+    struct = parser.get_structure("protein", pdb_path)
+    return list(struct.get_models())[model_id]
+
+
+def get_residues(chain, only_standard: bool = True):
+    """Return a list of standard amino acid residues from a chain."""
+    residues = []
+    for res in chain.get_residues():
+        if only_standard and not is_aa(res, standard=True):
+            continue
+        if res.get_id()[0] != ' ':  # skip HETATM
+            continue
+        residues.append(res)
+    return residues
+
+
+def get_backbone_coords(residues):
+    """
+    Extract backbone atom coordinates (N, CA, C, O) for each residue.
+    Returns: coords [N_res, 4, 3], mask [N_res] (True = all backbone atoms present)
+    """
+    N = len(residues)
+    coords = np.zeros((N, 4, 3), dtype=np.float32)
+    mask = np.zeros(N, dtype=bool)
+
+    for i, res in enumerate(residues):
+        try:
+            coords[i, 0] = res['N'].get_vector().get_array()
+            coords[i, 1] = res['CA'].get_vector().get_array()
+            coords[i, 2] = res['C'].get_vector().get_array()
+            if 'O' in res:
+                coords[i, 3] = res['O'].get_vector().get_array()
+            else:
+                # Estimate O position if missing
+                coords[i, 3] = coords[i, 2]
+            mask[i] = True
+        except KeyError:
+            pass
+    return coords, mask
+
+
+def get_aa_indices(residues):
+    """Return integer amino acid indices for each residue."""
+    return np.array([
+        AA3_TO_IDX.get(res.get_resname(), AA3_TO_IDX['UNK'])
+        for res in residues
+    ], dtype=np.int64)
+
+
+def compute_backbone_frames(coords, mask):
+    """
+    Compute SE(3)-equivariant backbone frames from N, CA, C atoms.
+    Frame: z-axis = CA->C, y-axis = component of CA->N perpendicular to z, x-axis = y x z.
+
+    Returns:
+        origins: [N, 3] = CA positions
+        rotations: [N, 3, 3] = rotation matrices (columns are x, y, z axes)
+    """
+    N_res = coords.shape[0]
+    origins = coords[:, 1, :]  # CA positions [N, 3]
+    rotations = np.zeros((N_res, 3, 3), dtype=np.float32)
+
+    for i in range(N_res):
+        if not mask[i]:
+            rotations[i] = np.eye(3)
+            continue
+        ca = coords[i, 1]
+        n = coords[i, 0]
+        c = coords[i, 2]
+
+        # z-axis: CA -> C
+        z = c - ca
+        z_norm = np.linalg.norm(z)
+        if z_norm < 1e-6:
+            rotations[i] = np.eye(3)
+            continue
+        z = z / z_norm
+
+        # y-axis: CA -> N, orthogonalized
+        y = n - ca
+        y = y - np.dot(y, z) * z
+        y_norm = np.linalg.norm(y)
+        if y_norm < 1e-6:
+            rotations[i] = np.eye(3)
+            continue
+        y = y / y_norm
+
+        # x-axis: y cross z
+        x = np.cross(y, z)
+
+        rotations[i] = np.stack([x, y, z], axis=-1)  # columns are axes
+
+    return origins, rotations
+
+
+def compute_torsion_angles(coords, mask):
+    """
+    Compute backbone torsion angles (phi, psi, omega) for each residue.
+    Returns sin/cos of each angle. [N, 6]
+    """
+    N = len(coords)
+    angles = np.zeros((N, 6), dtype=np.float32)
+
+    def dihedral(p0, p1, p2, p3):
+        """Praxelis dihedral angle computation."""
+        b1 = p1 - p0
+        b2 = p2 - p1
+        b3 = p3 - p2
+        n1 = np.cross(b1, b2)
+        n2 = np.cross(b2, b3)
+        n1_norm = np.linalg.norm(n1)
+        n2_norm = np.linalg.norm(n2)
+        if n1_norm < 1e-6 or n2_norm < 1e-6:
+            return 0.0
+        n1 = n1 / n1_norm
+        n2 = n2 / n2_norm
+        m1 = np.cross(n1, b2 / (np.linalg.norm(b2) + 1e-8))
+        cos_a = np.clip(np.dot(n1, n2), -1, 1)
+        sin_a = np.dot(m1, n2)
+        return np.arctan2(sin_a, cos_a)
+
+    for i in range(N):
+        if not mask[i]:
+            continue
+        ca_i = coords[i, 1]
+        n_i = coords[i, 0]
+        c_i = coords[i, 2]
+
+        # Phi: C_{i-1} - N_i - CA_i - C_i
+        if i > 0 and mask[i - 1]:
+            c_prev = coords[i - 1, 2]
+            phi = dihedral(c_prev, n_i, ca_i, c_i)
+            angles[i, 0] = np.sin(phi)
+            angles[i, 1] = np.cos(phi)
+
+        # Psi: N_i - CA_i - C_i - N_{i+1}
+        if i < N - 1 and mask[i + 1]:
+            n_next = coords[i + 1, 0]
+            psi = dihedral(n_i, ca_i, c_i, n_next)
+            angles[i, 2] = np.sin(psi)
+            angles[i, 3] = np.cos(psi)
+
+        # Omega: CA_{i-1} - C_{i-1} - N_i - CA_i
+        if i > 0 and mask[i - 1]:
+            ca_prev = coords[i - 1, 1]
+            c_prev = coords[i - 1, 2]
+            omega = dihedral(ca_prev, c_prev, n_i, ca_i)
+            angles[i, 4] = np.sin(omega)
+            angles[i, 5] = np.cos(omega)
+
+    return angles
+
+
+def get_interface_residues(rec_coords, binder_coords, rec_mask, binder_mask, cutoff: float = 8.0):
+    """
+    Find interface residues: receptor residues within cutoff of any binder Cα, and vice versa.
+    Uses CA-CA distances.
+
+    Returns:
+        rec_interface: bool array [N_rec]
+        binder_interface: bool array [N_binder]
+    """
+    rec_ca = rec_coords[:, 1, :]    # [N_rec, 3]
+    binder_ca = binder_coords[:, 1, :]  # [N_binder, 3]
+
+    # Pairwise CA-CA distances [N_rec, N_binder]
+    diff = rec_ca[:, None, :] - binder_ca[None, :, :]  # [N_rec, N_binder, 3]
+    dist = np.sqrt((diff ** 2).sum(axis=-1))  # [N_rec, N_binder]
+
+    # Mask out residues without coordinates
+    dist[~rec_mask, :] = np.inf
+    dist[:, ~binder_mask] = np.inf
+
+    rec_interface = (dist < cutoff).any(axis=1)
+    binder_interface = (dist < cutoff).any(axis=0)
+
+    return rec_interface, binder_interface
+
+
+def align_structures(mobile_ca, ref_ca, mobile_coords=None):
+    """
+    Kabsch alignment: align mobile to ref using CA positions.
+    Returns aligned CA coords and optionally full backbone coords.
+    """
+    assert mobile_ca.shape == ref_ca.shape, "Must have same number of residues"
+
+    # Center
+    mobile_center = mobile_ca.mean(axis=0)
+    ref_center = ref_ca.mean(axis=0)
+    m = mobile_ca - mobile_center
+    r = ref_ca - ref_center
+
+    # SVD
+    H = m.T @ r
+    U, S, Vt = np.linalg.svd(H)
+    d = np.sign(np.linalg.det(Vt.T @ U.T))
+    D = np.diag([1, 1, d])
+    R = Vt.T @ D @ U.T  # rotation matrix
+
+    mobile_ca_aligned = (m @ R.T) + ref_center
+
+    if mobile_coords is not None:
+        # Apply same rotation to full backbone
+        N_res, N_atoms, _ = mobile_coords.shape
+        flat = mobile_coords.reshape(-1, 3) - mobile_center
+        aligned_flat = (flat @ R.T) + ref_center
+        mobile_coords_aligned = aligned_flat.reshape(N_res, N_atoms, 3)
+        return mobile_ca_aligned, R, mobile_coords_aligned
+
+    return mobile_ca_aligned, R
+
+
+def compute_ca_rmsd(coords1, coords2, mask=None):
+    """Compute CA-RMSD between two sets of backbone coordinates."""
+    ca1 = coords1[:, 1, :]
+    ca2 = coords2[:, 1, :]
+    if mask is not None:
+        ca1 = ca1[mask]
+        ca2 = ca2[mask]
+    diff = ca1 - ca2
+    return np.sqrt((diff ** 2).sum(axis=-1).mean())
+
+
+def compute_fraction_native_contacts(
+    native_rec_ca, native_binder_ca,
+    model_rec_ca=None, model_binder_ca=None,
+    cutoff=8.0,
+    # Legacy 2-arg signature support
+    mask=None, delta=1.0,
+):
+    """
+    Compute fraction of native inter-chain contacts (fNAT).
+
+    fNAT = |recovered inter-chain contacts| / |native inter-chain contacts|
+
+    A native contact is a (receptor_i, binder_j) pair with CA-CA distance
+    < cutoff in the native complex.  A contact is "recovered" if the same
+    pair is < cutoff in the model complex.
+
+    Args:
+        native_rec_ca:    [N_rec, 3]   receptor CA coords in native complex
+        native_binder_ca: [N_bind, 3]  binder CA coords in native complex
+        model_rec_ca:     [N_rec, 3]   receptor CA in model (default: same as native)
+        model_binder_ca:  [N_bind, 3]  binder CA in model (default: same as native)
+        cutoff: contact distance threshold in Angstroms (default 8.0 for CA-CA)
+
+    Returns:
+        fNAT in [0, 1].  Returns 0.0 if no native contacts exist.
+    """
+    if model_rec_ca is None:
+        model_rec_ca = native_rec_ca
+    if model_binder_ca is None:
+        model_binder_ca = native_binder_ca
+
+    # Inter-chain distance matrices  [N_rec, N_bind]
+    native_dist = np.sqrt(
+        ((native_rec_ca[:, None, :] - native_binder_ca[None, :, :]) ** 2).sum(-1)
+    )
+    model_dist = np.sqrt(
+        ((model_rec_ca[:, None, :] - model_binder_ca[None, :, :]) ** 2).sum(-1)
+    )
+
+    native_contacts = native_dist < cutoff
+    recovered = native_contacts & (model_dist < cutoff)
+
+    n_native = native_contacts.sum()
+    if n_native == 0:
+        return 0.0
+    return float(recovered.sum()) / float(n_native)
+
+
+def rbf_encode(distances, d_min=0.0, d_max=20.0, n_bins=16):
+    """
+    RBF encoding of distances using Gaussian basis functions.
+    Returns: [*distances.shape, n_bins]
+    """
+    centers = np.linspace(d_min, d_max, n_bins)
+    sigma = (d_max - d_min) / (n_bins - 1)
+    encoded = np.exp(-((distances[..., None] - centers) ** 2) / (2 * sigma ** 2))
+    return encoded.astype(np.float32)
+
+
+# Candidate sidechain atoms for chi1 (first atom after CB)
+_CHI1_ATOMS = ['CG', 'CG1', 'OG', 'OG1', 'SG']
+# Candidate sidechain atoms for chi2 (second dihedral: CA-CB-XG-XD)
+_CHI2_ATOMS = ['CD', 'CD1', 'SD', 'OD1', 'ND1', 'CE', 'NE', 'OE1']
+
+
+def _dihedral_4pts(p0, p1, p2, p3):
+    """Compute dihedral angle between four 3D points (radians)."""
+    b1 = p1 - p0
+    b2 = p2 - p1
+    b3 = p3 - p2
+    n1 = np.cross(b1, b2)
+    n2 = np.cross(b2, b3)
+    n1_norm = np.linalg.norm(n1)
+    n2_norm = np.linalg.norm(n2)
+    if n1_norm < 1e-6 or n2_norm < 1e-6:
+        return 0.0
+    n1 = n1 / n1_norm
+    n2 = n2 / n2_norm
+    m1 = np.cross(n1, b2 / (np.linalg.norm(b2) + 1e-8))
+    return np.arctan2(np.dot(m1, n2), np.dot(n1, n2))
+
+
+def compute_chi_angles(residues, mask):
+    """
+    Compute chi1 and chi2 sidechain torsion angles for each residue.
+
+    Chi1: N - CA - CB - XG  (first sidechain dihedral)
+    Chi2: CA - CB - XG - XD (second sidechain dihedral)
+
+    For residues lacking the atoms (Gly, or missing coordinates), returns zeros.
+
+    Returns:
+        chi_feats: [N, 4]  (sin_chi1, cos_chi1, sin_chi2, cos_chi2)
+    """
+    N = len(residues)
+    chi_feats = np.zeros((N, 4), dtype=np.float32)
+
+    for i, res in enumerate(residues):
+        if not mask[i]:
+            continue
+        atoms = {atom.get_name(): atom.get_vector().get_array() for atom in res.get_atoms()
+                 if atom.get_name() in ('N', 'CA', 'CB') + tuple(_CHI1_ATOMS) + tuple(_CHI2_ATOMS)}
+
+        n_pos = atoms.get('N')
+        ca_pos = atoms.get('CA')
+        cb_pos = atoms.get('CB')
+
+        if n_pos is None or ca_pos is None or cb_pos is None:
+            continue
+
+        # Chi1: N - CA - CB - XG
+        xg_pos = None
+        for aname in _CHI1_ATOMS:
+            if aname in atoms:
+                xg_pos = atoms[aname]
+                break
+
+        if xg_pos is not None:
+            chi1 = _dihedral_4pts(np.array(n_pos), np.array(ca_pos),
+                                   np.array(cb_pos), np.array(xg_pos))
+            chi_feats[i, 0] = np.sin(chi1)
+            chi_feats[i, 1] = np.cos(chi1)
+
+            # Chi2: CA - CB - XG - XD
+            xd_pos = None
+            for aname in _CHI2_ATOMS:
+                if aname in atoms:
+                    xd_pos = atoms[aname]
+                    break
+
+            if xd_pos is not None:
+                chi2 = _dihedral_4pts(np.array(ca_pos), np.array(cb_pos),
+                                       np.array(xg_pos), np.array(xd_pos))
+                chi_feats[i, 2] = np.sin(chi2)
+                chi_feats[i, 3] = np.cos(chi2)
+
+    return chi_feats
+
+
+def get_cb_positions(residues, coords, mask):
+    """
+    Return CB positions for each residue (CA position for Gly or missing CB).
+
+    Returns:
+        cb_pos: [N, 3]
+    """
+    N = len(residues)
+    cb_pos = coords[:, 1, :].copy()  # default to CA
+
+    for i, res in enumerate(residues):
+        if not mask[i]:
+            continue
+        try:
+            cb_pos[i] = res['CB'].get_vector().get_array()
+        except KeyError:
+            pass  # Gly or missing CB: keep CA
+
+    return cb_pos.astype(np.float32)
+
+
+# Simplified hydrophobicity groups for contact energy
+_HYDROPHOBIC = {'ALA', 'VAL', 'ILE', 'LEU', 'MET', 'PHE', 'TRP', 'PRO', 'TYR'}
+_POS_CHARGED = {'ARG', 'LYS', 'HIS'}
+_NEG_CHARGED = {'ASP', 'GLU'}
+
+
+def _residue_group(resname):
+    if resname in _HYDROPHOBIC:
+        return 'H'
+    if resname in _POS_CHARGED:
+        return '+'
+    if resname in _NEG_CHARGED:
+        return '-'
+    return 'P'  # polar
+
+
+def compute_contact_energy(rec_residues, binder_residues,
+                           rec_cb, binder_cb,
+                           rec_mask, binder_mask,
+                           cutoff: float = 8.0):
+    """
+    Compute a simple CB-CB contact energy as a physics-based ddG proxy.
+
+    Uses a 4-group hydrophobicity potential:
+      HH: -1.0 (hydrophobic-hydrophobic, favorable)
+      +-: -0.5 (opposite charges, favorable)
+      H+/-: +0.3 (hydrophobic-charged, unfavorable)
+      else: 0.0
+
+    Returns a scalar in [0, 1] via sigmoid normalization.
+    """
+    n_rec = len(rec_residues)
+    n_binder = len(binder_residues)
+
+    # CB-CB distance matrix [n_rec, n_binder]
+    diff = rec_cb[:, None, :] - binder_cb[None, :, :]     # [n_rec, n_binder, 3]
+    dist = np.sqrt((diff ** 2).sum(axis=-1))               # [n_rec, n_binder]
+
+    # Mask invalid residues
+    dist[~rec_mask, :] = np.inf
+    dist[:, ~binder_mask] = np.inf
+
+    contact_mask = dist < cutoff
+
+    energy = 0.0
+    for i in range(n_rec):
+        for j in range(n_binder):
+            if not contact_mask[i, j]:
+                continue
+            gi = _residue_group(rec_residues[i].get_resname())
+            gj = _residue_group(binder_residues[j].get_resname())
+            if gi == 'H' and gj == 'H':
+                energy -= 1.0
+            elif (gi == '+' and gj == '-') or (gi == '-' and gj == '+'):
+                energy -= 0.5
+            elif (gi == 'H' and gj in ('+', '-')) or (gj == 'H' and gi in ('+', '-')):
+                energy += 0.3
+
+    # Normalize: sigmoid of (energy / 10) shifted so that 0 contacts → score 0.3
+    score = 1.0 / (1.0 + np.exp(-(energy - 5.0) / 5.0))
+    return float(score)

code/utils/sam.py

@@ -0,0 +1,54 @@
+"""
+Sharpness-Aware Minimization (SAM) optimizer wrapper.
+Seeks parameters in flatter minima for better OOD generalization.
+Reference: Foret et al., "Sharpness-Aware Minimization for Efficiently Improving Generalization" (ICLR 2021)
+"""
+
+import torch
+
+
+class SAM(torch.optim.Optimizer):
+    def __init__(self, params, base_optimizer, rho=0.05, **kwargs):
+        defaults = dict(rho=rho, **kwargs)
+        super().__init__(params, defaults)
+        self.base_optimizer = base_optimizer(self.param_groups, **kwargs)
+
+    @torch.no_grad()
+    def first_step(self):
+        grad_norm = self._grad_norm()
+        for group in self.param_groups:
+            scale = group['rho'] / (grad_norm + 1e-12)
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                e_w = p.grad * scale
+                p.add_(e_w)
+                self.state[p]['e_w'] = e_w
+
+    @torch.no_grad()
+    def second_step(self):
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                p.sub_(self.state[p]['e_w'])
+        self.base_optimizer.step()
+
+    def _grad_norm(self):
+        shared_device = self.param_groups[0]['params'][0].device
+        norm = torch.norm(
+            torch.stack([
+                p.grad.norm(p=2).to(shared_device)
+                for group in self.param_groups
+                for p in group['params']
+                if p.grad is not None
+            ]),
+            p=2,
+        )
+        return norm
+
+    def step(self, closure=None):
+        raise NotImplementedError("SAM requires manual first_step() and second_step() calls")
+
+    def zero_grad(self):
+        self.base_optimizer.zero_grad()

data/sample/README.md

@@ -0,0 +1,49 @@
+# Sample dataset (in-repo)
+
+This directory ships a single pre-built target — **`cam` (Calmodulin)** — so users
+can run a smoke test of the training and evaluation pipeline without first
+downloading the full multi-target dataset (~10 GB on Zenodo) or rebuilding
+from raw PDB files (~30 min per target).
+
+## Contents
+
+```
+sample/
+└── cam/
+    ├── train.pkl   # 84 paired holo/apo complex graphs   (~24 MB)
+    ├── val.pkl     # 12 validation graphs                (~1.3 MB)
+    └── test.pkl    # 96 held-out evaluation graphs       (~25 MB)
+```
+
+Each pickle is a list of dicts produced by `code/data/build_dataset.py`.
+Splits follow the family-stratified scheme used in the paper
+(equivalent to `data/processed_familysplit/cam/` train+val and
+`data/processed_familysplit_v5/cam/test.pkl` in the source tree).
+
+## Smoke test (1-epoch end-to-end)
+
+```bash
+# Train both phases for 1 epoch
+python code/scripts/train.py \
+    --target cam \
+    --phase both \
+    --data_dir data/sample \
+    --checkpoint_dir checkpoints_smoke \
+    --epochs 1 \
+    --no_wandb
+
+# Evaluate
+python code/scripts/evaluate.py \
+    --target cam \
+    --checkpoint checkpoints_smoke/best_phase2.pt \
+    --data_dir data/sample \
+    --outdir eval_smoke
+```
+
+Expected runtime: ~1 minute on a single GPU.
+
+## Want more data?
+
+- All 12 paper targets, pre-built: see `data/DOWNLOAD.md` for the Zenodo link.
+- Build from raw PDBs locally: `scripts/build_data.sh paper12`.
+- Per-target PDB lists and chain mappings: `data/target_lists/*.txt` (68 targets).

data/sample/cam/test.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:38e36d092bbcf4222c762e351fe305e8627f47c78c4acda74170c650ac09e1e8
+size 25608454

data/sample/esm2_embeddings/cam/1IWQ_A.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c52e5a467dc0e73ef7139475bdaafd05e6df8872c345e31fb3da1d35497c00a1
+size 712791

data/sample/esm2_embeddings/cam/1IWQ_B.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b7afc0e91e9095d3945adbf4fcf287fd167ac4fed51d10718db65ee80b89890e
+size 93271

data/sample/esm2_embeddings/cam/1K93_A.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ef76f07da93f60f77ee00fbd719fd81621e50492d245bf62f1a820140e853d61
+size 2484311

data/sample/esm2_embeddings/cam/1K93_B.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c9d65d07aaa8b1c219f5458d2a0502739f482b698c5579bbb4ebee681b5aecb
+size 2392151

data/sample/esm2_embeddings/cam/1NWD_A.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4138c902d0ca0e40ff0c8b0522f71974ba15d3ec1a11db2bb00f9fe8227339f9
+size 758871

data/sample/esm2_embeddings/cam/1NWD_B.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ee429a3b3cbaedf0471f02f7d636e4737f98630f655402317dd3438f8c69c30d
+size 144471

data/sample/esm2_embeddings/cam/1SY9_A.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8a7ca8538945c4b1e524ef3440c5aa7f73afd5d273034063518a0251e2a59f01
+size 758871

data/sample/esm2_embeddings/cam/1SY9_B.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b6258dc842103adc893aacee23c195e2b7da3038e4704e0c9166e1e5581ac784
+size 98391

data/sample/esm2_embeddings/cam/2BBM_A.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0630d3cbc267244f3a40ee83155373675f654cfd720aeaf03271c6438cb68b1d
+size 758871

data/sample/esm2_embeddings/cam/2BBM_B.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:485ff76a5d7ed5459035608a25322601d8fc3d1b73acd49a13d1dcb1fec22ac3
+size 134231

data/sample/esm2_embeddings/cam/2HQW_A.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e8b37dc4704f3c83e4c1d723da66d1873a8dbf27a6e9313450312b9e847f9f44
+size 707671

data/sample/esm2_embeddings/cam/2HQW_B.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:564220203f56b067cda236c35cd153d3b20dcaf0fb2080fadb664eed8f413607
+size 113751

data/sample/esm2_embeddings/cam/2O5G_A.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1187a8edd2d25b219c7ca4cc970a6a771fe8605ab2937163dc3ee595fad97bab
+size 753751

data/sample/esm2_embeddings/cam/2O5G_B.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:604f44b0738f0c9d0f3e1410187a94d95dd7002cead28e23635ddd33bf419b60
+size 98391

data/sample/esm2_embeddings/cam/3D33_A.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b4c553c3e7dfe746babf92dfcae09dab8ecbcf1490cf475aead72fc5f2d30a43
+size 472151