diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index f9e6856..de003c5 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -14,12 +14,6 @@ jobs: fail-fast: false matrix: include: - - os: ubuntu-latest - python: '3.9' - toxenv: py39 - - os: ubuntu-latest - python: '3.10' - toxenv: py310 - os: ubuntu-latest python: '3.11' toxenv: py311 @@ -59,12 +53,6 @@ jobs: fail-fast: false matrix: include: - - os: ubuntu-latest - python: '3.9' - toxenv: py39-all - - os: ubuntu-latest - python: '3.10' - toxenv: py310-all - os: ubuntu-latest python: '3.11' toxenv: py311-all diff --git a/.gitignore b/.gitignore index 0dafbe4..01d1dc9 100644 --- a/.gitignore +++ b/.gitignore @@ -135,6 +135,11 @@ dmypy.json # asyncflow related asyncflow.session.* +ddict_* + +# ROME runtime outputs +af_stats_*.csv +examples/protien_binding_usecase/logs/ # pdzbinder wf outputs af_pipeline_outputs_multi/ @@ -146,3 +151,5 @@ ddict* b0 slurm* *slurm +# scratch archives +arch/ diff --git a/examples/.gitignore b/examples/.gitignore new file mode 100644 index 0000000..31e9608 --- /dev/null +++ b/examples/.gitignore @@ -0,0 +1,14 @@ +# SLURM log output (all workflows) +*/logs/ + +# Per-pipeline task directories written to cwd +*/p*/ + +# Dragon telemetry session dirs +*/telemetry/ + +# Run metadata written to cwd +*/runinfo + +# Legacy output dirs +*/myoutputs/ diff --git a/examples/protein_binding/SKILL.md b/examples/protein_binding/SKILL.md new file mode 100644 index 0000000..e05135c --- /dev/null +++ b/examples/protein_binding/SKILL.md @@ -0,0 +1,131 @@ +# Protein Binding Use Case — SKILL.md + +## What it does + +5-step adaptive protein design pipeline running on Delta HPC (gpuA40x4): + +| Step | Task | Type | Notes | +|------|------|------|-------| +| s1 | ProteinMPNN — generate sequences | local (subprocess) | GPU via `CUDA_VISIBLE_DEVICES` | +| s2 | Sequence ranking | local (pure Python) | sorts by MPNN score | +| s3 | FASTA generation | local (pure Python) | writes `.fa` for each target | +| s4 | AlphaFold multimer prediction | local (subprocess) | GPU, runs via Apptainer | +| s4_post | Copy best models | local (pure Python) | glob → shutil.copy | +| s5 | pLDDT extraction | local (subprocess) | writes `af_stats__pass_.csv` | + +Passes repeat up to `max_passes` times. The adaptive function reads pLDDT CSVs and can spawn child pipelines for proteins that degrade. + +--- + +## Key files + +``` +protien_binding_usecase/ +├── run_protein_binding.py # entry point — backend, GPU policy, ImpressManager +├── delta_gpu_run.sh # SLURM script — Dragon launcher +├── delta_env_setup.sh # one-time venv setup (~/ve/impress/) +├── af2_multimer_reduced.sh # Apptainer wrapper for AlphaFold multimer +├── mpnn_wrapper.py # ProteinMPNN CLI wrapper +└── plddt_extract_pipeline.py # reads AF JSON → writes CSV +``` + +Pipeline implementation: `/scratch/bblj/mgoliyad1/IMPRESS/src/impress/pipelines/protein_binding.py` + +--- + +## Environment variables + +| Variable | Set in | Purpose | +|----------|--------|---------| +| `SCRATCH` | user env | base scratch path (`/scratch/bblj`) | +| `MPNN_PATH` | `delta_gpu_run.sh` | path to ProteinMPNN repo | +| `AF2_DATABASE` | `delta_gpu_run.sh` | AlphaFold database dir | +| `AF2_SIF` | `delta_gpu_run.sh` | AlphaFold Apptainer image | +| `IMPRESS_INPUT_DIR` | `delta_gpu_run.sh` | dir containing `{name}_in/` folders | +| `IMPRESS_OUTPUT_DIR` | `delta_gpu_run.sh` | output root (af_pipeline_outputs_multi/) | +| `IMPRESS_SCRIPTS_DIR` | `delta_gpu_run.sh` | dir with mpnn_wrapper.py, af2_multimer_reduced.sh | +| `SBATCH_ACCOUNT` | user env | SLURM account (e.g. `bblj-delta-gpu`) | + +`IMPRESS_PRE_EXEC` was removed — the venv is activated in the SLURM script before `dragon -s`, so all subprocesses inherit the environment. + +--- + +## Dragon integration + +**Backend**: `DragonExecutionBackendV2` (flow engine only — no tasks dispatched to Dragon workers). + +**Why not V3**: Dragon worker subprocesses silently hang when a subprocess uses GPU ops (`cupy`, `torch`, Apptainer). Same failure mode documented in `DeepDriveSim/workflows/miniapps_workflow/miniapps_workflow.py`. + +**Fix**: All pipeline steps are `local_task=True`. GPU steps (s1, s4) run as `asyncio.create_subprocess_shell()` on the Dragon head process, which has direct GPU access from SLURM allocation. Dragon manages the `ImpressManager`/`WorkflowEngine` loop; it never dispatches tasks to workers. + +**GPU affinity**: `_find_gpus()` + `_make_policy()` in `run_protein_binding.py` use `dragon.native.machine.System()` to discover GPUs and build a `dragon.infrastructure.policy.Policy` with `gpu_affinity`. The policy's GPU list is passed as `CUDA_VISIBLE_DEVICES` to each subprocess via `_gpu_env()` in the pipeline. + +```python +# To change GPU count per pipeline (default n_gpus=1 for MPNN): +policy = _make_policy(all_gpus, idx=0, n_gpus=1) +``` + +**Launcher** (`delta_gpu_run.sh`): +```bash +source ~/ve/impress/bin/activate +dragon-config add --ofi-runtime-lib="${FAB_LIB}" +rm -rf asyncflow.session.* +dragon -s run_protein_binding.py # -s = single-node +``` + +**SLURM resource request**: `--tasks-per-node=4`, `--gpus=4`, `--exclusive` (matches Dragon's single-node layout). + +--- + +## Key design decisions + +### Three separate path env vars +`base_path` was split into `IMPRESS_INPUT_DIR` / `IMPRESS_OUTPUT_DIR` / `IMPRESS_SCRIPTS_DIR` so input data, outputs, and scripts can live in different locations independently. + +### Output dirs auto-created in `__init__` +`os.makedirs(..., exist_ok=True)` is called in `ProteinBindingPipeline.__init__` for all required subdirs. No manual `mkdir` needed in the SLURM script. + +### `s4_post` replaces `post_exec` +`post_exec` in task descriptions is a RADICAL-Pilot artifact — `ConcurrentExecutionBackend` and Dragon never execute it. Replaced with a `local_task=True` coroutine using `glob` + `shutil.copy` to move `ranked_0.pdb` and `ranking_debug.json` to `best_models/` and `best_ptm/`. + +### AlphaFold container flags +- `--no-home`: prevents host's newer Biopython from shadowing the container's version (`SCOPData` import error) +- `--run_relax` removed: not supported by this container version +- Database bind: `/scratch/rhaas/SUP-5301/database:/database` +- Container: `/scratch/rhaas/SUP-5301/alphafold.sif` + +--- + +## One-time setup + +```bash +export SCRATCH=/scratch/bblj +cd /scratch/bblj/mgoliyad1/IMPRESS/examples/protien_binding_usecase +bash delta_env_setup.sh # creates ~/ve/impress/ + +# Unzip inputs (GNU tar, not zip): +cd /scratch/bblj/mgoliyad1/IMPRESS_inputs +tar -xf prod_in.tar # or whatever archive name +``` + +## Submitting a job + +```bash +export SCRATCH=/scratch/bblj +export SBATCH_ACCOUNT=bblj-delta-gpu +cd /scratch/bblj/mgoliyad1/IMPRESS/examples/protien_binding_usecase +sbatch delta_gpu_run.sh +``` + +## Checking results + +```bash +# pLDDT scores per pass: +cat af_stats_p1_pass_1.csv + +# AlphaFold best models: +ls IMPRESS_outputs/af_pipeline_outputs_multi/p1/af/prediction/best_models/ + +# Full logs: +tail -f logs/impress_.out +``` diff --git a/examples/protein_binding/af2_multimer_reduced.sh b/examples/protein_binding/af2_multimer_reduced.sh deleted file mode 100644 index ac686ce..0000000 --- a/examples/protein_binding/af2_multimer_reduced.sh +++ /dev/null @@ -1,44 +0,0 @@ -#!/bin/bash - -# have this *first* since they change the WORK env variable in jobs -# module reset -# module load cuda/12.3.0 - -set -e -set -x - -# work and upperdir need to be on same file system -WORK=/tmp/work -UPPER=/tmp/upper -mkdir -p $WORK $UPPER - -export XLA_PYTHON_CLIENT_PREALLOCATE="false" -export XLA_PYTHON_CLIENT_MEM_FRACTION=".75" -export XLA_PYTHON_CLIENT_ALLOCATOR="platform" - -#-B database.squashfs:/database:image-src=/ -INPUT_FASTA_FILE_DIR=$1 -INPUT_FASTA_FILE_NAME=$2 -OUTPUT_DATA_DIR=$3 - -apptainer run --nv \ - --bind $INPUT_FASTA_FILE_DIR:/fasta \ - --bind $OUTPUT_DATA_DIR:/dimer_models \ - --bind /anvil/datasets/alphafold/db_20230311:/database \ - /apps/biocontainers/images/tacc_alphafold:2.3.1.sif \ - --data_dir=/database \ - --uniref90_database_path=/database/uniref90/uniref90.fasta \ - --mgnify_database_path=/database/mgnify/mgy_clusters_2022_05.fa \ - --template_mmcif_dir=/database/pdb_mmcif/mmcif_files/ \ - --obsolete_pdbs_path=/database/pdb_mmcif/obsolete.dat \ - --fasta_paths=/fasta/$INPUT_FASTA_FILE_NAME \ - --output_dir=/dimer_models \ - --model_preset=multimer \ - --db_preset=reduced_dbs \ - --small_bfd_database_path=/database/small_bfd/bfd-first_non_consensus_sequences.fasta \ - --uniprot_database_path=/database/uniprot/uniprot.fasta \ - --pdb_seqres_database_path=/database/pdb_seqres/pdb_seqres.txt \ - --max_template_date=2020-12-01 \ - --use_gpu_relax=False \ - --num_multimer_predictions_per_model=1 \ - --run_relax=False \ No newline at end of file diff --git a/examples/protein_binding/delta_env_setup.sh b/examples/protein_binding/delta_env_setup.sh new file mode 100644 index 0000000..8a2d47f --- /dev/null +++ b/examples/protein_binding/delta_env_setup.sh @@ -0,0 +1,245 @@ +#!/bin/bash +# ============================================================================= +# IMPRESS Protein 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 via `module load python` (Delta default: 3.13+) +# ============================================================================= +if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then + set -euo pipefail +fi + +# ── Initialize lmod (needed when run as non-interactive bash script) ────────── +if ! declare -f module &>/dev/null; then + _lmod_init=/usr/share/lmod/lmod/init/bash + [ -f "${_lmod_init}" ] && source "${_lmod_init}" +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" + +echo "=================================================================" +echo " ENV_DIR = ${ENV_DIR}" +echo " IMPRESS_DIR = ${IMPRESS_DIR}" +echo "=================================================================" + +# ── 1. Create venv ──────────────────────────────────────────────────────────── +echo "" +echo "── Step 1: Creating venv ──" + +if [ -n "${BASE_PY_OVERRIDE}" ]; then + BASE_PY="${BASE_PY_OVERRIDE}" + echo "Using Python override: ${BASE_PY}" +else + # On Delta, `module load python` gives the default Python 3.13+. + module load python 2>/dev/null || true + BASE_PY=$(command -v python3 2>/dev/null || command -v python 2>/dev/null || true) + if [ -z "${BASE_PY}" ]; then + echo "ERROR: no Python found after 'module load python'." + echo " Pass an explicit interpreter: --python /path/to/python3" + exit 1 + fi + ver=$("${BASE_PY}" -c "import sys; v=sys.version_info; print(v.major*100+v.minor)") + if [ "${ver}" -lt 311 ]; then + echo "ERROR: ${BASE_PY} is Python ${ver} — need 3.11+." + echo " Pass an explicit interpreter: --python /path/to/python3.11" + 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) ─────────────────────────────────────────────────── +echo "" +echo "── Step 6: PyTorch (cu121) ──" +"${PIP}" install -q torch --index-url https://download.pytorch.org/whl/cu121 + +# ── 7. Additional dependencies ─────────────────────────────────────────────── +echo "" +echo "── Step 7: pandas + biopandas + matplotlib ──" +"${PIP}" install -q pandas biopandas matplotlib + +# ── 8. PyRosetta ───────────────────────────────────────────────────────────── +echo "" +echo "── Step 8: PyRosetta ──" +# pyrosetta_installer handles credential lookup internally. +# Activate the venv in the environment so its subprocess pip installs there. +export VIRTUAL_ENV="${ENV_DIR}" +export PATH="${ENV_DIR}/bin:${PATH}" +"${PIP}" install -q pyrosetta-installer +"${PY}" -c "import pyrosetta_installer; pyrosetta_installer.install_pyrosetta()" + +# ── 9. Boltz (structure prediction — separate conda env) ───────────────────── +echo "" +echo "── Step 9: Boltz (separate conda env) ──" +# boltz 2.x is kept in a separate conda env so its dependency pins (scipy, etc.) +# don't constrain the main IMPRESS venv. Boltz 2.0+ also installs fine via pip +# in Python 3.13, but we keep the separation to avoid pin conflicts. +# s4_boltz.sh activates BOLTZ_VENV (set in delta_gpu_run.sh) instead of VIRTUAL_ENV. +MINIFORGE="${MINIFORGE:-${SCRATCH}/${USER}/miniforge3}" +BOLTZ_ENV="${BOLTZ_ENV:-${HOME}/ve/boltz}" +# Cache lives in home dir — scratch inode quota can't hold the 45K CCD files. +BOLTZ_CACHE="${BOLTZ_CACHE:-${HOME}/boltz}" +if [ ! -x "${BOLTZ_ENV}/bin/python" ]; then + echo " Creating conda env (Python 3.11) at ${BOLTZ_ENV}" + # Use /tmp for package cache to avoid scratch quota exhaustion. + CONDA_PKGS_DIRS=/tmp/conda_pkgs "${MINIFORGE}/bin/conda" create -p "${BOLTZ_ENV}" python=3.11 -y -q +else + echo " boltz conda env already exists at ${BOLTZ_ENV}" +fi +echo " Installing boltz into ${BOLTZ_ENV}" +"${BOLTZ_ENV}/bin/pip" install -q boltz +# Pre-warm the boltz cache so compute nodes (no internet) find weights ready. +# boltz downloads mols.tar (45K CCD pkl files) + model weights on first run. +# Running predict on the login node (which has internet) pre-populates them. +echo " Pre-warming boltz cache at ${BOLTZ_CACHE}" +mkdir -p "${BOLTZ_CACHE}" +_WARM_FA="$(mktemp /tmp/boltz_warmup_XXXXXX.fasta)" +printf ">warmup|A\nGSSGSSGSS\n>warmup|B\nGSSGSSGSS\n" > "${_WARM_FA}" +BOLTZ_CACHE_DIR="${BOLTZ_CACHE}" "${BOLTZ_ENV}/bin/boltz" predict "${_WARM_FA}" \ + --out_dir "$(mktemp -d /tmp/boltz_warmup_out_XXXXXX)" \ + --cache "${BOLTZ_CACHE}" \ + --override 2>&1 | grep -E "Download|Extracting|Error|error" || true +rm -f "${_WARM_FA}" +echo " Boltz cache pre-warm done (model weights cached at ${BOLTZ_CACHE})" + +# Pre-compute MSAs for all input proteins using the MSA server (login node has internet). +# protein_binding.py s3() reads from BOLTZ_MSA_CACHE and embeds paths in the FASTA so +# compute nodes do not need internet access. Entity 0 = receptor, entity 1 = peptide. +BOLTZ_MSA_CACHE="${BOLTZ_CACHE}/msa_cache" +mkdir -p "${BOLTZ_MSA_CACHE}" +echo " Pre-computing MSAs into ${BOLTZ_MSA_CACHE}" +# IMPRESS_BASE_DIR = parent of prod_in/; IMPRESS_OUTPUT_DIR = parent of af_pipeline_outputs_multi/ +_scratch="${SCRATCH}" +_base_dir="${IMPRESS_BASE_DIR:-${_scratch}/IMPRESS_inputs}" +_out_dir="${IMPRESS_OUTPUT_DIR:-${_scratch}/IMPRESS_outputs}" +_msa_inputs_dir="${_base_dir}/prod_in" +if [ -d "${_msa_inputs_dir}" ]; then + for _pdb_dir in "${_msa_inputs_dir}"/p*_in; do + for _pdb in "${_pdb_dir}"/*.pdb; do + [ -f "${_pdb}" ] || continue + _stem="$(basename "${_pdb}" .pdb)" + _msa_csv="${BOLTZ_MSA_CACHE}/boltz_results_${_stem}/msa/${_stem}_0.csv" + if [ -f "${_msa_csv}" ]; then + echo " ${_stem}: MSA already cached, skipping" + continue + fi + echo " ${_stem}: generating MSA via server..." + _tmp_fa="$(mktemp /tmp/boltz_msa_XXXXXX.fasta)" + # Use FASTA from a prior run (sequences match the actual protein); fall back to + # a placeholder that will trigger MSA generation but gives a generic MSA. + _prior_fa="" + for _cand in "${_out_dir}"/af_pipeline_outputs_multi/*/af/fasta/"${_stem}.fa"; do + [ -f "${_cand}" ] && { _prior_fa="${_cand}"; break; } + done + if [ -n "${_prior_fa}" ]; then + cp "${_prior_fa}" "${_tmp_fa}" + else + printf ">pdz|protein\nGSSGSS\n>pep|protein\nGSSG\n" > "${_tmp_fa}" + fi + BOLTZ_CACHE_DIR="${BOLTZ_CACHE}" "${BOLTZ_ENV}/bin/boltz" predict "${_tmp_fa}" \ + --out_dir "${BOLTZ_MSA_CACHE}" \ + --use_msa_server \ + --cache "${BOLTZ_CACHE}" \ + --output_format pdb \ + --override 2>&1 | grep -E "MSA|Generat|Error|error|skip" || true + rm -f "${_tmp_fa}" + done + done +else + echo " prod_in not found at ${_msa_inputs_dir}, skipping MSA pre-compute" +fi +echo " MSA pre-compute done" + +# ── 10. Verify ─────────────────────────────────────────────────────────────── +echo "" +echo "── Step 10: 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 "pandas" "${PY}" -c "import pandas; print(pandas.__version__)" +BOLTZ_ENV="${BOLTZ_ENV:-${HOME}/ve/boltz}" +_check "boltz" "${BOLTZ_ENV}/bin/python" -c "import boltz; print('ok')" + +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/protein_binding" +echo " sbatch delta_gpu_run.sh" +echo "=================================================================" diff --git a/examples/protein_binding/delta_gpu_run.sh b/examples/protein_binding/delta_gpu_run.sh new file mode 100644 index 0000000..14161e9 --- /dev/null +++ b/examples/protein_binding/delta_gpu_run.sh @@ -0,0 +1,117 @@ +#!/bin/bash +# +# Protein 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/ +# +# Example: +# sbatch delta_gpu_run.sh +# +# 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=02:30:00 +#SBATCH --job-name=impress_protein +#SBATCH --mail-user= +#SBATCH --mail-type=ALL +#SBATCH --output=logs/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_A}" +unset SLURM_EXPORT_ENV +source "${IMPRESS_VENV}/bin/activate" +dragon-config add --ofi-runtime-lib="${FAB_LIB}" + +# ── Tool paths (adjust for your allocation) ─────────────────────────────────── +export MPNN_PATH="${MPNN_PATH:-${SCRATCH}/${USER}/ProteinMPNN}" +export AF2_DATABASE="${AF2_DATABASE:-${SCRATCH}/${USER}/alphafold_database}" +export AF2_SIF="${AF2_SIF:-${SCRATCH}/${USER}/alphafold.sif}" +# Boltz lives in a separate Python 3.12 conda env (boltz 2.x requires numpy<2.0, +# scipy==1.13.1 etc. which have no Python 3.13 wheels). +export BOLTZ_VENV="${BOLTZ_VENV:-${HOME}/ve/boltz}" +# Boltz model weight cache — kept in home dir; scratch inode quota can't hold +# the 45K CCD molecule files that boltz extracts from mols.tar on first run. +export BOLTZ_CACHE_DIR="${BOLTZ_CACHE_DIR:-${HOME}/boltz}" +mkdir -p "${BOLTZ_CACHE_DIR}" + +# ── IMPRESS paths ───────────────────────────────────────────────────────────── +export IMPRESS_SCRIPTS_DIR="${IMPRESS_SCRIPTS_DIR:-${SCRATCH}/${USER}/IMPRESS/examples/protein_binding}" +# IMPRESS_BASE_DIR: parent of prod_in/ — pipeline builds prod_in/_in from here +export IMPRESS_BASE_DIR="${IMPRESS_BASE_DIR:-${SCRATCH}/${USER}/IMPRESS_inputs}" +export IMPRESS_OUTPUT_DIR="${IMPRESS_OUTPUT_DIR:-${SCRATCH}/${USER}/IMPRESS_outputs}" + +# 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, max_passes=1, no child pipelines. +# Runs a single MPNN → score → AF2 cycle to verify end-to-end path. +# 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}" + +# ── Working directory ───────────────────────────────────────────────────────── +WORKDIR="${IMPRESS_SCRIPTS_DIR}" +cd "${WORKDIR}" +mkdir -p logs + +# 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:-${WORKDIR}/logs/sessions}" +mkdir -p "${IMPRESS_SESSION_DIR}" + +# ── 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 + +if [ "${IMPRESS_BACKEND}" = "dragon" ]; then + rm -f ddict_orc* + echo "Running: dragon ${DRAGON_MODE} run_protein_binding.py (nodes=${SLURM_NNODES:-1})" + dragon ${DRAGON_MODE} run_protein_binding.py +else + echo "Running: python3 run_protein_binding.py (backend=${IMPRESS_BACKEND})" + python3 run_protein_binding.py +fi + +echo "=== Protein Binding pipeline done: $(date) ===" diff --git a/examples/protein_binding/mpnn_wrapper.py b/examples/protein_binding/mpnn_wrapper.py index 1e8942d..811fd4d 100644 --- a/examples/protein_binding/mpnn_wrapper.py +++ b/examples/protein_binding/mpnn_wrapper.py @@ -1,4 +1,4 @@ -#!/bin/sh +#!/usr/bin/env python3 import argparse parser = argparse.ArgumentParser() import subprocess @@ -15,7 +15,7 @@ parser.add_argument("-homo", "--homo", help="Are your input files homomers? (0 or 1) Note: Overrides tied positions. If homomer is specified, all positions on designed chains will be tied. If positions are restricted, lists must be of same length. Example: -index='1 3 5 7, 1 3 5 7'", type=int) parser.add_argument("-bias_AA", "--bias_AA", help="For which amino acids would you like to install bias? Example: -bias_AA='D E H'", type=str) parser.add_argument("-bias_weight", "--bias_weight", help="What weights would you like to install for the biased amino acids? Lists must match in length. Example: -bias_weight='0.3 -0.3 0.5'", type=str) -parser.add_argument("-temp", "--temp", help="What temperature would you like to sample from? Example: 0.3", type=int, default=0.1) +parser.add_argument("-temp", "--temp", help="What temperature would you like to sample from? Example: 0.3", type=float, default=0.1) parser.add_argument("-inter", "--interface", help="Would you like to design the interface? Do not specify indices if designing the interface. (1 or 0)", type=int, default=0) @@ -27,12 +27,12 @@ mpnn_path=args.mpnn_path is_monomer=args.is_monomer #default is_monomer false chains=args.design_chains -if chains == None: +if chains is None: chains='A' #default design chain A index=args.index fix=args.fix #default false, specify non fixed seqs=args.seqs -if seqs == None: +if seqs is None: seqs=1 #default 1 design per structure tie=args.tie homo=args.homo diff --git a/examples/protein_binding/plddt_extract_pipeline.py b/examples/protein_binding/plddt_extract_pipeline.py index d1c8a7d..6dc57bd 100644 --- a/examples/protein_binding/plddt_extract_pipeline.py +++ b/examples/protein_binding/plddt_extract_pipeline.py @@ -59,4 +59,5 @@ print(f"Processed {len(rows)} structure(s)") df = pd.DataFrame(rows, columns=['ID', 'avg_plddt', 'ptm', 'avg_pae']) -df.to_csv('af_stats_' + args.out + '_pass_' + args.iter + '.csv', index=False) +csv_path = os.path.join(args.path, 'af_stats_' + args.out + '_pass_' + args.iter + '.csv') +df.to_csv(csv_path, index=False) diff --git a/examples/protein_binding/protein_binding.py b/examples/protein_binding/protein_binding.py index 64ca80d..32258ac 100644 --- a/examples/protein_binding/protein_binding.py +++ b/examples/protein_binding/protein_binding.py @@ -6,7 +6,7 @@ from impress.pipelines.impress_pipeline import ImpressBasePipeline -MPNN_PATH = f"/anvil/projects/x-nairr240405/mason/ProteinMPNN" +MPNN_PATH = os.environ.get("MPNN_PATH", "") _BOLTZ_CHAIN_MAP = {'pdz': 'A', 'pep': 'B'} @@ -20,6 +20,10 @@ def _copy_pdb_rename_chains(src, dst, chain_map=_BOLTZ_CHAIN_MAP): line = line[:21] + chain_map[chain] + line[24:] f_out.write(line) +# One semaphore per GPU shared across all pipeline instances — caps concurrent +# Boltz launches per GPU at 2 regardless of how many pipelines share that GPU. +_boltz_sem_per_gpu: dict = {} + class ProteinBindingPipeline(ImpressBasePipeline): def __init__(self, name, flow, configs=None, **kwargs): # Execution metadata @@ -34,7 +38,11 @@ def __init__(self, name, flow, configs=None, **kwargs): self.num_seqs = kwargs.get("num_seqs", 10) self.sub_order = kwargs.get("sub_order", 0) self.max_passes = kwargs.get("max_passes", 10) - self.mpnn_path = kwargs.get("mpnn_path", MPNN_PATH) + self.mpnn_path = kwargs.get("mpnn_path") or MPNN_PATH + if not self.mpnn_path: + raise ValueError("mpnn_path must be supplied via kwarg or MPNN_PATH env var") + self.peptide_seq: str = kwargs.get("peptide_seq", "EGYQDYEPEA") + self.gpu_id = kwargs.get("gpu_id", None) # Sequence and score state self.current_scores = {} @@ -46,12 +54,18 @@ def __init__(self, name, flow, configs=None, **kwargs): # Input-related self.fasta_list_2 = kwargs.get("fasta_list_2", []) self.base_path = kwargs.get("base_path", os.getcwd()) + # input_base_path is the parent of prod_in/; defaults to base_path so + # existing callers that keep input data next to scripts still work. + self.input_base_path = kwargs.get("input_base_path", self.base_path) + # output_base_path is where af_pipeline_outputs_multi/ is written; + # defaults to base_path so existing callers without IMPRESS_OUTPUT_DIR work. + self.output_base_path = kwargs.get("output_base_path", self.base_path) self.scripts_path = os.path.join(self.base_path, "scripts") - self.input_path = os.path.join(self.base_path, f"prod_in/{self.name}_in") + self.input_path = os.path.join(self.input_base_path, f"prod_in/{self.name}_in") # Output paths self.output_path = os.path.join( - self.base_path, "af_pipeline_outputs_multi", self.name + self.output_base_path, "af_pipeline_outputs_multi", self.name ) self.output_path_mpnn = os.path.join(self.output_path, "mpnn") self.output_path_af = os.path.join( @@ -65,12 +79,9 @@ def __init__(self, name, flow, configs=None, **kwargs): def set_up_new_pipeline_dirs(self, new_pipeline_name): base_output = os.path.join( - self.base_path, "af_pipeline_outputs_multi", new_pipeline_name + self.output_base_path, "af_pipeline_outputs_multi", new_pipeline_name ) - input_dir = os.path.join(self.base_path, f"prod_in/{new_pipeline_name}_in") - - if os.path.isdir(base_output): - return # already exists, nothing to do + input_dir = os.path.join(self.input_base_path, f"prod_in/{new_pipeline_name}_in") # all directories to create subdirs = [ @@ -95,15 +106,16 @@ def register_pipeline_tasks(self): """Register all pipeline tasks""" @self.auto_register_task(capture_stdio=True) # MPNN - async def s1(task_description={"gpus_per_rank": 1}): + async def s1(): # noqa: B006 self.step_id += 1 mpnn_script = os.path.join(self.base_path, "mpnn_wrapper.py") output_dir = os.path.join(self.output_path_mpnn, f"job_{self.passes}") + os.makedirs(output_dir, exist_ok=True) chain = "A" input_path = self.input_path if self.passes == 1 else self.output_path_af - return ( + cmd = ( f"bash {self.scripts_path}/s1_mpnn.sh " f"{mpnn_script} " f"{input_path} " @@ -112,6 +124,7 @@ async def s1(task_description={"gpus_per_rank": 1}): f"{self.num_seqs} " f"{chain}" ) + return cmd @self.auto_register_task(local_task=True) async def s2(): @@ -139,17 +152,38 @@ async def s2(): async def s3(): self.step_id += 1 output_dir = os.path.join(self.output_path, "af", "fasta") + # Pre-computed MSA cache (populated by delta_env_setup.sh on the login node, + # which has internet access). Entity 0 = receptor (PDZ), entity 1 = peptide. + msa_cache = os.path.expanduser("~/boltz/msa_cache") fasta_file_to_return = [] for fasta_file in self.fasta_list_2: base_name = fasta_file.split(".")[0] fasta_file_to_return.append(base_name) design_seq = self.iter_seqs[base_name][self.seq_rank][0] - pep_seq = "EGYQDYEPEA" + pep_seq = self.peptide_seq + + # Boltz embeds pre-computed MSAs in the FASTA header rather than + # running MSA search at prediction time (compute nodes have no internet). + # Naming convention: entity 0 = receptor (PDZ), entity 1 = peptide. + pdz_msa = os.path.join( + msa_cache, f"boltz_results_{base_name}", "msa", f"{base_name}_0.csv" + ) + pep_msa = os.path.join( + msa_cache, f"boltz_results_{base_name}", "msa", f"{base_name}_1.csv" + ) + if os.path.exists(pdz_msa) and os.path.exists(pep_msa): + # |path/to/msa.csv| tells Boltz to load the pre-computed MSA + pdz_tag = f">pdz|protein|{pdz_msa}" + pep_tag = f">pep|protein|{pep_msa}" + else: + # |empty| runs single-sequence mode — no MSA, slightly less accurate + pdz_tag = ">pdz|protein|empty" + pep_tag = ">pep|protein|empty" fasta_path = os.path.join(output_dir, f"{base_name}.fa") with open(fasta_path, "w") as f: - f.write(f">pdz|protein\n{design_seq}\n>pep|protein\n{pep_seq}\n") + f.write(f"{pdz_tag}\n{design_seq}\n{pep_tag}\n{pep_seq}\n") return fasta_file_to_return @@ -163,14 +197,14 @@ async def s3(): # ) @self.auto_register_task(capture_stdio=True) - async def s4(target_fasta, task_description={"gpus_per_rank": 1}): # noqa: B006 + async def s4(target_fasta): # noqa: B006 self.step_id += 1 cmd = ( f"bash {self.scripts_path}/s4_boltz.sh " f"{self.output_path}/af/fasta/{target_fasta}.fa " f"{self.output_path}/af/prediction/dimer_models/{target_fasta}" + + (f" {self.gpu_id}" if self.gpu_id is not None else "") ) - self.logger.pipeline_log(f"s4 command for {target_fasta}: {cmd}") return cmd @self.auto_register_task(local_task=True) @@ -206,7 +240,7 @@ async def s5(): self.step_id += 1 return ( f"bash {self.scripts_path}/s5_plddt_extract.sh " - f"{self.base_path} " + f"{self.output_base_path} " f"{self.passes} " f"{self.name}" ) @@ -217,15 +251,19 @@ async def get_scores_map(self): def finalize(self, sub_iter_seqs): # finalize the "cleanup" of the current pipeline + from pathlib import Path for a in sub_iter_seqs: self.fasta_list_2.remove(f"{a}.pdb") - os.unlink(f"{self.output_path_af}/{a}.pdb") - os.unlink(f"{self.output_path}/af/fasta/{a}.fa") + Path(f"{self.output_path_af}/{a}.pdb").unlink(missing_ok=True) + Path(f"{self.output_path}/af/fasta/{a}.fa").unlink(missing_ok=True) self.previous_scores = copy.deepcopy(self.current_scores) async def run(self): """Main execution logic""" + if self.gpu_id is not None: + self.logger.pipeline_log(f"gpu={self.gpu_id}") + self.logger.pipeline_log(f"Running for a maximum of {self.max_passes} passes") self.set_up_new_pipeline_dirs(self.name) @@ -241,7 +279,20 @@ async def run(self): else: self.logger.pipeline_log("Submitting MPNN task") - await self.s1() + try: + await self.s1() + except Exception as exc: + # The execution backend may report a spurious failure (TypeError/ + # 'NoneType' subscriptable, or ProcessGroup state error) even when + # MPNN completed successfully. Check for output before propagating. + seqs_dir = os.path.join( + self.output_path_mpnn, f"job_{self.passes}", "seqs" + ) + if not (os.path.isdir(seqs_dir) and os.listdir(seqs_dir)): + raise + self.logger.pipeline_log( + f"s1 raised {exc!r} but seqs output exists — treating as success" + ) self.logger.pipeline_log("MPNN task finished") self.logger.pipeline_log("Submitting sequence ranking task") @@ -255,6 +306,29 @@ async def run(self): alphafold_tasks = [] post_exec_tasks = [] + # Shared per-GPU semaphore caps concurrent Boltz launches at 2 per GPU + # across all pipeline instances pinned to the same GPU. + gpu_key = self.gpu_id if self.gpu_id is not None else "default" + if gpu_key not in _boltz_sem_per_gpu: + _boltz_sem_per_gpu[gpu_key] = asyncio.Semaphore(2) + _boltz_sem = _boltz_sem_per_gpu[gpu_key] + + async def _guarded_s4(target_fasta): + async with _boltz_sem: + try: + return await self.s4(target_fasta=target_fasta) + except Exception as exc: + # The execution backend may raise a spurious failure even when + # Boltz completed successfully. Check for the output PDB before propagating. + pred_dir = os.path.join( + self.output_path, "af", "prediction", "dimer_models", + target_fasta, f"boltz_results_{target_fasta}", + "predictions", target_fasta, + ) + if os.path.isfile(os.path.join(pred_dir, f"{target_fasta}_model_0.pdb")): + return None # output exists; treat as success + raise + for target_fasta in fasta_files: models_path = os.path.join( self.output_path, "af", "prediction", "dimer_models", target_fasta, @@ -282,8 +356,8 @@ async def run(self): f"{target_fasta}.pdb", ) - # launch coroutine without awaiting yet - alphafold_tasks.append(self.s4(target_fasta=target_fasta)) + # launch coroutine without awaiting yet (semaphore-gated) + alphafold_tasks.append(_guarded_s4(target_fasta)) post_exec_tasks.append( self.s4_post_exec( target_fasta=target_fasta, @@ -306,26 +380,40 @@ async def run(self): self.logger.pipeline_log(f"s4 DONE for {fasta_name}") s4_post_results = await asyncio.gather(*post_exec_tasks, return_exceptions=True) + any_s4_ok = False for fasta_name, result in zip(fasta_files, s4_post_results): if isinstance(result, Exception): self.logger.pipeline_log(f"s4_post_exec FAILED for {fasta_name}: {result}") else: self.logger.pipeline_log(f"s4_post_exec DONE for {fasta_name}") + any_s4_ok = True + + if not any_s4_ok: + raise RuntimeError("All s4 tasks failed — skipping pLDDT extraction") self.logger.pipeline_log("Submitting pLDTT extraction task") staged_file = f"af_stats_{self.name}_pass_{self.passes}.csv" - await self.s5( - task_description={ - "output_staging": [ - { - "source": f"task:///{staged_file}", - "target": f"client:///{staged_file}", - } - ], - } - ) + try: + await self.s5( + task_description={ + "output_staging": [ + { + "source": f"task:///{staged_file}", + "target": f"client:///{staged_file}", + } + ], + } + ) + except Exception as exc: + # Spurious backend failure: check if s5 wrote the CSV despite the error. + csv_path = os.path.join(self.output_base_path, staged_file) + if not os.path.isfile(csv_path): + raise + self.logger.pipeline_log( + f"s5 raised {exc!r} but CSV exists — treating as success" + ) self.logger.pipeline_log("pLDTT extract finished") await self.run_adaptive_step(wait=True) diff --git a/examples/protein_binding/protein_binding_run.py b/examples/protein_binding/protein_binding_run.py index fa9e4b9..bacbd3c 100644 --- a/examples/protein_binding/protein_binding_run.py +++ b/examples/protein_binding/protein_binding_run.py @@ -11,10 +11,9 @@ import numpy as np import pandas as pd -# from rhapsody.backends import DragonExecutionBackendV3 -from radical.asyncflow import LocalExecutionBackend +from rhapsody.backends import DragonExecutionBackend import rhapsody -rhapsody.enable_logging(level=logging.DEBUG) +rhapsody.enable_logging(level=logging.INFO) from impress import PipelineSetup from impress import ImpressManager @@ -326,8 +325,7 @@ async def adaptive_decision(pipeline: ProteinBindingPipeline) -> None: async def impress_protein_bind() -> None: """Execute protein binding analysis with LLM-driven adaptive optimization.""" -# backend = await DragonExecutionBackendV3() - backend = await LocalExecutionBackend(ProcessPoolExecutor()) + backend = await DragonExecutionBackend() manager: ImpressManager = ImpressManager(execution_backend=backend) diff --git a/examples/protein_binding/run_nonadaptive.py b/examples/protein_binding/run_nonadaptive.py index 6eea85e..6a97fe7 100644 --- a/examples/protein_binding/run_nonadaptive.py +++ b/examples/protein_binding/run_nonadaptive.py @@ -1,7 +1,7 @@ import asyncio from typing import List -from rhapsody.backends import DragonExecutionBackendV3 +from rhapsody.backends import DragonExecutionBackend from rhapsody.telemetry import define_event from impress import PipelineSetup @@ -9,7 +9,7 @@ from protein_binding import ProteinBindingPipeline import rhapsody, logging -rhapsody.enable_logging(level=logging.DEBUG) +rhapsody.enable_logging(level=logging.INFO) def _on_task_event(event) -> None: @@ -19,7 +19,7 @@ def _on_task_event(event) -> None: async def impress_protein_bind_nonadaptive() -> None: - backend = await DragonExecutionBackendV3() + backend = await DragonExecutionBackend() manager: ImpressManager = ImpressManager( execution_backend=backend, diff --git a/examples/protein_binding/run_protein_binding.py b/examples/protein_binding/run_protein_binding.py index cf66937..14a3e8a 100644 --- a/examples/protein_binding/run_protein_binding.py +++ b/examples/protein_binding/run_protein_binding.py @@ -1,20 +1,36 @@ import copy +import os import shutil import asyncio from typing import Dict, Any, Optional, List -from rhapsody.backends import DragonExecutionBackendV3 from rhapsody.telemetry import define_event from rhapsody.telemetry.events import make_event -from impress import PipelineSetup -from impress import ImpressManager +from impress import find_gpus, ImpressManager, PipelineSetup from protein_binding import ProteinBindingPipeline import rhapsody, logging rhapsody.enable_logging(level=logging.DEBUG) +# ── Backend / test mode ─────────────────────────────────────────────────── +# IMPRESS_BACKEND: "dragon" (default, multi-node HPC) or "local" (single-node, +# ProcessPoolExecutor — useful for development / non-Dragon clusters). +BACKEND = os.environ.get("IMPRESS_BACKEND", "dragon").lower() + +if BACKEND == "dragon": + from rhapsody.backends import DragonExecutionBackend +else: + from concurrent.futures import ProcessPoolExecutor + from rhapsody.backends import ConcurrentExecutionBackend +TEST_MODE = os.getenv("IMPRESS_TEST_MODE", "0") == "1" +N_PIPELINES = 4 if TEST_MODE else 16 +MAX_PASSES = 10 if TEST_MODE else 10 +MAX_SUB_PIPELINES_OVERRIDE = 3 if TEST_MODE else None # None = use inline default + +print(f"[INFO] IMPRESS_BACKEND={BACKEND} TEST_MODE={TEST_MODE} N_PIPELINES={N_PIPELINES} MAX_PASSES={MAX_PASSES} MAX_SUB_PIPELINES_OVERRIDE={MAX_SUB_PIPELINES_OVERRIDE}") + # --------------------------------------------------------------------------- # Custom application-level telemetry events # --------------------------------------------------------------------------- @@ -52,21 +68,22 @@ # --------------------------------------------------------------------------- def _on_task_event(event) -> None: - if event.event_type == "TaskFailed": - wid = getattr(event, "workflow_id", None) - print(f"[TELEMETRY] TaskFailed task={event.task_id} workflow={wid}") + pass # --------------------------------------------------------------------------- # Adaptive helpers # --------------------------------------------------------------------------- -async def adaptive_criteria(current_score: float, previous_score: float) -> bool: +def adaptive_criteria(current_score: float, previous_score: float) -> bool: return current_score > previous_score async def impress_protein_bind() -> None: - backend = await DragonExecutionBackendV3() + if BACKEND == "dragon": + backend = await DragonExecutionBackend() + else: + backend = await ConcurrentExecutionBackend.create(ProcessPoolExecutor()) manager: ImpressManager = ImpressManager( execution_backend=backend, @@ -80,12 +97,12 @@ async def impress_protein_bind() -> None: # adaptive_decision closes over `manager` so it can emit events via # manager.telemetry, which is set inside start() before any pipeline runs. async def adaptive_decision(pipeline: ProteinBindingPipeline) -> Optional[Dict[str, Any]]: - MAX_SUB_PIPELINES: int = 3 + MAX_SUB_PIPELINES: int = MAX_SUB_PIPELINES_OVERRIDE if MAX_SUB_PIPELINES_OVERRIDE is not None else 3 tel = manager.telemetry sid = tel.session_id if tel else None # Read current scores from CSV - file_name = f'af_stats_{pipeline.name}_pass_{pipeline.passes}.csv' + file_name = os.path.join(pipeline.output_base_path, f'af_stats_{pipeline.name}_pass_{pipeline.passes}.csv') with open(file_name) as fd: for line in fd.readlines()[1:]: line = line.strip() @@ -108,13 +125,14 @@ async def adaptive_decision(pipeline: ProteinBindingPipeline) -> Optional[Dict[s continue prev_score = pipeline.previous_scores[protein] - decision = await adaptive_criteria(curr_score, prev_score) + decision = adaptive_criteria(curr_score, prev_score) pipeline.logger.pipeline_log(f'Adaptive decision: {decision}') if tel: tel.emit(make_event( ProteinScore, session_id=sid, + backend="rhapsody", protein=protein, pipeline_name=pipeline.name, pass_num=pipeline.passes, @@ -135,13 +153,14 @@ async def adaptive_decision(pipeline: ProteinBindingPipeline) -> Optional[Dict[s for protein in sub_iter_seqs: src = f'{pipeline.output_path_af}/{protein}.pdb' - dst = f'{pipeline.base_path}/prod_in/{new_name}_in/{protein}.pdb' + dst = f'{pipeline.input_base_path}/prod_in/{new_name}_in/{protein}.pdb' shutil.copyfile(src, dst) if tel: tel.emit(make_event( ChildPipelineSpawned, session_id=sid, + backend="rhapsody", parent_name=pipeline.name, child_name=new_name, num_proteins=len(sub_iter_seqs), @@ -160,6 +179,9 @@ async def adaptive_decision(pipeline: ProteinBindingPipeline) -> Optional[Dict[s 'seq_rank': pipeline.seq_rank + 1, 'sub_order': pipeline.sub_order + 1, 'previous_scores': copy.deepcopy(pipeline.previous_scores), + 'input_base_path': pipeline.input_base_path, + 'output_base_path': pipeline.output_base_path, + 'gpu_id': pipeline.gpu_id, } } @@ -177,6 +199,7 @@ async def adaptive_decision(pipeline: ProteinBindingPipeline) -> Optional[Dict[s tel.emit(make_event( PassSummary, session_id=sid, + backend="rhapsody", pipeline_name=pipeline.name, pass_num=pipeline.passes, num_proteins=len(pipeline.current_scores), @@ -184,13 +207,38 @@ async def adaptive_decision(pipeline: ProteinBindingPipeline) -> Optional[Dict[s child_spawned=child_spawned, )) + # IMPRESS_SCRIPTS_DIR = protein_binding examples dir (scripts/, mpnn_wrapper.py) + # IMPRESS_BASE_DIR = parent of prod_in/ (input PDB files) + # IMPRESS_OUTPUT_DIR = where af_pipeline_outputs_multi/ is written + scripts_dir = os.environ.get( + "IMPRESS_SCRIPTS_DIR", os.path.dirname(os.path.abspath(__file__)) + ) + input_base_dir = os.environ.get("IMPRESS_BASE_DIR", scripts_dir) + output_base_dir = os.environ.get("IMPRESS_OUTPUT_DIR", scripts_dir) + os.makedirs(output_base_dir, exist_ok=True) + + all_gpus = find_gpus() + + if all_gpus: + print("[INFO] GPU assignment:") + for i in range(1, N_PIPELINES + 1): + gpu_id = all_gpus[(i - 1) % len(all_gpus)] + print(f"[INFO] p{i:>2} -> gpu={gpu_id}") + pipeline_setups: List[PipelineSetup] = [ PipelineSetup( name=f"p{str(i)}", type=ProteinBindingPipeline, - adaptive_fn=adaptive_decision + config={ + "base_path": scripts_dir, + "input_base_path": input_base_dir, + "output_base_path": output_base_dir, + "max_passes": MAX_PASSES, + **({"gpu_id": all_gpus[(i - 1) % len(all_gpus)]} if all_gpus else {}), + }, + adaptive_fn=adaptive_decision, ) - for i in range(1, 17) + for i in range(1, N_PIPELINES + 1) ] await manager.start(pipeline_setups=pipeline_setups) @@ -201,8 +249,8 @@ async def adaptive_decision(pipeline: ProteinBindingPipeline) -> Optional[Dict[s dur = summary.get("duration") if dur: print(f"[TELEMETRY] mean task time: {dur['mean_seconds'] * 1000:.1f} ms") + await manager.telemetry.stop() - await manager.flow.shutdown() if __name__ == "__main__": diff --git a/examples/protein_binding/scripts/s1_mpnn.sh b/examples/protein_binding/scripts/s1_mpnn.sh index 67f4110..377a38d 100755 --- a/examples/protein_binding/scripts/s1_mpnn.sh +++ b/examples/protein_binding/scripts/s1_mpnn.sh @@ -11,8 +11,8 @@ mpnn_path="$4" num_seqs="$5" chain="$6" -source /anvil/projects/x-nairr240405/mason/LigandMPNN/.venv/bin/activate -#source /ocean/projects/dmr170002p/hooten/LigandMPNN/.venv/bin/activate +# Re-activate the IMPRESS venv if running inside a subprocess (VIRTUAL_ENV is exported by sbatch). +[ -n "${VIRTUAL_ENV:-}" ] && source "${VIRTUAL_ENV}/bin/activate" python3 "$mpnn_script" \ -pdb="$input_path" \ diff --git a/examples/protein_binding/scripts/s4_alphafold.sh b/examples/protein_binding/scripts/s4_alphafold.sh index 29b2126..91fedc1 100755 --- a/examples/protein_binding/scripts/s4_alphafold.sh +++ b/examples/protein_binding/scripts/s4_alphafold.sh @@ -7,14 +7,32 @@ set -euo pipefail fasta_path="$1" output_dir="$2" -module load modtree/gpu -module load cuda/12.8.0 -module load gcc/11.2.0 -#source /anvil/scratch/x-mason/IMPRESS/.venv/bin/activate -source /ocean/projects/dmr170002p/hooten/IMPRESS/.venv/bin/activate +# Re-activate the IMPRESS venv if running inside a subprocess (VIRTUAL_ENV is exported by sbatch). +[ -n "${VIRTUAL_ENV:-}" ] && source "${VIRTUAL_ENV}/bin/activate" -pixi run --manifest-path /anvil/scratch/x-mason/localcolabfold \ - colabfold_batch \ +# ── Test-mode stub (IMPRESS_TEST_MODE=1) ────────────────────────────────── +# Write minimal Boltz-shaped output so plddt_extract_pipeline can run without +# invoking the real AlphaFold model. +if [ "${IMPRESS_TEST_MODE:-0}" = "1" ]; then + name="$(basename "${fasta_path}" .fa)" + pred_dir="${output_dir}/boltz_results_${name}/predictions/${name}" + mkdir -p "${pred_dir}" + python3 - "${pred_dir}" "${name}" <<'PYEOF' +import sys, json +import numpy as np +pred_dir, name = sys.argv[1], sys.argv[2] +n = 110 # 100 PDZ residues + 10 peptide (PEP_LEN=10 assumed by extractor) +np.savez(f"{pred_dir}/plddt_{name}_model_0.npz", plddt=np.full(n, 0.85)) +np.savez(f"{pred_dir}/pae_{name}_model_0.npz", pae=np.full((n, n), 2.0)) +with open(f"{pred_dir}/confidence_{name}_model_0.json", "w") as f: + json.dump({"iptm": 0.75, "ptm": 0.80}, f) +PYEOF + echo "[MOCK] s4_alphafold stub done for ${name}" + exit 0 +fi +# ── End test-mode stub ──────────────────────────────────────────────────── + +colabfold_batch \ --model-type alphafold2_multimer_v3 \ --max-template-date 2020-12-01 \ --rank multimer \ diff --git a/examples/protein_binding/scripts/s4_boltz.sh b/examples/protein_binding/scripts/s4_boltz.sh index bab328c..bb5b50f 100755 --- a/examples/protein_binding/scripts/s4_boltz.sh +++ b/examples/protein_binding/scripts/s4_boltz.sh @@ -2,30 +2,112 @@ set -e # Step 4: Structure prediction via Boltz -# Args: $1=fasta_path $2=output_dir +# Args: $1=fasta_path $2=output_dir $3=gpu_id (optional) fasta_path="$1" output_dir="$2" +# Optional GPU assignment passed by the caller so tasks spread across GPUs. +if [ -n "${3:-}" ]; then + export CUDA_VISIBLE_DEVICES="$3" +fi -#source /ocean/projects/dmr170002p/hooten/IMPRESS/.venv/bin/activate -source /anvil/projects/x-nairr240405/mason/IMPRESS/.venv/bin/activate -module load modtree/gpu +# Boltz requires Python <=3.12 (numpy<2.0 etc.) so it lives in its own env. +# BOLTZ_VENV may point to a conda env (no bin/activate) or a pip venv; prepend +# its bin/ to PATH so the correct python/boltz are found in either case. +_BOLTZ_ENV="${BOLTZ_VENV:-${VIRTUAL_ENV:-}}" +[ -n "${_BOLTZ_ENV}" ] && export PATH="${_BOLTZ_ENV}/bin:${PATH}" export SSL_CERT_FILE=/etc/pki/tls/certs/ca-bundle.crt -# Boltz caches MSA in boltz_results_/msa/ and reuses it across runs even with -# --override, so pass 2+ would fold new MPNN-designed sequences using the pass-1 MSA. -# Delete the stale MSA before each run to force recomputation for the current sequence. -fasta_stem=$(basename "${fasta_path}" .fa) -stale_msa="${output_dir}/boltz_results_${fasta_stem}/msa" -if [ -d "${stale_msa}" ]; then - rm -rf "${stale_msa}" +# Compute nodes typically have no internet access, so --use_msa_server is off +# by default. Set BOLTZ_USE_MSA_SERVER=1 to enable it on nodes with internet. +_MSA_FLAG="" +[ "${BOLTZ_USE_MSA_SERVER:-0}" = "1" ] && _MSA_FLAG="--use_msa_server" + +# ── Test-mode stub (IMPRESS_TEST_MODE=1) ────────────────────────────────── +# Write minimal Boltz-shaped output so downstream steps (plddt_extract_pipeline) +# can run without invoking the real Boltz model. numpy is available because +# BOLTZ_VENV/bin is already on PATH above. +if [ "${IMPRESS_TEST_MODE:-0}" = "1" ]; then + name="$(basename "${fasta_path}" .fa)" + pred_dir="${output_dir}/boltz_results_${name}/predictions/${name}" + mkdir -p "${pred_dir}" + python3 - "${pred_dir}" "${name}" <<'PYEOF' +import sys, json +import numpy as np +pred_dir, name = sys.argv[1], sys.argv[2] +n = 110 # 100 PDZ residues + 10 peptide (PEP_LEN=10 assumed by extractor) +np.savez(f"{pred_dir}/plddt_{name}_model_0.npz", plddt=np.full(n, 0.85)) +np.savez(f"{pred_dir}/pae_{name}_model_0.npz", pae=np.full((n, n), 2.0)) +with open(f"{pred_dir}/confidence_{name}_model_0.json", "w") as f: + json.dump({"iptm": 0.75, "ptm": 0.80}, f) +PYEOF + echo "[MOCK] s4_boltz stub done for ${name}" + exit 0 fi +# ── End test-mode stub ──────────────────────────────────────────────────── + +mkdir -p "${output_dir}" + +_boltz_cache_dir="${BOLTZ_CACHE_DIR:-${HOME}/.boltz}" + +# $_boltz_cache_dir is shared across concurrently-dispatched 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 task 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 (see examples/small_molecule_binding/scripts/boltz.sh +# for the same fix, ported here after this exact race killed 13/16 pipelines +# in a production run). +# +# 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 task 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 + 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" boltz predict \ "${fasta_path}" \ --out_dir "${output_dir}" \ - --use_msa_server \ - --cache /anvil/projects/x-nairr240405/mason/boltz \ + ${_MSA_FLAG} \ + --cache "$_boltz_cache_dir" \ --output_format pdb \ --write_full_pae \ + --no_kernels \ + --devices 1 \ --override diff --git a/examples/protein_binding/scripts/s5_plddt_extract.sh b/examples/protein_binding/scripts/s5_plddt_extract.sh index 38cb576..e4999ac 100755 --- a/examples/protein_binding/scripts/s5_plddt_extract.sh +++ b/examples/protein_binding/scripts/s5_plddt_extract.sh @@ -2,16 +2,19 @@ set -e # Step 5: pLDDT extraction -# Args: $1=base_path $2=iter $3=out_name +# Args: $1=output_base_path $2=iter $3=out_name -base_path="$1" +output_base_path="$1" iter="$2" out_name="$3" -source /anvil/projects/x-nairr240405/mason/IMPRESS/.venv/bin/activate -#source /ocean/projects/dmr170002p/hooten/IMPRESS/.venv/bin/activate +# plddt_extract_pipeline.py lives one level above this scripts/ directory. +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -python3 "$base_path/plddt_extract_pipeline.py" \ - --path="$base_path" \ - --iter="$iter" \ - --out="$out_name" +# Re-activate the IMPRESS venv if running inside a subprocess (VIRTUAL_ENV is exported by sbatch). +[ -n "${VIRTUAL_ENV:-}" ] && source "${VIRTUAL_ENV}/bin/activate" + +python3 "${SCRIPT_DIR}/../plddt_extract_pipeline.py" \ + --path="${output_base_path}" \ + --iter="${iter}" \ + --out="${out_name}" diff --git a/examples/small_molecule_binding/CLAUDE.md b/examples/small_molecule_binding/CLAUDE.md index aaa9f96..f0af22b 100644 --- a/examples/small_molecule_binding/CLAUDE.md +++ b/examples/small_molecule_binding/CLAUDE.md @@ -7,6 +7,12 @@ 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 | +| 2026-09-09 | — | Fixed guided-RFD3 ligand atom-name mismatch that crashed 4/4 pipelines in a real production run (job `21916521`) on their first guided-backbone-feedback attempt: `_normalize_ligand_id()` only rewrote the Boltz co-folded PDB's ligand *residue* name, not its Boltz-assigned *atom* names, so `select_exposed`/`select_buried` (copied verbatim from the base spec, keyed by canonical `.params` atom names) never matched and RFD3's validator rejected every guided run. New `_infer_ligand_atom_mapping()`/`_normalize_ligand_atom_names()` establish atom correspondence via element+connectivity graph isomorphism (rdkit) with a Kabsch-RMSD tie-break; `_write_guided_rfd3_json()` now also verifies atom-name coverage before writing. Adds an `rdkit` runtime dependency | +| 2026-09-09 | — | Fixed a second guided-RFD3 crash found in a real production run (job `21928556`, 3/3 pipelines that reached guided feedback crashed on their very first attempt): `_write_guided_rfd3_json()` copied the base spec's `partial.length` field verbatim into the guided JSON, but RFD3's `DesignInputSpecification` validator rejects `length` outright whenever `partial.input`/`partial_t` (partial diffusion) are set (`ValidationError: ... Length argument must not be provided during partial diffusion`) — length is inferred from the input structure in that mode. `_write_guided_rfd3_json()` now drops `length` from the guided spec | ## Context @@ -24,7 +30,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 +38,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 +48,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 +58,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,16 +104,47 @@ 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 `_normalize_ligand_atom_names()` to rewrite that same PDB's ligand *atom* names to the canonical `.params` names (see "Guided-RFD3 ligand atom-name mapping" below) — Boltz assigns its own arbitrary atom names during co-folding, unrelated to the params file, so this is a separate fix from step 2. +4. Calls `_write_guided_rfd3_json()` to copy the base `ALR_binder_design.json`'s `ligand`/`select_exposed`/`select_buried` fields verbatim (dropping `length` — RFD3's validator rejects it during partial diffusion, since length is inferred from the input structure) into a new spec with `input` pointed at the normalized PDB and `partial_t` set, writing `{taskdir}/in/guided_binder_design.json` — after first verifying every `select_exposed`/`select_buried` atom name is actually present in the normalized PDB. +5. Passes that guided JSON (instead of the base one) as `rfd3.sh`'s `inputs=` argument. If any of steps 2–4 fails (no ligand found in the Boltz model, no full atom-name mapping found, or the coverage check fails), falls back to the base, unguided JSON rather than erroring. + +### Guided-RFD3 ligand atom-name mapping + +Boltz-2's co-folded ligand output uses its own arbitrary atom names (e.g. `C41`, `O24`, ...) that have nothing to do with the canonical names in the ligand's `.params` file (e.g. `C18`, `O3`, ...). Since `select_exposed`/`select_buried` are copied verbatim from the base spec and are keyed by those canonical names, a guided PDB with unrenamed atoms fails RFD3's own input validation (`ComponentValidationError: Number of atoms must be a multiple of the requested names`) — this was the confirmed root cause of a real production run (job `21916521`) crashing all 4 pipeline instances on their first guided-feedback attempt. + +`_infer_ligand_atom_mapping()` establishes the correspondence via element + heavy-atom-connectivity graph isomorphism (rdkit), ignoring bond order throughout (the `.params` format has none): +- **Reference graph**: heavy atoms + bonds parsed directly from the `.params` file's `ATOM`/`BOND` records (`_params_heavy_atom_graph()`) — exact, no perception needed. Reference *coordinates* come from the real, correctly-named structure the base spec's own `partial.input` already points at (`_resolve_reference_pdb_path()`) — no synthetic conformer is built. +- **Query graph**: the Boltz ligand's connectivity, perceived from 3D distances via `rdkit.Chem.rdDetermineBonds.DetermineConnectivity()` (Boltz's output carries no CONECT records for the ligand). +- **Isomorphism + tie-break**: `GetSubstructMatches()` enumerates every graph-valid atom correspondence; local topological symmetry (e.g. a sulfonate's three interchangeable terminal oxygens) can yield more than one. Each candidate is Kabsch-superposed (reusing `_kabsch_rmsd()`) against the reference coordinates, and the lowest-RMSD mapping wins — grounded in real geometry rather than an arbitrary tiebreak. When more than one isomorphism exists, the best-vs-next-best RMSD gap is logged. + +This is **not a generic guarantee for every future ligand**: it's only provably safe for a ligand whose "sides" (whatever `select_exposed`/`select_buried` partition into) aren't themselves graph-isomorphic to each other — true for `ALR` (a monocyclic benzene-sulfonate ring isn't isomorphic to a fused naphthalene-sulfonate ring), but not checked automatically for `IND`/`RED`/`IAI` or any future ligand. Run `scripts/check_ligand_atom_mapping.py` against a new ligand's `.params` + reference PDB before trusting guided feedback with it. + +`_write_guided_rfd3_json()` adds a final defensive check: before writing, it verifies every `select_exposed`/`select_buried` atom name is present in the (now atom-renamed) guided PDB, returning `False` (fail safe, fall back to unguided) rather than reproducing RFD3's rejection in a new form if not. + +`scripts/check_ligand_atom_mapping.py` validates this mapping logic against real fixtures already in the repo (the job-`21916521` crash-artifact `guided_scaffold.pdb` files under `logs/p{1..4}/*_rfd3/in/`, plus `p1_in/ALR.params` + `p1_in/input_pdbs/scaffold-with-ALR.pdb` as ground truth) — not synthetic test data. `scripts/validate_run.py`'s check 8 (`check_guided_ligand_atom_names`) regression-tests the same invariant against any completed run's output tree. Ensemble similarity utilities (all in `small_molecule_binding.py`): - `_ca_rmsd(path1, path2)` — Kabsch-aligned CA RMSD between two PDB files +- `_kabsch_rmsd(coords1, coords2)` — the underlying generic Kabsch-alignment RMSD, also reused by `_infer_ligand_atom_mapping()`'s isomorphism tie-break - `_seq_identity(fasta1, fasta2)` — fraction matching residues over shorter sequence - `_ensemble_selective_avg(current, prior, sim_fn, similar_if_low)` — returns `(overall_avg, selective_avg, has_data)` for scores of entries whose similarity is on the "similar" side of the mean pairwise similarity @@ -121,10 +158,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 +181,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 +197,34 @@ 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 + +`rdkit` is a runtime dependency (installed by `delta_env_setup.sh`'s Step 10, pinned `2024.9.6`) used both by the offline `derive_ligand_smiles.py` tool above and, per-cycle, by `rfd3()`'s guided-backbone-feedback atom-name mapping (see "Guided-RFD3 ligand atom-name mapping" above). `scripts/check_ligand_atom_mapping.py` and `scripts/validate_run.py` (check 8) validate that mapping against real fixtures/completed runs, respectively. 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..10c60d6 --- /dev/null +++ b/examples/small_molecule_binding/delta_env_setup.sh @@ -0,0 +1,301 @@ +#!/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. rdkit — ligand atom-name graph-isomorphism mapping for guided RFD3 ──── +# +# Used by rfd3()'s guided-backbone-feedback path (_infer_ligand_atom_mapping +# in small_molecule_binding.py) to reconcile Boltz-2's arbitrary ligand atom +# names against the canonical names in the ligand's .params file, via +# element+connectivity graph isomorphism (rdDetermineBonds.DetermineConnectivity +# + GetSubstructMatches) with a Kabsch-RMSD tie-break. Without this, RFD3's +# input validator rejects every guided run (ComponentValidationError) -- +# confirmed as the root cause of 4/4 pipeline crashes in a real production +# run (job 21916521). +# +# Pinned to 2024.9.6, the same version already used by the offline +# scripts/derive_ligand_smiles.py tool in this repo (rdkit has no +# dependency on numpy/gemmi's own pins, so it should not disturb Step 7's +# numpy<2.0/gemmi==0.6.5 resolution -- `pip check` after this step should +# stay clean; re-investigate only if it doesn't). +# +echo "" +echo "── Step 10: rdkit ──" +"${PIP}" install -q "rdkit==2024.9.6" + +# ── 11. Additional dependencies ─────────────────────────────────────────────── +echo "" +echo "── Step 11: pandas + biopandas ──" +"${PIP}" install -q pandas biopandas + +# ── 12. PyRosetta ───────────────────────────────────────────────────────────── +echo "" +echo "── Step 12: PyRosetta ──" +"${PIP}" install -q pyrosetta-installer +"${PY}" -c "import pyrosetta_installer; pyrosetta_installer.install_pyrosetta()" + +# ── 13. 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 13: 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}" + +# ── 14. Verify ──────────────────────────────────────────────────────────────── +echo "" +echo "── Step 14: 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 "rdkit" "${PY}" -c "import rdkit; print(rdkit.__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..f6179d4 --- /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=mh1314@scarletmail.rutgers.edu +#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 13 (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..d7c2490 100644 --- a/examples/small_molecule_binding/run_small_molecule_binding.py +++ b/examples/small_molecule_binding/run_small_molecule_binding.py @@ -1,32 +1,95 @@ 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 + mpnn_ensemble_size: int # cycle-0 MPNN sequence candidates per backbone + rfd3_partial_t: float # RFD3 partial-diffusion noise (A) for guided backbone feedback -# ── 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) + +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, + mpnn_ensemble_size = 10, + 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, + mpnn_ensemble_size = 2, + # 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() + +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: @@ -42,10 +105,33 @@ def _prior(ttype): if step == 'backbone': if not passed: + # Safety net: a guided (partial-diffusion) backbone that keeps + # failing QC -- whether for real structural reasons or an + # unforeseen RFD3 metrics-schema gap -- would otherwise loop on + # the same rfd3_input_pdb forever, since only a *successful* fold + # ever clears it. Fall back to unguided regeneration after a few + # consecutive guided-mode failures instead of deadlocking. + if pipeline.state.get('rfd3_input_pdb') is not None: + count = pipeline.state.get('backbone_guided_fail_count', 0) + 1 + pipeline.state['backbone_guided_fail_count'] = count + if count >= 3: + pipeline.state['rfd3_input_pdb'] = None + pipeline.state['backbone_guided_fail_count'] = 0 + pipeline.logger.pipeline_log( + "[adaptive/backbone] guided backbone QC failed 3x in a " + "row -- abandoning guided mode, reverting to unguided " + "(scratch) RFD3 generation" + ) 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['backbone_guided_fail_count'] = 0 + 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: @@ -73,40 +159,113 @@ def _prior(ttype): pipeline.state['seq_retry_count'] = count if count >= 3: pipeline.state['seq_retry_count'] = 0 + # A guided backbone that fails downstream of backbone-QC + # gets only this one shot -- otherwise rfd3_input_pdb stays + # pinned to the same seed indefinitely (only 'backbone' QC + # failures and 'fold' decisions used to clear it), causing + # RFD3 to regenerate near-duplicate doomed backbones for + # dozens of cycles in a row (confirmed via job 21945304: + # 17 consecutive p2 rfd3 generations produced byte-identical + # guided_scaffold.pdb output from one stuck seed). + pipeline.state['rfd3_input_pdb'] = None pipeline.next_step = STEP_RFD3 + pipeline.logger.pipeline_log( + "[adaptive/sequence] sequence-similarity gate failed " + "3x in a row on this backbone -- escalating to a new, " + "unguided backbone (STEP_RFD3) instead of another " + "resequencing retry" + ) else: pipeline.next_step = STEP_RETRY_SEQ elif step == 'packmin': total_score = metrics.get('total_score') if total_score is not None and total_score > 0: + # See the sequence-stage comment above: any STEP_RFD3 escalation + # must clear a stuck guided seed, not just backbone-QC failures. + pipeline.state['rfd3_input_pdb'] = None pipeline.next_step = STEP_RFD3 # badly packed — restart backbone + pipeline.logger.pipeline_log( + f"[adaptive/packmin] total_score={total_score} > 0 -- badly " + "packed, escalating to a new, unguided backbone (STEP_RFD3)" + ) else: pipeline.next_step = STEP_MPNN 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: + reason = "safety cap (5 attempts)" if count >= 5 else "non-improving metrics" + pipeline.state['fastrelax_fail_count'] = 0 + pipeline.state['fastrelax_prev_metrics'] = None + # See the sequence-stage comment above: any STEP_RFD3 escalation + # must clear a stuck guided seed, not just backbone-QC failures. + pipeline.state['rfd3_input_pdb'] = None pipeline.next_step = STEP_RFD3 + pipeline.logger.pipeline_log( + f"[adaptive/fastrelax] escalating to a new, unguided " + f"backbone (STEP_RFD3) ({reason}); attempt={count} metrics={metrics}" + ) 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: + reason = "safety cap (5 attempts)" if count >= 5 else "non-improving metrics" + pipeline.state['interface_fail_count'] = 0 + pipeline.state['interface_prev_metrics'] = None + # See the sequence-stage comment above: any STEP_RFD3 escalation + # must clear a stuck guided seed, not just backbone-QC failures. + pipeline.state['rfd3_input_pdb'] = None pipeline.next_step = STEP_RFD3 + pipeline.logger.pipeline_log( + f"[adaptive/interface] escalating to a new, unguided " + f"backbone (STEP_RFD3) ({reason}); attempt={count} metrics={metrics}" + ) else: pipeline.next_step = STEP_MPNN @@ -115,16 +274,32 @@ def _prior(ttype): if not passed: # Failed fold — don't use this model as a backbone guide pipeline.state['rfd3_input_pdb'] = None + pipeline.logger.pipeline_log( + "[adaptive/fold] fold failed -- next backbone will be unguided (scratch)" + ) else: if not prior: pipeline.state['rfd3_input_pdb'] = None + pipeline.logger.pipeline_log( + "[adaptive/fold] fold passed but no prior fold history yet -- " + "next backbone will be unguided (scratch)" + ) else: overall, selective, has_data = _ensemble_selective_avg( current[3], prior, _ca_rmsd, similar_if_low=True) if has_data and selective is not None and selective > overall: pipeline.state['rfd3_input_pdb'] = current[3] # guided backbone + pipeline.logger.pipeline_log( + f"[adaptive/fold] similar-cluster avg ({selective:.2f}) > " + f"overall avg ({overall:.2f}) -- next backbone guided from {current[3]}" + ) else: pipeline.state['rfd3_input_pdb'] = None # scratch + pipeline.logger.pipeline_log( + f"[adaptive/fold] guided-feedback condition not met " + f"(has_data={has_data}, selective={selective}, overall={overall}) " + "-- next backbone will be unguided (scratch)" + ) pipeline.next_step = STEP_RFD3 else: @@ -139,32 +314,53 @@ 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. Each + # pipeline reads its own p{i}_in/ directory rather than sharing one. + + 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": os.path.join(examples_dir, f"p{i}_in"), + "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, + "mpnn_ensemble_size": cfg.mpnn_ensemble_size, + "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/check_ligand_atom_mapping.py b/examples/small_molecule_binding/scripts/check_ligand_atom_mapping.py new file mode 100644 index 0000000..89ccebe --- /dev/null +++ b/examples/small_molecule_binding/scripts/check_ligand_atom_mapping.py @@ -0,0 +1,210 @@ +#!/usr/bin/env python3 +"""Standalone validation for small_molecule_binding.py's RFD3 guided-input +ligand atom-name mapping (_infer_ligand_atom_mapping / +_normalize_ligand_atom_names), run against real fixtures already in this +repo -- not synthetic test data. + +Background: job 21916521 (a real ~3h production HPC run) crashed all 4 +pipeline instances on their first use of RFD3 guided-backbone feedback, +because Boltz-2's co-folded ligand output uses its own arbitrary atom names +that don't match the canonical params-file names baked into the base RFD3 +spec's select_exposed/select_buried fields. The crash-artifact guided PDBs +from that run (logs/p{1..4}/*_rfd3/in/guided_scaffold.pdb) are still on +disk and are the primary fixtures used here. + +Usage: + python scripts/check_ligand_atom_mapping.py + python scripts/check_ligand_atom_mapping.py --base-path logs + +Exits 0 if every check passes, nonzero otherwise. +""" + +import argparse +import glob +import json +import os +import sys + +_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) + +from small_molecule_binding import ( # noqa: E402 + _find_ligand_hetatm_residue, + _infer_ligand_atom_mapping, + _ligand_resname_from_params, + _resolve_reference_pdb_path, +) + +# Boltz atom -> canonical name pairs that must land in the same +# select_exposed/select_buried bucket for ALR specifically (its two ring +# systems -- a monocyclic benzene-sulfonate and a fused naphthalene-sulfonate +# -- are not graph-isomorphic to each other, so the *only* real ambiguity is +# each sulfonate's 3 interchangeable terminal oxygens; see CLAUDE.md's +# "Ensemble-guided backbone feedback" section for why this is ALR-specific, +# not a general guarantee for future ligands). +_ALR_BURIED_SULFONATE_OXYGENS = {"O5", "O7", "O8"} # bonded to S1 +_ALR_EXPOSED_SULFONATE_OXYGENS = {"O6", "O9", "O10"} # bonded to S2 + + +def _load_boltz_atoms(pdb_path: str, ligand_chain_id: str = "B"): + import gemmi + st = gemmi.read_structure(pdb_path) + res = _find_ligand_hetatm_residue(st, ligand_chain_id) + if res is None: + return None + return {atom.name: (atom.element.name, (atom.pos.x, atom.pos.y, atom.pos.z)) for atom in res} + + +def check_real_crash_artifacts(examples_dir: str, base_path: str): + """Run _infer_ligand_atom_mapping against every real *_rfd3/in/guided_scaffold.pdb + left on disk from job 21916521 (or any other run). For each: assert a + full bijection is found, every mapped pair is element-consistent, the + ALR select_exposed/select_buried atom-name lists are fully covered, and + each sulfonate's 3 oxygens land together in the correct bucket (a real + correctness check beyond "isomorphism exists").""" + params_path = os.path.join(examples_dir, "p1_in", "ALR.params") + base_json_path = os.path.join(examples_dir, "p1_in", "ALR_binder_design.json") + if not (os.path.isfile(params_path) and os.path.isfile(base_json_path)): + return [f"SKIPPED: {params_path} or {base_json_path} not found"] + reference_pdb = _resolve_reference_pdb_path(base_json_path) + if not os.path.isfile(reference_pdb): + return [f"reference pdb {reference_pdb} not found"] + + with open(base_json_path) as fh: + base_partial = json.load(fh)["partial"] + ligand_key = base_partial["ligand"] + expected_names = set() + for field in ("select_exposed", "select_buried"): + expected_names.update(base_partial[field][ligand_key].split(",")) + + candidates = sorted(glob.glob(os.path.join(base_path, "p*", "*_rfd3", "in", "guided_scaffold.pdb"))) + if not candidates: + return [f"SKIPPED: no '*_rfd3/in/guided_scaffold.pdb' files found under {base_path}"] + + failures = [] + for pdb_path in candidates: + boltz_atoms = _load_boltz_atoms(pdb_path) + if boltz_atoms is None: + failures.append(f"{pdb_path}: no ligand HETATM residue found") + continue + + mapping = _infer_ligand_atom_mapping(boltz_atoms, params_path, reference_pdb) + if mapping is None: + failures.append(f"{pdb_path}: _infer_ligand_atom_mapping returned None (no mapping found)") + continue + + if len(mapping) != len(boltz_atoms): + failures.append(f"{pdb_path}: mapping covers {len(mapping)}/{len(boltz_atoms)} atoms, not a full bijection") + + mapped_names = set(mapping.values()) + if not expected_names <= mapped_names: + failures.append( + f"{pdb_path}: select_exposed/select_buried coverage FAILED -- " + f"missing {sorted(expected_names - mapped_names)}" + ) + + if ligand_key == "A:R": + buried_group = {b for b, c in mapping.items() if c in _ALR_BURIED_SULFONATE_OXYGENS} + exposed_group = {b for b, c in mapping.items() if c in _ALR_EXPOSED_SULFONATE_OXYGENS} + if len(buried_group) != 3 or len(exposed_group) != 3: + failures.append( + f"{pdb_path}: sulfonate oxygen grouping broken -- " + f"buried={buried_group} exposed={exposed_group} (expected 3 each)" + ) + + return failures + + +def check_negative_paths(examples_dir: str, base_path: str): + """Corrupt a real fixture (drop an atom; swap an element to one absent + from ALR) and confirm _infer_ligand_atom_mapping fails safe (returns + None) rather than raising.""" + params_path = os.path.join(examples_dir, "p1_in", "ALR.params") + base_json_path = os.path.join(examples_dir, "p1_in", "ALR_binder_design.json") + reference_pdb = _resolve_reference_pdb_path(base_json_path) + + fixture = None + for pdb_path in sorted(glob.glob(os.path.join(base_path, "p*", "*_rfd3", "in", "guided_scaffold.pdb"))): + boltz_atoms = _load_boltz_atoms(pdb_path) + if boltz_atoms: + fixture = boltz_atoms + break + if fixture is None: + return [f"SKIPPED: no usable '*_rfd3/in/guided_scaffold.pdb' fixture found under {base_path}"] + + failures = [] + + missing_atom = dict(fixture) + missing_atom.pop(next(iter(missing_atom))) + try: + result = _infer_ligand_atom_mapping(missing_atom, params_path, reference_pdb) + except Exception as e: + failures.append(f"missing-atom case raised {e!r} instead of returning None") + else: + if result is not None: + failures.append("missing-atom case returned a mapping instead of None") + + bad_element = dict(fixture) + name0 = next(iter(bad_element)) + _, xyz = bad_element[name0] + bad_element[name0] = ("Cl", xyz) # not present in ALR at all + try: + result2 = _infer_ligand_atom_mapping(bad_element, params_path, reference_pdb) + except Exception as e: + failures.append(f"bad-element case raised {e!r} instead of returning None") + else: + if result2 is not None: + failures.append("bad-element case returned a mapping instead of None") + + return failures + + +def main(argv=None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--base-path", default=os.path.join(_EXAMPLE_DIR, "logs"), + help="Directory containing p1/, p2/, ... pipeline output trees (default: examples_dir/logs)", + ) + args = parser.parse_args(argv) + + checks = [ + ("1. Real crash-artifact mapping correctness", + lambda: check_real_crash_artifacts(_EXAMPLE_DIR, args.base_path)), + ("2. Negative-path fail-safe behavior", + lambda: check_negative_paths(_EXAMPLE_DIR, args.base_path)), + ] + + print("=" * 72) + print("LIGAND ATOM-NAME MAPPING VALIDATION") + print("=" * 72) + + any_failed = False + for name, fn in checks: + try: + failures = fn() + except Exception as e: + 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/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..c651bcb --- /dev/null +++ b/examples/small_molecule_binding/scripts/validate_run.py @@ -0,0 +1,629 @@ +#!/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 13/14). +_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, partial.ligand/select_exposed/ + select_buried must match the base ALR_binder_design.json verbatim (only + input/partial_t may legitimately differ), and partial.length must be + absent (RFD3 rejects it during partial diffusion; it's only valid for + the base spec's from-scratch diffusion) -- 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})" + ) + + if "length" in partial: + failures.append( + f"{gj}: partial.length={partial['length']!r} must not be present -- " + f"RFD3 rejects 'length' during partial diffusion" + ) + + for field in ("ligand", "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 + + +# ── Check 8: guided ligand atom names cover select_exposed/select_buried ─── + +def check_guided_ligand_atom_names(base_path: str, pipeline_name: str, pipeline_inputs: str): + """Every */_rfd3/in/guided_scaffold.pdb's ligand HETATM atom names must be + a superset of its sibling guided_binder_design.json's select_exposed + + select_buried atom-name lists -- this is exactly the invariant RFD3's + input validator itself enforces (rejecting the guided spec with + 'ComponentValidationError: Number of atoms must be a multiple of the + requested names' when it doesn't hold). Regression test for the + guided-input ligand atom-name mapping fix; see + scripts/check_ligand_atom_mapping.py for the deeper unit-level check of + the mapping logic itself against real fixtures.""" + pipeline_dir = os.path.join(base_path, pipeline_name) + pairs = [] + for guided_pdb in sorted(glob.glob(os.path.join(pipeline_dir, "*_rfd3", "in", "guided_scaffold.pdb"))): + guided_json = os.path.join(os.path.dirname(guided_pdb), "guided_binder_design.json") + if os.path.isfile(guided_json): + pairs.append((guided_pdb, guided_json)) + if not pairs: + # No guided rfd3 runs have happened yet -- not a failure by itself. + return [] + + failures = [] + for guided_pdb, guided_json in pairs: + try: + with open(guided_json) as fh: + guided = json.load(fh) + except (OSError, json.JSONDecodeError) as e: + failures.append(f"{guided_json}: failed to parse as JSON: {e}") + continue + + partial = guided.get("partial", {}) + ligand_key = partial.get("ligand") + expected_names = set() + for field in ("select_exposed", "select_buried"): + names_csv = partial.get(field, {}).get(ligand_key, "") + expected_names.update(n for n in names_csv.split(",") if n) + if not expected_names: + failures.append(f"{guided_json}: no select_exposed/select_buried atom names found for ligand {ligand_key!r}") + continue + + present_names = set() + with open(guided_pdb, errors="replace") as fh: + for line in fh: + if line.startswith("HETATM") and line[17:20].strip() == ligand_key: + present_names.add(line[12:16].strip()) + + missing = expected_names - present_names + if missing: + failures.append( + f"{guided_pdb}: missing {sorted(missing)} from select_exposed/select_buried " + f"(this is exactly what RFD3's own validator would reject with " + f"ComponentValidationError)" + ) + 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)), + ("8. Guided ligand atom names cover select_exposed/select_buried", + lambda: check_guided_ligand_atom_names(base_path, pipeline_name, pipeline_inputs)), + ] + + 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..2f1d1cd 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,445 @@ 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 _find_ligand_hetatm_residue(st, ligand_chain_id: str = "B"): + """Locates the ligand HETATM residue in a gemmi Structure. Tries + 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. + Returns None if no HETATM residue exists at all. Shared by + _normalize_ligand_id and _normalize_ligand_atom_names so both agree on + exactly which residue is "the ligand".""" + for model in st: + for chain in model: + if chain.name != ligand_chain_id: + continue + for res in chain: + if res.het_flag == 'H': + return res + + for model in st: + for chain in model: + for res in chain: + if res.het_flag == 'H': + return res + + return None + + +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. 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. + + NOTE: this only fixes the residue name. Boltz also assigns its own, + unrelated atom names within that residue -- see _normalize_ligand_atom_names + for why those need fixing too before the guided spec is usable.""" + import gemmi + st = gemmi.read_structure(fold_pdb_path) + + target_res = _find_ligand_hetatm_residue(st, ligand_chain_id) + if target_res is None: + return False + + target_res.name = ligand_resname + st.write_pdb(out_pdb_path) + return True + + +# Fallback table keyed on Rosetta atom TYPE prefix, used only when a params +# ATOM name's leading alphabetic run doesn't parse to a valid element symbol. +# Mirrors scripts/derive_ligand_smiles.py's _ROSETTA_TYPE_ELEMENT_FALLBACK +# (duplicated, not imported -- see _params_heavy_atom_graph). Verified only +# against ALR's C/N/O/S/H atom set; extend if a future ligand needs 2-letter +# elements (Cl/Br/Zn, etc.). +_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", +} + + +def _infer_ligand_element(atom_name: str, rosetta_type: str) -> str: + """Infers an element symbol from a params ATOM record's name/type. Primary + rule: the atom NAME's leading alphabetic run is the element itself (e.g. + 'C18' -> 'C', 'N11' -> 'N'); falls back to _ROSETTA_TYPE_ELEMENT_FALLBACK + keyed on the Rosetta TYPE prefix. Mirrors + scripts/derive_ligand_smiles.py's _infer_element (duplicated, not + imported -- that script is a standalone offline tool, not part of this + per-cycle pipeline path; scripts/ isn't an importable package).""" + import re + from rdkit import Chem + periodic_table = Chem.GetPeriodicTable() + + match = re.match(r'[A-Za-z]+', atom_name) + if match: + candidate = match.group(0) + for length in (2, 1): + if len(candidate) >= length: + symbol = candidate[:length].capitalize() + if periodic_table.GetAtomicNumber(symbol) > 0: + return symbol + + 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 _params_heavy_atom_graph(params_path: str): + """Parses a Rosetta .params file's ATOM/BOND records into a heavy-atom-only + connectivity graph: ({atom_name: element}, [(atom1, atom2), ...]) where + both bond endpoints are heavy atoms. Hydrogens are dropped entirely -- + unlike scripts/derive_ligand_smiles.py (which needs them to resolve bond + order via DetermineBondOrders), _infer_ligand_atom_mapping only needs + connectivity for graph-isomorphism matching, so no bond-order solving or + placeholder hydrogen placement is needed here.""" + atom_types = {} + bonds = [] + with open(params_path) as fh: + for line in fh: + fields = line.split() + if not fields: + continue + record = fields[0] + if record == 'ATOM': + atom_types[fields[1]] = fields[2] + elif record in ('BOND', 'BOND_TYPE'): + bonds.append((fields[1], fields[2])) + + elements = {name: _infer_ligand_element(name, rtype) for name, rtype in atom_types.items()} + heavy_names = {name for name, el in elements.items() if el != 'H'} + heavy_elements = {name: elements[name] for name in heavy_names} + heavy_bonds = [(a, b) for a, b in bonds if a in heavy_names and b in heavy_names] + return heavy_elements, heavy_bonds + + +def _resolve_reference_pdb_path(base_json_path: str) -> str: + """Reads partial.input from the base RFD3 design spec and resolves it + relative to the spec's own directory (always pipeline_inputs). This is + the correctly-named, real-coordinate reference ligand structure already + shipped alongside every pipeline's inputs -- used as ground truth for + atom identity/connectivity/geometry in _infer_ligand_atom_mapping.""" + with open(base_json_path) as fh: + base = json.load(fh) + ref = base['partial']['input'] + if os.path.isabs(ref): + return ref + return os.path.join(os.path.dirname(base_json_path), ref) + + +def _iter_ligand_atom_names(pdb_path: str, resname: str): + """Atom names (fixed-column PDB parsing) of every HETATM record in + pdb_path whose resname matches exactly. Fixed-column parsing (not a + whitespace split) is required because names like 'A:R' contain a colon + -- see _ligand_resname_from_params's docstring.""" + names = [] + with open(pdb_path) as fh: + for line in fh: + if line.startswith('HETATM') and line[17:20].strip() == resname: + names.append(line[12:16].strip()) + return names + + +def _load_reference_ligand_coords(pdb_path: str, resname: str): + """Reads {atom_name: (x, y, z)} for the first HETATM residue named + exactly resname in pdb_path. Mirrors + scripts/derive_ligand_smiles.py's _load_reference_coords (duplicated, + not imported -- see _params_heavy_atom_graph); uses the same + fixed-column parsing for the same reason (resnames like 'A:R' contain a + colon a whitespace split would mangle).""" + coords = {} + target_resseq = None + with open(pdb_path) as fh: + for line in fh: + if not (line.startswith('HETATM') or line.startswith('ATOM ')): + continue + if line[17:20].strip() != resname: + if coords: + break # moved past the matching residue's contiguous block + continue + resseq = line[22:26].strip() + if target_resseq is None: + target_resseq = resseq + elif resseq != target_resseq: + break # a different residue instance with the same name + atom_name = line[12:16].strip() + coords[atom_name] = (float(line[30:38]), float(line[38:46]), float(line[46:54])) + return coords + + +def _infer_ligand_atom_mapping(boltz_atoms: dict, params_path: str, reference_pdb_path: str): + """Maps Boltz's arbitrarily-named ligand atom names onto the canonical + names read from params_path, via element+connectivity graph isomorphism + with a Kabsch-RMSD tie-break. boltz_atoms is + {boltz_atom_name: (element, (x, y, z))} for one ligand residue. + + Boltz co-folding assigns its own atom names to the ligand (unrelated to + the params file's canonical names), so select_exposed/select_buried + (copied verbatim from the base RFD3 spec, keyed by canonical names) never + match a Boltz-derived PDB's atom names without this step. Bond order is + ignored throughout (the .params file has none) -- only element identity + and heavy-atom connectivity establish correspondence: + - reference graph: heavy atoms + bonds parsed straight from + params_path's ATOM/BOND records (exact, no perception needed) + - reference coordinates: the real 3D structure at reference_pdb_path + (already correctly named -- see _resolve_reference_pdb_path) + - query graph: boltz_atoms' connectivity, perceived from 3D distances + via rdkit's DetermineConnectivity (Boltz's ligand output carries no + CONECT records) + Local topological symmetry (e.g. a sulfonate's three interchangeable + terminal oxygens) can produce more than one graph-valid isomorphism; both + structures carry real, roughly comparable 3D coordinates for the same + ligand pose, so each candidate mapping is Kabsch-superposed against the + reference and the lowest-RMSD one wins -- deterministic, and grounded in + actual geometry rather than an arbitrary tiebreak. When more than one + isomorphism exists, the best-vs-next-best RMSD gap is logged so a + suspiciously close tie (a symmetry case this heuristic can't actually + distinguish) is visible after the fact rather than silently accepted. + + Returns {boltz_name: canonical_name}, or None (never raises) if: heavy + atom counts or element multisets differ, the reference PDB is missing + coordinates for a params heavy atom, Boltz connectivity perception fails + or yields a disconnected graph, or no isomorphism exists at all -- + callers must treat None exactly like _normalize_ligand_id returning + False (fall back to unguided diffusion).""" + from rdkit import Chem + from rdkit.Chem import rdDetermineBonds + from rdkit.Geometry import Point3D + + ref_elements, ref_bonds = _params_heavy_atom_graph(params_path) + ref_resname = _ligand_resname_from_params(params_path) + ref_coords = _load_reference_ligand_coords(reference_pdb_path, ref_resname) + ref_names = [name for name in ref_elements if name in ref_coords] + if len(ref_names) != len(ref_elements): + return None # reference PDB is missing coordinates for a params heavy atom + + if len(boltz_atoms) != len(ref_names): + return None + if sorted(element for element, _ in boltz_atoms.values()) != sorted(ref_elements[n] for n in ref_names): + return None + + ref_mol = Chem.RWMol() + ref_idx = {} + for name in ref_names: + ref_idx[name] = ref_mol.AddAtom(Chem.Atom(ref_elements[name])) + for a, b in ref_bonds: + if a in ref_idx and b in ref_idx: + i, j = ref_idx[a], ref_idx[b] + if ref_mol.GetBondBetweenAtoms(i, j) is None: + ref_mol.AddBond(i, j, Chem.BondType.SINGLE) + ref_conf = Chem.Conformer(ref_mol.GetNumAtoms()) + for name, idx in ref_idx.items(): + ref_conf.SetAtomPosition(idx, Point3D(*ref_coords[name])) + ref_mol.AddConformer(ref_conf, assignId=True) + Chem.SanitizeMol(ref_mol, sanitizeOps=Chem.SanitizeFlags.SANITIZE_NONE) + + boltz_names = list(boltz_atoms) + boltz_mol = Chem.RWMol() + boltz_idx = {} + for name in boltz_names: + element, _ = boltz_atoms[name] + boltz_idx[name] = boltz_mol.AddAtom(Chem.Atom(element)) + boltz_conf = Chem.Conformer(boltz_mol.GetNumAtoms()) + for name, idx in boltz_idx.items(): + _, xyz = boltz_atoms[name] + boltz_conf.SetAtomPosition(idx, Point3D(*xyz)) + boltz_mol.AddConformer(boltz_conf, assignId=True) + Chem.SanitizeMol(boltz_mol, sanitizeOps=Chem.SanitizeFlags.SANITIZE_NONE) + + try: + rdDetermineBonds.DetermineConnectivity(boltz_mol) + except Exception: + return None + if len(Chem.GetMolFrags(boltz_mol)) != 1: + return None + + matches = boltz_mol.GetSubstructMatches(ref_mol, uniquify=False, useChirality=False, maxMatches=10000) + if not matches: + return None + + ref_coord_list = [ref_coords[name] for name in ref_names] # order == ref_mol atom order + best_mapping, best_rmsd, second_best_rmsd = None, None, None + for match in matches: + # match[i] is the boltz_mol atom index matched to ref_mol atom i (ref_names[i]). + query_coord_list = [boltz_atoms[boltz_names[qi]][1] for qi in match] + rmsd = _kabsch_rmsd(ref_coord_list, query_coord_list) + if best_rmsd is None or rmsd < best_rmsd: + second_best_rmsd = best_rmsd + best_rmsd = rmsd + best_mapping = {boltz_names[qi]: ref_names[i] for i, qi in enumerate(match)} + elif second_best_rmsd is None or rmsd < second_best_rmsd: + second_best_rmsd = rmsd + + if len(matches) > 1: + gap = (second_best_rmsd - best_rmsd) if second_best_rmsd is not None else float('inf') + print( + f"[rfd3 guided] ligand atom-name mapping: {len(matches)} candidate " + f"isomorphisms, best RMSD={best_rmsd:.4f} vs next-best={second_best_rmsd:.4f} " + f"(gap={gap:.4f}) -- a small gap means the tie-break may not be decisive" + ) + + return best_mapping + + +def _normalize_ligand_atom_names(pdb_path: str, params_path: str, base_json_path: str, + out_pdb_path: str, ligand_chain_id: str = "B") -> bool: + """Rewrites a resname-normalized guided PDB's ligand HETATM *atom* names + (not just its residue name -- see _normalize_ligand_id) to match the + canonical names read from params_path, via _infer_ligand_atom_mapping. + Required because select_exposed/select_buried in the guided RFD3 spec are + copied verbatim from the base spec and are keyed by those canonical + names, but Boltz assigns its own arbitrary atom names during co-folding + -- without this, RFD3's input validator rejects every guided run. + + Writes out_pdb_path (may be the same path as pdb_path) and returns True + on success. Returns False (writing nothing) if no ligand residue is + found, or no full atom-name mapping could be established, so the caller + falls back to unguided diffusion instead of producing a guided spec RFD3 + will reject.""" + import gemmi + st = gemmi.read_structure(pdb_path) + + target_res = _find_ligand_hetatm_residue(st, ligand_chain_id) + if target_res is None: + return False + + boltz_atoms = { + atom.name: (atom.element.name, (atom.pos.x, atom.pos.y, atom.pos.z)) + for atom in target_res + } + reference_pdb_path = _resolve_reference_pdb_path(base_json_path) + mapping = _infer_ligand_atom_mapping(boltz_atoms, params_path, reference_pdb_path) + if mapping is None: + return False + + for atom in target_res: + atom.name = mapping[atom.name] + 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) -> bool: + """Loads the base per-pipeline RFD3 InputSpecification JSON, copies its + ligand/select_exposed/select_buried fields verbatim, drops 'length', + 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 (besides the dropped 'length'). + + 'length' is dropped because RFD3's DesignInputSpecification validator + rejects it outright when partial.input/partial_t (partial diffusion) are + set -- length is inferred from the input structure in that mode. The + base file's 'length' is only valid for from-scratch (non-partial) + diffusion. + + Before writing, verifies every atom name referenced by select_exposed/ + select_buried is actually present in guided_pdb_path's ligand residue -- + fails safe (returns False, writes nothing) on a stale/mismatched base + spec or an atom-mapping bug, rather than reproducing RFD3's + ComponentValidationError in a new form. Returns True on success.""" + with open(base_json_path) as fh: + base = json.load(fh) + partial = dict(base.get('partial', {})) + partial.pop('length', None) + + ligand_key = partial.get('ligand') + expected_names = set() + for field in ('select_exposed', 'select_buried'): + names_csv = partial.get(field, {}).get(ligand_key, '') + expected_names.update(name for name in names_csv.split(',') if name) + present_names = set(_iter_ligand_atom_names(guided_pdb_path, ligand_key)) + if not expected_names <= present_names: + return False + + 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) + return True + + +def _prepare_guided_rfd3_inputs(base_json_path: str, fold_pdb_path: str, ligand_resname: str, + params_path: str, partial_t: float, taskdir: str): + """Orchestrates _normalize_ligand_id + _normalize_ligand_atom_names + + _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 any step failed (no ligand found in fold_pdb_path, no full + atom-name mapping could be established, or the select_exposed/ + select_buried coverage check failed) -- 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 + if not _normalize_ligand_atom_names(guided_pdb, params_path, base_json_path, guided_pdb): + return None + if not _write_guided_rfd3_json(base_json_path, guided_pdb, partial_t, guided_json): + return None + return guided_json + + class SmallMoleculeBindingPipeline(ImpressBasePipeline): def __init__(self, name, flow, configs=None, **kwargs): if configs is None: @@ -140,18 +580,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 +618,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 +636,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 +660,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 +669,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" - - input_pdb = self.state.get('rfd3_input_pdb') - scaffold_arg = f"scaffoldguided.target_pdb={input_pdb}" if input_pdb else "" + base_inputs = f"{self.pipeline_inputs}/ALR_binder_design.json" + output_dir = f"{taskdir}/out" + + fold_pdb = self.state.get('rfd3_input_pdb') + inputs = base_inputs + if fold_pdb: + params_path = f"{self.pipeline_inputs}/{self.ligand_params}" + ligand_resname = _ligand_resname_from_params(params_path) + guided_json = _prepare_guided_rfd3_inputs( + base_json_path=base_inputs, + fold_pdb_path=fold_pdb, + ligand_resname=ligand_resname, + params_path=params_path, + 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(): @@ -232,15 +709,29 @@ async def analysis_backbone(): for jf in json_files: with open(f"{out_dir}/{jf}") as fh: data = json.load(fh) - m = data.get('metrics', {}) - clashes = m.get('n_clashing.ligand_clashes', float('inf')) - dev = m.get('max_ca_deviation', float('inf')) - ss = m.get('helix_fraction', 0) + m.get('sheet_fraction', 0) + m = data.get('metrics', {}) + # RFD3's partial-diffusion (guided-backbone-feedback) mode never + # computes ligand-clash or secondary-structure metrics -- the + # 'n_clashing.ligand_clashes' key is absent entirely and + # helix_fraction/sheet_fraction come back as literal JSON NaN + # (so ss = NaN + NaN, and NaN > threshold is always False in + # Python) -- only max_ca_deviation and the interresidue clash + # counts are populated there. Unguided mode has all of these. + guided = 'n_clashing.ligand_clashes' not in m + if guided: + clashes = ( + m.get('n_clashing.interresidue_clashes_w_sidechain', float('inf')) + + m.get('n_clashing.interresidue_clashes_w_backbone', float('inf')) + ) + else: + clashes = m.get('n_clashing.ligand_clashes', float('inf')) + dev = m.get('max_ca_deviation', float('inf')) + ss = m.get('helix_fraction', 0) + m.get('sheet_fraction', 0) if best is None or clashes < best['clashes'] or ( clashes == best['clashes'] and dev < best['dev'] ): - best = {'file': jf, 'clashes': clashes, 'dev': dev, 'ss': ss} + best = {'file': jf, 'clashes': clashes, 'dev': dev, 'ss': ss, 'guided': guided} if best is None: self.state.update({ @@ -257,13 +748,16 @@ async def analysis_backbone(): passed = ( best['clashes'] == 0 and best['dev'] < self.backbone_max_ca_deviation - and best['ss'] > self.backbone_min_ss_fraction + # SS fraction isn't computable in guided mode (see above) -- + # skip that check there rather than fail unconditionally. + and (best['guided'] or best['ss'] > self.backbone_min_ss_fraction) ) self.state.update({ 'last_analysis_step': 'backbone', 'last_analysis_metrics': { 'pass': passed, 'best_model': best['file'], + 'guided': best['guided'], 'ligand_clashes': best['clashes'], 'max_ca_deviation': best['dev'], 'ss_fraction': best['ss'], @@ -273,7 +767,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 +794,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 +820,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 +841,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 +897,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 +914,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 +942,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 +955,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 +998,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 +1007,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 +1049,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 +1153,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 +1161,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 +1195,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 +1231,9 @@ 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.state.setdefault('backbone_guided_fail_count', 0) self.logger.pipeline_log("SmallMoleculeBindingPipeline starting (state machine)") while self.next_step != STEP_DONE: @@ -678,9 +1272,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() diff --git a/examples/small_molecule_binding/validation-plan-postfix-atomname.md b/examples/small_molecule_binding/validation-plan-postfix-atomname.md new file mode 100644 index 0000000..d4c9a5b --- /dev/null +++ b/examples/small_molecule_binding/validation-plan-postfix-atomname.md @@ -0,0 +1,96 @@ +# Validation plan: RFD3 guided-input atom-name fix + +Reference this after the next full HPC production run, once the atom-name mapping fix +(graph isomorphism + Kabsch tie-break in `_infer_ligand_atom_mapping` / +`_normalize_ligand_atom_names`, `small_molecule_binding.py`) has been deployed. Goal: confirm +the fix actually resolved job 21916521's crash pattern, and specifically probe the residual +risks flagged in planning that couldn't be closed by code review alone. + +## 1. Did the crash pattern actually go away? + +- `grep -c "ComponentValidationError" impress_.out` — expect **0** occurrences (job + 21916521 had one per pipeline, 4/4, each fatal). +- `grep -c "Pipeline FAILED" impress_.out` — any hits need individual triage; none + should trace back to a guided-RFD3 atom-name mismatch anymore. If any pipeline still dies + on its first guided-feedback attempt, the fix did not work and needs re-examination before + trusting anything else in this file. +- Compare `rfd3` attempt counts and final `ensemble=N` sizes per pipeline against job + 21916521's baseline (p1: 17 rfd3 / 69 ensemble; p2: 14 / 61; p3: 9 / 42; p4: 20 / 96, all + crashed short of `max_tasks=300`). Pipelines should now run substantially longer and/or + reach `max_tasks` — if they still terminate early, check whether it's a *new* failure mode + (see §4) rather than the same one recurring. + +## 2. Did the guided inputs actually validate correctly this time? + +- Run `python scripts/validate_run.py logs p1` (and p2–p4) against the new run's output — + the new check 8 (`check_guided_ligand_atom_names`) should pass for every + `*_rfd3/in/guided_scaffold.pdb` found. If the new run reused this same validate_run.py + against the *old* job 21916521 logs first, confirm it correctly flagged those as broken + (proves the check itself discriminates, not just rubber-stamps). +- Spot-check at least one fresh guided pair directly: confirm every name in + `guided_binder_design.json`'s `select_exposed`/`select_buried` is present among the ligand + HETATM atom names in the sibling `guided_scaffold.pdb`. + +## 3. Residual risk: `DetermineConnectivity`'s distance tolerance on real Boltz geometry + +This was flagged as unvalidated against actual Boltz-predicted (not crystallographic) ligand +geometry at plan time. + +- For each fresh `guided_scaffold.pdb` produced this run, check the perceived bond/degree + distribution against `ALR.params`'s (or whichever ligand's) known valences (e.g. each S + should have exactly 4 heavy neighbors). A mismatch here — even if the run didn't crash — + signals `covFactor=1.3` may be silently producing a technically-valid-but-wrong isomorphism + that happened not to trip the coverage check. +- If any pipeline used a different ligand than `ALR` this run (`IND`, `RED`, or `IAI` via + `sm_binder_design.json`), re-run this check per ligand — the tolerance was only validated + against `ALR`'s real crash artifacts during implementation. + +## 4. Residual risk: near-tie isomorphism candidates (symmetry ambiguity) + +Recommended addition during implementation: log the RMSD gap between the winning isomorphism +candidate and the next-best one every time `_infer_ligand_atom_mapping` succeeds. If that +logging was added: + +- `grep` the new run's logs for these RMSD-gap lines. A small gap (near-tie) on any accepted + mapping is a signal the tie-break may have picked arbitrarily between two graph-valid but + possibly semantically different atom assignments — flag any such case for manual review + even if the run didn't crash, since a wrong-but-plausible mapping degrades guidance quality + silently rather than failing loudly. +- If this logging was *not* added during implementation, add it now before trusting any + further runs that touch a ligand other than `ALR` — `ALR` was proven benign (its two ring + systems aren't isomorphic to each other, and its only symmetric atoms — each sulfonate's 3 + terminal oxygens — always land in the same exposed/buried bucket), but that guarantee does + **not** extend to `IND`/`RED`/`IAI` or any future ligand without checking their own + topology. + +## 5. Residual risk: `rdkit` dependency / env pin compatibility + +- Confirm `pip check` is clean in the run's actual venv (no regressions from adding `rdkit` + alongside the existing `boltz`/`gemmi==0.6.5`/`numpy` pins). +- Confirm `boltz` and `gemmi`-dependent steps (co-folding, `mpnn()`'s CIF.GZ→PDB conversion) + still function normally elsewhere in the same run — a silent pin downgrade caused by + `rdkit`'s install could show up as an unrelated failure downstream, not necessarily at the + `rdkit` import site itself. + +## 6. Residual risk: `gemmi` atom-name rewrite correctness + +`_normalize_ligand_id`'s residue-name rewrite via `gemmi` was already proven correct in +production; atom-name rewriting via the same API was new and only spot-checked (write, +re-read, confirm column alignment) during implementation, not exercised against a real HPC +run until now. + +- Re-read a fresh `guided_scaffold.pdb`'s ligand HETATM block and confirm fixed-column PDB + parsing (as `_load_reference_coords`/`_iter_hetatm_resnames`-style code depends on) still + finds the residue and every atom name correctly — no column misalignment, no truncated + names, especially for the 1-character element names (`S1`, `S2`, `O3`, etc.) that weren't + present in Boltz's own longer names (`S43`, `O24`). + +## 7. New-ligand onboarding checklist (for whenever this comes up next) + +Not specific to this run, but worth attaching here since it's the same unresolved gap: before +trusting guided feedback for a ligand other than `ALR`, manually verify its two "sides" +(whatever `select_exposed`/`select_buried` partition into) aren't graph-isomorphic to each +other — if they are, the isomorphism step could in principle map the wrong side onto the +wrong bucket without tripping any automated check. `scripts/check_ligand_atom_mapping.py` +(added with this fix) is the tool to run manually against the new ligand's `.params` + +reference PDB before it's ever used in a live guided run. diff --git a/src/impress/__init__.py b/src/impress/__init__.py index 89dd8d3..6a37de7 100644 --- a/src/impress/__init__.py +++ b/src/impress/__init__.py @@ -1,10 +1,10 @@ -from __future__ import annotations - +from impress.gpu import find_gpus from impress.impress_manager import ImpressManager from impress.pipelines.impress_pipeline import ImpressBasePipeline from impress.pipelines.setup import PipelineSetup __all__ = [ + "find_gpus", "ImpressManager", "ImpressBasePipeline", "PipelineSetup", diff --git a/src/impress/gpu.py b/src/impress/gpu.py new file mode 100644 index 0000000..42583d3 --- /dev/null +++ b/src/impress/gpu.py @@ -0,0 +1,32 @@ +import os +import subprocess + + +def find_gpus() -> list[int]: + """Return GPU IDs available to this process. + + Checks CUDA_VISIBLE_DEVICES first, then nvidia-smi. + Falls back to an empty list when neither yields results. + """ + val = os.environ.get("CUDA_VISIBLE_DEVICES", "") + ids = [int(g) for g in val.split(",") if g.strip().isdigit()] + if ids: + return ids + + try: + out = subprocess.run( + ["nvidia-smi", "--query-gpu=index", "--format=csv,noheader"], + capture_output=True, + text=True, + timeout=5, + ) + if out.returncode == 0: + return [ + int(ln.strip()) + for ln in out.stdout.splitlines() + if ln.strip().isdigit() + ] + except Exception: + pass + + return [] diff --git a/src/impress/impress_manager.py b/src/impress/impress_manager.py index 5fac46e..b9446f7 100644 --- a/src/impress/impress_manager.py +++ b/src/impress/impress_manager.py @@ -1,6 +1,8 @@ import asyncio -from collections.abc import Awaitable -from typing import Any, Callable, Optional, Union +import os +import tempfile +from collections.abc import Awaitable, Callable +from typing import Any, Optional, Union from radical.asyncflow import WorkflowEngine @@ -43,6 +45,7 @@ def __init__( self._telemetry_config: dict[str, Any] = telemetry_config or {} self._telemetry_subscribers: list[Callable] = telemetry_subscribers or [] self.telemetry: Any = None + self.flow: Optional[WorkflowEngine] = None def _normalize_pipeline_setup( self, setup: Union[dict[str, Any], PipelineSetup] @@ -77,6 +80,10 @@ def submit_new_pipelines( ValueError: If pipeline type is not a subclass of ImpressBasePipeline """ + if self.flow is None: + raise RuntimeError( + "ImpressManager.start() must be called before submit_new_pipelines()" + ) for setup_input in pipeline_setups: # Normalize to PipelineSetup object setup = self._normalize_pipeline_setup(setup_input) @@ -136,113 +143,139 @@ async def start( """ self.logger.separator("IMPRESS MANAGER STARTING") - self.flow: WorkflowEngine = await WorkflowEngine.create( - backend=self.execution_backend + # Write asyncflow session dirs to /tmp (node-local, no quota) instead of + # cwd on scratch, which exhausts inodes over many runs. + _session_base = os.environ.get("IMPRESS_SESSION_DIR", tempfile.gettempdir()) + self.flow = await WorkflowEngine.create( + backend=self.execution_backend, + work_dir=_session_base, ) - if self._telemetry_config: - self.telemetry = await self.flow.start_telemetry(**self._telemetry_config) - for fn in self._telemetry_subscribers: - self.telemetry.subscribe(fn) - - self.logger.manager_starting(len(pipeline_setups)) - - self.submit_new_pipelines(pipeline_setups) - - while True: - any_activity: bool = False - completed_pipelines: list[ImpressBasePipeline] = [] - - for pipeline, pipeline_future in list(self.pipeline_tasks.items()): - # Check if pipeline needs adaptive step and isn't already running one - if ( - getattr(pipeline, "invoke_adaptive_step", False) - and pipeline not in self.adaptive_tasks - ): - adaptive_task: asyncio.Task = asyncio.create_task( - self._run_adaptive_fn(pipeline) + try: + if self._telemetry_config: + self.telemetry = await self.flow.start_telemetry( + **self._telemetry_config + ) + for fn in self._telemetry_subscribers: + self.telemetry.subscribe(fn) + + self.logger.manager_starting(len(pipeline_setups)) + + self.submit_new_pipelines(pipeline_setups) + + while True: + any_activity: bool = False + completed_pipelines: list[tuple] = [] + + for pipeline, pipeline_future in list(self.pipeline_tasks.items()): + # Check if pipeline needs adaptive step and isn't running one yet + if ( + getattr(pipeline, "invoke_adaptive_step", False) + and pipeline not in self.adaptive_tasks + ): + adaptive_task: asyncio.Task = asyncio.create_task( + self._run_adaptive_fn(pipeline) + ) + self.adaptive_tasks[pipeline] = adaptive_task + any_activity = True + + # Check if pipeline has new config ready + config: Optional[dict[str, Any]] = ( + pipeline.get_child_pipeline_request() ) - self.adaptive_tasks[pipeline] = adaptive_task - any_activity = True - # Check if pipeline has new config ready - config: Optional[dict[str, Any]] = pipeline.get_child_pipeline_request() - - if config: - self.logger.child_pipeline_submitted(config["name"], pipeline.name) - # Convert dict to PipelineSetup for consistency - child_setup = PipelineSetup.from_dict(config) - self.new_pipeline_buffer.append(child_setup) - any_activity = True + if config: + self.logger.child_pipeline_submitted( + config["name"], pipeline.name + ) + # Convert dict to PipelineSetup for consistency + child_setup = PipelineSetup.from_dict(config) + self.new_pipeline_buffer.append(child_setup) + any_activity = True + + # Check if parent should be killed + if getattr(pipeline, "kill_parent", False): + self.logger.pipeline_killed(pipeline.name) + pipeline_future.cancel() + completed_pipelines.append((pipeline, pipeline_future)) + continue - # Check if parent should be killed - if getattr(pipeline, "kill_parent", False): - self.logger.pipeline_killed(pipeline.name) - pipeline_future.cancel() - completed_pipelines.append(pipeline) - continue - - # Check if pipeline is done - but only mark as completed - # if adaptive task is also done - if pipeline_future.done(): - # If there's an adaptive task running, don't mark as completed yet + # Check if pipeline is done - but only mark as completed + # if adaptive task is also done + if pipeline_future.done(): + # Adaptive task still running — wait before marking completed + if pipeline in self.adaptive_tasks: + adaptive_task = self.adaptive_tasks[pipeline] + if not adaptive_task.done(): + continue + + completed_pipelines.append((pipeline, pipeline_future)) + + # Clean up completed pipelines - but only if their + # adaptive tasks are also done + actually_completed: list[ImpressBasePipeline] = [] + for pipeline, future in completed_pipelines: + # Double-check: only clean up if adaptive task is + # done or doesn't exist if pipeline in self.adaptive_tasks: adaptive_task = self.adaptive_tasks[pipeline] if not adaptive_task.done(): continue + self.adaptive_tasks.pop(pipeline) + + self.pipeline_tasks.pop(pipeline, None) + exc = None + if future.done() and not future.cancelled(): + try: + exc = future.exception() + except Exception: + pass + if exc is not None: + self.logger.pipeline_failed(pipeline.name, exc) + else: + self.logger.pipeline_completed(pipeline.name) + actually_completed.append(pipeline) + + completed_pipelines = actually_completed + + # Clean up completed adaptive tasks + completed_adaptive: list[ImpressBasePipeline] = [] + for pipeline, adaptive_task in list(self.adaptive_tasks.items()): + if adaptive_task.done(): + completed_adaptive.append(pipeline) + + for pipeline in completed_adaptive: + self.adaptive_tasks.pop(pipeline, None) + + # Submit new pipelines; capture count before clearing so + # activity_summary reports the real number submitted. + if self.new_pipeline_buffer: + buffered_count = len(self.new_pipeline_buffer) + self.submit_new_pipelines(self.new_pipeline_buffer) + self.new_pipeline_buffer.clear() + any_activity = True + else: + buffered_count = 0 + + # Log activity summary periodically + if any_activity: + self.logger.activity_summary( + len(self.pipeline_tasks), + len(self.adaptive_tasks), + buffered_count, + ) - completed_pipelines.append(pipeline) - - # Clean up completed pipelines - but only if their - # adaptive tasks are also done - actually_completed: list[ImpressBasePipeline] = [] - for pipeline in completed_pipelines: - # Double-check: only clean up if adaptive task is - # done or doesn't exist - if pipeline in self.adaptive_tasks: - adaptive_task = self.adaptive_tasks[pipeline] - if not adaptive_task.done(): - continue - self.adaptive_tasks.pop(pipeline) - - self.pipeline_tasks.pop(pipeline, None) - self.logger.pipeline_completed(pipeline.name) - actually_completed.append(pipeline) - - completed_pipelines = actually_completed - - # Clean up completed adaptive tasks - completed_adaptive: list[ImpressBasePipeline] = [] - for pipeline, adaptive_task in list(self.adaptive_tasks.items()): - if adaptive_task.done(): - completed_adaptive.append(pipeline) - - for pipeline in completed_adaptive: - self.adaptive_tasks.pop(pipeline, None) - - # Submit new pipelines - if self.new_pipeline_buffer: - self.submit_new_pipelines(self.new_pipeline_buffer) - self.new_pipeline_buffer.clear() - any_activity = True - - # Log activity summary periodically - if any_activity: - self.logger.activity_summary( - len(self.pipeline_tasks), - len(self.adaptive_tasks), - len(self.new_pipeline_buffer), - ) + # Exit condition + if ( + not self.pipeline_tasks + and not self.new_pipeline_buffer + and not self.adaptive_tasks + ): + self.logger.manager_exiting() + self.logger.separator("IMPRESS MANAGER FINISHED") + break - # Exit condition - if ( - not self.pipeline_tasks - and not self.new_pipeline_buffer - and not self.adaptive_tasks - ): - self.logger.manager_exiting() - self.logger.separator("IMPRESS MANAGER FINISHED") - break - - if not any_activity: - await asyncio.sleep(0.5) + if not any_activity: + await asyncio.sleep(0.5) + finally: + await self.flow.shutdown() diff --git a/src/impress/pipelines/impress_pipeline.py b/src/impress/pipelines/impress_pipeline.py index f856fab..5279384 100644 --- a/src/impress/pipelines/impress_pipeline.py +++ b/src/impress/pipelines/impress_pipeline.py @@ -95,8 +95,7 @@ async def get_scores_map(self): """Optional: Return scores mapping""" return {} - @abstractmethod - async def finalize(self): + async def finalize(self): # noqa: B027 """Optional: Cleanup or finalization logic""" pass diff --git a/src/impress/pipelines/setup.py b/src/impress/pipelines/setup.py index e0b63e0..4e10036 100644 --- a/src/impress/pipelines/setup.py +++ b/src/impress/pipelines/setup.py @@ -1,5 +1,5 @@ -from collections.abc import Awaitable -from typing import Annotated, Any, Callable, Optional +from collections.abc import Awaitable, Callable +from typing import Annotated, Any, Optional from pydantic import BaseModel, Field, field_validator diff --git a/src/impress/utils/logger.py b/src/impress/utils/logger.py index 642d699..2919c8f 100644 --- a/src/impress/utils/logger.py +++ b/src/impress/utils/logger.py @@ -34,10 +34,25 @@ class LogLevel(Enum): class ImpressLogger: - def __init__(self, name="ImpressManager", use_colors=True, output_stream=None): + _LEVEL_ORDER = [ + LogLevel.DEBUG, + LogLevel.INFO, + LogLevel.WARNING, + LogLevel.ERROR, + LogLevel.CRITICAL, + ] + + def __init__( + self, + name="ImpressManager", + use_colors=True, + output_stream=None, + min_level: LogLevel = LogLevel.DEBUG, + ): self.name = name self.use_colors = use_colors self.output_stream = output_stream or sys.stdout + self.min_level = min_level self.level_colors = { LogLevel.DEBUG: Colors.BRIGHT_BLACK, @@ -91,40 +106,52 @@ def _format_message(self, level, component, message, pipeline_name=None): f"{timestamp} {colored_level} {colored_component}{pipeline_part} {message}" ) - def _write_log(self, message, to_stderr=False): - stream = sys.stderr if to_stderr else self.output_stream - stream.write(message + "\n") - stream.flush() + def _is_enabled(self, level: LogLevel) -> bool: + return self._LEVEL_ORDER.index(level) >= self._LEVEL_ORDER.index(self.min_level) + + def _write_log(self, message): + self.output_stream.write(message + "\n") + self.output_stream.flush() def debug(self, message, component="manager", pipeline_name=None): + if not self._is_enabled(LogLevel.DEBUG): + return formatted = self._format_message( LogLevel.DEBUG, component, message, pipeline_name ) self._write_log(formatted) def info(self, message, component="manager", pipeline_name=None): + if not self._is_enabled(LogLevel.INFO): + return formatted = self._format_message( LogLevel.INFO, component, message, pipeline_name ) self._write_log(formatted) def warning(self, message, component="manager", pipeline_name=None): + if not self._is_enabled(LogLevel.WARNING): + return formatted = self._format_message( LogLevel.WARNING, component, message, pipeline_name ) self._write_log(formatted) def error(self, message, component="manager", pipeline_name=None): + if not self._is_enabled(LogLevel.ERROR): + return formatted = self._format_message( LogLevel.ERROR, component, message, pipeline_name ) - self._write_log(formatted, to_stderr=True) + self._write_log(formatted) def critical(self, message, component="manager", pipeline_name=None): + if not self._is_enabled(LogLevel.CRITICAL): + return formatted = self._format_message( LogLevel.CRITICAL, component, message, pipeline_name ) - self._write_log(formatted, to_stderr=True) + self._write_log(formatted) def pipeline_started(self, pipeline_name): colored_name = self._colorize(pipeline_name, Colors.BRIGHT_WHITE) @@ -136,6 +163,11 @@ def pipeline_completed(self, pipeline_name): message = f"Pipeline completed: {colored_name}" self.info(message, "manager") + def pipeline_failed(self, pipeline_name, exc): + colored_name = self._colorize(pipeline_name, Colors.BRIGHT_WHITE) + message = f"Pipeline FAILED: {colored_name} — {exc}" + self.error(message, "manager") + def pipeline_killed(self, pipeline_name): colored_name = self._colorize(pipeline_name, Colors.BRIGHT_WHITE) message = f"Pipeline killed: {colored_name}" @@ -182,10 +214,11 @@ def activity_summary(self, active_pipelines, active_adaptive, buffered_pipelines self.debug(summary, "manager") def pipeline_log(self, message, level=LogLevel.INFO): + if not self._is_enabled(level): + return pipeline_component = f"PIPELINE-{self.name.upper()}" formatted = self._format_message(level, pipeline_component, message) - stderr_levels = [LogLevel.ERROR, LogLevel.CRITICAL] - self._write_log(formatted, to_stderr=level in stderr_levels) + self._write_log(formatted) def separator(self, title=None): if title: diff --git a/tests/conftest.py b/tests/conftest.py index c35a0f5..cb6ad28 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,9 +1,10 @@ # tests/conftest.py -import pytest import shutil from pathlib import Path - from unittest.mock import Mock + +import pytest + from impress import ImpressManager diff --git a/tests/unit/test_logger.py b/tests/unit/test_logger.py new file mode 100644 index 0000000..930d5b1 --- /dev/null +++ b/tests/unit/test_logger.py @@ -0,0 +1,278 @@ +import io +import sys + +from impress.utils.logger import ImpressLogger, LogLevel + + +class TestImpressLoggerInit: + def test_default_init(self): + logger = ImpressLogger() + assert logger.name == "ImpressManager" + assert logger.use_colors is True + assert logger.output_stream is sys.stdout + assert logger.min_level == LogLevel.DEBUG + + def test_custom_stream(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream) + assert logger.output_stream is stream + + def test_custom_min_level(self): + logger = ImpressLogger(min_level=LogLevel.WARNING) + assert logger.min_level == LogLevel.WARNING + + def test_custom_name(self): + logger = ImpressLogger(name="my_pipeline") + assert logger.name == "my_pipeline" + + +class TestLogLevelFiltering: + def test_is_enabled_at_exact_level(self): + logger = ImpressLogger(min_level=LogLevel.INFO) + assert logger._is_enabled(LogLevel.INFO) is True + + def test_is_enabled_above_min(self): + logger = ImpressLogger(min_level=LogLevel.INFO) + assert logger._is_enabled(LogLevel.WARNING) is True + assert logger._is_enabled(LogLevel.ERROR) is True + assert logger._is_enabled(LogLevel.CRITICAL) is True + + def test_is_disabled_below_min(self): + logger = ImpressLogger(min_level=LogLevel.INFO) + assert logger._is_enabled(LogLevel.DEBUG) is False + + def test_debug_suppressed_at_info_level(self): + stream = io.StringIO() + logger = ImpressLogger( + output_stream=stream, min_level=LogLevel.INFO, use_colors=False + ) + logger.debug("should not appear") + assert stream.getvalue() == "" + + def test_info_written_at_info_level(self): + stream = io.StringIO() + logger = ImpressLogger( + output_stream=stream, min_level=LogLevel.INFO, use_colors=False + ) + logger.info("should appear") + assert "should appear" in stream.getvalue() + + def test_warning_suppressed_below_min(self): + stream = io.StringIO() + logger = ImpressLogger( + output_stream=stream, min_level=LogLevel.ERROR, use_colors=False + ) + logger.warning("should not appear") + assert stream.getvalue() == "" + + def test_all_levels_write_at_debug_min(self): + stream = io.StringIO() + logger = ImpressLogger( + output_stream=stream, min_level=LogLevel.DEBUG, use_colors=False + ) + logger.debug("d") + logger.info("i") + logger.warning("w") + logger.error("e") + logger.critical("c") + output = stream.getvalue() + assert "d" in output + assert "i" in output + assert "w" in output + assert "e" in output + assert "c" in output + + def test_critical_only_at_critical_min(self): + stream = io.StringIO() + logger = ImpressLogger( + output_stream=stream, min_level=LogLevel.CRITICAL, use_colors=False + ) + logger.debug("d") + logger.info("i") + logger.warning("w") + logger.error("e") + logger.critical("c") + output = stream.getvalue() + assert "d" not in output + assert "i" not in output + assert "w" not in output + assert "e" not in output + assert "c" in output + + +class TestOutputStreamRespected: + def test_error_writes_to_output_stream_not_stderr(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=False) + logger.error("an error") + assert "an error" in stream.getvalue() + + def test_critical_writes_to_output_stream(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=False) + logger.critical("critical msg") + assert "critical msg" in stream.getvalue() + + def test_custom_stream_receives_all_output(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=False) + logger.info("info line") + logger.error("error line") + logger.warning("warn line") + output = stream.getvalue() + assert "info line" in output + assert "error line" in output + assert "warn line" in output + + +class TestColors: + def test_no_ansi_when_colors_off(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=False) + logger.info("hello") + assert "\033[" not in stream.getvalue() + + def test_ansi_present_when_colors_on(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=True) + logger.info("hello") + assert "\033[" in stream.getvalue() + + def test_colorize_returns_plain_when_off(self): + logger = ImpressLogger(use_colors=False) + result = logger._colorize("text", "\033[31m") + assert result == "text" + + def test_colorize_wraps_when_on(self): + logger = ImpressLogger(use_colors=True) + result = logger._colorize("text", "\033[31m") + assert result.startswith("\033[31m") + assert "text" in result + + +class TestHighLevelMethods: + def _make_logger(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=False) + return logger, stream + + def test_pipeline_started(self): + logger, stream = self._make_logger() + logger.pipeline_started("my_pipe") + assert "my_pipe" in stream.getvalue() + + def test_pipeline_completed(self): + logger, stream = self._make_logger() + logger.pipeline_completed("my_pipe") + assert "my_pipe" in stream.getvalue() + + def test_pipeline_failed(self): + logger, stream = self._make_logger() + logger.pipeline_failed("bad_pipe", ValueError("oops")) + out = stream.getvalue() + assert "bad_pipe" in out + assert "oops" in out + + def test_pipeline_killed(self): + logger, stream = self._make_logger() + logger.pipeline_killed("dead_pipe") + assert "dead_pipe" in stream.getvalue() + + def test_adaptive_started(self): + logger, stream = self._make_logger() + logger.adaptive_started("p1") + assert "p1" in stream.getvalue() + + def test_adaptive_completed(self): + logger, stream = self._make_logger() + logger.adaptive_completed("p1") + assert "p1" in stream.getvalue() + + def test_adaptive_failed(self): + logger, stream = self._make_logger() + logger.adaptive_failed("p1", "bad fn") + out = stream.getvalue() + assert "p1" in out + assert "bad fn" in out + + def test_child_pipeline_submitted(self): + logger, stream = self._make_logger() + logger.child_pipeline_submitted("child", "parent") + out = stream.getvalue() + assert "child" in out + assert "parent" in out + + def test_manager_starting(self): + logger, stream = self._make_logger() + logger.manager_starting(5) + assert "5" in stream.getvalue() + + def test_manager_exiting(self): + logger, stream = self._make_logger() + logger.manager_exiting() + assert stream.getvalue() != "" + + def test_activity_summary_suppressed_at_info(self): + stream = io.StringIO() + logger = ImpressLogger( + output_stream=stream, use_colors=False, min_level=LogLevel.INFO + ) + logger.activity_summary(3, 1, 2) + assert stream.getvalue() == "" # activity_summary is DEBUG level + + def test_activity_summary_written_at_debug(self): + stream = io.StringIO() + logger = ImpressLogger( + output_stream=stream, use_colors=False, min_level=LogLevel.DEBUG + ) + logger.activity_summary(3, 1, 2) + out = stream.getvalue() + assert "3" in out + + +class TestPipelineLog: + def test_pipeline_log_default_info(self): + stream = io.StringIO() + logger = ImpressLogger("pipe1", output_stream=stream, use_colors=False) + logger.pipeline_log("step done") + assert "step done" in stream.getvalue() + + def test_pipeline_log_suppressed_below_min(self): + stream = io.StringIO() + logger = ImpressLogger( + "pipe1", output_stream=stream, use_colors=False, min_level=LogLevel.WARNING + ) + logger.pipeline_log("step done", level=LogLevel.INFO) + assert stream.getvalue() == "" + + def test_pipeline_log_debug_level(self): + stream = io.StringIO() + logger = ImpressLogger("pipe1", output_stream=stream, use_colors=False) + logger.pipeline_log("debug step", level=LogLevel.DEBUG) + assert "debug step" in stream.getvalue() + + def test_pipeline_log_includes_pipeline_name_component(self): + stream = io.StringIO() + logger = ImpressLogger("mypipe", output_stream=stream, use_colors=False) + logger.pipeline_log("event") + assert "PIPELINE-MYPIPE" in stream.getvalue() + + +class TestSeparator: + def test_separator_no_title(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=False) + logger.separator() + assert "=" in stream.getvalue() + + def test_separator_with_title(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=False) + logger.separator("HELLO WORLD") + assert "HELLO WORLD" in stream.getvalue() + + def test_separator_ends_with_newline(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=False) + logger.separator() + assert stream.getvalue().endswith("\n") diff --git a/tests/unit/test_manager_core.py b/tests/unit/test_manager_core.py index d35495c..883e912 100644 --- a/tests/unit/test_manager_core.py +++ b/tests/unit/test_manager_core.py @@ -3,8 +3,7 @@ import pytest # Import the classes we're testing -from impress import ImpressBasePipeline, PipelineSetup -from impress import ImpressManager +from impress import ImpressBasePipeline, ImpressManager, PipelineSetup class MockPipeline(ImpressBasePipeline): diff --git a/tests/unit/test_manager_life_cycle.py b/tests/unit/test_manager_life_cycle.py index 9783467..bd110b3 100644 --- a/tests/unit/test_manager_life_cycle.py +++ b/tests/unit/test_manager_life_cycle.py @@ -12,9 +12,12 @@ class MockWorkflowEngine: """Mock workflow engine""" @classmethod - async def create(cls, backend=None): + async def create(cls, backend=None, **kwargs): return cls() + async def shutdown(self): + pass + class TestManagerLifecycle: @pytest.mark.asyncio @@ -113,3 +116,18 @@ async def run(self): # Should take at least 0.15 seconds due to slow adaptive function assert end_time - start_time >= 0.15 + + @pytest.mark.asyncio + @patch("impress.impress_manager.WorkflowEngine", MockWorkflowEngine) + async def test_start_exception_still_shuts_down_engine(self, impress_manager): + """WorkflowEngine.shutdown() is called even when a pipeline raises mid-run""" + + class ExplodingPipeline(MockPipeline): + async def run(self): + raise RuntimeError("pipeline exploded") + + await impress_manager.start( + [{"name": "boom", "type": ExplodingPipeline, "config": {}, "kwargs": {}}] + ) + # If shutdown() raises AttributeError, the try/finally fix is broken + assert len(impress_manager.pipeline_tasks) == 0 diff --git a/tests/unit/test_manager_pipeline_submission.py b/tests/unit/test_manager_pipeline_submission.py index e390b93..92bbdfa 100644 --- a/tests/unit/test_manager_pipeline_submission.py +++ b/tests/unit/test_manager_pipeline_submission.py @@ -1,5 +1,7 @@ from unittest.mock import Mock, patch +import pytest + # Import the classes we're testing from impress import PipelineSetup @@ -64,3 +66,10 @@ def test_submit_multiple_pipelines(self, mock_create_task, impress_manager): assert len(impress_manager.pipeline_tasks) == 2 assert mock_create_task.call_count == 2 + + def test_submit_before_start_raises(self, impress_manager): + """submit_new_pipelines raises RuntimeError when called before start()""" + with pytest.raises(RuntimeError, match="start\\(\\) must be called"): + impress_manager.submit_new_pipelines( + [{"name": "p", "type": MockPipeline, "config": {}, "kwargs": {}}] + ) diff --git a/tests/unit/test_pipeline_base.py b/tests/unit/test_pipeline_base.py new file mode 100644 index 0000000..928ea2d --- /dev/null +++ b/tests/unit/test_pipeline_base.py @@ -0,0 +1,157 @@ +import asyncio + +import pytest + +from impress.pipelines.impress_pipeline import ImpressBasePipeline + + +class MinimalPipeline(ImpressBasePipeline): + """Minimal concrete subclass for testing the base class.""" + + async def run(self): + pass + + def register_pipeline_tasks(self): + pass + + +class TestImpressBasePipelineInit: + def test_name_set(self): + p = MinimalPipeline(name="p1") + assert p.name == "p1" + + def test_flow_defaults_to_none(self): + p = MinimalPipeline(name="p1") + assert p.flow is None + + def test_flow_passed_through(self): + sentinel = object() + p = MinimalPipeline(name="p1", flow=sentinel) + assert p.flow is sentinel + + def test_state_is_empty_dict(self): + p = MinimalPipeline(name="p1") + assert p.state == {} + + def test_kill_parent_false(self): + p = MinimalPipeline(name="p1") + assert p.kill_parent is False + + def test_invoke_adaptive_step_false(self): + p = MinimalPipeline(name="p1") + assert p.invoke_adaptive_step is False + + def test_adaptive_barrier_is_event(self): + p = MinimalPipeline(name="p1") + assert isinstance(p._adaptive_barrier, asyncio.Event) + + def test_incoming_child_pipeline_request_empty(self): + p = MinimalPipeline(name="p1") + assert not p.incoming_child_pipeline_request + + def test_kwargs_stored_in_config(self): + p = MinimalPipeline(name="p1", foo="bar", baz=42) + assert p.config["foo"] == "bar" + assert p.config["baz"] == 42 + + +class TestChildPipelineRequest: + def test_submit_sets_request(self): + p = MinimalPipeline(name="p1") + config = {"name": "child", "type": MinimalPipeline} + p.submit_child_pipeline_request(config) + assert p.incoming_child_pipeline_request == config + + def test_get_returns_config_then_clears(self): + p = MinimalPipeline(name="p1") + config = {"name": "child", "type": MinimalPipeline} + p.submit_child_pipeline_request(config) + + result1 = p.get_child_pipeline_request() + assert result1 == config + + result2 = p.get_child_pipeline_request() + assert result2 is None + + def test_get_returns_none_when_nothing_pending(self): + p = MinimalPipeline(name="p1") + assert p.get_child_pipeline_request() is None + + def test_get_clears_after_retrieval(self): + p = MinimalPipeline(name="p1") + p.submit_child_pipeline_request({"name": "c"}) + p.get_child_pipeline_request() + assert not p.incoming_child_pipeline_request + + +class TestAdaptiveStep: + @pytest.mark.asyncio + async def test_run_adaptive_step_sets_flag(self): + p = MinimalPipeline(name="p1") + # wait=False: flag is set without blocking on the barrier + await p.run_adaptive_step(wait=False) + assert p.invoke_adaptive_step is True + + @pytest.mark.asyncio + async def test_run_adaptive_step_no_wait_does_not_hang(self): + p = MinimalPipeline(name="p1") + # _adaptive_barrier is clear — with wait=False this must return immediately + await asyncio.wait_for(p.run_adaptive_step(wait=False), timeout=1.0) + assert p.invoke_adaptive_step is True + + def test_set_adaptive_flag_true_clears_barrier(self): + p = MinimalPipeline(name="p1") + p._adaptive_barrier.set() + p._set_adaptive_flag(True) + assert p.invoke_adaptive_step is True + assert not p._adaptive_barrier.is_set() + + def test_set_adaptive_flag_false_does_not_touch_barrier(self): + p = MinimalPipeline(name="p1") + p._adaptive_barrier.set() + p._set_adaptive_flag(False) + assert p.invoke_adaptive_step is False + assert p._adaptive_barrier.is_set() # barrier unchanged + + +class TestOptionalMethods: + @pytest.mark.asyncio + async def test_finalize_is_noop(self): + p = MinimalPipeline(name="p1") + result = await p.finalize() + assert result is None + + @pytest.mark.asyncio + async def test_get_scores_map_returns_empty_dict(self): + p = MinimalPipeline(name="p1") + scores = await p.get_scores_map() + assert scores == {} + + def test_get_current_config_has_name_and_type(self): + p = MinimalPipeline(name="p1") + cfg = p.get_current_config_for_next_pipeline() + assert "name" in cfg + assert "type" in cfg + + def test_get_current_config_type_is_class(self): + p = MinimalPipeline(name="p1") + cfg = p.get_current_config_for_next_pipeline() + assert cfg["type"] is MinimalPipeline + + +class TestAbstractMethods: + def test_cannot_instantiate_without_run(self): + class NoRun(ImpressBasePipeline): + def register_pipeline_tasks(self): + pass + + with pytest.raises(TypeError): + NoRun(name="x") + + def test_cannot_instantiate_without_register(self): + class NoRegister(ImpressBasePipeline): + async def run(self): + pass + + with pytest.raises(TypeError): + NoRegister(name="x") diff --git a/tests/unit/test_pipeline_management.py b/tests/unit/test_pipeline_management.py index af35ae9..593327f 100644 --- a/tests/unit/test_pipeline_management.py +++ b/tests/unit/test_pipeline_management.py @@ -82,7 +82,7 @@ async def run(self): # Run the manager with timeout try: await asyncio.wait_for(impress_manager.start([pipeline_setup]), timeout=3.0) - except asyncio.TimeoutError: + except TimeoutError: # Print debug info if it times out print(f"Completed pipelines: {completed_pipelines}") print(f"Pipeline tasks: {len(impress_manager.pipeline_tasks)}") diff --git a/tests/unit/test_pipeline_setup.py b/tests/unit/test_pipeline_setup.py new file mode 100644 index 0000000..c56d398 --- /dev/null +++ b/tests/unit/test_pipeline_setup.py @@ -0,0 +1,141 @@ +import pytest +from pydantic import ValidationError + +from impress import PipelineSetup + +from .test_manager_core import MockPipeline + + +class TestPipelineSetupConstruction: + def test_all_fields(self): + async def fn(p): + pass + + setup = PipelineSetup( + name="my_pipe", + type=MockPipeline, + config={"a": 1}, + kwargs={"b": 2}, + adaptive_fn=fn, + ) + assert setup.name == "my_pipe" + assert setup.type is MockPipeline + assert setup.config == {"a": 1} + assert setup.kwargs == {"b": 2} + assert setup.adaptive_fn is fn + + def test_defaults(self): + setup = PipelineSetup(name="p", type=MockPipeline) + assert setup.config == {} + assert setup.kwargs == {} + assert setup.adaptive_fn is None + + def test_validate_type_rejects_non_subclass(self): + with pytest.raises(ValidationError): + PipelineSetup(name="p", type=str) + + def test_validate_type_rejects_non_type(self): + with pytest.raises(ValidationError): + PipelineSetup(name="p", type="not_a_class") + + def test_validate_type_accepts_subclass(self): + class Sub(MockPipeline): + pass + + setup = PipelineSetup(name="p", type=Sub) + assert setup.type is Sub + + +class TestFromDict: + def test_known_fields_separated(self): + async def fn(p): + pass + + data = { + "name": "p1", + "type": MockPipeline, + "config": {"x": 1}, + "adaptive_fn": fn, + } + setup = PipelineSetup.from_dict(data) + assert setup.name == "p1" + assert setup.type is MockPipeline + assert setup.config == {"x": 1} + assert setup.adaptive_fn is fn + assert setup.kwargs == {} + + def test_extra_keys_go_to_kwargs(self): + data = { + "name": "p1", + "type": MockPipeline, + "foo": "bar", + "baz": 42, + } + setup = PipelineSetup.from_dict(data) + assert setup.kwargs == {"foo": "bar", "baz": 42} + + def test_kwargs_key_in_dict_lands_in_kwargs(self): + # "kwargs" is not a known field, so it ends up nested inside kwargs + data = { + "name": "p1", + "type": MockPipeline, + "kwargs": {"inner": "value"}, + } + setup = PipelineSetup.from_dict(data) + assert setup.kwargs == {"kwargs": {"inner": "value"}} + + def test_minimal_dict(self): + setup = PipelineSetup.from_dict({"name": "p", "type": MockPipeline}) + assert setup.name == "p" + assert setup.config == {} + assert setup.kwargs == {} + assert setup.adaptive_fn is None + + +class TestToDict: + def test_basic_structure(self): + setup = PipelineSetup( + name="p1", + type=MockPipeline, + config={"c": 1}, + ) + d = setup.to_dict() + assert d["name"] == "p1" + assert d["type"] is MockPipeline + assert d["config"] == {"c": 1} + + def test_adaptive_fn_omitted_when_none(self): + setup = PipelineSetup(name="p", type=MockPipeline, adaptive_fn=None) + d = setup.to_dict() + assert "adaptive_fn" not in d + + def test_adaptive_fn_included_when_set(self): + async def fn(p): + pass + + setup = PipelineSetup(name="p", type=MockPipeline, adaptive_fn=fn) + d = setup.to_dict() + assert d["adaptive_fn"] is fn + + def test_kwargs_spread_into_result(self): + setup = PipelineSetup( + name="p", + type=MockPipeline, + kwargs={"foo": "bar", "num": 7}, + ) + d = setup.to_dict() + assert d["foo"] == "bar" + assert d["num"] == 7 + + def test_roundtrip_from_dict(self): + original = { + "name": "rt", + "type": MockPipeline, + "config": {"k": "v"}, + "extra_param": 99, + } + setup = PipelineSetup.from_dict(original) + d = setup.to_dict() + assert d["name"] == "rt" + assert d["type"] is MockPipeline + assert d["extra_param"] == 99