diff --git a/examples/small_molecule_binding/CLAUDE.md b/examples/small_molecule_binding/CLAUDE.md index aaa9f96..248c09e 100644 --- a/examples/small_molecule_binding/CLAUDE.md +++ b/examples/small_molecule_binding/CLAUDE.md @@ -7,6 +7,10 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co | Date | Commit | Notes | |---|---|---| | 2026-04-06 | 3390b61 | change log added | +| 2026-09-01 | — | RunConfig dataclass; PROD/TEST named configs replace flat if/else constants | +| 2026-09-09 | — | Replaced broken RFD3 `scaffoldguided.target_pdb` scaffold feedback with real RFD3 partial-diffusion guidance (`partial.input`/`partial_t`); replaced AlphaFold2/ColabFold fold-validation step with Boltz-2 (protein+ligand co-folding) | +| 2026-09-09 | — | Fixed `analysis_sequence()` silently never comparing MPNN candidates (it only ever read each `.fa` file's first line — the un-designed template record — since real LigandMPNN writes multiple candidates into one file, not one file per candidate); now parses every candidate record and picks the true highest-confidence one | +| 2026-09-09 | — | Added a metric-agnostic non-improvement short-circuit to `fastrelax`/`interface` retry logic — escalates to a new backbone (`STEP_RFD3`) as soon as a retry fails to improve on the previous attempt, instead of always exhausting 5 resequencing retries on backbones that real data showed never recover | ## Context @@ -24,7 +28,7 @@ pip install . python run_small_molecule_binding.py ``` -Before running on HPC, edit the path constants at the top of `run_small_molecule_binding.py` and the `__init__` kwargs in `SmallMoleculeBindingPipeline` (`foundry_sif_path`, `colabfold_path`, `mpnn_dir`, `ligand_params`, etc.) to match the target system. +Before running on HPC, edit the path constants at the top of `run_small_molecule_binding.py` and the `__init__` kwargs in `SmallMoleculeBindingPipeline` (`foundry_sif_path`, `boltz_cache_path`, `mpnn_dir`, `ligand_params`, etc.) to match the target system. ## Architecture @@ -32,7 +36,7 @@ Before running on HPC, edit the path constants at the top of `run_small_molecule - **`small_molecule_binding.py`** — defines `SmallMoleculeBindingPipeline(ImpressBasePipeline)`, all step constants, ensemble utility functions (`_ca_rmsd`, `_seq_identity`, `_ensemble_selective_avg`), and the inner `_run_refine_cycle()` loop. All pipeline tasks (HPC and local analysis) are registered via `@self.auto_register_task()` inside `_register_real_tasks()`. The `run()` method drives a state-machine loop; `_run_refine_cycle()` handles the MPNN+PackMin inner loop with per-cycle sequence retry support. -- **`run_small_molecule_binding.py`** — entry point. Sets threshold constants, defines the `adaptive_decision()` callback, creates an `ImpressManager`, and launches via `manager.start(pipeline_setups=[...])`. +- **`run_small_molecule_binding.py`** — entry point. Defines the `RunConfig` dataclass and two named instances (`PROD`, `TEST`); selects between them via `IMPRESS_TEST_MODE`; defines the `adaptive_decision()` callback; creates an `ImpressManager` and launches via `manager.start(pipeline_setups=[...])`. ### Step constants (state-machine constants in `small_molecule_binding.py`) @@ -42,8 +46,8 @@ Before running on HPC, edit the path constants at the top of `run_small_molecule | `STEP_RFD3` | 1 | backbone diffusion | | `STEP_MPNN` | 2 | MPNN + PackMin refinement cycle | | `STEP_FASTRELAX` | 3 | Rosetta FastRelax | -| `STEP_INTERFACE` | 4 | filter_shape (PyRosetta, gates AF2) | -| `STEP_AF2` | 5 | fold prediction | +| `STEP_INTERFACE` | 4 | filter_shape (PyRosetta, gates fold prediction) | +| `STEP_AF2` | 5 | fold prediction — backed by Boltz-2 co-folding (constant name kept as `STEP_AF2` for compatibility; it no longer runs AlphaFold2) | | `STEP_RETRY_SEQ` | 6 | internal: retry sequence prediction without backbone restart | ### Pipeline tasks and scripts @@ -52,16 +56,16 @@ Before running on HPC, edit the path constants at the top of `run_small_molecule |---|---|---|---| | `rfd3` | HPC | `scripts/rfd3.sh` (RFDiffusion3 via `apptainer exec`) | GPU | | `analysis_backbone` | local | reads JSON metrics from `rfd3` output dir | CPU | -| `mpnn` | HPC | `scripts/mpnn.sh` → `scripts/mpnn_wrapper.sh` (LigandMPNN) | CPU | -| `analysis_sequence` | local | reads `.fa` headers from MPNN `seqs/` output | CPU | +| `mpnn` | HPC | `scripts/mpnn.sh` → `mpnn_run.py` (LigandMPNN) | CPU | +| `analysis_sequence` | local | parses every record in MPNN's `seqs/*.fa` output (LigandMPNN writes one file per input structure containing a template record plus `batch_size` designed candidates — not one candidate per file) and selects the highest-`overall_confidence` candidate | CPU | | `packmin` | HPC | `scripts/packmin.sh` → `scripts/packmin.py` (PyRosetta pack+minimize) | CPU | | `analysis_packmin` | local | reads `_packmin_score.json` from packmin output | CPU | | `fastrelax` | HPC | `scripts/fastrelax.sh` → `scripts/fastrelax.py` (Rosetta FastRelax) | CPU | | `analysis_fastrelax` | local | reads `.fasc` score file from fastrelax output | CPU | | `filter_shape` | HPC | `scripts/filter_shape.sh` → `scripts/filter_shape.py` (PyRosetta shape complementarity) | CPU | | `analysis_interface` | local | reads `shape_complementarity_values.txt` | CPU | -| `af2` | HPC | `scripts/af2.sh` (ColabFold/LocalColabFold) | GPU | -| `analysis_fold` | local | reads ColabFold `_scores.json` files | CPU | +| `boltz` (dispatched from the `STEP_AF2` state, whose constant name is kept for compatibility) | HPC | `scripts/boltz.sh` (Boltz-2, pip-installed CLI, co-folds protein+ligand — no container) | GPU | +| `analysis_fold` | local | reads Boltz `confidence_*.json` files under `predictions/boltz_input/` | CPU | | `filter_energy` | HPC | `scripts/filter_energy.sh` → `scripts/filter_energy.py` (ligand energy filter) | CPU | ### State-machine execution flow @@ -98,13 +102,27 @@ After a successful fold, `adaptive_decision()` always returns to `STEP_RFD3` for | `backbone` | no ligand clashes, `max_ca_deviation < threshold`, sufficient secondary structure | `STEP_MPNN` (with ensemble similarity gating) | `STEP_RFD3` | | `sequence` | ensemble similarity check (sequence identity) | `STEP_MPNN` | `STEP_RETRY_SEQ` (up to 3x), then `STEP_RFD3` | | `packmin` | always passes | `STEP_MPNN` | — | -| `fastrelax` | interaction energy, total score, fa_rep below thresholds | `STEP_INTERFACE` | `STEP_MPNN` | -| `interface` | shape complementarity `max_sc >= interface_min_sc` | `STEP_AF2` | `STEP_MPNN` (up to 5x), then `STEP_RFD3` | -| `fold` | mean pLDDT `>= fold_min_plddt` | sets `rfd3_input_pdb` for guided backbone → `STEP_RFD3` | clears `rfd3_input_pdb` → `STEP_RFD3` | +| `fastrelax` | interaction energy, total score, fa_rep below thresholds | `STEP_INTERFACE` | `STEP_MPNN`, unless none of the failing metrics improved vs. the previous attempt on this backbone (see below), then `STEP_RFD3` | +| `interface` | shape complementarity `max_sc >= interface_min_sc` | `STEP_AF2` | `STEP_MPNN`, unless `max_sc` didn't improve vs. the previous attempt (see below), then `STEP_RFD3`; `STEP_RFD3` regardless after 5x (safety cap) | +| `fold` | Boltz `complex_plddt * 100 >= fold_min_plddt`, and (if set) `ligand_iptm >= fold_min_ligand_iptm` | sets `rfd3_input_pdb` for guided backbone → `STEP_RFD3` | clears `rfd3_input_pdb` → `STEP_RFD3` | + +### fastrelax / interface non-improvement short-circuit + +`fastrelax` and `interface` failures used to always retry via `STEP_MPNN` up to a flat 5x cap before escalating to `STEP_RFD3`. Real HPC data showed this was frequently wasteful: a backbone whose fastrelax metrics (some combination of `interact`/`total_score`/`fa_rep`) or interface shape complementarity are failing for backbone-level structural reasons (packing, energetics, surface complementarity) doesn't improve no matter which MPNN-designed sequence is tried — only regenerating the backbone (`STEP_RFD3`) can help, so burning through all 5 resequencing attempts wastes significant HPC time (confirmed: ~15-20 min per doomed backbone). Observed on a single real run (job `21913252`): three distinct failure-mode combinations across `p3`'s first three backbones (`interact`+`total_score`-only, `fa_rep`-only, and interface shape-complementarity), each flat/non-improving across every attempt, zero eventual recoveries. + +Both branches now use `_stage_metrics_improving()` (small_molecule_binding.py) to compare the current failure's metrics against the previous attempt on the *same* backbone (tracked in `fastrelax_prev_metrics`/`interface_prev_metrics`, reset whenever a new backbone starts). A metric counts as "improved" if its gap to threshold shrank by more than 5% of the previous gap; metrics already passing on the previous attempt aren't considered. If **none** of the currently-failing metrics improved, the pipeline escalates straight to `STEP_RFD3` instead of retrying — effectively a retry cap of 1 (the very first attempt always gets one retry, since there's nothing to compare against yet; the second non-improving attempt escalates). The original 5x counter (`fastrelax_fail_count`/`interface_fail_count`) remains as an outer safety net. ### Ensemble-guided backbone feedback -After a successful fold prediction, `adaptive_decision()` computes CA-RMSD between the current AF2 model and all prior fold ensemble entries. If the selective average score (for structurally similar models) exceeds the overall average, the current AF2 model is fed back as `rfd3_input_pdb` for the next RFDiffusion run (`scaffoldguided.target_pdb`), biasing the next backbone toward successful structural motifs. +After a successful fold prediction, `adaptive_decision()` computes CA-RMSD between the current Boltz model and all prior fold ensemble entries. If the selective average score (for structurally similar models) exceeds the overall average, the current Boltz model is fed back as `rfd3_input_pdb`. + +RFD3 has no `scaffoldguided.*`-style CLI override (that was a leftover from an older RFDiffusion version and doesn't exist in RFD3 — the pipeline's `rfd3()` task no longer attempts one). Guidance is expressed entirely through RFD3's `InputSpecification` JSON, via **partial diffusion**: the `partial.input` field points at a real structure and `partial.partial_t` (Å of noise added before re-denoising; `rfd3_partial_t` kwarg, default `10.0`) controls how closely the result stays to it. + +When `rfd3_input_pdb` is set, `rfd3()`: +1. Reads the ligand's literal residue name from `ligand_params`'s `NAME` record via `_ligand_resname_from_params()` — this is **not** always the params filename stem (e.g. `ALR.params`'s `NAME` is `A:R`, not `ALR`; the colon is a deliberate workaround for RFD3 misresolving the bare `"ALR"` literal — never hardcode or "clean up" this value). +2. Calls `_normalize_ligand_id()` to rewrite the Boltz model's ligand HETATM residue name (via gemmi, no coordinate transform — Boltz already places the ligand correctly relative to the protein it just co-folded) to match that literal, writing `{taskdir}/in/guided_scaffold.pdb`. +3. Calls `_write_guided_rfd3_json()` to copy the base `ALR_binder_design.json`'s `ligand`/`length`/`select_exposed`/`select_buried` fields verbatim into a new spec with `input` pointed at the normalized PDB and `partial_t` set, writing `{taskdir}/in/guided_binder_design.json`. +4. Passes that guided JSON (instead of the base one) as `rfd3.sh`'s `inputs=` argument. If normalization fails (no ligand found in the Boltz model), falls back to the base, unguided JSON rather than erroring. Ensemble similarity utilities (all in `small_molecule_binding.py`): - `_ca_rmsd(path1, path2)` — Kabsch-aligned CA RMSD between two PDB files @@ -121,10 +139,13 @@ Ensemble similarity utilities (all in `small_molecule_binding.py`): | `fastrelax_max_total_score` | 0.0 | total Rosetta score (REU) | | `fastrelax_max_fa_rep` | 150.0 | fa_rep repulsion energy (REU) | | `interface_min_sc` | 0.5 | minimum shape complementarity score | -| `fold_min_plddt` | 70.0 | minimum mean pLDDT | +| `fold_min_plddt` | 70.0 | minimum Boltz `complex_plddt` (rescaled ×100, so this stays on the same 0–100 scale as the old AlphaFold2 pLDDT) | +| `fold_min_ligand_iptm` | `None` | minimum Boltz `ligand_iptm` (protein-ligand interface confidence, 0–1 scale); `None` disables this gate — a new capability plain AlphaFold2 couldn't provide since it never folded the ligand | | `max_tasks` | 300 | maximum ensemble entries before stopping | -Threshold constants in `run_small_molecule_binding.py` override these defaults at `PipelineSetup` construction. +Also configurable, not a pass/fail threshold: `rfd3_partial_t` (default `10.0`, Å of noise added during RFD3 partial diffusion — see "Ensemble-guided backbone feedback" below). + +The `PROD` config in `run_small_molecule_binding.py` overrides these class defaults at `PipelineSetup` construction (e.g. `backbone_max_ca_deviation=1.0`, `fastrelax_max_interact=-8.0`, `fold_min_plddt=75.0`). The `TEST` config sets inert thresholds (everything passes) and `max_tasks=10` for integration testing. ### Output directory structure @@ -141,12 +162,13 @@ Each HPC task creates its working directory as `{base_path}/{name}/{taskcount}_{ ... N_fastrelax/out/ # FastRelax PDB + .fasc score file N+1_filter_shape/out/ - N+2_alphafold/out/ + N+2_boltz/out/boltz_results_boltz_input/predictions/boltz_input/ # Boltz-2 PDBs + confidence_*.json + # (boltz nests its own out_dir/boltz_results_/ automatically) ``` Mock mode (`mock=True`, `mock.py`) mirrors this same `{base_path}/{name}/{taskcount}_{taskname}/...` layout with hardcoded fixture outputs, so the two modes stay directly comparable. -MPNN copies the input backbone to a short fixed filename (`binder.cif.gz` or `binder.`) in `{taskdir}/in/` each cycle to avoid 255-character filename limits in AF2 result archives. +MPNN copies the input backbone to a short fixed filename (`binder.cif.gz` or `binder.`) in `{taskdir}/in/` each cycle to avoid 255-character filename limits in fold-prediction result archives. ### Inter-step state passing @@ -156,29 +178,32 @@ Steps communicate via `self.state`: - `best_backbone_path` — path to best `.cif.gz` from `rfd3` (set by `analysis_backbone`) - `best_packed_pdb` — path to best packed PDB (set by `analysis_sequence`, updated by `packmin`) - `last_seq_fasta` — path to best FASTA from MPNN (set by `analysis_sequence`) -- `best_af2_model` — path to best AF2 PDB (set by `analysis_fold`) +- `best_fold_model` — path to best Boltz-2 co-folded PDB (protein+ligand; set by `analysis_fold`) - `last_analysis_step` — `'backbone'` / `'sequence'` / `'packmin'` / `'fastrelax'` / `'interface'` / `'fold'` - `last_analysis_metrics` — dict with `pass` bool and step-specific score fields - `ensemble` — list of `(etype, score, input_path, output_path)` tuples **Set by `adaptive_decision`:** -- `rfd3_input_pdb` — if set, passed to `rfd3` as `scaffoldguided.target_pdb` for guided diffusion +- `rfd3_input_pdb` — if set, the Boltz co-folded model `rfd3()` normalizes and feeds into RFD3 as partial-diffusion `input` (see "Ensemble-guided backbone feedback") - `seq_retry_count` — retry counter for sequence stage (reset on new backbone or successful sequence) - `interface_fail_count` — retry counter for interface stage (reset on pass or after 5 failures) +- `fastrelax_prev_metrics` / `interface_prev_metrics` — the previous attempt's `last_analysis_metrics` for the current backbone, used by the non-improvement short-circuit (see above); `None` when there's no prior attempt to compare against, reset on pass or new backbone **Set at run start (`setdefault`):** - `ensemble` — initialized to `[]` - `rfd3_input_pdb` — initialized to `None` - `seq_retry_count` — initialized to `0` - `last_seq_fasta` — initialized to `None` +- `fastrelax_prev_metrics` / `interface_prev_metrics` — initialized to `None` ### Execution backends -`run_small_molecule_binding.py` has `LocalExecutionBackend(ProcessPoolExecutor())` active by default. `DragonExecutionBackendV3()` is commented out — swap it in for HPC production runs. +`run_small_molecule_binding.py` uses `DragonExecutionBackend` for HPC production runs. `LocalExecutionBackend(ProcessPoolExecutor())` can be swapped in for local testing. ### Pipeline inputs Each pipeline instance (named e.g. `p1`) expects a `{name}_in/` directory containing: - `ALR_binder_design.json` — RFDiffusion3 input spec (contig, ligand, scaffold args) - `.params` — Rosetta ligand params file (default `ALR.params`) +- `.smiles` — ligand SMILES string, read by the `boltz` task to build its co-folding input; derive it once with `scripts/derive_ligand_smiles.py .params .pdb` (RDKit bond-order perception from the params file's exact connectivity + the reference structure's 3D coordinates — there is no SMILES in a Rosetta `.params` file itself) - Optionally `common_filenames.txt` — used by `filter_energy` for cross-filtering diff --git a/examples/small_molecule_binding/README.md b/examples/small_molecule_binding/README.md index e0d9956..c2bed53 100644 --- a/examples/small_molecule_binding/README.md +++ b/examples/small_molecule_binding/README.md @@ -5,6 +5,9 @@ | Date | Commit | Notes | |---|---|---| | 2026-04-06 | 3390b61 | change log added | +| 2026-09-09 | — | Replaced RFDiffusion3's broken `scaffoldguided.target_pdb` scaffold feedback with real RFD3 partial-diffusion guidance (`partial.input`/`partial_t`); replaced AlphaFold2/ColabFold fold-validation with Boltz-2 (protein+ligand co-folding, adds a real ligand-binding-confidence signal AF2 never had) | +| 2026-09-09 | — | Fixed `analysis_sequence()` silently never comparing MPNN candidates against each other (it only ever read a `.fa` file's first line — an un-designed template record — since real LigandMPNN writes every candidate for a design into one file, not one file per candidate) | +| 2026-09-09 | — | Added a metric-agnostic short-circuit to the `fastrelax`/`interface` retry loops — escalates to a new backbone as soon as a resequencing retry stops improving, instead of always exhausting a flat 5x retry cap on backbones that real data showed never recover | --- @@ -34,8 +37,8 @@ rfd3 ──► analysis_backbone ──► [adaptive] analysis_interface ──► [adaptive] │ pass ▼ - af2 - analysis_fold ──► [adaptive] ──► rfd3 (loop) + boltz + analysis_fold ──► [adaptive] ──► rfd3 (loop, optionally guided) ``` Each `[adaptive]` call invokes `adaptive_decision` (defined in `run_small_molecule_binding.py`), which reads the analysis metrics and the growing ensemble history to decide the next step. The pipeline loops until the task budget is exhausted (`max_tasks` ensemble entries). @@ -45,18 +48,20 @@ Each `[adaptive]` call invokes `adaptive_decision` (defined in `run_small_molecu ## Transformation Tasks ### `rfd3` — Backbone Diffusion -Generates a new protein backbone scaffold conditioned on the ligand binding site using RFdiffusion3 (via Apptainer). On the first iteration or after a failed fold, generation starts from scratch. After a successful fold that lands in a high-scoring neighbourhood (see adaptive rules below), the previous fold decoy is passed as `scaffoldguided.target_pdb=` to bias sampling toward that region. +Generates a new protein backbone scaffold conditioned on the ligand binding site using RFDiffusion3 (via Apptainer). RFD3 has no `scaffoldguided.*`-style CLI override for scaffold guidance — that was a leftover from an older RFDiffusion version. Guidance is expressed entirely through RFD3's `InputSpecification` JSON via **partial diffusion**: when a previous fold decoy is set as the guide (see adaptive rules below), the task normalizes that Boltz-2 model's ligand identity to match the pipeline's ligand `.params` file (Boltz assigns its own placeholder residue name; RFD3's `ligand`/`select_exposed`/`select_buried` selectors need the pipeline's literal name), then writes a guided copy of the base JSON spec with `partial.input` pointed at the normalized model and `partial.partial_t` (Å of diffusion noise, `rfd3_partial_t` kwarg) set. If normalization fails (no ligand found), falls back to the unguided base spec rather than erroring. -- **Input**: `_in/ALR_binder_design.json` (diffusion config), optionally a scaffold PDB from the previous fold +- **Input**: `_in/ALR_binder_design.json` (diffusion config); when guided, `_rfd3/in/guided_scaffold.pdb` + `guided_binder_design.json` (generated by this task, not user-supplied) - **Output**: `_rfd3/out/.cif.gz` + `.json` (per-model metrics) - **HPC**: 1 GPU per rank ### `mpnn` + `analysis_sequence` — Sequence Design Runs LigandMPNN to design amino acid sequences for the current backbone. On cycle 0, `mpnn_ensemble_size` independent sequence batches are generated from the backbone; on subsequent cycles within the same refinement loop, 1 batch is generated from the best packed structure so far. Side-chain packing is performed alongside sequence design. +LigandMPNN writes **one file per input structure** (`seqs/binder.fa`), containing a template record (an echo of the input sequence, no confidence fields) followed by `batch_size` real designed candidates (`id=1`..`id=N`, each with `overall_confidence`/`ligand_confidence`). `analysis_sequence` parses every candidate record across every `.fa` file and selects the single highest-`overall_confidence` candidate, writing its sequence to a clean single-record `seqs/best_candidate.fa` and pointing `best_packed_pdb` at that candidate's actual packed structure (`packed/binder_packed__1.pdb`). + - **Input**: backbone PDB (cycle 0) or best packed PDB (cycle > 0); optional `fixed_residues.txt` -- **Output**: `_mpnn/out/seqs/*.fa` (FASTA with confidence scores in header), `_mpnn/out/packed/*.pdb` (packed structures) -- **Scores extracted**: `overall_confidence` (0–1), `ligand_confidence` (0–1) from FASTA header +- **Output**: `_mpnn/out/seqs/binder.fa` (all candidates), `_mpnn/out/seqs/best_candidate.fa` (winning candidate only), `_mpnn/out/packed/*.pdb` (packed structures, one per candidate) +- **Scores extracted**: `overall_confidence` (0–1), `ligand_confidence` (0–1) of the winning candidate ### `packmin` + `analysis_packmin` — Side-Chain Pack & Minimize PyRosetta script that repacks side chains and performs energy minimization on the best-confidence sequence. Used between MPNN cycles to propagate structural improvements. @@ -71,40 +76,42 @@ Full backbone + side-chain relaxation of the best packed structure using Rosetta - **Input**: best packed PDB; ligand `.params` file - **Output**: `_fastrelax/out/_relaxed_0001.pdb`, `_relaxed.fasc` - **Scores extracted**: `total_score` (REU), `interaction_energy` (REU, protein–ligand interaction), `fa_rep` (REU, Lennard-Jones repulsion), `rmsd` (Å, deviation from input) +- **Retry behavior**: a failure only retries via `STEP_MPNN` if at least one failing metric improved over the previous attempt on the same backbone (see [Adaptive Decision Rules](#adaptive-decision-rules)) — a flat/non-improving backbone escalates to a fresh `rfd3` call immediately rather than burning through resequencing attempts that real data showed never help ### `filter_shape` + `analysis_interface` — Shape Complementarity -PyRosetta script computing shape complementarity (SC) and interface energetics between the designed protein and ligand. Gates progression to fold prediction. +PyRosetta script computing shape complementarity (SC) and interface energetics between the designed protein and ligand. Gates progression to fold prediction. Uses the same non-improvement short-circuit as `fastrelax`. -- **Input**: directory of relaxed PDB files from the previous FastRelax step; ligand directory under `_in/` +- **Input**: directory of relaxed PDB files from the previous FastRelax step; ligand params path (`_in/`, stem resolved by the underlying PyRosetta script) - **Output**: `_filter_shape/out/shape_complementarity_values.txt` (SC per model), `interface_values.txt` (full interface metrics CSV) - **Scores extracted**: `max_sc` — maximum SC value across all models in the batch (0–1 scale) -### `af2` + `analysis_fold` — AlphaFold2 Fold Prediction -ColabFold (AlphaFold2 multimer) predicts the fold of the best-confidence sequence to assess structural self-consistency between the diffused backbone and the designed sequence. +### `boltz` + `analysis_fold` — Boltz-2 Co-Folding +Boltz-2 (pip-installed, no container — see [`scripts/boltz.sh`](scripts/boltz.sh)) co-folds the best-confidence designed sequence **together with the ligand** (specified by SMILES — see [`.smiles`](#user-inputs) below) to assess structural self-consistency and predict the bound complex directly. This is a real capability upgrade over the AlphaFold2 step it replaces: AF2 ran in single-sequence mode with zero ligand awareness (a bare "does this sequence fold" sanity check), while Boltz-2 predicts the actual complex and reports a ligand-binding-confidence metric (`ligand_iptm`) AF2 could never provide. -- **Input**: FASTA of best-confidence sequence (`state['last_seq_fasta']`) -- **Output**: `_alphafold/out/rank_*.pdb`, `rank_*_scores.json` -- **Scores extracted**: `best_mean_plddt` — mean per-residue pLDDT (0–100) of the highest-scoring ranked model +- **Input**: `_boltz/in/boltz_input.yaml` (built from `state['last_seq_fasta']` + `.smiles`) +- **Output**: `_boltz/out/boltz_results_boltz_input/predictions/boltz_input/boltz_input_model_0.pdb` + `confidence_boltz_input_model_0.json` — note Boltz nests its own output one level deeper than its `--out_dir` argument (`boltz_results_/`), confirmed against Boltz's own source, not just its docs +- **Scores extracted**: `complex_plddt` (0–1, rescaled ×100 for `best_complex_plddt` to match the old AF2 pLDDT's 0–100 scale), `ligand_iptm` (0–1, protein–ligand interface confidence) - **HPC**: 1 GPU per rank --- ## Scores and Quality Thresholds -All thresholds are configurable at pipeline construction time (see [Configurable Parameters](#configurable-parameters)). +All thresholds are configurable at pipeline construction time (see [Configurable Parameters](#configurable-parameters)); values below are the class defaults. The production run (`run_small_molecule_binding.py`'s `PROD` config) overrides several of these — see [Usage](#usage). -| Analysis step | Score | Threshold (default) | Meaning | +| Analysis step | Score | Threshold (class default) | Meaning | |---|---|---|---| | `backbone` | `ligand_clashes` | must be `== 0` | No ligand atom clashes in backbone | | `backbone` | `max_ca_deviation` | `< 2.0 Å` | Backbone stays close to diffusion target | | `backbone` | `ss_fraction` | `> 0.2` | At least 20% secondary structure (helix + sheet) | -| `fastrelax` | `interaction_energy` | `< 0.0 REU` | Favourable protein–ligand interaction energy | +| `fastrelax` | `interact` | `< 0.0 REU` | Favourable protein–ligand interaction energy | | `fastrelax` | `total_score` | `< 0.0 REU` | Net favourable total Rosetta energy | | `fastrelax` | `fa_rep` | `< 150.0 REU` | Low steric clash energy after relaxation | | `interface` | `max_sc` | `>= 0.5` | Shape complementarity at ligand interface | -| `fold` | `best_mean_plddt` | `>= 70.0` | AlphaFold2 confidence in predicted structure | +| `fold` | `best_complex_plddt` | `>= 70.0` | Boltz-2 confidence in predicted complex (0–100 scale) | +| `fold` | `ligand_iptm` | `>= None` (off) | Boltz-2 protein–ligand interface confidence, when `fold_min_ligand_iptm` is set | -Sequence analysis (`analysis_sequence`) always sets `pass=True`; routing is handled entirely by the ensemble comparison (see below). +Sequence analysis (`analysis_sequence`) always sets `pass=True`; routing is handled entirely by the ensemble comparison (see below). `packmin` gates only on `total_score > 0` (badly packed → restart backbone), not a fixed threshold. --- @@ -119,8 +126,8 @@ Every analysis task appends a tuple `(type, score, input_path, output_path)` to | Type | Score | Input | Output | |---|---|---|---| | `generate backbone` | `ss_fraction` (0–1) | scaffold PDB or `None` | backbone `.cif.gz` | -| `predict sequence` | `overall_confidence` (0–1) | backbone path | FASTA path | -| `fold decoy` | `best_mean_plddt` (0–100) | FASTA path | fold PDB path | +| `predict sequence` | `overall_confidence` (0–1) | backbone path | best-candidate FASTA path | +| `fold decoy` | `best_complex_plddt` (0–100) | FASTA path | Boltz co-folded PDB path | ### Selective average check @@ -130,9 +137,15 @@ For a given entry type, the function computes: The current result is considered to be in a **productive neighbourhood** when `selective_avg > overall_avg`, i.e. the entries most similar to the current result score better than average. -- **Backbone similarity**: Kabsch-aligned CA-RMSD (lower = more similar). Falls back to simple pass/fail when RMSD data is unavailable (RFdiffusion outputs `.cif.gz`, not `.pdb`). +- **Backbone similarity**: Kabsch-aligned CA-RMSD (lower = more similar). Falls back to simple pass/fail when RMSD data is unavailable (RFDiffusion outputs `.cif.gz`, not `.pdb`). - **Sequence similarity**: per-position identity fraction (higher = more similar). -- **Fold similarity**: Kabsch-aligned CA-RMSD between ColabFold PDB outputs. +- **Fold similarity**: Kabsch-aligned CA-RMSD between Boltz co-folded PDB outputs. + +### Non-improvement short-circuit (fastrelax / interface) + +A `fastrelax` or `interface` failure only retries via `STEP_MPNN` if at least one currently-failing metric improved (closed its gap to threshold by more than 5% of the previous gap) compared to the previous attempt on the *same* backbone. The comparison state (`fastrelax_prev_metrics` / `interface_prev_metrics`) resets whenever a new backbone starts. If nothing improved, the pipeline escalates straight to `STEP_RFD3` — effectively a retry cap of 1 (the very first failure always gets one retry, since there's nothing to compare against yet). The original flat 5x counter (`fastrelax_fail_count` / `interface_fail_count`) remains as an outer safety net in case metrics oscillate rather than genuinely plateauing. + +This exists because some fastrelax/interface failures are driven by the backbone itself (packing, energetics, surface shape) rather than by sequence choice — no amount of resequencing fixes them, only a new backbone can. Confirmed against real HPC data: three different failure-mode combinations, each flat across every resequencing attempt observed, zero eventual recoveries. ### Step-by-step routing @@ -147,15 +160,18 @@ The current result is considered to be in a **productive neighbourhood** when `s | `sequence` | `selective_avg > overall_avg` | `STEP_MPNN` (reset retry count) | | `sequence` | `selective_avg ≤ overall_avg`, retry < 3 | `STEP_RETRY_SEQ` — re-run MPNN same cycle | | `sequence` | `selective_avg ≤ overall_avg`, retry = 3 | `STEP_RFD3` — abandon backbone, start over | -| `packmin` | always | `STEP_MPNN` — continue refinement cycle | +| `packmin` | `total_score ≤ 0` (or unavailable) | `STEP_MPNN` — continue refinement cycle | +| `packmin` | `total_score > 0` | `STEP_RFD3` — badly packed, restart backbone | | `fastrelax` | pass | `STEP_INTERFACE` — run shape complementarity | -| `fastrelax` | fail | `STEP_MPNN` — retry sequence design | -| `interface` | pass | `STEP_AF2` — run fold prediction | -| `interface` | fail | `STEP_MPNN` — retry sequence design | +| `fastrelax` | fail, improving | `STEP_MPNN` — retry sequence design | +| `fastrelax` | fail, not improving (or 5x safety cap) | `STEP_RFD3` — short-circuit to new backbone | +| `interface` | pass | `STEP_AF2` (dispatches to `boltz`) — run fold prediction | +| `interface` | fail, improving | `STEP_MPNN` — retry sequence design | +| `interface` | fail, not improving (or 5x safety cap) | `STEP_RFD3` — short-circuit to new backbone | | `fold` | `selective_avg > overall_avg` | `STEP_RFD3` with `rfd3_input_pdb` set to current fold decoy (guided diffusion) | | `fold` | otherwise | `STEP_RFD3` with `rfd3_input_pdb = None` (scratch) | -Note: fold analysis never sets `STEP_DONE`. The pipeline terminates exclusively via the task budget check. +Note: fold analysis never sets `STEP_DONE`. The pipeline terminates exclusively via the task budget check. `STEP_AF2`'s constant name is kept for compatibility even though it now dispatches to `boltz()`, not AlphaFold2. --- @@ -165,12 +181,14 @@ Place all input files in `/_in/` (default: `p1_in/`). | File | Required | Description | |---|---|---| -| `ALR_binder_design.json` | Yes | RFdiffusion3 design config (target structure, hotspot residues, diffusion settings) | +| `ALR_binder_design.json` | Yes | RFDiffusion3 design config (target structure, hotspot residues, diffusion settings) | | `fixed_residues.txt` | Yes | Space-separated residue indices to hold fixed during MPNN sequence design | | `.params` | Yes | Rosetta ligand parameter file (e.g. `ALR.params`); filename must match `ligand_params` kwarg | -| `/` | Yes | Directory containing the ligand PDB/SDF files for shape complementarity analysis (named by `ligand_name`, default `ALR`) | +| `.smiles` | Yes | Ligand SMILES string, read by the `boltz` task to build its co-folding input. There is no SMILES in a Rosetta `.params` file — derive it once with `scripts/derive_ligand_smiles.py .params .pdb` (RDKit bond-order perception from the params file's exact atom/bond connectivity plus the reference structure's 3D coordinates) | +| `input_pdbs/` | Yes | Target/scaffold PDB referenced by the diffusion config (`ALR_binder_design.json`'s `partial.input`) | | `common_filenames.txt` | Yes (filter_energy) | List of accepted filenames for ligand energy filtering | -| `input_pdbs/` | Optional | Target PDB files referenced by the diffusion config | + +**Important**: whatever literal residue name a ligand's `.params` file declares in its `NAME` record is not necessarily the filename stem — e.g. `ALR.params`'s `NAME` is `A:R`, not `ALR` (the colon is a deliberate workaround for RFD3 misresolving the bare `"ALR"` literal). Always resolve the real name from the `.params` file; never hardcode or "clean up" it. --- @@ -183,19 +201,22 @@ All parameters are passed as `kwargs` to `PipelineSetup`: | Parameter | Default | Description | |---|---|---| | `base_path` | `os.getcwd()` | Root directory for all task subdirectories and input files | -| `mpnn_dir` | `/ocean/projects/dmr170002p/hooten/LigandMPNN` | Path to LigandMPNN repository checkout | -| `foundry_sif_path` | `/ocean/projects/dmr170002p/hooten/foundry_medprec.sif` | Apptainer SIF image containing RFdiffusion3 | -| `colabfold_path` | `/ocean/projects/dmr170002p/hooten/localcolabfold` | LocalColabFold installation (pixi manifest path) | +| `mpnn_dir` | env var `MPNN_DIR` (required) | Path to LigandMPNN repository checkout | +| `foundry_sif_path` | env var `FOUNDRY_SIF_PATH` (required) | Apptainer sandbox/`.sif` containing RFDiffusion3 | +| `boltz_cache_path` | env var `BOLTZ_CACHE` (required) | Boltz-2 model-weights cache directory (pip-installed CLI, no container — see `scripts/boltz.sh`) | | `ligand_params` | `ALR.params` | Ligand parameter filename (relative to `_in/`) | +`mpnn_dir`, `foundry_sif_path`, and `boltz_cache_path` all raise `ValueError` at construction time if neither the kwarg nor the corresponding environment variable is set — there is no silent path default. On Delta HPC, `delta_gpu_run.sh` sets all three env vars before launching (see [Usage](#usage)). + ### Pipeline behaviour | Parameter | Default | Description | |---|---|---| | `mock` | `False` | Run with lightweight mock tasks (no HPC tools required) | | `num_refine_cycles` | `3` | Number of MPNN → PackMin cycles per backbone attempt | -| `mpnn_ensemble_size` | `10` | Number of independent sequence batches on cycle 0 | -| `diffusion_batch_size` | `1` | Number of backbone models to generate per RFdiffusion3 call | +| `mpnn_ensemble_size` | `1` | Number of independent sequence batches on cycle 0 | +| `diffusion_batch_size` | `2` | Number of backbone models to generate per RFDiffusion3 call | +| `rfd3_partial_t` | `10.0` Å | RFD3 partial-diffusion noise level for guided backbone feedback (lower = stays closer to the guide structure) | | `max_tasks` | `300` | Total ensemble entries before stopping (counts every backbone, sequence, and fold entry, pass and fail) | ### Quality thresholds @@ -208,58 +229,45 @@ All parameters are passed as `kwargs` to `PipelineSetup`: | `fastrelax_max_total_score` | `0.0` REU | Maximum total Rosetta score after relaxation | | `fastrelax_max_fa_rep` | `150.0` REU | Maximum Lennard-Jones repulsion after relaxation | | `interface_min_sc` | `0.5` | Minimum shape complementarity score | -| `fold_min_plddt` | `70.0` | Minimum mean pLDDT from ColabFold | +| `fold_min_plddt` | `70.0` | Minimum Boltz-2 `complex_plddt`, rescaled ×100 | +| `fold_min_ligand_iptm` | `None` | Minimum Boltz-2 `ligand_iptm`; `None` disables this gate (off by default so it doesn't silently make existing configs stricter) | --- ## Output Structure -Each task creates a numbered directory `_/` under `base_path`. The counter `N` increments with every HPC task (rfd3, mpnn, packmin, fastrelax, af2); analysis tasks share the counter with the preceding HPC task. +Each task creates a numbered directory `_/` under `base_path`. The counter `N` increments with every HPC task (rfd3, mpnn, packmin, fastrelax, filter_shape, boltz); analysis tasks share the counter with the preceding HPC task. ``` / - _in/ # user inputs - 1_rfd3/out/ # backbone diffusion outputs - 2_mpnn/out/seqs/ # FASTA files with confidence scores - 2_mpnn/out/packed/ # packed PDB structures - 3_packmin/out/ # minimized PDB + score JSON + _in/ # user inputs + 1_rfd3/out/ # backbone diffusion outputs + 2_mpnn/out/seqs/ # binder.fa (all candidates) + best_candidate.fa + 2_mpnn/out/packed/ # packed PDB structures, one per candidate + 3_packmin/out/ # minimized PDB + score JSON ... - _alphafold/out/ # ColabFold rank PDBs + score JSONs + N_fastrelax/out/ # relaxed PDB + .fasc score file + N+1_filter_shape/out/ + N+2_boltz/out/boltz_results_boltz_input/predictions/boltz_input/ # co-folded PDB + confidence JSON ``` +Mock mode (`mock=True`, `mock.py`) mirrors this same layout with hardcoded fixture outputs, matching the real multi-candidate MPNN shape and the real nested Boltz output path, so the two modes stay directly comparable. + --- ## Usage -### Production run (HPC) +### Production run (Delta HPC) -Edit the threshold constants and tool paths in `run_small_molecule_binding.py`, then: +1. One-time environment setup: `bash delta_env_setup.sh` (creates the venv, installs all dependencies including Boltz-2 and PyRosetta, warms the Boltz weights cache). See that script's header comment for required env vars (`SCRATCH`, etc.). +2. Derive each ligand's SMILES once: `python scripts/derive_ligand_smiles.py .params .pdb`, then copy the resulting `.smiles` file into every `_in/` directory that uses that ligand. +3. Adjust `PROD`/`TEST` in `run_small_molecule_binding.py` if the default thresholds don't fit your target (class defaults above are permissive; `PROD` is tuned tighter — e.g. `fastrelax_max_fa_rep=100.0`, `interface_min_sc=0.55`, `fold_min_plddt=75.0`). +4. Submit: `sbatch delta_gpu_run.sh` (sets `MPNN_DIR`/`FOUNDRY_SIF_PATH`/`BOLTZ_CACHE` from `SCRATCH`-relative defaults, or export them yourself beforehand to override — see that script's header comment). Pass `IMPRESS_TEST_MODE=1` before `sbatch` to run the inert `TEST` config instead of `PROD`. -```bash -cd examples/small_molecule_binding -python run_small_molecule_binding.py -``` +After a run, validate the output against expected invariants (ligand identity preserved through the guided-RFD3 path, Boltz output shape, no regression to rejected design states, etc.): -Key variables to set before running: - -```python -# run_small_molecule_binding.py -BACKBONE_MAX_CA_DEVIATION = 2.0 -BACKBONE_MIN_SS_FRACTION = 0.2 -FASTRELAX_MAX_FA_REP = 10.0 -FASTRELAX_MAX_SCORE = 0.0 -INTERFACE_MIN_SC = 0.5 -FOLD_MIN_PLDDT = 70.0 -``` - -And in the pipeline kwargs: - -```python -"foundry_sif_path": "/path/to/foundry.sif", # overrides default -"colabfold_path": "/path/to/localcolabfold", # overrides default -"mpnn_dir": "/path/to/LigandMPNN", # overrides default -"ligand_params": "YOURLIGAND.params", -"max_tasks": 300, +```bash +python scripts/validate_run.py ``` ### Mock / dry run (no HPC required) @@ -269,7 +277,7 @@ cd examples/small_molecule_binding python run_test_small_molecule_binding.py ``` -Mock tasks write placeholder files and hardcode passing metrics, so the full orchestration and adaptive routing logic can be exercised without any external tools. The mock run terminates after 100 ensemble entries (`max_tasks=100`). +Mock tasks write placeholder files and hardcode passing metrics (matching the real multi-candidate MPNN and nested-Boltz-output shapes), so the full orchestration and adaptive routing logic can be exercised without any external tools. The mock run terminates after 100 ensemble entries (`max_tasks=100`). A handful of standalone regression checks (MPNN candidate selection, Boltz filename derivation, the fastrelax/interface short-circuit) run before the mock pipeline itself. --- @@ -282,14 +290,14 @@ The following values are embedded in the code and not exposed as constructor kwa | `mpnn` task | `--seed 111` | Fixed random seed for LigandMPNN | | `mpnn` task | `--temperature 0.1` | Sampling temperature for sequence design | | `mpnn` task | `--number_of_packs_per_design 1` | Side-chain packs per sequence | -| `af2` task | `--random-seed 999` | Fixed random seed for ColabFold | -| `af2` task | `--model-type alphafold2 --rank multimer` | AlphaFold2 multimer ranking | +| `boltz` task | `--diffusion_samples 1` (via `BOLTZ_DIFFUSION_SAMPLES` env var, default `1`) | Number of Boltz-2 structural samples per prediction | +| `boltz` task | `--output_format pdb` | Structure output format | | `fastrelax` task | `-n 1` | One FastRelax round | | `_parse_pdb_ca_coords` | `lru_cache(maxsize=512)` | Max cached PDB files for RMSD re-use | -| `adaptive_decision` | retry count `>= 3` | Retries before abandoning backbone on sequence plateau | -| `PYROSETTA_PRE_EXEC` | `source /anvil/scratch/x-mason/env_pyrosetta` | PyRosetta environment activation (Anvil-specific) | -| `AF2_PRE_EXEC` | CUDA + pixi PATH setup | GPU environment for ColabFold (Anvil-specific) | +| `adaptive_decision` | sequence retry count `>= 3` | Retries before abandoning backbone on sequence plateau | +| `adaptive_decision` | fastrelax/interface `rel_tolerance=0.05` | Relative-improvement threshold for the non-improvement short-circuit (see above) | +| `adaptive_decision` | fastrelax/interface safety cap `>= 5` | Outer retry cap regardless of the short-circuit, in case metrics oscillate | ### Execution backend -`run_small_molecule_binding.py` has `LocalExecutionBackend(ProcessPoolExecutor())` active by default. `DragonExecutionBackendV3()` is commented out — swap it in for HPC production runs. +`run_small_molecule_binding.py` uses `DragonExecutionBackend` for HPC production runs (`IMPRESS_BACKEND=dragon`, the default). Set `IMPRESS_BACKEND=local` before running to use `ConcurrentExecutionBackend(ProcessPoolExecutor())` instead for single-node/non-Dragon development. diff --git a/examples/small_molecule_binding/delta_env_setup.sh b/examples/small_molecule_binding/delta_env_setup.sh new file mode 100755 index 0000000..406f8ff --- /dev/null +++ b/examples/small_molecule_binding/delta_env_setup.sh @@ -0,0 +1,279 @@ +#!/bin/bash +# ============================================================================= +# IMPRESS Small Molecule Binding environment setup — Delta HPC (NCSA) +# +# Creates a Python 3.11+ venv and installs all dependencies. +# +# Usage: +# export SCRATCH=/scratch/ +# bash delta_env_setup.sh [--env-dir DIR] [--impress-dir DIR] [--python PATH] +# +# Defaults: +# ENV_DIR = /u/$USER/ve/impress +# IMPRESS_DIR = $SCRATCH/$USER/IMPRESS +# python = auto-detected (python/3.11, cray-python/3.11.7, anaconda3) +# +# Tool directories (cloned by this script if absent): +# MPNN_DIR = $SCRATCH/$USER/LigandMPNN +# BOLTZ_CACHE = $SCRATCH/$USER/.cache/boltz (model weights cache) +# +# Foundry container (RFD3 backbone diffusion) is managed separately: +# Run pull_foundry.sh to build the sandbox tarball; delta_gpu_run.sh unpacks +# it to /tmp at job start. +# ============================================================================= +if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then + set -euo pipefail +fi + +# ── Require SCRATCH ─────────────────────────────────────────────────────────── +if [[ -z "${SCRATCH:-}" ]]; then + echo "ERROR: set the SCRATCH env var to your allocation scratch root, e.g.:" + echo " export SCRATCH=/scratch/" + echo " bash delta_env_setup.sh" + exit 1 +fi + +# ── Defaults / arg parsing ──────────────────────────────────────────────────── +ENV_DIR="${ENV_DIR:-/u/${USER}/ve/impress}" +IMPRESS_DIR="${IMPRESS_DIR:-${SCRATCH}/${USER}/IMPRESS}" +BASE_PY_OVERRIDE="" + +while [[ $# -gt 0 ]]; do + case $1 in + --env-dir) ENV_DIR="$2"; shift 2 ;; + --impress-dir) IMPRESS_DIR="$2"; shift 2 ;; + --python) BASE_PY_OVERRIDE="$2"; shift 2 ;; + *) echo "Unknown argument: $1"; exit 1 ;; + esac +done + +PY="${ENV_DIR}/bin/python" +PIP="${ENV_DIR}/bin/pip" + +MPNN_DIR="${MPNN_DIR:-${SCRATCH}/${USER}/LigandMPNN}" +BOLTZ_CACHE="${BOLTZ_CACHE:-${SCRATCH}/${USER}/.cache/boltz}" + +echo "=================================================================" +echo " ENV_DIR = ${ENV_DIR}" +echo " IMPRESS_DIR = ${IMPRESS_DIR}" +echo " MPNN_DIR = ${MPNN_DIR}" +echo " BOLTZ_CACHE = ${BOLTZ_CACHE}" +echo "=================================================================" + +# ── 1. Create venv ──────────────────────────────────────────────────────────── +echo "" +echo "── Step 1: Creating venv ──" + +_find_python() { + for candidate in python3.12 python3.11 python3 python; do + local p + p=$(command -v "${candidate}" 2>/dev/null) || continue + local ver + ver=$("${p}" -c "import sys; v=sys.version_info; print(v.major*100+v.minor)" 2>/dev/null) || continue + [ "${ver}" -ge 311 ] && echo "${p}" && return 0 + done + return 1 +} + +if [ -n "${BASE_PY_OVERRIDE}" ]; then + BASE_PY="${BASE_PY_OVERRIDE}" + echo "Using Python override: ${BASE_PY}" +else + BASE_PY=$(_find_python || true) + if [ -z "${BASE_PY}" ]; then + echo "python3.11+ not in PATH — trying modules..." + for mod in python/3.13.5-gcc13.3.1 cray-python/3.12.12 anaconda3; do + module load "${mod}" 2>/dev/null || true + BASE_PY=$(_find_python || true) + [ -n "${BASE_PY}" ] && echo " loaded module: ${mod}" && break + done + fi + if [ -z "${BASE_PY}" ]; then + echo "ERROR: no Python 3.11+ interpreter found." + echo " Pass an explicit interpreter: --python /path/to/python3.11" + echo " Or load a module manually before running this script." + exit 1 + fi +fi +echo "Using Python: ${BASE_PY} ($(${BASE_PY} --version))" + +if [ ! -x "${PY}" ]; then + "${BASE_PY}" -m venv "${ENV_DIR}" +else + echo "venv already exists at ${ENV_DIR}" +fi + +echo "Python: $("${PY}" --version)" + +# ── 2. Bootstrap pip ────────────────────────────────────────────────────────── +echo "" +echo "── Step 2: Bootstrapping pip ──" +"${PY}" -m pip install -q --upgrade pip wheel +"${PIP}" install -q --force-reinstall "setuptools<71" + +# ── 3. radical.asyncflow (PyPI) ────────────────────────────────────────────── +echo "" +echo "── Step 3: radical-asyncflow (PyPI) ──" +"${PIP}" install -q radical-asyncflow + +# ── 4. rhapsody-py (PyPI) ──────────────────────────────────────────────────── +echo "" +echo "── Step 4: rhapsody-py[dragon] (PyPI) ──" +"${PIP}" install -q "rhapsody-py[dragon,telemetry]" + +# ── 5. IMPRESS (local editable) ─────────────────────────────────────────────── +echo "" +echo "── Step 5: IMPRESS (editable) ──" +"${PIP}" install -q -e "${IMPRESS_DIR}" + +# ── 6. PyTorch (CUDA 12.1) — required by LigandMPNN ───────────────────────── +echo "" +echo "── Step 6: PyTorch (cu121) ──" +"${PIP}" install -q torch --index-url https://download.pytorch.org/whl/cu121 + +# ── 7. Boltz-2 ──────────────────────────────────────────────────────────────── +# +# EMPIRICALLY CONFIRMED: `pip install "boltz[cuda]"` (with or without `-U`) +# thrashes pip's resolver for a very long time (observed: 28GB+ pip cache, +# 60-75+ pip-metadata/pip-unpack temp dirs, no completion after ~1hr each +# attempt) -- boltz pins several dependencies (`numpy<2.0`, `gemmi==0.6.5`, +# `pytorch-lightning==2.5.0`, etc.) that genuinely conflict with what's +# already installed in this venv (numpy 2.x from other packages, gemmi 0.7.5, +# etc. -- see this repo's other steps). Resolving a real, deep conflict like +# this is inherently slow/combinatorial for pip's resolver, `-U` or not. +# +# WORKING APPROACH (installs cleanly in seconds instead of hanging): +# install boltz with --no-deps, then install its actually-imported runtime +# dependencies individually, also with --no-deps, accepting the versions +# already present rather than forcing boltz's exact pins. Verified working: +# boltz 2.2.1 imports and `boltz predict --help` runs correctly against +# numpy 1.26.4 (downgraded from whatever was there before -- re-verify +# pyrosetta/ProDy/impress/asyncflow/rhapsody still import after this step, +# they were confirmed OK against numpy 1.26.4 during initial validation) and +# gemmi 0.6.5 (downgraded from 0.7.5). `pip check` will still report several +# cosmetic mismatches (pytorch-lightning, cuequivariance-ops-torch-cu12, +# colabfold/ml-dtypes leftovers from before ColabFold was removed, torch's +# own sympy/triton/nvidia-cublas sub-pins) -- none of these broke any actual +# import in testing; only re-investigate if a real runtime failure surfaces. +# +echo "" +echo "── Step 7: Boltz-2 ──" +"${PIP}" install -q --no-deps "boltz[cuda]" +"${PIP}" install -q --no-deps \ + pytorch_lightning torchmetrics fairscale einops einx mashumaro modelcif \ + wandb dm-tree chembl_structure_pipeline \ + cuequivariance_ops_cu12 cuequivariance_ops_torch_cu12 +"${PY}" -c "import boltz; import torch; print('boltz', boltz.__version__ if hasattr(boltz, '__version__') else '(no __version__)', '+ torch', torch.__version__, 'import OK')" + +# ── 8. LigandMPNN ───────────────────────────────────────────────────────────── +# +# LigandMPNN is run directly from its source tree (no package install). +# This step clones the repo; all Python dependencies (torch, ProDy, biopython, +# numpy) are already satisfied by the venv above. +# LigandMPNN's own requirements.txt pins older torch/cudnn versions — do NOT +# install it into this venv; the newer versions here are compatible at runtime. +# +echo "" +echo "── Step 8: LigandMPNN (clone) ──" +if [ ! -d "${MPNN_DIR}" ]; then + echo " Cloning LigandMPNN to ${MPNN_DIR}" + git clone https://github.com/dauparas/LigandMPNN "${MPNN_DIR}" +else + echo " LigandMPNN already at ${MPNN_DIR}, pulling latest" + git -C "${MPNN_DIR}" pull --ff-only || echo " (pull skipped — non-fast-forward or detached HEAD)" +fi +# Install ProDy and biopython (needed by LigandMPNN; torch already installed). +"${PIP}" install -q ProDy biopython + +# ── 9. gemmi — CIF.GZ parsing for backbone conversion ──────────────────────── +# +# Pinned to 0.6.5, NOT latest: Step 7 installs boltz, which pins gemmi==0.6.5 +# exactly. An unpinned `pip install gemmi` here would silently upgrade to +# latest and re-break that pin (this happened during initial validation). +# Verified empirically that 0.6.5 has everything this pipeline's gemmi usage +# needs: mpnn()'s CIF.GZ->PDB conversion (gemmi.cif.read_string, +# make_structure_from_block, write_pdb) and rfd3()'s ligand-normalization +# helper (read_structure, res.het_flag, mutable res.name, write_pdb) both +# round-trip correctly against 0.6.5. +# +echo "" +echo "── Step 9: gemmi ──" +"${PIP}" install -q "gemmi==0.6.5" + +# ── 10. Additional dependencies ─────────────────────────────────────────────── +echo "" +echo "── Step 10: pandas + biopandas ──" +"${PIP}" install -q pandas biopandas + +# ── 11. PyRosetta ───────────────────────────────────────────────────────────── +echo "" +echo "── Step 11: PyRosetta ──" +"${PIP}" install -q pyrosetta-installer +"${PY}" -c "import pyrosetta_installer; pyrosetta_installer.install_pyrosetta()" + +# ── 12. Boltz-2 model weights ───────────────────────────────────────────────── +# +# Boltz has no dedicated "download weights" subcommand — weights auto-download +# on first `boltz predict` call. Warm the cache with a trivial CPU prediction +# on a login node so compute nodes (no internet) find them already present at +# BOLTZ_CACHE. +# +echo "" +echo "── Step 12: Boltz-2 model weights (cache warm-up) ──" +BOLTZ_CACHE="${BOLTZ_CACHE:-${SCRATCH}/${USER}/.cache/boltz}" +mkdir -p "${BOLTZ_CACHE}" +_WARM_DIR=$(mktemp -d) +cat > "${_WARM_DIR}/warm.yaml" <<'YAML' +version: 1 +sequences: + - protein: + id: [A] + sequence: MAAAAAAAAAAAAAAAAAAA + msa: empty +YAML +"${ENV_DIR}/bin/boltz" predict "${_WARM_DIR}/warm.yaml" \ + --out_dir "${_WARM_DIR}/out" --cache "${BOLTZ_CACHE}" \ + --devices 1 --accelerator cpu --output_format pdb \ + || echo "WARNING: boltz cache warm-up failed — check login-node internet access" +rm -rf "${_WARM_DIR}" + +# ── 13. Verify ──────────────────────────────────────────────────────────────── +echo "" +echo "── Step 13: Verifying installation ──" +_check() { + local label="$1"; shift + if out=$("$@" 2>&1); then + echo " ${label}: OK (${out})" + else + echo " WARNING: ${label} failed" + echo " ${out}" | head -3 + fi +} + +_check "radical.asyncflow" "${PY}" -c "import radical.asyncflow; print(radical.asyncflow.__version__)" +_check "rhapsody-py" "${PY}" -c "import rhapsody; print('ok')" +_check "impress" "${PY}" -c "import impress; print('ok')" +_check "torch" "${PY}" -c "import torch; print(torch.__version__)" +_check "boltz" "${PY}" -c "import boltz; print('ok')" +_check "gemmi" "${PY}" -c "import gemmi; print(gemmi.__version__)" +_check "pyrosetta" "${PY}" -c "import pyrosetta; print('ok')" +_check "ProDy" "${PY}" -c "import prody; print(prody.__version__)" +_check "LigandMPNN" test -d "${MPNN_DIR}" && echo "present" +_check "boltz weights" test -f "${BOLTZ_CACHE}/boltz2_conf.ckpt" && echo "present" + +echo "" +echo "=================================================================" +echo "Setup complete." +echo "" +echo "Activate with:" +echo " source ${ENV_DIR}/bin/activate" +echo "" +echo "Run the pipeline:" +echo " export SCRATCH=${SCRATCH}" +echo " export SBATCH_ACCOUNT=bblj-delta-gpu" +echo " cd ${IMPRESS_DIR}/examples/small_molecule_binding" +echo " sbatch delta_gpu_run.sh" +echo "" +echo "Note: Foundry container (RFD3) is managed separately." +echo " Build once with: sbatch pull_foundry.sh" +echo "=================================================================" diff --git a/examples/small_molecule_binding/delta_gpu_run.sh b/examples/small_molecule_binding/delta_gpu_run.sh new file mode 100644 index 0000000..f952e25 --- /dev/null +++ b/examples/small_molecule_binding/delta_gpu_run.sh @@ -0,0 +1,167 @@ +#!/bin/bash +# +# Small Molecule Binding Pipeline — SLURM batch script (Delta HPC / GPU) +# +# Set before calling sbatch (only SBATCH_ACCOUNT and SCRATCH are required; +# the rest default to standard Delta locations): +# export SBATCH_ACCOUNT=-delta-gpu +# export SCRATCH=/scratch/ +# +# Optional overrides (all have defaults based on SCRATCH/$USER): +# export MPNN_DIR=/path/to/LigandMPNN +# export BOLTZ_CACHE=/path/to/boltz_cache +# +# Foundry container (RFD3): +# The foundry sandbox is stored as a .tar.gz on scratch (built by pull_foundry.sh). +# This script extracts it to /tmp at job start (no scratch quota cost) and removes +# it on exit. Override FOUNDRY_TAR to point to a different archive, or set +# FOUNDRY_SIF_PATH directly to skip extraction entirely (e.g. a pre-extracted dir). +# +# Example: +# sbatch delta_gpu_run.sh +# sbatch delta_gpu_run.sh run_nonadaptive.py # non-adaptive runner +# +# Account: set SBATCH_ACCOUNT=-delta-gpu before calling sbatch +#SBATCH --partition=gpuA40x4 +#SBATCH --nodes=1 +#SBATCH --tasks-per-node=1 +#SBATCH --cpus-per-task=16 +#SBATCH --gpus-per-node=4 +#SBATCH --mem=220G +#SBATCH --time=04:00:00 +#SBATCH --job-name=impress_sm_binding +#SBATCH --mail-user= +#SBATCH --mail-type=ALL +#SBATCH --output=impress_%j.out +##SBATCH --error=logs/impress_%j.err +# NOTE: logs/ must exist before sbatch is called. Create it once with: +# mkdir -p /logs +# NOTE: IMPRESS log output (including errors) goes to .out, not .err. +# On failure, check logs/impress_.out — the .err file will only +# contain Python interpreter crashes or output from non-IMPRESS processes. + +set -e + +# ── Sanity checks ───────────────────────────────────────────────────────────── +if [ -z "${SBATCH_ACCOUNT:-}${SLURM_JOB_ACCOUNT:-}" ]; then + echo "WARNING: SBATCH_ACCOUNT is not set — job may be charged to default account." +fi +echo "Account: ${SLURM_JOB_ACCOUNT:-unknown}" + +if [ -z "${SCRATCH:-}" ]; then + echo "ERROR: SCRATCH is not set." + echo " export SCRATCH=/scratch/ && sbatch delta_gpu_run.sh" + exit 1 +fi + +# ── System library paths (Delta-specific, required by Dragon) ───────────────── +export CUDA_HOME=/opt/nvidia/hpc_sdk/Linux_x86_64/25.3/cuda/12.8 +export MPI_LIB=/opt/cray/pe/mpich/8.1.32/ofi/gnu/11.2/lib-abi-mpich +export FAB_LIB=/opt/cray/libfabric/1.22.0/lib64 +export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${MPI_LIB}:${FAB_LIB}:${LD_LIBRARY_PATH:-} + +# ── Environment ─────────────────────────────────────────────────────────────── +IMPRESS_VENV="${IMPRESS_VENV:-${HOME}/ve/impress}" +unset SLURM_EXPORT_ENV +source "${IMPRESS_VENV}/bin/activate" +dragon-config add --ofi-runtime-lib="${FAB_LIB}" + +# ── Tool paths (read by SmallMoleculeBindingPipeline via env vars) ───────────── +# These are picked up by the pipeline's __init__ when not passed as kwargs. +export MPNN_DIR="${MPNN_DIR:-${SCRATCH}/${USER}/LigandMPNN}" + +# Boltz-2 model weights cache — kept on scratch to avoid home quota exhaustion. +# Pre-warm once on a login node via delta_env_setup.sh's Step 12 (boltz has no +# dedicated "download weights" subcommand; weights auto-download on first +# `boltz predict` call). +export BOLTZ_CACHE="${BOLTZ_CACHE:-${SCRATCH}/${USER}/.cache/boltz}" +mkdir -p "${BOLTZ_CACHE}" + +# ── Foundry sandbox: extract to /tmp at job start, clean up on exit ─────────── +# Extracting to /tmp avoids the scratch quota. Compute nodes have ample /tmp +# space that is not quota-counted. If FOUNDRY_SIF_PATH is already set (e.g. +# a pre-built .sif or a persistent sandbox on a large allocation), extraction +# is skipped entirely. +if [ -z "${FOUNDRY_SIF_PATH:-}" ] && [ -f "${SCRATCH}/foundry.sif" ]; then + export FOUNDRY_SIF_PATH="${SCRATCH}/foundry.sif" +fi +if [ -z "${FOUNDRY_SIF_PATH:-}" ]; then + FOUNDRY_TAR="${FOUNDRY_TAR:-${SCRATCH}/${USER}/foundry_sandbox.tar.gz}" + if [ ! -f "${FOUNDRY_TAR}" ]; then + echo "ERROR: foundry sandbox tarball not found: ${FOUNDRY_TAR}" + echo " Build it first: sbatch pull_foundry.sh" + echo " (or set FOUNDRY_SIF_PATH to an existing .sif/sandbox)" + exit 1 + fi + _FOUNDRY_TMP="/tmp/foundry_${SLURM_JOB_ID:-$$}" + echo "Extracting foundry sandbox from ${FOUNDRY_TAR} to ${_FOUNDRY_TMP} ..." + mkdir -p "${_FOUNDRY_TMP}" + tar -xzf "${FOUNDRY_TAR}" -C "${_FOUNDRY_TMP}" --strip-components=1 + export FOUNDRY_SIF_PATH="${_FOUNDRY_TMP}" + # shellcheck disable=SC2064 + trap "echo 'Removing ${_FOUNDRY_TMP}'; rm -rf '${_FOUNDRY_TMP}'" EXIT +fi + +echo "MPNN_DIR: ${MPNN_DIR}" +echo "FOUNDRY_SIF_PATH: ${FOUNDRY_SIF_PATH}" +echo "BOLTZ_CACHE: ${BOLTZ_CACHE}" + +# ── Tool existence checks ────────────────────────────────────────────────────── +if [ ! -d "${MPNN_DIR}" ]; then + echo "ERROR: MPNN_DIR does not exist: ${MPNN_DIR}" + echo " Clone LigandMPNN: git clone https://github.com/dauparas/LigandMPNN ${MPNN_DIR}" + exit 1 +fi + +# ── Working directory ───────────────────────────────────────────────────────── +#WORKDIR="${IMPRESS_SCRIPTS_DIR:-${SCRATCH}/${USER}/IMPRESS/examples/small_molecule_binding}" +WORKDIR="${IMPRESS_SCRIPTS_DIR:-${SCRATCH}/IMPRESS/examples/small_molecule_binding}" +cd "${WORKDIR}" +mkdir -p logs + +# IMPRESS_WORK_DIR: where pipeline task dirs (p1/, p2/, …) are written. +# Defaults to logs/ so all run artifacts stay out of the source tree and are +# covered by .gitignore. Override to write outputs elsewhere. +export IMPRESS_WORK_DIR="${IMPRESS_WORK_DIR:-${WORKDIR}/logs}" +mkdir -p "${IMPRESS_WORK_DIR}" + +# IMPRESS_SESSION_DIR: asyncflow session dir — runinfo, captured task +# stdout/stderr (.stdout/.stderr per task UID). Must be on Lustre so files +# survive the job and can be reviewed after failures. +export IMPRESS_SESSION_DIR="${IMPRESS_SESSION_DIR:-${IMPRESS_WORK_DIR}/sessions}" +mkdir -p "${IMPRESS_SESSION_DIR}" + +# IMPRESS_BACKEND: "dragon" (default, multi-node HPC) or "local" (single-node, +# ProcessPoolExecutor — useful for development / non-Dragon clusters). +# Set before sbatch: IMPRESS_BACKEND=local sbatch delta_gpu_run.sh +export IMPRESS_BACKEND="${IMPRESS_BACKEND:-dragon}" +echo "IMPRESS_BACKEND: ${IMPRESS_BACKEND}" + +# IMPRESS_TEST_MODE=1: 2 pipelines, inert thresholds, max_tasks=10. +# Runs one full rfd3→mpnn→fastrelax→filter_shape→af2 cycle to verify the +# end-to-end path without looping. Set before sbatch: +# IMPRESS_TEST_MODE=1 sbatch delta_gpu_run.sh +export IMPRESS_TEST_MODE="${IMPRESS_TEST_MODE:-0}" +echo "TEST_MODE: ${IMPRESS_TEST_MODE}" + +# ── Run ─────────────────────────────────────────────────────────────────────── + +# -s = single-node Dragon runtime; -m = multi-node (uses MPI/OFI fabric). +if [ "${SLURM_NNODES:-1}" -gt 1 ]; then + DRAGON_MODE="-m" +else + DRAGON_MODE="-s" +fi + +RUNNER="${1:-run_small_molecule_binding.py}" + +if [ "${IMPRESS_BACKEND}" = "dragon" ]; then + rm -f ddict_orc* + echo "Running: dragon ${DRAGON_MODE} ${RUNNER} (nodes=${SLURM_NNODES:-1})" + dragon ${DRAGON_MODE} "${RUNNER}" +else + echo "Running: python3 ${RUNNER} (backend=${IMPRESS_BACKEND})" + python3 "${RUNNER}" +fi + +echo "=== Small Molecule Binding pipeline done: $(date) ===" diff --git a/examples/small_molecule_binding/mock.py b/examples/small_molecule_binding/mock.py index cb7bc07..f6c97dc 100644 --- a/examples/small_molecule_binding/mock.py +++ b/examples/small_molecule_binding/mock.py @@ -74,6 +74,17 @@ async def rfd3(task_description=None, **kwargs): with open(f"{taskdir}/out/{model_name}.pdb", "w") as fh: fh.write(_synthetic_ca_pdb(seed=pipeline.taskcount)) + rfd3_input_pdb = pipeline.state.get('rfd3_input_pdb') + if rfd3_input_pdb: + # Marker file so tests can assert the guided-diffusion path fired, + # without simulating the real gemmi-based ligand-normalization + # logic (mock's synthetic CA trace has no ligand HETATM block). + with open(f"{taskdir}/in/guided_binder_design.json", "w") as fh: + json.dump({ + "guided_source": rfd3_input_pdb, + "partial_t": pipeline.rfd3_partial_t, + }, fh) + @pipeline.auto_register_task(local_task=True) async def analysis_backbone(task_description=None, **kwargs): taskdir = f"{pipeline.base_path}/{pipeline.name}/{pipeline.taskcount}_rfd3" @@ -124,18 +135,40 @@ async def mpnn(task_description=None, **kwargs): with open(f"{taskdir}/in/{short_name}", "w") as fh: fh.write("REMARK mock mpnn input copy\n") - with open(f"{taskdir}/out/packed/binder_rank_001_packed_1_1.pdb", "w") as fh: - fh.write("REMARK mock mpnn output\nEND\n") - - sequence = _synthetic_sequence( + # Mirror real LigandMPNN's actual output shape (confirmed against live + # HPC runs): ONE file (seqs/binder.fa) containing a template record (no + # 'id='/confidence fields, an echo of the input) followed by several + # designed candidate records ('id=1'..'id=N', each with + # overall_confidence/ligand_confidence) -- NOT one file per candidate. + # id=3 is deliberately the highest-confidence candidate here (not id=1, + # the first) so a regression test can tell "picks the best" apart from + # "picks the first"/"picks the last". + candidate_base_conf = {"1": 0.35, "2": 0.40, "3": 0.55, "4": 0.38} + for cand_id in ("1", "2", "3", "4"): + with open(f"{taskdir}/out/packed/binder_packed_{cand_id}_1.pdb", "w") as fh: + fh.write("REMARK mock mpnn output\nEND\n") + + template_seq = _synthetic_sequence( seed=pipeline.taskcount, base_seq="MAGICKSEQUENCEALPHA", ) - with open(f"{taskdir}/out/seqs/binder_rank_001.fa", "w") as fh: + with open(f"{taskdir}/out/seqs/binder.fa", "w") as fh: fh.write( - ">binder_rank_001, T=0.1, seed=111, overall_confidence=0.85, " - "ligand_confidence=0.75, seq_rec=0.90\n" - f"{sequence}\n" + f">binder, T=0.1, seed=111, num_res={len(template_seq)}, " + "num_ligand_res=10\n" + f"{template_seq}\n" ) + for cand_id, base_conf in candidate_base_conf.items(): + conf = _synthetic_jitter(pipeline.taskcount * 10 + int(cand_id), base=base_conf, jitter=0.02) + lig_conf = _synthetic_jitter(pipeline.taskcount * 10 + int(cand_id) + 400_000, base=base_conf - 0.05, jitter=0.02) + cand_seq = _synthetic_sequence( + seed=pipeline.taskcount * 10 + int(cand_id), base_seq=template_seq, + ) + fh.write( + f">binder, id={cand_id}, T=0.1, seed=111, " + f"overall_confidence={conf:.4f}, ligand_confidence={lig_conf:.4f}, " + "seq_rec=0.5000\n" + f"{cand_seq}\n" + ) @pipeline.auto_register_task(local_task=True) async def analysis_sequence(task_description=None, **kwargs): @@ -143,18 +176,48 @@ async def analysis_sequence(task_description=None, **kwargs): seqs_dir = f"{taskdir}/out/seqs" pipeline.state['last_mpnn_seqs_dir'] = seqs_dir pipeline.state['last_analysis_step'] = 'sequence' - pipeline.state['best_packed_pdb'] = ( - f"{taskdir}/out/packed/binder_rank_001_packed_1_1.pdb" - ) + + # Mirrors the real analysis_sequence()'s parsing: evaluate every id= + # record across every .fa file, skip the template record. + best_conf, best_lig_conf, best_id, best_seq = -1.0, 0.0, None, None + for fa_file in os.listdir(seqs_dir): + if not fa_file.endswith('.fa'): + continue + with open(f"{seqs_dir}/{fa_file}") as fh: + content = fh.read() + for record in content.split('>')[1:]: + lines = record.splitlines() + if not lines: + continue + header, seq = lines[0], ''.join(lines[1:]).strip() + parts = { + kv.split('=')[0].strip(): kv.split('=')[1].strip() + for kv in header.split(',') if '=' in kv + } + if 'id' not in parts: + continue + conf = float(parts.get('overall_confidence', 0)) + lig_conf = float(parts.get('ligand_confidence', 0)) + if conf > best_conf: + best_conf, best_lig_conf, best_id, best_seq = conf, lig_conf, parts['id'], seq + + if best_id is not None: + pipeline.state['best_packed_pdb'] = f"{taskdir}/out/packed/binder_packed_{best_id}_1.pdb" + fasta_path = f"{seqs_dir}/best_candidate.fa" + with open(fasta_path, "w") as fh: + fh.write(f">binder_id_{best_id}\n{best_seq}\n") + pipeline.state['last_seq_fasta'] = fasta_path + else: + pipeline.state['last_seq_fasta'] = None + pipeline.state['last_analysis_metrics'] = { 'pass': True, - 'best_overall_confidence': 0.85, - 'best_ligand_confidence': 0.75, + 'best_overall_confidence': best_conf, + 'best_ligand_confidence': best_lig_conf, } - fasta_path = f"{seqs_dir}/binder_rank_001.fa" - pipeline.state['last_seq_fasta'] = fasta_path pipeline.state['ensemble'].append(( - ETYPE_SEQUENCE, 0.85, pipeline.state.get('best_backbone_path'), fasta_path, + ETYPE_SEQUENCE, best_conf, pipeline.state.get('best_backbone_path'), + pipeline.state.get('last_seq_fasta'), )) @pipeline.auto_register_task(local_task=True) @@ -244,35 +307,80 @@ async def analysis_interface(task_description=None, **kwargs): } @pipeline.auto_register_task(local_task=True) - async def af2(task_description=None, **kwargs): + async def boltz(task_description=None, **kwargs): pipeline.taskcount += 1 - taskname = "alphafold" + taskname = "boltz" pipeline.previous_task = taskname taskdir = f"{pipeline.base_path}/{pipeline.name}/{pipeline.taskcount}_{taskname}" - os.makedirs(f"{taskdir}/in", exist_ok=True) - os.makedirs(f"{taskdir}/out", exist_ok=True) - - for rank in range(1, 6): - with open(f"{taskdir}/out/rank_{rank:03d}.pdb", "w") as fh: - fh.write(_synthetic_ca_pdb(seed=pipeline.taskcount * 100 + rank)) - with open(f"{taskdir}/out/rank_{rank:03d}_scores.json", "w") as fh: - json.dump({"plddt": [85.0 + rank] * 50, "max_pae": 5.0}, fh) + # Mirrors real boltz's own out_dir/boltz_results_/predictions/ + # nesting (see small_molecule_binding.py's analysis_fold() comment). + pred_dir = f"{taskdir}/out/boltz_results_boltz_input/predictions/boltz_input" + os.makedirs(f"{taskdir}/in", exist_ok=True) + os.makedirs(pred_dir, exist_ok=True) + + for model_i in range(5): + seed = pipeline.taskcount * 100 + model_i + with open(f"{pred_dir}/boltz_input_model_{model_i}.pdb", "w") as fh: + fh.write(_synthetic_ca_pdb(seed=seed)) + complex_plddt = _synthetic_jitter(seed, base=0.90, jitter=0.03) + ligand_iptm = _synthetic_jitter(seed + 500_000, base=0.70, jitter=0.05) + with open(f"{pred_dir}/confidence_boltz_input_model_{model_i}.json", "w") as fh: + json.dump({ + "complex_plddt": complex_plddt, + "ligand_iptm": ligand_iptm, + "confidence_score": complex_plddt, + "ptm": _synthetic_jitter(seed + 600_000, base=0.75, jitter=0.05), + "iptm": _synthetic_jitter(seed + 700_000, base=0.70, jitter=0.05), + "protein_iptm": _synthetic_jitter(seed + 800_000, base=0.72, jitter=0.05), + }, fh) @pipeline.auto_register_task(local_task=True) async def analysis_fold(task_description=None, **kwargs): - taskdir = f"{pipeline.base_path}/{pipeline.name}/{pipeline.taskcount}_alphafold" - best_model = f"{taskdir}/out/rank_005.pdb" - best_mean_plddt = _synthetic_jitter(pipeline.taskcount + 900_000, base=90.0, jitter=3.0) - pipeline.state['best_af2_model'] = best_model + pred_dir = ( + f"{pipeline.base_path}/{pipeline.name}/{pipeline.taskcount}_boltz/out/" + "boltz_results_boltz_input/predictions/boltz_input" + ) + conf_files = [ + f for f in os.listdir(pred_dir) + if f.startswith('confidence_') and f.endswith('.json') + ] if os.path.isdir(pred_dir) else [] + + best_complex_plddt = -1.0 + best_model = None + best_ligand_iptm = None + for cf in conf_files: + with open(f"{pred_dir}/{cf}") as fh: + data = json.load(fh) + score = data.get('complex_plddt', 0.0) + if score > best_complex_plddt: + best_complex_plddt = score + best_model = cf.replace('confidence_', '', 1).replace('.json', '.pdb') + best_ligand_iptm = data.get('ligand_iptm') + + # Rescale 0-1 -> 0-100 to preserve fold_min_plddt's existing semantics + # (mirrors the real analysis_fold()). + best_plddt_100 = best_complex_plddt * 100.0 + passed = best_plddt_100 >= pipeline.fold_min_plddt + if pipeline.fold_min_ligand_iptm is not None: + passed = passed and ( + best_ligand_iptm is not None + and best_ligand_iptm >= pipeline.fold_min_ligand_iptm + ) + + if best_model: + full_model_path = f"{pred_dir}/{best_model}" + pipeline.state['best_fold_model'] = full_model_path + pipeline.state['ensemble'].append(( + ETYPE_FOLD, best_plddt_100, pipeline.state.get('last_seq_fasta'), full_model_path, + )) + pipeline.state['last_analysis_step'] = 'fold' pipeline.state['last_analysis_metrics'] = { - 'pass': best_mean_plddt >= pipeline.fold_min_plddt, - 'best_mean_plddt': best_mean_plddt, - 'best_model': best_model, + 'pass': passed, + 'best_complex_plddt': best_plddt_100, + 'best_ligand_iptm': best_ligand_iptm, + 'best_model': best_model, } - pipeline.state['ensemble'].append(( - ETYPE_FOLD, best_mean_plddt, pipeline.state.get('last_seq_fasta'), best_model, - )) @pipeline.auto_register_task(local_task=True) async def filter_energy(ligand_name="ALR", task_description=None, **kwargs): diff --git a/examples/small_molecule_binding/p1_in/ALR.smiles b/examples/small_molecule_binding/p1_in/ALR.smiles new file mode 100644 index 0000000..438a090 --- /dev/null +++ b/examples/small_molecule_binding/p1_in/ALR.smiles @@ -0,0 +1 @@ +COc1cc(S(=O)(=O)[O-])c(C)cc1/N=N/c1c(O)ccc2cc(S(=O)(=O)[O-])ccc12 diff --git a/examples/small_molecule_binding/p2_in/ALR.smiles b/examples/small_molecule_binding/p2_in/ALR.smiles new file mode 100644 index 0000000..438a090 --- /dev/null +++ b/examples/small_molecule_binding/p2_in/ALR.smiles @@ -0,0 +1 @@ +COc1cc(S(=O)(=O)[O-])c(C)cc1/N=N/c1c(O)ccc2cc(S(=O)(=O)[O-])ccc12 diff --git a/examples/small_molecule_binding/p3_in/ALR.smiles b/examples/small_molecule_binding/p3_in/ALR.smiles new file mode 100644 index 0000000..438a090 --- /dev/null +++ b/examples/small_molecule_binding/p3_in/ALR.smiles @@ -0,0 +1 @@ +COc1cc(S(=O)(=O)[O-])c(C)cc1/N=N/c1c(O)ccc2cc(S(=O)(=O)[O-])ccc12 diff --git a/examples/small_molecule_binding/p4_in/ALR.smiles b/examples/small_molecule_binding/p4_in/ALR.smiles new file mode 100644 index 0000000..438a090 --- /dev/null +++ b/examples/small_molecule_binding/p4_in/ALR.smiles @@ -0,0 +1 @@ +COc1cc(S(=O)(=O)[O-])c(C)cc1/N=N/c1c(O)ccc2cc(S(=O)(=O)[O-])ccc12 diff --git a/examples/small_molecule_binding/p5_in/ALR.smiles b/examples/small_molecule_binding/p5_in/ALR.smiles new file mode 100644 index 0000000..438a090 --- /dev/null +++ b/examples/small_molecule_binding/p5_in/ALR.smiles @@ -0,0 +1 @@ +COc1cc(S(=O)(=O)[O-])c(C)cc1/N=N/c1c(O)ccc2cc(S(=O)(=O)[O-])ccc12 diff --git a/examples/small_molecule_binding/p6_in/ALR.smiles b/examples/small_molecule_binding/p6_in/ALR.smiles new file mode 100644 index 0000000..438a090 --- /dev/null +++ b/examples/small_molecule_binding/p6_in/ALR.smiles @@ -0,0 +1 @@ +COc1cc(S(=O)(=O)[O-])c(C)cc1/N=N/c1c(O)ccc2cc(S(=O)(=O)[O-])ccc12 diff --git a/examples/small_molecule_binding/p7_in/ALR.smiles b/examples/small_molecule_binding/p7_in/ALR.smiles new file mode 100644 index 0000000..438a090 --- /dev/null +++ b/examples/small_molecule_binding/p7_in/ALR.smiles @@ -0,0 +1 @@ +COc1cc(S(=O)(=O)[O-])c(C)cc1/N=N/c1c(O)ccc2cc(S(=O)(=O)[O-])ccc12 diff --git a/examples/small_molecule_binding/p8_in/ALR.smiles b/examples/small_molecule_binding/p8_in/ALR.smiles new file mode 100644 index 0000000..438a090 --- /dev/null +++ b/examples/small_molecule_binding/p8_in/ALR.smiles @@ -0,0 +1 @@ +COc1cc(S(=O)(=O)[O-])c(C)cc1/N=N/c1c(O)ccc2cc(S(=O)(=O)[O-])ccc12 diff --git a/examples/small_molecule_binding/pull_foundry.sh b/examples/small_molecule_binding/pull_foundry.sh new file mode 100644 index 0000000..b883023 --- /dev/null +++ b/examples/small_molecule_binding/pull_foundry.sh @@ -0,0 +1,35 @@ +#!/bin/bash +#SBATCH --account=bblj-delta-gpu +#SBATCH --partition=gpuA40x4 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=128G +#SBATCH --gpus-per-node=1 +#SBATCH --time=00:30:00 +#SBATCH --job-name=pull_foundry +#SBATCH --output=slurm-%j.out +#SBATCH --error=slurm-%j.err + +set -euo pipefail +ulimit -c 0 # disable core dumps + +export APPTAINER_CACHEDIR="${APPTAINER_CACHEDIR:-${SCRATCH:?SCRATCH must be set}/${USER}/.apptainer_cache}" +export APPTAINER_TMPDIR=/tmp/apptainer_$$ +mkdir -p "$APPTAINER_CACHEDIR" "$APPTAINER_TMPDIR" + +# Build as sandbox (directory) to /tmp — no mksquashfs involved. +# Then tar to a single archive on scratch for storage. +SANDBOX=/tmp/foundry_sandbox_$$ +DEST_TAR="${FOUNDRY_SANDBOX_TAR:-${SCRATCH:?SCRATCH must be set}/${USER}/foundry_sandbox.tar.gz}" + +echo "=== Building sandbox to /tmp ===" +apptainer build --sandbox "$SANDBOX" docker://rosettacommons/foundry + +echo "=== Compressing sandbox to scratch ===" +tar -czf "$DEST_TAR" -C /tmp "foundry_sandbox_$$" + +echo "Done: $(ls -lh "$DEST_TAR")" +rm -rf "$SANDBOX" + +rm -rf "$APPTAINER_TMPDIR" diff --git a/examples/small_molecule_binding/run_nonadaptive.py b/examples/small_molecule_binding/run_nonadaptive.py index 75f56fd..2f87e6e 100644 --- a/examples/small_molecule_binding/run_nonadaptive.py +++ b/examples/small_molecule_binding/run_nonadaptive.py @@ -1,9 +1,8 @@ import asyncio +import os from typing import List -from radical.asyncflow import LocalExecutionBackend -from concurrent.futures import ProcessPoolExecutor -from rhapsody.backends import DragonExecutionBackendV3 +from rhapsody.backends import DragonExecutionBackend from impress import ImpressManager, PipelineSetup from small_molecule_binding import ( @@ -13,7 +12,7 @@ import logging import rhapsody -rhapsody.enable_logging(level=logging.DEBUG) +rhapsody.enable_logging(level=logging.INFO) # ── Per-step quality thresholds ──────────────────────────────────────────── # These are passed to pipeline analysis tasks for metric logging but are not @@ -51,10 +50,21 @@ async def nonadaptive_decision(pipeline: SmallMoleculeBindingPipeline) -> None: ) +# Indices of pipelines to run — corresponds to the protein IDs selected for +# this campaign (non-contiguous because some were dropped after earlier runs). +PIPELINE_INDICES = [1, 2, 4, 6, 7, 8, 10, 11, 12, 13, 14, 15, 16, 18, 19, 20, 23, 26, 27, 30, 32] + + async def impress_smallmol_nonadaptive() -> None: """Execute the small-molecule binding pipeline without adaptive routing.""" - #backend = await LocalExecutionBackend(ProcessPoolExecutor()) - backend = await DragonExecutionBackendV3() + examples_dir = os.path.dirname(os.path.abspath(__file__)) + work_dir = os.environ.get( + "IMPRESS_WORK_DIR", os.path.join(examples_dir, "logs") + ) + os.makedirs(work_dir, exist_ok=True) + input_dir = os.path.join(examples_dir, "p1_in") + + backend = await DragonExecutionBackend() manager: ImpressManager = ImpressManager(execution_backend=backend) pipeline_setups: List[PipelineSetup] = [ @@ -63,6 +73,9 @@ async def impress_smallmol_nonadaptive() -> None: type=SmallMoleculeBindingPipeline, adaptive_fn=nonadaptive_decision, kwargs={ + "base_path": work_dir, + "scripts_path": os.path.join(examples_dir, "scripts"), + "input_dir": input_dir, "backbone_max_ca_deviation": BACKBONE_MAX_CA_DEVIATION, "backbone_min_ss_fraction": BACKBONE_MIN_SS_FRACTION, "fastrelax_max_fa_rep": FASTRELAX_MAX_FA_REP, @@ -74,7 +87,7 @@ async def impress_smallmol_nonadaptive() -> None: "num_refine_cycles": 2, } ) - for i in [1,2,4,6,7,8,10,11,12,13,14,15,16,18,19,20,23,26,27,30,32] + for i in PIPELINE_INDICES ] await manager.start(pipeline_setups=pipeline_setups) diff --git a/examples/small_molecule_binding/run_small_molecule_binding.py b/examples/small_molecule_binding/run_small_molecule_binding.py index f39f6d7..9c27937 100644 --- a/examples/small_molecule_binding/run_small_molecule_binding.py +++ b/examples/small_molecule_binding/run_small_molecule_binding.py @@ -1,32 +1,92 @@ import asyncio import os -from concurrent.futures import ThreadPoolExecutor,ProcessPoolExecutor +from dataclasses import dataclass from typing import List -from radical.asyncflow import LocalExecutionBackend -from rhapsody.backends import DragonExecutionBackendV3 - -from impress import ImpressManager, PipelineSetup +from impress import find_gpus, ImpressManager, PipelineSetup from small_molecule_binding import ( SmallMoleculeBindingPipeline, STEP_DONE, STEP_RFD3, STEP_MPNN, STEP_FASTRELAX, STEP_INTERFACE, STEP_AF2, STEP_RETRY_SEQ, ETYPE_BACKBONE, ETYPE_SEQUENCE, ETYPE_FOLD, - _ca_rmsd, _seq_identity, _ensemble_selective_avg, + _ca_rmsd, _seq_identity, _ensemble_selective_avg, _stage_metrics_improving, ) import logging import rhapsody -rhapsody.enable_logging(level=logging.DEBUG) +rhapsody.enable_logging(level=logging.INFO) + + +@dataclass +class RunConfig: + n_pipelines: int + max_tasks: int + # backbone + backbone_max_ca_deviation: float + backbone_min_ss_fraction: float + # fastrelax + fastrelax_max_fa_rep: float + fastrelax_max_score: float # total_score REU + fastrelax_max_interact: float # interaction energy REU + # interface + interface_min_sc: float # shape complementarity + # fold + fold_min_plddt: float # mean pLDDT; init is -1.0 so -1.0 = always pass + fold_min_ligand_iptm: float | None # Boltz-2 protein-ligand interface confidence; None = off + # diffusion / refinement + diffusion_batch_size: int + num_refine_cycles: int + rfd3_partial_t: float # RFD3 partial-diffusion noise (A) for guided backbone feedback + + +PROD = RunConfig( + n_pipelines = 4, + max_tasks = 300, + backbone_max_ca_deviation = 1.0, + backbone_min_ss_fraction = 0.5, + fastrelax_max_fa_rep = 100.0, + fastrelax_max_score = -250.0, # data range -193 to -510 + fastrelax_max_interact = -8.0, # p75 = -8.8 + interface_min_sc = 0.55, + fold_min_plddt = 75.0, + # Off by default so switching to Boltz-2 doesn't silently make PROD stricter. + # 0.5 is a reasonable starting point if ligand-binding-confidence gating is + # wanted later — a signal plain AlphaFold2 could never provide since it + # never folded the ligand at all. + fold_min_ligand_iptm = None, + diffusion_batch_size = 4, + num_refine_cycles = 2, + rfd3_partial_t = 10.0, +) + +# Inert thresholds — everything passes; low task budget for one full cycle. +TEST = RunConfig( + n_pipelines = 2, + max_tasks = 10, + backbone_max_ca_deviation = 9999.0, + backbone_min_ss_fraction = 0.0, + fastrelax_max_fa_rep = 9999.0, + fastrelax_max_score = 9999.0, + fastrelax_max_interact = 9999.0, + interface_min_sc = 0.0, + fold_min_plddt = -1.0, + fold_min_ligand_iptm = None, + diffusion_batch_size = 1, + num_refine_cycles = 1, + # Not a pass/fail threshold like the fields above — a diffusion-noise + # parameter, so kept at a sane real value rather than an inert extreme. + rfd3_partial_t = 10.0, +) + +BACKEND = os.environ.get("IMPRESS_BACKEND", "dragon").lower() -# ── Per-step quality thresholds ──────────────────────────────────────────── -BACKBONE_MAX_CA_DEVIATION = 1.0 -BACKBONE_MIN_SS_FRACTION = 0.5 -FASTRELAX_MAX_FA_REP = 100.0 # fa_rep REU -FASTRELAX_MAX_SCORE = -250.0 # total_score REU (was 0.0 — inert; data range -193 to -510) -FASTRELAX_MAX_INTERACT = -8.0 # interaction energy REU (was absent/0.0 — inert; p75=-8.8) -INTERFACE_MIN_SC = 0.55 # shape complementarity (was 0.35 — passed 100%; data min=0.37) -FOLD_MIN_PLDDT = 75.0 # mean pLDDT (was 70.0; ~70% of AF2 runs exceed 75) +if BACKEND == "dragon": + from rhapsody.backends import DragonExecutionBackend +else: + from concurrent.futures import ProcessPoolExecutor + from rhapsody.backends import ConcurrentExecutionBackend + +cfg = TEST if os.getenv("IMPRESS_TEST_MODE", "0") == "1" else PROD async def adaptive_decision(pipeline: SmallMoleculeBindingPipeline) -> None: @@ -45,7 +105,12 @@ def _prior(ttype): pipeline.next_step = STEP_RFD3 else: current, prior = _prior(ETYPE_BACKBONE) - pipeline.state['seq_retry_count'] = 0 # reset on any new backbone + # Reset on any new backbone -- these are per-backbone retry state, + # not per-pipeline, and must not leak into the next backbone's + # first fastrelax/interface attempt (see _stage_metrics_improving). + pipeline.state['seq_retry_count'] = 0 + pipeline.state['fastrelax_prev_metrics'] = None + pipeline.state['interface_prev_metrics'] = None if not prior: pipeline.next_step = STEP_MPNN else: @@ -86,26 +151,61 @@ def _prior(ttype): elif step == 'fastrelax': if passed: - pipeline.state['fastrelax_fail_count'] = 0 + pipeline.state['fastrelax_fail_count'] = 0 + pipeline.state['fastrelax_prev_metrics'] = None pipeline.next_step = STEP_INTERFACE else: + # Metric-agnostic short-circuit: a backbone whose fastrelax metrics + # (interact/total_score/fa_rep, whichever combination is failing) + # aren't improving attempt-over-attempt is backbone-driven, not + # sequence-driven -- resequencing (STEP_MPNN) can't fix it, only a + # new backbone (STEP_RFD3) can. Confirmed against real HPC data + # (job 21913252): 15 fastrelax attempts across 3 backbones, zero + # improvement and zero eventual recoveries once a backbone's + # metrics went flat. Retry cap = 1 (escalate on the first + # non-improving retry) per that evidence; the count>=5 check + # remains as an outer safety net in case metrics oscillate. + prev = pipeline.state.get('fastrelax_prev_metrics') + specs = [ + ('interact', True, pipeline.fastrelax_max_interact), + ('total_score', True, pipeline.fastrelax_max_total_score), + ('fa_rep', True, pipeline.fastrelax_max_fa_rep), + ] + improving = _stage_metrics_improving(metrics, prev, specs) + pipeline.state['fastrelax_prev_metrics'] = dict(metrics) + count = pipeline.state.get('fastrelax_fail_count', 0) + 1 pipeline.state['fastrelax_fail_count'] = count - if count >= 5: - pipeline.state['fastrelax_fail_count'] = 0 + + if (prev is not None and not improving) or count >= 5: + pipeline.state['fastrelax_fail_count'] = 0 + pipeline.state['fastrelax_prev_metrics'] = None pipeline.next_step = STEP_RFD3 else: pipeline.next_step = STEP_MPNN elif step == 'interface': if passed: - pipeline.state['interface_fail_count'] = 0 + pipeline.state['interface_fail_count'] = 0 + pipeline.state['interface_prev_metrics'] = None pipeline.next_step = STEP_AF2 else: + # Same metric-agnostic short-circuit as 'fastrelax' above, applied + # to shape complementarity -- confirmed the identical wasteful + # pattern occurs here too (job 21913252: 5 interface attempts on + # one backbone, shape complementarity flat within ~0.02, never + # nearing threshold). + prev = pipeline.state.get('interface_prev_metrics') + specs = [('max_sc', False, pipeline.interface_min_sc)] # higher is better + improving = _stage_metrics_improving(metrics, prev, specs) + pipeline.state['interface_prev_metrics'] = dict(metrics) + count = pipeline.state.get('interface_fail_count', 0) + 1 pipeline.state['interface_fail_count'] = count - if count >= 5: - pipeline.state['interface_fail_count'] = 0 + + if (prev is not None and not improving) or count >= 5: + pipeline.state['interface_fail_count'] = 0 + pipeline.state['interface_prev_metrics'] = None pipeline.next_step = STEP_RFD3 else: pipeline.next_step = STEP_MPNN @@ -139,32 +239,52 @@ def _prior(ttype): async def impress_smallmol_bind() -> None: """Execute the small-molecule binding pipeline.""" - #backend = await LocalExecutionBackend(ProcessPoolExecutor()) - backend = await DragonExecutionBackendV3() + # Resolve paths before launching Dragon (os.getcwd() is the examples dir). + examples_dir = os.path.dirname(os.path.abspath(__file__)) + work_dir = os.environ.get( + "IMPRESS_WORK_DIR", os.path.join(examples_dir, "logs") + ) + os.makedirs(work_dir, exist_ok=True) + # Input data lives in the source tree; pass as absolute so it resolves + # correctly regardless of what base_path / work_dir is set to. + input_dir = os.path.join(examples_dir, "p1_in") + + if BACKEND == "dragon": + backend = await DragonExecutionBackend() + else: + backend = await ConcurrentExecutionBackend.create(ProcessPoolExecutor()) manager: ImpressManager = ImpressManager(execution_backend=backend) + all_gpus = find_gpus() + pipeline_setups: List[PipelineSetup] = [ PipelineSetup( name=f"p{str(i)}", type=SmallMoleculeBindingPipeline, adaptive_fn=adaptive_decision, kwargs={ - "backbone_max_ca_deviation": BACKBONE_MAX_CA_DEVIATION, - "backbone_min_ss_fraction": BACKBONE_MIN_SS_FRACTION, - "fastrelax_max_fa_rep": FASTRELAX_MAX_FA_REP, - "fastrelax_max_total_score": FASTRELAX_MAX_SCORE, - "fastrelax_max_interact": FASTRELAX_MAX_INTERACT, - "interface_min_sc": INTERFACE_MIN_SC, - "fold_min_plddt": FOLD_MIN_PLDDT, - "diffusion_batch_size": 4, - "num_refine_cycles": 2, + "base_path": work_dir, + "scripts_path": os.path.join(examples_dir, "scripts"), + "input_dir": input_dir, + "backbone_max_ca_deviation": cfg.backbone_max_ca_deviation, + "backbone_min_ss_fraction": cfg.backbone_min_ss_fraction, + "fastrelax_max_fa_rep": cfg.fastrelax_max_fa_rep, + "fastrelax_max_total_score": cfg.fastrelax_max_score, + "fastrelax_max_interact": cfg.fastrelax_max_interact, + "interface_min_sc": cfg.interface_min_sc, + "fold_min_plddt": cfg.fold_min_plddt, + "fold_min_ligand_iptm": cfg.fold_min_ligand_iptm, + "diffusion_batch_size": cfg.diffusion_batch_size, + "num_refine_cycles": cfg.num_refine_cycles, + "rfd3_partial_t": cfg.rfd3_partial_t, + "max_tasks": cfg.max_tasks, + **({"gpu_id": all_gpus[(i - 1) % len(all_gpus)]} if all_gpus else {}), } ) - for i in range(1,9) + for i in range(1, cfg.n_pipelines + 1) ] await manager.start(pipeline_setups=pipeline_setups) - await manager.flow.shutdown() if __name__ == "__main__": diff --git a/examples/small_molecule_binding/run_test_small_molecule_binding.py b/examples/small_molecule_binding/run_test_small_molecule_binding.py index a780ee7..7142e21 100644 --- a/examples/small_molecule_binding/run_test_small_molecule_binding.py +++ b/examples/small_molecule_binding/run_test_small_molecule_binding.py @@ -16,7 +16,7 @@ from concurrent.futures import ThreadPoolExecutor from typing import List -from radical.asyncflow import ConcurrentExecutionBackend +from radical.asyncflow import LocalExecutionBackend from impress import ImpressManager, PipelineSetup from small_molecule_binding import SmallMoleculeBindingPipeline @@ -43,29 +43,112 @@ def setup_mock_inputs(pipeline_name: str) -> None: fh.write(f"# mock placeholder: {fname}\n") -def check_af2_filename_derivation() -> None: - """Regression check for analysis_fold()'s scores.json -> unrelaxed.pdb - filename derivation, against real ColabFold (colabfold_batch) naming.""" +def check_boltz_filename_derivation() -> None: + """Regression check for analysis_fold()'s confidence_*.json -> *.pdb + filename derivation, against real Boltz-2 (`boltz predict`) naming.""" cases = [ ( - "binder_scores_rank_001_alphafold2_model_3_seed_999.json", - "binder_unrelaxed_rank_001_alphafold2_model_3_seed_999.pdb", + "confidence_boltz_input_model_0.json", + "boltz_input_model_0.pdb", ), ( - "binder_scores_rank_005_alphafold2_ptm_model_1_seed_000.json", - "binder_unrelaxed_rank_005_alphafold2_ptm_model_1_seed_000.pdb", + "confidence_boltz_input_model_4.json", + "boltz_input_model_4.pdb", ), ] - for sf, expected in cases: - derived = sf.replace('_scores_', '_unrelaxed_').replace('.json', '.pdb') - assert derived == expected, f"{sf!r} -> {derived!r}, expected {expected!r}" + for cf, expected in cases: + derived = cf.replace('confidence_', '', 1).replace('.json', '.pdb') + assert derived == expected, f"{cf!r} -> {derived!r}, expected {expected!r}" + + +def check_mpnn_candidate_selection() -> None: + """Regression check for analysis_sequence()'s candidate-selection logic + against real LigandMPNN's actual output shape: ONE file per input + structure containing a template record (no 'id=') followed by several + designed candidate records ('id=1'..'id=N', each with overall_confidence). + A prior version of this logic only ever read a file's first line (the + template), so it silently always "selected" the template's defaulted + 0.0 confidence and never compared real candidates -- this asserts the + fix actually distinguishes and picks the highest-confidence candidate, + not the template and not simply the first candidate in the file.""" + fixture = ( + ">binder, T=0.1, seed=111, num_res=94, num_ligand_res=39\n" + "TEMPLATESEQUENCE\n" + ">binder, id=1, T=0.1, seed=111, overall_confidence=0.4167, " + "ligand_confidence=0.4290, seq_rec=0.5000\n" + "CANDIDATEONE\n" + ">binder, id=2, T=0.1, seed=111, overall_confidence=0.4092, " + "ligand_confidence=0.4448, seq_rec=0.4574\n" + "CANDIDATETWO\n" + ">binder, id=3, T=0.1, seed=111, overall_confidence=0.4048, " + "ligand_confidence=0.4235, seq_rec=0.4574\n" + "CANDIDATETHREE\n" + ">binder, id=4, T=0.1, seed=111, overall_confidence=0.4257, " + "ligand_confidence=0.4414, seq_rec=0.5319\n" + "CANDIDATEFOUR\n" + ) + # Mirrors analysis_sequence()'s parsing exactly (small_molecule_binding.py). + best_conf, best_id, best_seq = -1.0, None, None + for record in fixture.split('>')[1:]: + lines = record.splitlines() + header, seq = lines[0], ''.join(lines[1:]).strip() + parts = { + kv.split('=')[0].strip(): kv.split('=')[1].strip() + for kv in header.split(',') if '=' in kv + } + if 'id' not in parts: + continue + conf = float(parts.get('overall_confidence', 0)) + if conf > best_conf: + best_conf, best_id, best_seq = conf, parts['id'], seq + + assert best_id == '4', f"expected id=4 (highest overall_confidence), got id={best_id!r}" + assert best_seq == 'CANDIDATEFOUR', f"expected candidate 4's sequence, got {best_seq!r}" + assert abs(best_conf - 0.4257) < 1e-6, f"expected conf=0.4257, got {best_conf}" + + +def check_fastrelax_interface_shortcircuit() -> None: + """Regression check for _stage_metrics_improving(), the metric-agnostic + fastrelax/interface short-circuit (see plan-shortcircuit-farep-loop.md). + Validated against real HPC data (job 21913252, all three of p3's + backbones) before landing -- these fixtures are that same real data.""" + from small_molecule_binding import _stage_metrics_improving + + fastrelax_specs = [ + ('interact', True, -8.0), + ('total_score', True, -250.0), + ('fa_rep', True, 100.0), + ] + + # First attempt on a backbone: nothing to compare against yet -- always retry. + assert _stage_metrics_improving({'interact': -7.5, 'total_score': -200.0, 'fa_rep': 56.0}, None, fastrelax_specs) is True + + # fa_rep-only failure, flat across attempts (real data: p3 backbone 3, + # attempts 1->2) -- must be detected as NOT improving. + prev = {'interact': -20.17, 'total_score': -413.98, 'fa_rep': 103.81} + cur = {'interact': -17.94, 'total_score': -409.12, 'fa_rep': 104.03} + assert _stage_metrics_improving(cur, prev, fastrelax_specs) is False, \ + "flat fa_rep-only failure should not be read as improving" + + # A real, meaningful improvement should still be allowed to retry: fa_rep + # starts above threshold (failing, 110.0 > 100.0) and drops well under it. + prev = {'interact': -20.0, 'total_score': -400.0, 'fa_rep': 110.0} + cur = {'interact': -20.0, 'total_score': -400.0, 'fa_rep': 60.0} # fa_rep way down + assert _stage_metrics_improving(cur, prev, fastrelax_specs) is True, \ + "a real fa_rep improvement should be read as improving" + + # interface (higher-is-better) uses the same function with lower_is_better=False. + interface_specs = [('max_sc', False, 0.55)] + prev = {'max_sc': 0.5179} + cur = {'max_sc': 0.5148} # real data: p3 backbone 2, attempts 3->5 direction + assert _stage_metrics_improving(cur, prev, interface_specs) is False async def run_mock_test() -> None: pipeline_name = "p1" setup_mock_inputs(pipeline_name) - backend = await ConcurrentExecutionBackend(ThreadPoolExecutor()) + backend = await LocalExecutionBackend(ThreadPoolExecutor()) manager: ImpressManager = ImpressManager(execution_backend=backend) pipeline_setups: List[PipelineSetup] = [ @@ -88,5 +171,7 @@ async def run_mock_test() -> None: if __name__ == "__main__": - check_af2_filename_derivation() + check_boltz_filename_derivation() + check_mpnn_candidate_selection() + check_fastrelax_interface_shortcircuit() asyncio.run(run_mock_test()) diff --git a/examples/small_molecule_binding/scripts/af2.sh b/examples/small_molecule_binding/scripts/af2.sh deleted file mode 100755 index ede53a8..0000000 --- a/examples/small_molecule_binding/scripts/af2.sh +++ /dev/null @@ -1,23 +0,0 @@ -#!/bin/bash -set -euo pipefail - -# AlphaFold2 structure prediction via LocalColabFold (pixi) -# Args: $1=colabfold_path $2=short_fasta $3=output_dir - -colabfold_path="$1" -short_fasta="$2" -output_dir="$3" - -module load gcc/11.2.0 -module load cuda -source /anvil/projects/x-nairr240405/mason/IMPRESS/.venv/bin/activate - -pixi run --manifest-path "$colabfold_path" \ - colabfold_batch \ - --model-type alphafold2 \ - --rank auto \ - --random-seed 999 \ - --save-all \ - --debug-logging \ - "$short_fasta" \ - "$output_dir" diff --git a/examples/small_molecule_binding/scripts/boltz.sh b/examples/small_molecule_binding/scripts/boltz.sh new file mode 100755 index 0000000..6afb7e7 --- /dev/null +++ b/examples/small_molecule_binding/scripts/boltz.sh @@ -0,0 +1,76 @@ +#!/bin/bash +set -euo pipefail +# Protein+ligand co-folding via Boltz-2 (pip CLI, no container) +# Args: $1=input_yaml $2=output_dir $3=boltz_cache_dir +input_yaml="$1"; output_dir="$2"; boltz_cache_dir="$3" + +# $boltz_cache_dir is shared across concurrently-running pipelines. boltz's own +# download_boltz2() checks `mols.exists()` (directory presence), not +# completeness, before skipping extraction — tarfile.extractall() creates the +# "mols" directory entry immediately, so a *second* concurrent pipeline calling +# download_boltz2() while a first one is still mid-extract sees mols/ already +# existing and skips extraction outright, then reads a half-populated +# directory and fails with "CCD component not found!" for whatever +# hasn't been extracted yet. A prior fix here only checked mols.tar's archive +# integrity under a lock, which doesn't guard this race at all (mols.tar can +# be perfectly valid while mols/ is still being extracted from it elsewhere). +# +# Fix: hold the lock for the entire check-and-repair, verify mols/ actually +# contains every file mols.tar lists (not just that the directory exists), +# and if not, delete and re-extract *inside* the lock via boltz's own +# download_boltz2() so no other concurrent pipeline can observe a +# partially-populated mols/ while this one repairs it. A `.mols_complete` +# marker (written only after a verified-complete extraction) lets later +# invocations skip the O(45k) file-count re-check once warmed. +mkdir -p "$boltz_cache_dir" +( + flock -x 200 + + tar_ok=false + if tar -tf "$boltz_cache_dir/mols.tar" >/dev/null 2>&1; then + tar_ok=true + fi + + mols_complete=false + if $tar_ok && [ -f "$boltz_cache_dir/.mols_complete" ]; then + mols_complete=true + elif $tar_ok && [ -d "$boltz_cache_dir/mols" ]; then + # No marker yet (cache predates this fix, or a previous repair was + # interrupted) -- verify extraction actually completed rather than + # trusting mere directory existence. + expected=$(tar -tf "$boltz_cache_dir/mols.tar" | grep -vc '/$') + actual=$(find "$boltz_cache_dir/mols" -maxdepth 1 -type f | wc -l) + if [ "$actual" -eq "$expected" ]; then + mols_complete=true + touch "$boltz_cache_dir/.mols_complete" + fi + fi + + if ! $mols_complete; then + rm -rf "$boltz_cache_dir/mols.tar" "$boltz_cache_dir/mols" "$boltz_cache_dir/.mols_complete" + BOLTZ_CACHE_DIR_FOR_PY="$boltz_cache_dir" python -c " +import os +from pathlib import Path +from boltz.main import download_boltz2 +download_boltz2(Path(os.environ['BOLTZ_CACHE_DIR_FOR_PY'])) +" + touch "$boltz_cache_dir/.mols_complete" + fi +) 200>"$boltz_cache_dir/.download.lock" + +# --no_kernels: cuequivariance_ops_torch's compiled kernel (used for the fused +# triangular-multiplication op) requires cublasGemmGroupedBatchedEx, which is +# absent from nvidia-cublas-cu12==12.1.3.1 (the exact version torch==2.5.1+cu121 +# pins and loads first via its own RPATH) and only present from 12.5.3.2+ -- +# verified by inspecting both wheels' libcublas.so.12 with `nm -D`. That ABI +# mismatch makes the kernel import fail every time (reproduced with no GPU +# present: `python -c "import torch; import cuequivariance_ops_torch"`), not +# just intermittently, so --no_kernels (falls back to plain PyTorch ops) is +# required until torch's pinned nvidia-cublas-cu12 and cuequivariance-ops-cu12 +# are reconciled -- don't remove this thinking it's a leftover. +boltz predict "$input_yaml" \ + --out_dir "$output_dir" --cache "$boltz_cache_dir" \ + --devices 1 --accelerator gpu \ + --diffusion_samples "${BOLTZ_DIFFUSION_SAMPLES:-1}" \ + --output_format pdb \ + --no_kernels diff --git a/examples/small_molecule_binding/scripts/derive_ligand_smiles.py b/examples/small_molecule_binding/scripts/derive_ligand_smiles.py new file mode 100644 index 0000000..9b51239 --- /dev/null +++ b/examples/small_molecule_binding/scripts/derive_ligand_smiles.py @@ -0,0 +1,377 @@ +"""One-time SMILES derivation for a Rosetta-params-defined ligand. + +Rosetta `.params` files give exact atom identity and bond *connectivity* +(which atoms are bonded to which) but no bond order and no explicit +hydrogens beyond whatever is literally listed as an ATOM record. To hand a +ligand to Boltz-2 (which wants a SMILES string) we need real bond orders, +aromaticity, and formal charges -- RDKit's `rdDetermineBonds.DetermineBondOrders` +is built exactly for this: given known connectivity plus a 3D conformer, it +solves for a chemically sensible Lewis structure. + +This is a one-time, offline tool -- not part of the per-cycle pipeline. + +Usage: + derive_ligand_smiles.py .params .pdb [--charge N] [--out PATH] + +Implementation note (found empirically against ALR.params + scaffold-with-ALR.pdb, +a fused-bicyclic (naphthalene) + monocyclic aromatic azo-linked bis-sulfonate): +DetermineBondOrders on a *heavy-atom-only* skeleton (no explicit hydrogens at +all, relying on RDKit's implicit-H valence fill) reliably failed or returned +chemically nonsensical structures (cumulated double/triple bonds, absurd +formal charges) across a wide charge sweep -- for both the full molecule and +the fused-ring "core" with the sulfonate groups excluded. Adding placeholder +explicit hydrogens (approximate, not chemically precise, positions -- one per +params BOND record referencing an atom absent from the reference PDB) gave +DetermineBondOrders a fully-specified valence at every atom and immediately +produced the correct aromatic, charge-balanced structure. This script +therefore reconstructs those hydrogens with approximate 3D placeholder +coordinates (never claimed to be chemically accurate bond geometry -- just +distinct, non-degenerate positions) rather than dropping them outright. +""" + +import argparse +import pathlib +import re + +import numpy as np +from rdkit import Chem +from rdkit.Chem import rdDetermineBonds +from rdkit.Geometry import Point3D + + +# Fallback table keyed on Rosetta atom TYPE prefix, used only when the atom +# NAME's leading alphabetic run doesn't parse to a valid element symbol. +# Verified only against ALR's atom set (C/N/O/S/H). NOT verified for 2-letter +# elements (Cl/Br/Zn, etc.) -- extend this table if IND/RED (or any future +# ligand) ever need it. Nothing ALR-specific is hardcoded into the matching +# logic itself, only into the table's contents. +_ROSETTA_TYPE_ELEMENT_FALLBACK = { + "Nhis": "N", "Nlys": "N", + "CH1": "C", "CH2": "C", "CH3": "C", "COO": "C", "aroC": "C", + "OH": "O", "OOC": "O", "ONH2": "O", + "Hapo": "H", "Hpol": "H", + "S": "S", "SH1": "S", +} + +_PERIODIC_TABLE = Chem.GetPeriodicTable() + + +def _parse_params(params_path): + """Parse a Rosetta .params file. + + Returns (resname, {atom_name: rosetta_type}, [(atom1, atom2), ...]) from + the file's NAME, ATOM, and BOND records. Includes hydrogens (both in the + atom-type map and the bond list) -- callers that want a heavy-atom-only + view filter separately. + """ + resname = None + atom_types = {} + bonds = [] + + with open(params_path) as fh: + for line in fh: + line = line.rstrip("\n") + if not line.strip(): + continue + fields = line.split() + record = fields[0] + + if record == "NAME": + resname = fields[1] + elif record == "ATOM": + # ATOM + atom_name, rosetta_type = fields[1], fields[2] + atom_types[atom_name] = rosetta_type + elif record == "BOND" or record == "BOND_TYPE": + # BOND [bond order, for BOND_TYPE] + a1, a2 = fields[1], fields[2] + bonds.append((a1, a2)) + + if resname is None: + raise ValueError(f"{params_path}: no NAME record found") + if not atom_types: + raise ValueError(f"{params_path}: no ATOM records found") + + return resname, atom_types, bonds + + +def _infer_element(atom_name, rosetta_type): + """Infer the element symbol for a params ATOM entry. + + Primary rule: the leading alphabetic run of the atom NAME field is + literally the PDB-style element-derived atom name (e.g. "N11" -> "N", + "C13" -> "C", "S1" -> "S"). If that run isn't a valid element symbol + (e.g. it's empty, or the atom-naming convention doesn't follow this + pattern), fall back to a table keyed on the Rosetta atom TYPE prefix. + + Verified only for ALR's C/N/O/S/H atom set. NOT verified for 2-letter + elements (Cl/Br/Zn, ...) -- IND/RED ligands exist in this repo but have + no reference PDB, so they're out of scope until one is added; extending + this function for them should not require touching ALR's behavior. + """ + match = re.match(r"[A-Za-z]+", atom_name) + if match: + candidate = match.group(0) + # Try progressively shorter prefixes (handles e.g. "Cl1" correctly + # while still falling back cleanly for made-up multi-letter runs). + for length in (2, 1): + if len(candidate) >= length: + symbol = candidate[:length].capitalize() + if _PERIODIC_TABLE.GetAtomicNumber(symbol) > 0: + return symbol + + # Name-based inference failed -- fall back to the Rosetta TYPE prefix table. + for prefix, element in _ROSETTA_TYPE_ELEMENT_FALLBACK.items(): + if rosetta_type.startswith(prefix): + return element + + raise ValueError( + f"could not infer element for atom name={atom_name!r} type={rosetta_type!r}" + ) + + +def _load_reference_coords(pdb_path, resname): + """Read {atom_name: (x, y, z)} from the first HETATM residue in pdb_path + whose residue name matches `resname`. + + PDB fixed-column parsing is used for the residue-name field (columns + 18-20) since names like "A:R" contain a colon that a naive whitespace + split would otherwise mangle. + """ + coords = {} + found_residue = False + target_resseq = None + + with open(pdb_path) as fh: + for line in fh: + if not (line.startswith("HETATM") or line.startswith("ATOM ")): + continue + + line_resname = line[17:20].strip() + if line_resname != resname: + if found_residue: + # We've moved past the matching residue's contiguous block. + break + continue + + resseq = line[22:26].strip() + if target_resseq is None: + target_resseq = resseq + elif resseq != target_resseq: + # A different residue instance with the same name -- stop at + # the first one, per the docstring contract. + break + + found_residue = True + atom_name = line[12:16].strip() + x = float(line[30:38]) + y = float(line[38:46]) + z = float(line[46:54]) + coords[atom_name] = (x, y, z) + + if not coords: + raise ValueError(f"{pdb_path}: no HETATM residue named {resname!r} found") + + return coords + + +def _placeholder_h_coords(h_names, atom_names_by_parent, ref_coords, bonds): + """Approximate (not chemically precise) 3D positions for hydrogens absent + from the reference PDB, so DetermineBondOrders sees a fully-specified + valence at every heavy atom instead of guessing implicit H counts. + + Each H is placed near its (single) bonded heavy-atom parent, offset in a + direction generally pointing away from that parent's other heavy + neighbors -- distinct per-H when several hydrogens share one parent + (e.g. a methyl group) by fanning out around an arbitrary perpendicular + axis. Precise bond lengths/angles are not the goal (empirically, + DetermineBondOrders's bond-order search doesn't depend on them once + connectivity is fixed) -- only non-degenerate, distinguishable positions. + """ + # parent heavy atom for each H (H atoms have exactly one bond in a + # correctly-formed params file) + parent_of = {} + for a, b in bonds: + if a in h_names and b in ref_coords: + parent_of[a] = b + elif b in h_names and a in ref_coords: + parent_of[b] = a + + h_coords = {} + for parent, h_list in atom_names_by_parent.items(): + parent_pos = np.array(ref_coords[parent], dtype=float) + + other_heavy_neighbors = [ + ref_coords[nb] for a, b in bonds + for nb in ((b,) if a == parent else (a,) if b == parent else ()) + if nb in ref_coords and nb != parent + ] + if other_heavy_neighbors: + centroid = np.mean(np.array(other_heavy_neighbors, dtype=float), axis=0) + base_dir = parent_pos - centroid + else: + base_dir = np.array([1.0, 0.0, 0.0]) + norm = np.linalg.norm(base_dir) + base_dir = base_dir / norm if norm > 1e-6 else np.array([1.0, 0.0, 0.0]) + + # arbitrary axis perpendicular to base_dir, for fanning out multiple H's + arbitrary = np.array([0.0, 0.0, 1.0]) if abs(base_dir[2]) < 0.9 else np.array([0.0, 1.0, 0.0]) + perp_axis = np.cross(base_dir, arbitrary) + perp_axis /= np.linalg.norm(perp_axis) + + n_h = len(h_list) + for i, h_name in enumerate(h_list): + angle = np.radians((360.0 / n_h) * i) if n_h > 1 else 0.0 + # Rodrigues' rotation of base_dir around perp_axis by `angle` + rotated = ( + base_dir * np.cos(angle) + + np.cross(perp_axis, base_dir) * np.sin(angle) + + perp_axis * np.dot(perp_axis, base_dir) * (1 - np.cos(angle)) + ) + pos = parent_pos + rotated * 1.0 # arbitrary ~1 Angstrom offset + h_coords[h_name] = tuple(pos) + + return h_coords + + +def derive_smiles(params_path, reference_pdb_path, net_charge=0): + """Derive a SMILES string for the ligand described by params_path, + using 3D coordinates from reference_pdb_path (heavy atoms) plus + reconstructed placeholder coordinates for hydrogens absent from that PDB + to resolve bond orders. + + Atoms: every params heavy atom found in the reference PDB, plus every + params hydrogen bonded to one of those heavy atoms (given a placeholder + position -- see _placeholder_h_coords). Params ATOM/BOND entries for + atoms that are neither in the PDB nor a hydrogen bonded to a kept heavy + atom are dropped. + + Bonds are wired as single bonds initially; rdDetermineBonds.DetermineBondOrders + then infers real bond order/aromaticity/formal charges from connectivity + and valence. The hydrogens are stripped from the final returned molecule + (Chem.RemoveHs) so the SMILES reflects only the ligand's heavy-atom + skeleton, same as a normal canonical SMILES. + + net_charge defaults to 0 because ALR.params's per-atom partial charges + happen to sum to roughly zero -- this is a convenient heuristic, not a + rigorous formal-charge derivation. If charge=0 fails, +1/-1 are tried + next; if those also fail, the search widens further (+/-2, +/-3, +/-4) + since a real bis-sulfonic-acid ligand is, in practice, virtually always + doubly deprotonated (net charge -2) at neutral pH -- found empirically + for ALR, not assumed a priori. + """ + resname, atom_types, bond_pairs = _parse_params(params_path) + ref_coords = _load_reference_coords(reference_pdb_path, resname) + + heavy_names = [name for name in atom_types if name in ref_coords] + if not heavy_names: + raise ValueError( + f"no overlap between params atoms and reference PDB atoms for {resname!r}" + ) + heavy_set = set(heavy_names) + + # Hydrogens (or any other atom not in the PDB) bonded to a kept heavy atom. + h_names = [ + name for name in atom_types + if name not in heavy_set + and any(name in pair and (pair[0] in heavy_set or pair[1] in heavy_set) for pair in bond_pairs) + ] + h_set = set(h_names) + + kept_names = heavy_names + h_names + kept_set = set(kept_names) + kept_bonds = [ + (a1, a2) for (a1, a2) in bond_pairs + if a1 in kept_set and a2 in kept_set + ] + + # Group H names by their heavy-atom parent, for placeholder placement. + atom_names_by_parent = {} + for a, b in kept_bonds: + if a in h_set and b in heavy_set: + atom_names_by_parent.setdefault(b, []).append(a) + elif b in h_set and a in heavy_set: + atom_names_by_parent.setdefault(a, []).append(b) + + h_coords = _placeholder_h_coords(h_set, atom_names_by_parent, ref_coords, kept_bonds) + all_coords = {**ref_coords, **h_coords} + + # Build the RWMol: atoms first (recording an index map), then bonds. + mol = Chem.RWMol() + name_to_idx = {} + for name in kept_names: + element = _infer_element(name, atom_types[name]) + atom = Chem.Atom(element) + idx = mol.AddAtom(atom) + name_to_idx[name] = idx + + for a1, a2 in kept_bonds: + i, j = name_to_idx[a1], name_to_idx[a2] + if mol.GetBondBetweenAtoms(i, j) is None: + mol.AddBond(i, j, Chem.BondType.SINGLE) + + # Attach the 3D conformer (real PDB coords for heavy atoms, placeholder + # coords for reconstructed hydrogens). + conformer = Chem.Conformer(mol.GetNumAtoms()) + for name, idx in name_to_idx.items(): + x, y, z = all_coords[name] + conformer.SetAtomPosition(idx, Point3D(x, y, z)) + mol.AddConformer(conformer, assignId=True) + + # Sanitize NONE first -- DetermineBondOrders operates on the raw graph + # and will itself figure out valence/order/aromaticity/charges. + Chem.SanitizeMol(mol, sanitizeOps=Chem.SanitizeFlags.SANITIZE_NONE) + + charge_ladder = [net_charge, net_charge + 1, net_charge - 1, + net_charge + 2, net_charge - 2, + net_charge + 3, net_charge - 3, + net_charge + 4, net_charge - 4] + last_error = None + for charge in charge_ladder: + trial_mol = Chem.RWMol(mol) + try: + rdDetermineBonds.DetermineBondOrders(trial_mol, charge=charge) + Chem.SanitizeMol(trial_mol) + except Exception as exc: # noqa: BLE001 -- want to try every charge in the ladder + last_error = exc + continue + heavy_only = Chem.RemoveHs(trial_mol) + return Chem.MolToSmiles(heavy_only) + + raise RuntimeError( + f"DetermineBondOrders failed across charge ladder {charge_ladder}; " + f"last error: {last_error}" + ) + + +def main(): + parser = argparse.ArgumentParser( + description="Derive a SMILES string for a Rosetta-params ligand from its " + "params connectivity + a reference PDB's 3D coordinates." + ) + parser.add_argument("params_path", help="Rosetta .params file (e.g. ALR.params)") + parser.add_argument("reference_pdb_path", help="Reference PDB containing the ligand's HETATM block") + parser.add_argument("--charge", type=int, default=0, help="Net formal charge to try first (default 0)") + parser.add_argument("--out", default=None, help="Output .smiles path (default: /.smiles)") + args = parser.parse_args() + + params_path = pathlib.Path(args.params_path) + out_path = pathlib.Path(args.out) if args.out else params_path.with_suffix(".smiles") + + smiles = derive_smiles(str(params_path), args.reference_pdb_path, net_charge=args.charge) + + # Round-trip through Chem.MolFromSmiles() before writing -- abort if it + # doesn't parse back to a valid molecule. + round_trip_mol = Chem.MolFromSmiles(smiles) + if round_trip_mol is None: + raise RuntimeError( + f"derived SMILES failed to round-trip through Chem.MolFromSmiles(): {smiles!r}" + ) + + out_path.write_text(smiles + "\n") + print(smiles) + print(f"wrote {out_path}") + + +if __name__ == "__main__": + main() diff --git a/examples/small_molecule_binding/scripts/fastrelax.py b/examples/small_molecule_binding/scripts/fastrelax.py index b899106..cfd78be 100755 --- a/examples/small_molecule_binding/scripts/fastrelax.py +++ b/examples/small_molecule_binding/scripts/fastrelax.py @@ -100,7 +100,7 @@ def main(args): if __name__ == '__main__': args = parse_args() - opts = '-ex1 -ex2 -use_input_sc -flip_HNQ -no_optH false -out:level 500' + opts = '-ex1 -ex2 -use_input_sc -flip_HNQ -no_optH false -mute all' if args.constraints: opts += ' -enzdes::cstfile {}'.format(args.constraints) opts += ' -run:preserve_header' diff --git a/examples/small_molecule_binding/scripts/fastrelax.sh b/examples/small_molecule_binding/scripts/fastrelax.sh index 57a3987..b69403d 100755 --- a/examples/small_molecule_binding/scripts/fastrelax.sh +++ b/examples/small_molecule_binding/scripts/fastrelax.sh @@ -8,9 +8,9 @@ pdb_path="$1" lig_path="$2" output_dir="$3" -SCRIPT_DIR="$(dirname $0)" +SCRIPT_DIR="$(dirname "$0")" -source /anvil/projects/x-nairr240405/mason/IMPRESS/.venv/bin/activate +[ -n "${VIRTUAL_ENV:-}" ] && source "${VIRTUAL_ENV}/bin/activate" python "$SCRIPT_DIR/fastrelax.py" \ "$pdb_path" \ diff --git a/examples/small_molecule_binding/scripts/filter_energy.py b/examples/small_molecule_binding/scripts/filter_energy.py index 4acb46d..a570cd8 100644 --- a/examples/small_molecule_binding/scripts/filter_energy.py +++ b/examples/small_molecule_binding/scripts/filter_energy.py @@ -42,8 +42,6 @@ parts = line.split() ligand_energy = float(parts[-1]) # The last element is the total energy - print(f"Processed: {pdb_file}, Ligand Energy ({ligand_name}): {ligand_energy}") - # Check if the ligand energy is negative if ligand_energy < 0: with open(output_file, 'a') as of: diff --git a/examples/small_molecule_binding/scripts/filter_shape.py b/examples/small_molecule_binding/scripts/filter_shape.py index 87278a5..cfb8ab7 100644 --- a/examples/small_molecule_binding/scripts/filter_shape.py +++ b/examples/small_molecule_binding/scripts/filter_shape.py @@ -1,30 +1,37 @@ +import argparse import os import re import pyrosetta -import sys -pdb_directory = sys.argv[1] # '/WWW/PDB_Files' -SC_output_file = sys.argv[2] # 'shape_complementarity_values.txt' -ligand_name = sys.argv[3] # 'ALR' -gen_output_file = sys.argv[4] # 'interface_values.txt' -pyrosetta.init(f"-ignore_unrecognized_res -ignore_zero_occupancy --extra_res_fa {ligand_name}.params -corrections::beta_nov16 true") +def parse_args(): + parser = argparse.ArgumentParser(description="Rosetta shape-complementarity/interface analysis") + parser.add_argument("pdb_directory", help="Directory of PDB files to analyse") + parser.add_argument("SC_output_file", help="Output file for shape-complementarity values") + parser.add_argument("ligand_name", help="Ligand residue name (e.g. ALR)") + parser.add_argument("gen_output_file", help="Output CSV for interface metrics") + return parser.parse_args() -# Define directories and files -#pdb_directory = '/WWW/PDB_Files' -#SC_output_file = 'shape_complementarity_values.txt' -# Set up the general output file which will have metrics that look at the interface -with open(gen_output_file, 'w') as genout: - # Write the shape complementarity value to the output file - genout.write(f"FileName,Shape Complementarity,ddg,contact molecular surf,SASA,Very buried unsat hbond,Surface unsat hbond,SAP SCORE\n") - genout.close() +def main(): + args = parse_args() -# Create list of PDB files to analyze -pdb_files = [f for f in os.listdir(pdb_directory) if f.endswith('.pdb')] + pyrosetta.init( + f"-ignore_unrecognized_res -ignore_zero_occupancy --extra_res_fa {args.ligand_name}.params" + f" -corrections::beta_nov16 true -mute all" + ) -################################################################################################################################################################################################################################### -protocol = pyrosetta.rosetta.protocols.rosetta_scripts.XmlObjects().create_from_string( + # Set up the general output file which will have metrics that look at the interface + with open(args.gen_output_file, 'w') as genout: + genout.write( + "FileName,Shape Complementarity,ddg,contact molecular surf,SASA," + "Very buried unsat hbond,Surface unsat hbond,SAP SCORE\n" + ) + + # Create list of PDB files to analyze + pdb_files = [f for f in os.listdir(args.pdb_directory) if f.endswith('.pdb')] + + protocol = pyrosetta.rosetta.protocols.rosetta_scripts.XmlObjects().create_from_string( """ @@ -34,7 +41,7 @@ - + @@ -93,41 +100,25 @@ """).get_mover("ParsedProtocol") -#################################################################################################################################################################################################################################### - - -##print(f""" -##shape complementarity : {pose.scores['sc2']} -##ddg : {pose.scores['ddg']} -##contact molecular surf : {pose.scores['cms']} -##SASA : {pose.scores['IA_dSASA_int']} -##Very buried unsat hbond: {pose.scores['vbuns']} -##Surface unsat hbond : {pose.scores['sbuns']} -##SAP SCORE : {pose.scores['sap_score']} -##""") - -# Analyze the selected PDB files -for pdb_file in pdb_files: - full_path = os.path.join(pdb_directory, pdb_file) - - # Initialize variables for ligand energy - pose = pyrosetta.pose_from_pdb(full_path) - protocol.apply(pose) + # Analyze the selected PDB files + for pdb_file in pdb_files: + full_path = os.path.join(args.pdb_directory, pdb_file) - # Open the SC output file - with open(SC_output_file, 'a') as SCout: - - # Write the shape complementarity value to the output file - SCout.write(f"{pdb_file}\tShape Complementarity: {pose.scores['sc2']}\n") - SCout.close() + pose = pyrosetta.pose_from_pdb(full_path) + protocol.apply(pose) - # Open the general output file - with open(gen_output_file, 'a') as genout: + with open(args.SC_output_file, 'a') as SCout: + SCout.write(f"{pdb_file}\tShape Complementarity: {pose.scores['sc2']}\n") - # Write the shape complementarity value to the output file - genout.write(f"{pdb_file},{pose.scores['sc2']},{pose.scores['ddg']},{pose.scores['cms']},{pose.scores['IA_dSASA_int']},{pose.scores['vbuns']},{pose.scores['sbuns']},{pose.scores['sap_score']}\n") - genout.close() + with open(args.gen_output_file, 'a') as genout: + genout.write( + f"{pdb_file},{pose.scores['sc2']},{pose.scores['ddg']}," + f"{pose.scores['cms']},{pose.scores['IA_dSASA_int']}," + f"{pose.scores['vbuns']},{pose.scores['sbuns']},{pose.scores['sap_score']}\n" + ) -print(f"Shape complementarity values have been written to {SC_output_file}.") + print(f"Shape complementarity values have been written to {args.SC_output_file}.") +if __name__ == "__main__": + main() diff --git a/examples/small_molecule_binding/scripts/filter_shape.sh b/examples/small_molecule_binding/scripts/filter_shape.sh index 2208834..48979c0 100755 --- a/examples/small_molecule_binding/scripts/filter_shape.sh +++ b/examples/small_molecule_binding/scripts/filter_shape.sh @@ -11,7 +11,7 @@ interface_values_output="$4" SCRIPT_DIR="$(dirname "$0")" -source /anvil/projects/x-nairr240405/mason/IMPRESS/.venv/bin/activate +[ -n "${VIRTUAL_ENV:-}" ] && source "${VIRTUAL_ENV}/bin/activate" python "$SCRIPT_DIR/filter_shape.py" \ "$pdb_directory" \ diff --git a/examples/small_molecule_binding/scripts/mpnn.sh b/examples/small_molecule_binding/scripts/mpnn.sh index c5aaf1d..29378eb 100755 --- a/examples/small_molecule_binding/scripts/mpnn.sh +++ b/examples/small_molecule_binding/scripts/mpnn.sh @@ -11,9 +11,11 @@ n_batches="$4" batch_size="$5" fixed_residues="${6:-}" -source /anvil/projects/x-nairr240405/mason/LigandMPNN/.venv/bin/activate +SCRIPT_DIR="$(dirname "$0")" -python "$mpnn_dir/run.py" \ +# mpnn_run.py restores numpy deprecated aliases (np.int/np.bool/np.object) +# removed in NumPy 1.24+ that LigandMPNN's bundled openfold still uses. +python "$SCRIPT_DIR/mpnn_run.py" "$mpnn_dir" \ --model_type "ligand_mpnn" \ --checkpoint_path_sc "$mpnn_dir/model_params/ligandmpnn_sc_v_32_002_16.pt" \ --checkpoint_ligand_mpnn "$mpnn_dir/model_params/ligandmpnn_v_32_010_25.pt" \ diff --git a/examples/small_molecule_binding/scripts/mpnn_run.py b/examples/small_molecule_binding/scripts/mpnn_run.py new file mode 100644 index 0000000..4c7cfd2 --- /dev/null +++ b/examples/small_molecule_binding/scripts/mpnn_run.py @@ -0,0 +1,25 @@ +#!/usr/bin/env python +""" +Wrapper for LigandMPNN run.py that restores deprecated numpy aliases removed +in NumPy 1.24+. LigandMPNN's bundled openfold uses np.int, np.object, np.bool. + +Usage (from mpnn.sh): + python mpnn_run.py [run.py args...] +The mpnn_dir is stripped from sys.argv before run.py sees it. +""" +import sys +import os + +import numpy as np +for _alias, _builtin in [('int', int), ('float', float), ('bool', bool), + ('complex', complex), ('object', object), ('str', str)]: + if not hasattr(np, _alias): + setattr(np, _alias, _builtin) + +mpnn_dir = sys.argv[1] +sys.argv = [os.path.join(mpnn_dir, 'run.py')] + sys.argv[2:] +sys.path.insert(0, mpnn_dir) +os.chdir(mpnn_dir) + +import runpy +runpy.run_path(os.path.join(mpnn_dir, 'run.py'), run_name='__main__') diff --git a/examples/small_molecule_binding/scripts/mpnn_wrapper.sh b/examples/small_molecule_binding/scripts/mpnn_wrapper.sh deleted file mode 100755 index 545901b..0000000 --- a/examples/small_molecule_binding/scripts/mpnn_wrapper.sh +++ /dev/null @@ -1,28 +0,0 @@ -#!/bin/bash -# run ligandmpnn - -pdb_path=$1 -output_dir=$2 -lmpnn_dir=$3 - -echo "A16" > fixed_residues.txt -FIXED=`cat fixed_residues.txt` -set -x -python $lmpnn_dir/run.py \ - --model_type "ligand_mpnn" \ - --checkpoint_path_sc $lmpnn_dir/model_params/ligandmpnn_sc_v_32_002_16.pt \ - --checkpoint_ligand_mpnn $lmpnn_dir/model_params/ligandmpnn_v_32_010_25.pt \ - --seed 111 \ - --pdb_path $pdb_path \ - --out_folder $output_dir \ - --pack_side_chains 1 \ - --number_of_batches 1 \ - --batch_size 1 \ - --number_of_packs_per_design 1 \ - --pack_with_ligand_context 1 \ - --fixed_residues "$FIXED" \ - --repack_everything 1 \ - --temperature 0.1 -# --bias_AA "A:10.0" - - diff --git a/examples/small_molecule_binding/scripts/packmin.py b/examples/small_molecule_binding/scripts/packmin.py index 8dcbf97..d32bbdd 100644 --- a/examples/small_molecule_binding/scripts/packmin.py +++ b/examples/small_molecule_binding/scripts/packmin.py @@ -7,31 +7,12 @@ from pyrosetta.teaching import * #Core Includes -from rosetta.core.kinematics import MoveMap -from rosetta.core.kinematics import FoldTree -from rosetta.core.pack.task import TaskFactory -from rosetta.core.pack.task import operation -from rosetta.core.simple_metrics import metrics -from rosetta.core.select import residue_selector as selections -from rosetta.core import select -from rosetta.core.select.movemap import * +from pyrosetta.rosetta.core.pack.task import operation +from pyrosetta.rosetta.core.select.movemap import * #Protocol Includes -from rosetta.protocols import minimization_packing as pack_min -from rosetta.protocols import relax as rel -from rosetta.protocols.antibody.residue_selector import CDRResidueSelector -from rosetta.protocols.antibody import * -from rosetta.protocols.loops import * -''' -When downloading a new PDB file, do a pack_min minimization with coordinate constraints and a ligand. - -Requires a PDB file input. - -Options: -Name (-n, string): change the output PDB name from [original_name]_relaxed.pdb -Score function (-sf, string): change the score function from the default of ref2015_cst -Catalytic residues (-cat, int, multiple accepted): list residues that should not be moved -''' +from pyrosetta.rosetta.protocols import minimization_packing as pack_min + def parse_args(): parser = argparse.ArgumentParser() @@ -93,18 +74,12 @@ def main(args): # Packer tasks with -ex1 and -ex2 tf = ut.make_task_factory() - #From pack_min tutorial from jupyter notebooks - ###tf = TaskFactory() tf.push_back(operation.InitializeFromCommandline()) tf.push_back(operation.RestrictToRepacking()) packer = pack_min.PackRotamersMover() packer.task_factory(tf) - #This line is from khare lab relax code, not sure if this will be necessary: - #pp = Pose(pose) - - #Run the packer. (Note this may take a few minutes) packer.apply(pose) # Write Rosetta score JSON alongside the output PDB @@ -115,12 +90,11 @@ def main(args): #Dump the PDB pose.dump_pdb(out_name) - ##pose.dump_pdb('/outputs/2r0l_all_repack.pdb') if __name__ == '__main__': args = parse_args() - opts = '-ex1 -ex2 -use_input_sc -flip_HNQ -no_optH false' + opts = '-ex1 -ex2 -use_input_sc -flip_HNQ -no_optH false -mute all' if args.constraints: opts += ' -enzdes::cstfile {}'.format(args.constraints) opts += ' -run:preserve_header' diff --git a/examples/small_molecule_binding/scripts/packmin.sh b/examples/small_molecule_binding/scripts/packmin.sh index 5ce3be7..5be9629 100755 --- a/examples/small_molecule_binding/scripts/packmin.sh +++ b/examples/small_molecule_binding/scripts/packmin.sh @@ -10,7 +10,7 @@ output_dir="$3" SCRIPT_DIR="$(dirname "$0")" -source /anvil/projects/x-nairr240405/mason/IMPRESS/.venv/bin/activate +[ -n "${VIRTUAL_ENV:-}" ] && source "${VIRTUAL_ENV}/bin/activate" python "$SCRIPT_DIR/packmin.py" \ "$pdb_path" \ diff --git a/examples/small_molecule_binding/scripts/rfd3.sh b/examples/small_molecule_binding/scripts/rfd3.sh index 33b27e3..8761792 100755 --- a/examples/small_molecule_binding/scripts/rfd3.sh +++ b/examples/small_molecule_binding/scripts/rfd3.sh @@ -2,26 +2,29 @@ set -euo pipefail # Backbone generation via RFDiffusion3 (apptainer) -# Args: $1=foundry_sif_path $2=output_dir $3=inputs $4=scaffold_arg $5=diffusion_batch_size -# scaffold_arg: "scaffoldguided.target_pdb=" or "" if unused +# Args: $1=foundry_sif_path $2=output_dir $3=inputs $4=diffusion_batch_size +# +# RFD3 has no scaffold/guidance CLI override (no `scaffoldguided.*` namespace, +# unlike older RFDiffusion versions). Scaffold/motif guidance is expressed +# entirely inside the InputSpecification JSON passed as `inputs=` (see +# `input`/`partial_t` fields) -- guided vs. unguided diffusion is selected by +# which JSON file the caller points `inputs` at, never by an extra CLI arg. foundry_sif_path="$1" output_dir="$2" inputs="$3" diffusion_batch_size="$4" -if [ $# -eq 5 ]; then - scaffold_arg="$5" -else - scaffold_arg="" -fi +# Prevent host ~/.local Python packages from contaminating the container +# (apptainer mounts $HOME by default; PYTHONNOUSERSITE must be SET, not unset). +unset PYTHONPATH PYTHONUSERBASE PYTHONDONTWRITEBYTECODE +export PYTHONNOUSERSITE=1 -apptainer exec --nv "$foundry_sif_path" rfd3 design \ +apptainer exec --nv --writable-tmpfs ${SCRATCH:+--bind "${SCRATCH}:${SCRATCH}"} "$foundry_sif_path" rfd3 design \ out_dir="$output_dir" \ inputs="$inputs" \ skip_existing=False \ dump_trajectories=True \ prevalidate_inputs=True \ - diffusion_batch_size="$diffusion_batch_size" \ - ${scaffold_arg:+$scaffold_arg} + diffusion_batch_size="$diffusion_batch_size" diff --git a/examples/small_molecule_binding/scripts/validate_run.py b/examples/small_molecule_binding/scripts/validate_run.py new file mode 100644 index 0000000..b045060 --- /dev/null +++ b/examples/small_molecule_binding/scripts/validate_run.py @@ -0,0 +1,562 @@ +#!/usr/bin/env python3 +"""Post-HPC-run validation for the small_molecule_binding example pipeline. + +Standalone CLI (not a pipeline task): checks a completed or in-progress +`IMPRESS_TEST_MODE=1` (or production) HPC run's output tree against the +invariants the Boltz-2 / RFD3-guided-scaffold rewrite depends on, so a +successful-looking run is actually verified rather than assumed. See the +"Post-HPC-run validation" section of the design plan for the full rationale +behind each check. + +Usage: + python scripts/validate_run.py + python scripts/validate_run.py logs p1 + python scripts/validate_run.py logs p1 --python /path/to/venv/bin/python + +Exits 0 if every check passes (SKIPPED checks do not count as failures), +nonzero if any check fails. +""" + +import argparse +import glob +import json +import os +import subprocess +import sys + +# ── Reuse small_molecule_binding.py's ligand-resname parsing logic ───────── +# +# This script lives at examples/small_molecule_binding/scripts/validate_run.py, +# one level below small_molecule_binding.py, so the example directory can +# reasonably be added to sys.path. Importing the real module is preferred +# (single source of truth for the "never hardcode a ligand resname" rule) but +# small_molecule_binding.py imports the `impress` package at module scope, so +# it only works in an environment that has the framework installed. Fall back +# to a duplicated minimal implementation so this validator still works when +# run standalone (the task description explicitly anticipates this). + +_EXAMPLE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if _EXAMPLE_DIR not in sys.path: + sys.path.insert(0, _EXAMPLE_DIR) + +try: + from small_molecule_binding import _ligand_resname_from_params +except Exception: + def _ligand_resname_from_params(params_path: str) -> str: + """Fallback duplicate of small_molecule_binding._ligand_resname_from_params, + used only when importing the real module fails (e.g. `impress` is not + installed in this environment). Keep this in sync with the original -- + it reads the exact literal from the params file's NAME record and must + NEVER hardcode a resname (e.g. 'ALR'); see that function's docstring.""" + with open(params_path) as fh: + for line in fh: + parts = line.split() + if len(parts) >= 2 and parts[0] == "NAME": + return parts[1] + raise ValueError(f"no NAME record found in {params_path}") + + +# The exact literal that must never reappear as a HETATM resname (RFD3 +# misresolves the bare "ALR" -- see ALR.params's NAME record / design plan). +# Never treated as "the" expected resname -- always compared against whatever +# _ligand_resname_from_params() reads live from ALR.params. +_REGRESSION_RESNAME = "ALR" + +# State key that was designed, then explicitly rejected, during the RFD3 +# guided-scaffold redesign (an earlier Kabsch-superposition ligand-grafting +# approach needed it; Boltz's joint co-folding made it unnecessary). Its +# reappearance anywhere in run output would mean a regression to that +# rejected approach. +_REJECTED_STATE_KEY = "rfd3_guide_ligand_pdb" + +# At least one of these must exist somewhere under $BOLTZ_CACHE for the cache +# to be considered "warmed" (see plan's delta_env_setup.sh Step 12/13). +_BOLTZ_CACHE_MARKERS = ("boltz2_conf.ckpt", "boltz2_aff.ckpt", "mols.tar", "ccd.pkl") + +# Skip these extensions when grepping run output for the rejected state key -- +# they're binary/compressed and can be large; a literal ASCII key string +# would not usefully appear in them the way it would in JSON/log/PDB text. +_GREP_SKIP_EXTENSIONS = (".gz", ".png", ".npz", ".pt", ".ckpt", ".pdb.gz", ".cif.gz") +_GREP_MAX_BYTES = 50 * 1024 * 1024 # don't slurp huge files into memory + + +def _resolve_pipeline_inputs(base_path: str, pipeline_name: str) -> str: + """Resolve the pipeline_inputs directory the way this actually plays out in + production, not SmallMoleculeBindingPipeline.__init__'s bare fallback. + + __init__ only falls back to `{base_path}/{name}_in` when no `input_dir` + kwarg is given at all. In practice, run_small_molecule_binding.py always + passes an explicit `input_dir` that is a *sibling* of base_path (both + `logs/` and `p1_in/` live directly under the examples directory) -- see + `examples_dir`/`work_dir`/`input_dir` in that file's `impress_smallmol_bind()`. + We reproduce that sibling convention first (matches real runs), falling + back to the bare class-default location if the sibling doesn't exist. + """ + base_path = os.path.abspath(base_path) + sibling = os.path.join(os.path.dirname(base_path), f"{pipeline_name}_in") + if os.path.isdir(sibling): + return sibling + return os.path.join(base_path, f"{pipeline_name}_in") + + +def _iter_hetatm_resnames(pdb_path: str): + """Return the list of resname strings (PDB fixed-width columns 18-20) for + every HETATM record in pdb_path.""" + names = [] + with open(pdb_path, errors="replace") as fh: + for line in fh: + if line.startswith("HETATM"): + names.append(line[17:20].strip()) + return names + + +# ── Check 1: ligand identity preserved end-to-end ─────────────────────────── + +def check_ligand_identity(base_path: str, pipeline_name: str, pipeline_inputs: str): + """Every Boltz model PDB and every guided_scaffold.pdb must contain a + HETATM residue named exactly the literal read live from ALR.params's NAME + record. Fails loudly (and specifically) if the bad literal "ALR" shows up + instead -- that exact regression must never reappear.""" + params_path = os.path.join(pipeline_inputs, "ALR.params") + if not os.path.isfile(params_path): + return [f"cannot check ligand identity: {params_path} not found"] + try: + expected = _ligand_resname_from_params(params_path) + except Exception as e: + return [f"cannot parse NAME record from {params_path}: {e}"] + + pipeline_dir = os.path.join(base_path, pipeline_name) + # NOTE: boltz nests its own output under out_dir/boltz_results_/ + # before the predictions// layout (confirmed against boltz's + # source: `out_dir = out_dir / f"boltz_results_{data.stem}"` in main.py, + # and empirically against a real `boltz predict` run) -- this is easy to + # miss from the docs alone. + candidates = sorted( + glob.glob(os.path.join(pipeline_dir, "*_boltz", "out", "boltz_results_boltz_input", + "predictions", "boltz_input", "*_model_*.pdb")) + + glob.glob(os.path.join(pipeline_dir, "*_rfd3", "in", "guided_scaffold.pdb")) + ) + if not candidates: + return [ + f"no '*_boltz/out/boltz_results_boltz_input/predictions/boltz_input/*_model_*.pdb' " + f"or '*_rfd3/in/guided_scaffold.pdb' files found under {pipeline_dir} " + f"-- nothing to check (has this run produced any boltz/guided-rfd3 output yet?)" + ] + + failures = [] + for pdb_path in candidates: + resnames = _iter_hetatm_resnames(pdb_path) + if not resnames: + failures.append( + f"{pdb_path}: no HETATM records found at all " + f"(expected ligand resname {expected!r})" + ) + continue + if expected in resnames: + continue + if _REGRESSION_RESNAME in resnames and expected != _REGRESSION_RESNAME: + failures.append( + f"{pdb_path}: REGRESSION -- found literal {_REGRESSION_RESNAME!r} instead of " + f"expected {expected!r}. RFD3 misresolves the bare {_REGRESSION_RESNAME!r} " + f"literal; the colon in {expected!r} is a deliberate workaround (see " + f"{params_path}'s NAME record), not a typo. This must never reappear." + ) + else: + failures.append( + f"{pdb_path}: expected ligand resname {expected!r} not found among HETATM " + f"residues found: {sorted(set(resnames))!r}" + ) + return failures + + +# ── Check 2: Boltz output shape ───────────────────────────────────────────── + +def check_boltz_output_shape(base_path: str, pipeline_name: str): + """Every */_boltz/out/ dir must have boltz_results_boltz_input/predictions/ + boltz_input/confidence_boltz_input_model_*.json files that parse as JSON + and carry complex_plddt (numeric, ~0-1) and a ligand_iptm key (value may + be null).""" + pipeline_dir = os.path.join(base_path, pipeline_name) + boltz_out_dirs = sorted(glob.glob(os.path.join(pipeline_dir, "*_boltz", "out"))) + if not boltz_out_dirs: + return [f"no '*_boltz/out' directories found under {pipeline_dir}"] + + failures = [] + for out_dir in boltz_out_dirs: + pred_dir = os.path.join(out_dir, "boltz_results_boltz_input", "predictions", "boltz_input") + if not os.path.isdir(pred_dir): + failures.append(f"{out_dir}: missing boltz_results_boltz_input/predictions/boltz_input/ directory") + continue + conf_files = sorted( + glob.glob(os.path.join(pred_dir, "confidence_boltz_input_model_*.json")) + ) + if not conf_files: + failures.append( + f"{pred_dir}: no 'confidence_boltz_input_model_*.json' files found" + ) + continue + for cf in conf_files: + try: + with open(cf) as fh: + data = json.load(fh) + except (OSError, json.JSONDecodeError) as e: + failures.append(f"{cf}: failed to parse as JSON: {e}") + continue + if "complex_plddt" not in data: + failures.append(f"{cf}: missing 'complex_plddt' key") + else: + v = data["complex_plddt"] + if not isinstance(v, (int, float)) or isinstance(v, bool): + failures.append(f"{cf}: complex_plddt is not numeric: {v!r}") + elif not (-0.01 <= v <= 1.05): + failures.append( + f"{cf}: complex_plddt={v!r} is outside the expected ~0-1 range" + ) + if "ligand_iptm" not in data: + failures.append( + f"{cf}: missing 'ligand_iptm' key (value may legitimately be null)" + ) + return failures + + +# ── Check 3: guided JSON correctness ──────────────────────────────────────── + +def check_guided_json_correctness(base_path: str, pipeline_name: str, pipeline_inputs: str): + """Every */_rfd3/in/guided_binder_design.json must parse, its partial.input + must point at a file that exists, and partial.ligand/length/select_exposed/ + select_buried must match the base ALR_binder_design.json verbatim (only + input/partial_t may legitimately differ) -- mirrors _write_guided_rfd3_json.""" + pipeline_dir = os.path.join(base_path, pipeline_name) + guided_jsons = sorted( + glob.glob(os.path.join(pipeline_dir, "*_rfd3", "in", "guided_binder_design.json")) + ) + if not guided_jsons: + # No guided rfd3 runs have happened yet (e.g. very first backbone, or no + # fold has passed similarity gating yet) -- not a failure by itself. + return [] + + base_json_path = os.path.join(pipeline_inputs, "ALR_binder_design.json") + if not os.path.isfile(base_json_path): + return [f"cannot check guided JSONs: base spec {base_json_path} not found"] + try: + with open(base_json_path) as fh: + base = json.load(fh) + except (OSError, json.JSONDecodeError) as e: + return [f"{base_json_path}: failed to parse as JSON: {e}"] + base_partial = base.get("partial", {}) + + failures = [] + for gj in guided_jsons: + try: + with open(gj) as fh: + guided = json.load(fh) + except (OSError, json.JSONDecodeError) as e: + failures.append(f"{gj}: failed to parse as JSON: {e}") + continue + + partial = guided.get("partial") + if not isinstance(partial, dict): + failures.append(f"{gj}: missing/invalid top-level 'partial' object") + continue + + input_path = partial.get("input") + if not input_path: + failures.append(f"{gj}: 'partial.input' is missing") + else: + resolved = ( + input_path if os.path.isabs(input_path) + else os.path.join(os.path.dirname(gj), input_path) + ) + if not os.path.isfile(resolved): + failures.append( + f"{gj}: partial.input={input_path!r} does not point at an existing file " + f"(resolved: {resolved})" + ) + + for field in ("ligand", "length", "select_exposed", "select_buried"): + expected = base_partial.get(field) + actual = partial.get(field) + if actual != expected: + failures.append( + f"{gj}: partial.{field} differs from base {base_json_path} -- " + f"expected {expected!r}, found {actual!r} (only 'input'/'partial_t' " + f"may legitimately differ)" + ) + return failures + + +# ── Check 4: RFD3 didn't silently no-op ───────────────────────────────────── + +def check_rfd3_no_op(base_path: str, pipeline_name: str): + """Every */_rfd3/out/ dir must contain at least one '*_model_*.json' file + (mirrors analysis_backbone()'s own file-discovery: f.endswith('.json') and + '_model_' in f). Best-effort: if any captured task stdout/stderr log is + discoverable, grep it for TypeError/RFD3InferenceConfig crash signatures. + + NOTE on the log-discovery part (see report for detail): capture_stdio=True + task logs are written by the asyncflow/dragon execution backend to a + per-*session* work_dir (IMPRESS_SESSION_DIR env var, else + tempfile.gettempdir()), under a randomly-generated + 'asyncflow.session.' subdirectory that is unrelated to base_path, + and filenames use an internal task uid ('task.NNNNNN') that has no + relationship to this pipeline's '{taskcount}_{taskname}' directory + convention. There is therefore no reliable static way to tie a captured + log file back to one specific rfd3 taskdir. This check does a best-effort + scan of the taskdir itself (in case a future version copies logs there) + plus IMPRESS_SESSION_DIR if set, and flags a crash signature generically + if found, without claiming it belongs to any particular rfd3 invocation. + """ + pipeline_dir = os.path.join(base_path, pipeline_name) + rfd3_out_dirs = sorted(glob.glob(os.path.join(pipeline_dir, "*_rfd3", "out"))) + if not rfd3_out_dirs: + return [f"no '*_rfd3/out' directories found under {pipeline_dir}"] + + failures = [] + log_files = [] + for out_dir in rfd3_out_dirs: + model_jsons = [ + f for f in os.listdir(out_dir) if f.endswith(".json") and "_model_" in f + ] if os.path.isdir(out_dir) else [] + if not model_jsons: + failures.append( + f"{out_dir}: no '*_model_*.json' files found -- rfd3 may have silently " + f"produced no output" + ) + taskdir = os.path.dirname(out_dir) + for pattern in ("*.stdout", "*.stderr", "*.log"): + log_files.extend(glob.glob(os.path.join(taskdir, pattern))) + log_files.extend(glob.glob(os.path.join(taskdir, "*", pattern))) + + session_dir = os.environ.get("IMPRESS_SESSION_DIR") + if session_dir and os.path.isdir(session_dir): + for pattern in ("*.stdout", "*.stderr"): + log_files.extend( + glob.glob(os.path.join(session_dir, "**", pattern), recursive=True) + ) + + seen = set() + for lf in log_files: + if lf in seen or not os.path.isfile(lf): + continue + seen.add(lf) + try: + if os.path.getsize(lf) > _GREP_MAX_BYTES: + continue + with open(lf, errors="replace") as fh: + content = fh.read() + except OSError: + continue + if "TypeError" in content or "RFD3InferenceConfig" in content: + failures.append( + f"{lf}: contains 'TypeError' or 'RFD3InferenceConfig' -- possible RFD3 crash " + f"signature (the exact error the old scaffoldguided.target_pdb bug produced). " + f"NOTE: log discovery is best-effort and not reliably tied to a specific rfd3 " + f"taskdir -- see check_rfd3_no_op()'s docstring." + ) + return failures + + +# ── Check 5: state-key regression guard ───────────────────────────────────── + +def check_no_rejected_state_key(base_path: str, pipeline_name: str): + """Grep all files under base_path/pipeline_name (and IMPRESS_SESSION_DIR + logs, if discoverable) for the literal 'rfd3_guide_ligand_pdb' -- a state + key that was designed then explicitly rejected in favor of Boltz's joint + co-folding. Its reappearance anywhere means the rejected + Kabsch-superposition ligand-grafting approach crept back in.""" + pipeline_dir = os.path.join(base_path, pipeline_name) + if not os.path.isdir(pipeline_dir): + return [f"{pipeline_dir} does not exist -- nothing to grep"] + + failures = [] + key_bytes = _REJECTED_STATE_KEY.encode() + + def _grep_tree(root_dir): + for root, _dirs, files in os.walk(root_dir): + for fname in files: + if fname.endswith(_GREP_SKIP_EXTENSIONS): + continue + fpath = os.path.join(root, fname) + try: + if os.path.getsize(fpath) > _GREP_MAX_BYTES: + continue + with open(fpath, "rb") as fh: + chunk = fh.read() + except OSError: + continue + if key_bytes in chunk: + failures.append( + f"{fpath}: contains rejected state key " + f"{_REJECTED_STATE_KEY!r} -- regression to the rejected " + f"Kabsch-superposition ligand-grafting design" + ) + + _grep_tree(pipeline_dir) + + session_dir = os.environ.get("IMPRESS_SESSION_DIR") + if session_dir and os.path.isdir(session_dir): + _grep_tree(session_dir) + + return failures + + +# ── Check 6: ensemble sanity ───────────────────────────────────────────────── + +def check_ensemble_sanity(base_path: str, pipeline_name: str): + """Ensemble state (self.state['ensemble']) lives only in the pipeline's + in-memory process state -- ImpressBasePipeline and + SmallMoleculeBindingPipeline have no checkpoint/state-dump-to-disk + convention as of this writing (confirmed by reading + src/impress/pipelines/impress_pipeline.py and small_molecule_binding.py: + no pickle/json state-dump call anywhere in either). There is therefore + nothing on disk to check monotonic-taskcount / no-duplicate-tuple + invariants against. Rather than fabricate a check against directory + counts that don't actually reconstruct the ensemble list, this is an + explicit no-op.""" + return [ + "SKIPPED: no on-disk ensemble state found, cannot verify (ensemble lives in " + "in-memory self.state, not persisted to disk by this framework)" + ] + + +# ── Check 7: env sanity ────────────────────────────────────────────────────── + +def check_env_sanity(python_exe: str): + """$BOLTZ_CACHE must exist and contain at least one known weight/cache + marker file; `import boltz` must succeed under the given interpreter.""" + failures = [] + + boltz_cache = os.environ.get("BOLTZ_CACHE") + if not boltz_cache: + failures.append("BOLTZ_CACHE environment variable is not set") + elif not os.path.isdir(boltz_cache): + failures.append(f"BOLTZ_CACHE={boltz_cache!r} is not a directory") + else: + found = None + for root, _dirs, files in os.walk(boltz_cache): + for marker in _BOLTZ_CACHE_MARKERS: + if marker in files: + found = os.path.join(root, marker) + break + if found: + break + if not found: + failures.append( + f"BOLTZ_CACHE={boltz_cache!r} does not contain any of " + f"{_BOLTZ_CACHE_MARKERS!r} -- weights may not have been " + f"downloaded / cache-warmed yet" + ) + + try: + result = subprocess.run( + [python_exe, "-c", "import boltz"], + capture_output=True, text=True, timeout=120, + ) + except (OSError, subprocess.TimeoutExpired) as e: + failures.append(f"failed to run {python_exe!r} to check `import boltz`: {e}") + else: + if result.returncode != 0: + stderr_tail = result.stderr.strip()[-500:] + failures.append( + f"`{python_exe} -c 'import boltz'` failed (exit {result.returncode}): " + f"{stderr_tail}" + ) + return failures + + +# ── main ───────────────────────────────────────────────────────────────────── + +def main(argv=None) -> int: + parser = argparse.ArgumentParser( + description=( + "Validate a completed (or in-progress) small_molecule_binding HPC run's " + "output tree against the Boltz-2 / RFD3-guided-scaffold invariants." + ) + ) + parser.add_argument( + "base_path", + help="Pipeline base_path, e.g. 'logs' (same value passed as SmallMoleculeBindingPipeline's base_path kwarg)", + ) + parser.add_argument("pipeline_name", help="Pipeline name, e.g. 'p1'") + parser.add_argument( + "--python", + default=sys.executable, + help="Python interpreter to check `import boltz` with (default: the interpreter running this script)", + ) + args = parser.parse_args(argv) + + base_path = os.path.abspath(args.base_path) + pipeline_name = args.pipeline_name + pipeline_dir = os.path.join(base_path, pipeline_name) + pipeline_inputs = _resolve_pipeline_inputs(base_path, pipeline_name) + + print(f"Validating run: base_path={base_path} pipeline_name={pipeline_name}") + print(f"Resolved pipeline_dir: {pipeline_dir}") + print(f"Resolved pipeline_inputs: {pipeline_inputs}") + print() + + if not os.path.isdir(base_path): + print(f"FAIL: base_path {base_path!r} does not exist") + return 1 + if not os.path.isdir(pipeline_dir): + print( + f"FAIL: pipeline directory {pipeline_dir!r} does not exist -- " + f"has pipeline {pipeline_name!r} run yet under this base_path?" + ) + return 1 + if not os.path.isdir(pipeline_inputs): + print( + f"WARNING: pipeline_inputs directory {pipeline_inputs!r} does not exist -- " + f"checks 1 and 3 (which need ALR.params / ALR_binder_design.json) will fail\n" + ) + + checks = [ + ("1. Ligand identity preserved end-to-end", + lambda: check_ligand_identity(base_path, pipeline_name, pipeline_inputs)), + ("2. Boltz output shape", + lambda: check_boltz_output_shape(base_path, pipeline_name)), + ("3. Guided JSON correctness", + lambda: check_guided_json_correctness(base_path, pipeline_name, pipeline_inputs)), + ("4. RFD3 didn't silently no-op", + lambda: check_rfd3_no_op(base_path, pipeline_name)), + ("5. State-key regression guard (rfd3_guide_ligand_pdb)", + lambda: check_no_rejected_state_key(base_path, pipeline_name)), + ("6. Ensemble sanity", + lambda: check_ensemble_sanity(base_path, pipeline_name)), + ("7. Env sanity (BOLTZ_CACHE / import boltz)", + lambda: check_env_sanity(args.python)), + ] + + print("=" * 72) + print("VALIDATION RESULTS") + print("=" * 72) + + any_failed = False + for name, fn in checks: + try: + failures = fn() + except Exception as e: # a check itself must never crash the whole run + failures = [f"check raised an unexpected exception: {e!r}"] + + if failures and all(f.startswith("SKIPPED:") for f in failures): + print(f"[SKIP] {name}") + for f in failures: + print(f" {f}") + elif not failures: + print(f"[PASS] {name}") + else: + any_failed = True + print(f"[FAIL] {name} -- {len(failures)} issue(s)") + for f in failures: + print(f" - {f}") + + print("=" * 72) + if any_failed: + print("RESULT: FAIL") + return 1 + print("RESULT: PASS") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/small_molecule_binding/small_molecule_binding.py b/examples/small_molecule_binding/small_molecule_binding.py index e8fc3f2..8373470 100644 --- a/examples/small_molecule_binding/small_molecule_binding.py +++ b/examples/small_molecule_binding/small_molecule_binding.py @@ -1,6 +1,7 @@ import asyncio import copy +import gzip import json import os import pathlib @@ -15,8 +16,8 @@ STEP_RFD3 = 1 # backbone diffusion STEP_MPNN = 2 # mpnn + packmin refinement cycle STEP_FASTRELAX = 3 # Rosetta FastRelax -STEP_INTERFACE = 4 # filter_shape (PyRosetta, gates af2) -STEP_AF2 = 5 # fold prediction +STEP_INTERFACE = 4 # filter_shape (PyRosetta, gates fold prediction) +STEP_AF2 = 5 # fold prediction (Boltz-2 co-folding; name kept for compatibility) STEP_RETRY_SEQ = 6 # internal: retry sequence prediction without backbone restart # Ensemble transformation type labels @@ -119,6 +120,138 @@ def _ensemble_selective_avg( return overall_avg, sum(sel_scores) / len(sel_scores), True +def _stage_metrics_improving( + current: dict, + previous: dict | None, + specs: list, + rel_tolerance: float = 0.05, +) -> bool: + """Metric-agnostic "is this retry attempt actually helping" check, shared by + adaptive_decision()'s 'fastrelax' and 'interface' short-circuit logic. + + `specs` is a list of (metric_key, lower_is_better, threshold) triples. For + each metric present in both `current` and `previous` that was still failing + on the previous attempt (gap > 0), checks whether its gap-to-threshold + shrank by more than `rel_tolerance` of the previous gap. Returns True (an + improvement was found) if any tracked metric improved; False if every + still-failing metric stayed flat or got worse. + + Returns True (never short-circuit) when `previous` is None/empty -- there is + no prior attempt on this backbone to compare against yet, so the first + failure always gets one retry regardless of this check.""" + if not previous: + return True + for key, lower_is_better, threshold in specs: + cur_val, prev_val = current.get(key), previous.get(key) + if cur_val is None or prev_val is None: + continue + cur_gap = (cur_val - threshold) if lower_is_better else (threshold - cur_val) + prev_gap = (prev_val - threshold) if lower_is_better else (threshold - prev_val) + if prev_gap <= 0: + continue # this metric already passed on the previous attempt + if (prev_gap - cur_gap) > rel_tolerance * prev_gap: + return True + return False + + +# ── RFD3 guided-input utilities ──────────────────────────────────────────── + +def _ligand_resname_from_params(params_path: str) -> str: + """Reads the 'NAME ' record from a Rosetta .params file and returns + the exact literal residue name (e.g. 'A:R' for ALR.params). This is the + literal PDB/Rosetta residue name used in HETATM records for this ligand and + is NOT always equal to the params filename stem -- never assume otherwise, + and never hardcode a specific ligand's resname here or in any caller.""" + with open(params_path) as fh: + for line in fh: + parts = line.split() + if len(parts) >= 2 and parts[0] == 'NAME': + return parts[1] + raise ValueError(f"no NAME record found in {params_path}") + + +def _normalize_ligand_id(fold_pdb_path: str, ligand_resname: str, out_pdb_path: str, + ligand_chain_id: str = "B") -> bool: + """Rewrites a Boltz co-folded PDB's ligand HETATM residue name (via gemmi) to + match ligand_resname, so RFD3's ligand/select_exposed/select_buried selectors + (which key off the literal resname) resolve against it. Purely a string edit + -- no coordinate transform, since Boltz already places the ligand correctly + relative to the protein it just folded. Locates the ligand residue by HETATM + on ligand_chain_id first; if Boltz didn't honor the requested chain letter, + falls back to the first HETATM residue found anywhere in the structure. + Writes out_pdb_path and returns True on success. Returns False (writing + nothing) if no HETATM residue is found at all, so the caller can fall back + to unguided diffusion instead of crashing.""" + import gemmi + st = gemmi.read_structure(fold_pdb_path) + + target_res = None + for model in st: + for chain in model: + if chain.name != ligand_chain_id: + continue + for res in chain: + if res.het_flag == 'H': + target_res = res + break + if target_res: + break + if target_res: + break + + if target_res is None: + # Fallback: Boltz may not have honored the requested chain id. + for model in st: + for chain in model: + for res in chain: + if res.het_flag == 'H': + target_res = res + break + if target_res: + break + if target_res: + break + + if target_res is None: + return False + + target_res.name = ligand_resname + st.write_pdb(out_pdb_path) + return True + + +def _write_guided_rfd3_json(base_json_path: str, guided_pdb_path: str, partial_t: float, + out_json_path: str) -> None: + """Loads the base per-pipeline RFD3 InputSpecification JSON, copies its + ligand/length/select_exposed/select_buried fields verbatim, replaces 'input' + with guided_pdb_path, adds partial_t, and writes the result to + out_json_path. Only 'input'/'partial_t' differ from the base file.""" + with open(base_json_path) as fh: + base = json.load(fh) + partial = dict(base.get('partial', {})) + partial['input'] = guided_pdb_path + partial['partial_t'] = partial_t + guided = dict(base) + guided['partial'] = partial + with open(out_json_path, 'w') as fh: + json.dump(guided, fh, indent=4) + + +def _prepare_guided_rfd3_inputs(base_json_path: str, fold_pdb_path: str, ligand_resname: str, + partial_t: float, taskdir: str): + """Orchestrates _normalize_ligand_id + _write_guided_rfd3_json: writes + {taskdir}/in/guided_scaffold.pdb and {taskdir}/in/guided_binder_design.json. + Returns the guided JSON path, or None if ligand normalization failed (e.g. + no HETATM residue found in fold_pdb_path) -- callers should fall back to + unguided diffusion in that case.""" + guided_pdb = f"{taskdir}/in/guided_scaffold.pdb" + guided_json = f"{taskdir}/in/guided_binder_design.json" + if not _normalize_ligand_id(fold_pdb_path, ligand_resname, guided_pdb): + return None + _write_guided_rfd3_json(base_json_path, guided_pdb, partial_t, guided_json) + return guided_json + + class SmallMoleculeBindingPipeline(ImpressBasePipeline): def __init__(self, name, flow, configs=None, **kwargs): if configs is None: @@ -140,18 +273,35 @@ def __init__(self, name, flow, configs=None, **kwargs): super().__init__(name, flow, **configs, **kwargs) # Paths - self.base_path = kwargs.get("base_path", os.getcwd()) - self.scripts_path = os.path.join(self.base_path, "scripts") - self.pipeline_inputs = os.path.join(self.base_path, f"{self.name}_in") - self.mpnn_dir = kwargs.get("mpnn_dir", f"/anvil/projects/x-nairr240405/mason/LigandMPNN") + self.base_path = kwargs.get("base_path", os.getcwd()) + self.scripts_path = kwargs.get( + "scripts_path", os.path.join(self.base_path, "scripts") + ) + # input_dir: explicit input directory name (relative to base_path) or + # absolute path. Falls back to "_in" when not provided. + _input_dir = kwargs.get("input_dir", "") + if _input_dir and os.path.isabs(_input_dir): + self.pipeline_inputs = _input_dir + elif _input_dir: + self.pipeline_inputs = os.path.join(self.base_path, _input_dir) + else: + self.pipeline_inputs = os.path.join(self.base_path, f"{self.name}_in") + self.mpnn_dir = kwargs.get("mpnn_dir") or os.environ.get("MPNN_DIR") + if not self.mpnn_dir: + raise ValueError("mpnn_dir must be supplied via kwarg or MPNN_DIR env var") # Configurable tool paths and ensemble sizes - self.foundry_sif_path = kwargs.get("foundry_sif_path", "/anvil/projects/x-nairr240405/mason/foundry.sif") - self.colabfold_path = kwargs.get("colabfold_path", "/anvil/projects/x-nairr240405/mason/localcolabfold") + self.foundry_sif_path = kwargs.get("foundry_sif_path") or os.environ.get("FOUNDRY_SIF_PATH") + if not self.foundry_sif_path: + raise ValueError("foundry_sif_path must be supplied via kwarg or FOUNDRY_SIF_PATH env var") + self.boltz_cache_path = kwargs.get("boltz_cache_path") or os.environ.get("BOLTZ_CACHE") + if not self.boltz_cache_path: + raise ValueError("boltz_cache_path must be supplied via kwarg or BOLTZ_CACHE env var") self.ligand_params = kwargs.get("ligand_params", "ALR.params") self.mpnn_ensemble_size = kwargs.get("mpnn_ensemble_size", 1) self.num_refine_cycles = kwargs.get("num_refine_cycles", 3) self.diffusion_batch_size = kwargs.get("diffusion_batch_size", 2) + self.rfd3_partial_t = kwargs.get("rfd3_partial_t", 10.0) # Quality thresholds (overridable at construction time) self.backbone_max_ca_deviation = kwargs.get("backbone_max_ca_deviation", 2.0) @@ -161,7 +311,9 @@ def __init__(self, name, flow, configs=None, **kwargs): self.fastrelax_max_fa_rep = kwargs.get("fastrelax_max_fa_rep", 150.0) self.interface_min_sc = kwargs.get("interface_min_sc", 0.5) self.fold_min_plddt = kwargs.get("fold_min_plddt", 70.0) + self.fold_min_ligand_iptm = kwargs.get("fold_min_ligand_iptm", None) self.max_tasks = kwargs.get("max_tasks", 300) + self.gpu_id = kwargs.get("gpu_id", None) # Output paths (legacy) self.output_path = os.path.join(self.base_path, "myoutputs", self.name) @@ -177,6 +329,12 @@ def __init__(self, name, flow, configs=None, **kwargs): self.next_step = STEP_RFD3 self._current_cycle_i = 0 # set by run() before each mpnn call + def _gpu_env(self) -> dict: + env = {**os.environ} + if self.gpu_id is not None: + env["CUDA_VISIBLE_DEVICES"] = str(self.gpu_id) + return env + # ── Task registration ────────────────────────────────────────────────── def register_pipeline_tasks(self): @@ -195,9 +353,8 @@ def _register_mock_tasks(self): def _register_real_tasks(self): """Register real HPC tasks that return shell command strings.""" - @self.auto_register_task(capture_stdio=True) - async def rfd3(task_description={"gpus_per_rank": 1}): + async def rfd3(): self.taskcount += 1 taskname = "rfd3" self.previous_task = taskname @@ -205,20 +362,33 @@ async def rfd3(task_description={"gpus_per_rank": 1}): os.makedirs(f"{taskdir}/in", exist_ok=True) os.makedirs(f"{taskdir}/out", exist_ok=True) - inputs = f"{self.pipeline_inputs}/ALR_binder_design.json" - output_dir = f"{taskdir}/out" + base_inputs = f"{self.pipeline_inputs}/ALR_binder_design.json" + output_dir = f"{taskdir}/out" - input_pdb = self.state.get('rfd3_input_pdb') - scaffold_arg = f"scaffoldguided.target_pdb={input_pdb}" if input_pdb else "" + fold_pdb = self.state.get('rfd3_input_pdb') + inputs = base_inputs + if fold_pdb: + ligand_resname = _ligand_resname_from_params( + f"{self.pipeline_inputs}/{self.ligand_params}" + ) + guided_json = _prepare_guided_rfd3_inputs( + base_json_path=base_inputs, + fold_pdb_path=fold_pdb, + ligand_resname=ligand_resname, + partial_t=self.rfd3_partial_t, + taskdir=taskdir, + ) + if guided_json: + inputs = guided_json - return ( + cmd = ( f"bash {self.scripts_path}/rfd3.sh" f" {self.foundry_sif_path}" f" {output_dir}" f" {inputs}" f" {self.diffusion_batch_size}" - f" {scaffold_arg}" ) + return cmd @self.auto_register_task(local_task=True) async def analysis_backbone(): @@ -273,7 +443,7 @@ async def analysis_backbone(): ETYPE_BACKBONE, best['ss'], self.state.get('rfd3_input_pdb'), backbone_path, )) - @self.auto_register_task(capture_stdio=True) + @self.auto_register_task(local_task=True) async def mpnn( fixed_residues_file: str | None = None): self.taskcount += 1 @@ -300,13 +470,24 @@ async def mpnn( shutil.copy(pdb_path_orig, short_pdb) pdb_path = short_pdb + # LigandMPNN (ProDy parsePDB) only reads PDB format; convert CIF.GZ. + if pdb_path.endswith('.cif.gz'): + import gemmi as _gemmi + pdb_for_mpnn = f"{taskdir}/in/binder.pdb" + with gzip.open(pdb_path, 'rb') as _f: + _cif_data = _f.read().decode() + _doc = _gemmi.cif.read_string(_cif_data) + _st = _gemmi.make_structure_from_block(_doc.sole_block()) + _st.write_pdb(pdb_for_mpnn) + pdb_path = pdb_for_mpnn + if fixed_residues_file: with open(fixed_residues_file) as f: fixed_residues = f.read().strip() else: fixed_residues = "" - return ( + cmd = ( f"bash {self.scripts_path}/mpnn.sh" f" {self.mpnn_dir}" f" {pdb_path}" @@ -315,6 +496,17 @@ async def mpnn( f" {batch_size}" f' "{fixed_residues}"' ) + log_file = f"{taskdir}/mpnn.log" + with open(log_file, "wb") as _lf: + proc = await asyncio.create_subprocess_shell( + cmd, + stdout=_lf, + stderr=asyncio.subprocess.STDOUT, + env=self._gpu_env(), + ) + await proc.wait() + if proc.returncode != 0: + raise RuntimeError(f"mpnn failed with exit code {proc.returncode}\nSee {log_file}") @self.auto_register_task(local_task=True) async def analysis_sequence(): @@ -325,36 +517,51 @@ async def analysis_sequence(): best_conf = -1.0 best_lig_conf = 0.0 - best_seq_name = None - best_fa_file = None + best_id = None + best_seq = None for fa_file in [f for f in os.listdir(seqs_dir) if f.endswith('.fa')]: with open(f"{seqs_dir}/{fa_file}") as fh: - header = fh.readline().strip() - try: - parts = { - kv.split('=')[0].strip(): kv.split('=')[1].strip() - for kv in header.lstrip('>').split(',') - if '=' in kv - } - conf = float(parts.get('overall_confidence', 0)) - lig_conf = float(parts.get('ligand_confidence', 0)) - name = header.lstrip('>').split(',')[0].strip() - except (ValueError, IndexError): - continue - - if conf > best_conf: - best_conf = conf - best_lig_conf = lig_conf - best_seq_name = name - best_fa_file = fa_file + content = fh.read() + # LigandMPNN writes ONE file per input structure containing MULTIPLE + # records: a template record (echo of the input, no 'id=' field) + # followed by 'batch_size' real designed candidates ('id=1'..'id=N', + # each with overall_confidence/ligand_confidence). Evaluate every + # id= record across every file -- do not assume one candidate per file. + for record in content.split('>')[1:]: + lines = record.splitlines() + if not lines: + continue + header, seq = lines[0], ''.join(lines[1:]).strip() + try: + parts = { + kv.split('=')[0].strip(): kv.split('=')[1].strip() + for kv in header.split(',') + if '=' in kv + } + if 'id' not in parts: + continue # template record, not a real candidate + cand_id = parts['id'] + conf = float(parts.get('overall_confidence', 0)) + lig_conf = float(parts.get('ligand_confidence', 0)) + except (ValueError, IndexError): + continue + + if conf > best_conf: + best_conf = conf + best_lig_conf = lig_conf + best_id = cand_id + best_seq = seq # Always update best_packed_pdb so packmin reads the current mpnn output - if best_seq_name: - self.state['best_packed_pdb'] = f"{out_dir}/packed/{best_seq_name}_packed_1_1.pdb" - self.state['last_seq_fasta'] = f"{seqs_dir}/{best_fa_file}" + if best_id is not None: + self.state['best_packed_pdb'] = f"{out_dir}/packed/binder_packed_{best_id}_1.pdb" + best_fasta_path = f"{seqs_dir}/best_candidate.fa" + with open(best_fasta_path, 'w') as fh: + fh.write(f">binder_id_{best_id}\n{best_seq}\n") + self.state['last_seq_fasta'] = best_fasta_path else: - self.state['last_seq_fasta'] = None + self.state['last_seq_fasta'] = None self.state['last_analysis_metrics'] = { 'pass': True, @@ -366,7 +573,7 @@ async def analysis_sequence(): self.state.get('best_backbone_path'), self.state.get('last_seq_fasta'), )) - @self.auto_register_task(capture_stdio=True) + @self.auto_register_task(local_task=True) async def packmin(): self.taskcount += 1 taskname = "packmin" @@ -383,12 +590,20 @@ async def packmin(): # Predict output path so the next mpnn can read from it self.state['best_packed_pdb'] = f"{output_dir}/{pdb_stem}_minimized.pdb" - return ( + cmd = ( f"bash {self.scripts_path}/packmin.sh" f" {pdb_path}" f" {lig_path}" f" {output_dir}" ) + log_file = f"{taskdir}/packmin.log" + with open(log_file, "wb") as _lf: + proc = await asyncio.create_subprocess_shell( + cmd, stdout=_lf, stderr=asyncio.subprocess.STDOUT, env=self._gpu_env(), + ) + await proc.wait() + if proc.returncode != 0: + raise RuntimeError(f"packmin failed with exit code {proc.returncode}\nSee {log_file}") @self.auto_register_task(local_task=True) async def analysis_packmin(): @@ -403,7 +618,7 @@ async def analysis_packmin(): self.state['last_analysis_step'] = 'packmin' self.state['last_analysis_metrics'] = {'pass': True, 'total_score': total_score} - @self.auto_register_task(capture_stdio=True) + @self.auto_register_task(local_task=True) async def fastrelax(): self.taskcount += 1 taskname = "fastrelax" @@ -416,13 +631,20 @@ async def fastrelax(): lig_path = f"{self.pipeline_inputs}/{self.ligand_params}" output_dir = f"{taskdir}/out" - return ( + cmd = ( f"bash {self.scripts_path}/fastrelax.sh" f" {pdb_path}" f" {lig_path}" f" {output_dir}" - f" > fastrelax.txt" ) + log_file = f"{taskdir}/fastrelax.log" + with open(log_file, "wb") as _lf: + proc = await asyncio.create_subprocess_shell( + cmd, stdout=_lf, stderr=asyncio.subprocess.STDOUT, env=self._gpu_env(), + ) + await proc.wait() + if proc.returncode != 0: + raise RuntimeError(f"fastrelax failed with exit code {proc.returncode}\nSee {log_file}") @self.auto_register_task(local_task=True) async def analysis_fastrelax(): @@ -452,7 +674,7 @@ async def analysis_fastrelax(): 'rmsd': rmsd, } - @self.auto_register_task(capture_stdio=True) + @self.auto_register_task(local_task=True) async def filter_shape(ligand_name: str = "ALR"): taskname = "filter_shape" taskdir = f"{self.base_path}/{self.name}/{self.taskcount}_{taskname}" @@ -461,13 +683,21 @@ async def filter_shape(ligand_name: str = "ALR"): pdb_directory = f"{self.base_path}/{self.name}/{self.taskcount}_fastrelax/out" - return ( + cmd = ( f"bash {self.scripts_path}/filter_shape.sh" f" {pdb_directory}" f" {taskdir}/out/shape_complementarity_values.txt" f" {self.pipeline_inputs}/{ligand_name}" f" {taskdir}/out/interface_values.txt" ) + log_file = f"{taskdir}/filter_shape.log" + with open(log_file, "wb") as _lf: + proc = await asyncio.create_subprocess_shell( + cmd, stdout=_lf, stderr=asyncio.subprocess.STDOUT, env=self._gpu_env(), + ) + await proc.wait() + if proc.returncode != 0: + raise RuntimeError(f"filter_shape failed with exit code {proc.returncode}\nSee {log_file}") @self.auto_register_task(local_task=True) async def analysis_interface(): @@ -495,68 +725,98 @@ async def analysis_interface(): } @self.auto_register_task(capture_stdio=True) - async def af2(task_description={"gpus_per_rank": 1}): + async def boltz(): self.taskcount += 1 - taskname = "alphafold" + taskname = "boltz" self.previous_task = taskname taskdir = f"{self.base_path}/{self.name}/{self.taskcount}_{taskname}" os.makedirs(f"{taskdir}/in", exist_ok=True) os.makedirs(f"{taskdir}/out", exist_ok=True) - src_fasta = self.state['last_seq_fasta'] - short_fasta = f"{taskdir}/in/binder.fa" - seq_lines = [] - with open(src_fasta) as fh: - for line in fh: - if not line.startswith('>'): - seq_lines.append(line) - with open(short_fasta, 'w') as fh: - fh.write('>binder\n') - fh.writelines(seq_lines) + seq = _read_fasta_seq(self.state['last_seq_fasta']) + if not seq: + raise RuntimeError(f"boltz: no usable sequence in {self.state['last_seq_fasta']}") + + ligand_stem = pathlib.Path(self.ligand_params).stem + with open(f"{self.pipeline_inputs}/{ligand_stem}.smiles") as fh: + ligand_smiles = fh.read().strip() + + yaml_path = f"{taskdir}/in/boltz_input.yaml" + with open(yaml_path, "w") as fh: + fh.write( + "version: 1\nsequences:\n - protein:\n id: [A]\n" + f" sequence: {seq}\n msa: empty\n" + " - ligand:\n id: [B]\n" + f" smiles: '{ligand_smiles}'\n" + ) output_dir = f"{taskdir}/out" - - return ( - f"bash {self.scripts_path}/af2.sh" - f" {self.colabfold_path}" - f" {short_fasta}" + cmd = ( + f"bash {self.scripts_path}/boltz.sh" + f" {yaml_path}" f" {output_dir}" + f" {self.boltz_cache_path}" ) + return cmd @self.auto_register_task(local_task=True) async def analysis_fold(): - out_dir = f"{self.base_path}/{self.name}/{self.taskcount}_alphafold/out" - score_files = [ - f for f in os.listdir(out_dir) - if 'scores' in f and f.endswith('.json') - ] + # Boltz nests its own output under out_dir/boltz_results_/ + # (see boltz/main.py: `out_dir = out_dir / f"boltz_results_{data.stem}"`) + # before the documented predictions// layout -- confirmed + # empirically against a real `boltz predict` run, not just the docs. + pred_dir = ( + f"{self.base_path}/{self.name}/{self.taskcount}_boltz/out/" + "boltz_results_boltz_input/predictions/boltz_input" + ) + conf_files = [ + f for f in os.listdir(pred_dir) + if f.startswith('confidence_') and f.endswith('.json') + ] if os.path.isdir(pred_dir) else [] + + if not conf_files: + raise RuntimeError( + f"boltz produced no confidence files in {pred_dir} — " + "GPU/predict failure" + ) - best_plddt = -1.0 - best_model = None - for sf in score_files: - with open(f"{out_dir}/{sf}") as fh: - arr = json.load(fh).get('plddt', []) - if arr: - mean_plddt = sum(arr) / len(arr) - if mean_plddt > best_plddt: - best_plddt = mean_plddt - best_model = sf.replace('_scores_', '_unrelaxed_').replace('.json', '.pdb') + best_complex_plddt = -1.0 + best_model = None + best_ligand_iptm = None + for cf in conf_files: + with open(f"{pred_dir}/{cf}") as fh: + data = json.load(fh) + score = data.get('complex_plddt', 0.0) + if score > best_complex_plddt: + best_complex_plddt = score + best_model = cf.replace('confidence_', '', 1).replace('.json', '.pdb') + best_ligand_iptm = data.get('ligand_iptm') + + # Rescale 0-1 -> 0-100 to preserve fold_min_plddt's existing semantics. + best_plddt_100 = best_complex_plddt * 100.0 + passed = best_plddt_100 >= self.fold_min_plddt + if self.fold_min_ligand_iptm is not None: + passed = passed and ( + best_ligand_iptm is not None + and best_ligand_iptm >= self.fold_min_ligand_iptm + ) if best_model: - full_model_path = f"{out_dir}/{best_model}" - self.state['best_af2_model'] = full_model_path + full_model_path = f"{pred_dir}/{best_model}" + self.state['best_fold_model'] = full_model_path self.state['ensemble'].append(( - ETYPE_FOLD, best_plddt, self.state.get('last_seq_fasta'), full_model_path, + ETYPE_FOLD, best_plddt_100, self.state.get('last_seq_fasta'), full_model_path, )) self.state['last_analysis_step'] = 'fold' self.state['last_analysis_metrics'] = { - 'pass': best_plddt >= self.fold_min_plddt, - 'best_mean_plddt': best_plddt, - 'best_model': best_model, + 'pass': passed, + 'best_complex_plddt': best_plddt_100, + 'best_ligand_iptm': best_ligand_iptm, + 'best_model': best_model, } - @self.auto_register_task(capture_stdio=True) + @self.auto_register_task(local_task=True) async def filter_energy(ligand_name: str = "ALR"): taskname = "filter_energy" taskdir = f"{self.base_path}/{self.name}/{self.taskcount}_{taskname}" @@ -569,7 +829,7 @@ async def filter_energy(ligand_name: str = "ALR"): output_energy_file = f"{outputs_dir}/negative_ligand_energies.txt" common_filenames_file = f"{self.pipeline_inputs}/common_filenames.txt" - return ( + cmd = ( f"bash {self.scripts_path}/filter_energy.sh" f" {pdb_directory}" f" {output_file}" @@ -577,6 +837,14 @@ async def filter_energy(ligand_name: str = "ALR"): f" {common_filenames_file}" f" {ligand_name}" ) + log_file = f"{taskdir}/filter_energy.log" + with open(log_file, "wb") as _lf: + proc = await asyncio.create_subprocess_shell( + cmd, stdout=_lf, stderr=asyncio.subprocess.STDOUT, env=self._gpu_env(), + ) + await proc.wait() + if proc.returncode != 0: + raise RuntimeError(f"filter_energy failed with exit code {proc.returncode}\nSee {log_file}") # ── Score utils ──────────────────────────────────────────────────────── @@ -603,7 +871,6 @@ async def _run_refine_cycle(self): self.logger.pipeline_log(f"running mpnn [cycle {cycle_i}]") await self.mpnn() -# fixed_residues_file=f"{self.pipeline_inputs}/fixed_residues.txt" ) self.logger.pipeline_log(f"mpnn [cycle {cycle_i}] finished") await self.analysis_sequence() await self.run_adaptive_step() @@ -640,6 +907,8 @@ async def run(self): self.state.setdefault('rfd3_input_pdb', None) self.state.setdefault('seq_retry_count', 0) self.state.setdefault('last_seq_fasta', None) + self.state.setdefault('fastrelax_prev_metrics', None) + self.state.setdefault('interface_prev_metrics', None) self.logger.pipeline_log("SmallMoleculeBindingPipeline starting (state machine)") while self.next_step != STEP_DONE: @@ -678,9 +947,9 @@ async def run(self): await self.run_adaptive_step() elif self.next_step == STEP_AF2: - self.logger.pipeline_log("running af2") - await self.af2() - self.logger.pipeline_log("af2 finished") + self.logger.pipeline_log("running boltz") + await self.boltz() + self.logger.pipeline_log("boltz finished") await self.analysis_fold() await self.run_adaptive_step()