""" 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()