From 250250281f90b5364615aac959038f1f169f267b Mon Sep 17 00:00:00 2001 From: AJPreto Date: Fri, 12 Jun 2026 11:13:24 +0000 Subject: [PATCH 1/2] Update with GNINA, Karmadock fix, multi chain --- .dockerignore | 1 + .gitignore | 18 + .pre-commit-config.yaml | 4 +- Dockerfile | 43 +- Dockerfile.gnina-bundle | 55 + Makefile | 77 +- README.md | 192 ++- UPGRADE_SUMMARY.md | 146 ++ docs/adfrsuite_class.py | 91 +- guild/bulk.py | 1343 +++++++++++++---- guild/constants/bulk.py | 46 +- guild/constants/gnina.py | 31 + guild/constants/guild.py | 33 +- guild/constants/p2rank.py | 12 +- guild/constants/system.py | 10 +- guild/constants/visualization.py | 10 +- guild/docking/boltz.py | 309 +++- guild/docking/diffdock.py | 283 +++- guild/docking/gnina.py | 352 +++++ guild/docking/karmadock.py | 198 ++- guild/docking/vina.py | 318 +--- guild/run.py | 279 +++- guild/tools/bulk.py | 83 +- guild/tools/preparation.py | 71 +- guild/tools/scores.py | 2 +- guild/tools/subprocess_log.py | 68 + guild/tools/utils.py | 2 + guild/transformers/chembl.py | 14 +- guild/transformers/converters.py | 19 + guild/transformers/msa.py | 28 +- guild/transformers/pdb.py | 163 +- pyproject.toml | 9 +- scripts/apply_karmadock_patches.py | 131 ++ scripts/run_guild.py | 178 ++- scripts/vina_rescore_all.py | 11 +- tests/bulk_run/bulk_run_imports.py | 0 tests/bulk_run/test_diffdock_resume.py | 124 ++ tests/bulk_run/test_gnina_orchestration.py | 88 ++ .../test_prepare_docking_multi_protein.py | 160 ++ tests/bulk_run/test_progress_logging.py | 133 ++ tests/bulk_run/test_run_docking_order.py | 120 ++ tests/bulk_run/test_subprocess_log.py | 79 + tests/bulk_run/test_vina_rescore_split.py | 130 ++ ...test_multimethod_underscore_combination.py | 107 ++ tests/single_run/single_run_imports.py | 0 tests/single_run/test_gnina_input_mode.py | 200 +++ tests/single_run/test_gnina_scoring.py | 65 + tests/single_run/test_multi_chain_support.py | 111 ++ .../test_pocket_contacts_from_box.py | 95 ++ uv.lock | 2 +- 50 files changed, 5150 insertions(+), 894 deletions(-) create mode 100644 Dockerfile.gnina-bundle create mode 100644 UPGRADE_SUMMARY.md create mode 100644 guild/constants/gnina.py create mode 100644 guild/docking/gnina.py create mode 100644 guild/tools/subprocess_log.py create mode 100644 scripts/apply_karmadock_patches.py create mode 100644 tests/bulk_run/bulk_run_imports.py create mode 100644 tests/bulk_run/test_diffdock_resume.py create mode 100644 tests/bulk_run/test_gnina_orchestration.py create mode 100644 tests/bulk_run/test_prepare_docking_multi_protein.py create mode 100644 tests/bulk_run/test_progress_logging.py create mode 100644 tests/bulk_run/test_run_docking_order.py create mode 100644 tests/bulk_run/test_subprocess_log.py create mode 100644 tests/bulk_run/test_vina_rescore_split.py create mode 100644 tests/scores/test_multimethod_underscore_combination.py create mode 100644 tests/single_run/single_run_imports.py create mode 100644 tests/single_run/test_gnina_input_mode.py create mode 100644 tests/single_run/test_gnina_scoring.py create mode 100644 tests/single_run/test_multi_chain_support.py create mode 100644 tests/single_run/test_pocket_contacts_from_box.py diff --git a/.dockerignore b/.dockerignore index 19e3949..425e12b 100644 --- a/.dockerignore +++ b/.dockerignore @@ -7,6 +7,7 @@ !README.md !guild !tests +!scripts/apply_karmadock_patches.py # Ignore things that might be inside the folders above **/*~ diff --git a/.gitignore b/.gitignore index eed5aa6..92355b3 100644 --- a/.gitignore +++ b/.gitignore @@ -55,11 +55,13 @@ DATABASE_SCHEMA.md /.databrickscfg /.data /.local +/.azure # Chache /.pytest_cache /.ruff_cache /.cache +/_receptor_pdbqt_cache # Software vina @@ -73,3 +75,19 @@ openbabel/* pdbfixer/* plip/* KarmaDock + +# Local-only / scratch areas +/notebooks/debug/*.py +/notebooks/database +/guild/OLD +/temp_data + +# scripts/ — only the load-bearing scripts are tracked. +# run_guild.py is the Makefile entry point; apply_karmadock_patches.py is +# COPY'd by the Dockerfile during base-build; vina_rescore_all.py is a +# published utility. Anything else under scripts/ is treated as a local +# experiment and should not be committed (avoids leaking machine-specific paths). +scripts/* +!scripts/run_guild.py +!scripts/apply_karmadock_patches.py +!scripts/vina_rescore_all.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index df10de5..b92623d 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -4,7 +4,7 @@ repos: hooks: - id: nbstripout - repo: https://github.com/astral-sh/ruff-pre-commit - rev: 'v0.15.11' + rev: 'v0.15.15' hooks: - id: ruff args: [--fix, --exit-non-zero-on-fix] @@ -31,6 +31,6 @@ repos: - id: debug-statements - repo: https://github.com/astral-sh/uv-pre-commit - rev: 0.11.7 + rev: 0.11.18 hooks: - id: uv-lock diff --git a/Dockerfile b/Dockerfile index 46b361e..41a4b42 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,4 +1,22 @@ # syntax=docker/dockerfile:1 + +#################################################################################################### +# GNINA — pulled as a curated ~1.4 GB bundle (binary + the .so deps it loads +# at runtime). Build the bundle separately with Dockerfile.gnina-bundle and +# tag it (locally or in your own registry), then point GNINA_BUNDLE_IMAGE at +# that tag: +# +# docker build -f Dockerfile.gnina-bundle -t gnina-bundle:local . +# docker build --build-arg GNINA_BUNDLE_IMAGE=gnina-bundle:local -t guild . +# +# We do NOT consume gnina/gnina:latest directly here: pulling the full ~9 GB +# upstream image into the build graph caused disk exhaustion in CI. The bundle +# is a `FROM scratch` image whose only contents are /export/{bin,lib}, so it +# stays off the main build's dependency graph until the final runtime stage. +ARG GNINA_BUNDLE_IMAGE=gnina-bundle:latest +FROM ${GNINA_BUNDLE_IMAGE} AS gnina-source + +#################################################################################################### FROM python:3.10-slim-bookworm AS base ARG APP_NAME="guild" @@ -48,6 +66,13 @@ RUN wget https://github.com/ccsb-scripps/AutoDock-Vina/releases/download/v1.2.5/ -O /usr/local/bin/vina && \ chmod a+x /usr/local/bin/vina +# NOTE: GNINA is intentionally not copied into base-build. The gnina-source +# stage pulls a large upstream image and would force every target that +# transitively depends on base-build (including `test`) to materialize it on +# disk. The gnina bundle is only needed at runtime, so the COPY lives in the +# final `docker` stage instead — buildkit then skips gnina-source entirely for +# the `test` target. + # Python build tools RUN python -m pip install --upgrade pip setuptools wheel @@ -150,6 +175,14 @@ ENV LD_LIBRARY_PATH=/app/.venv/lib/python3.10/site-packages/nvidia/cu13/lib:/opt # Vina COPY --from=base-build /usr/local/bin/vina /usr/local/bin/vina +# GNINA — curated bundle (binary + isolated lib tree). Pulled directly from +# the gnina-source stage so base-build (and therefore the `test` target) does +# not have to materialize the large upstream image. Invoked with +# LD_LIBRARY_PATH=/opt/gnina/lib so its torch/openbabel/boost don't conflict +# with the venv-managed torch and the system openbabel. +COPY --from=gnina-source /export /opt/gnina +RUN chmod a+x /opt/gnina/bin/gnina + # LocalColabFold COPY --from=base-build /opt/localcolabfold /opt/localcolabfold ENV COLABFOLD_BIN="/opt/localcolabfold/.pixi/envs/default/bin" @@ -177,4 +210,12 @@ COPY --chown=appuser:appuser guild /app/guild RUN git clone https://github.com/schrojunzhang/KarmaDock.git /app/KarmaDock && \ git -C /app/KarmaDock checkout 9a35d0cb7caaa1a4d0a61f6ea96821dc1edefa81 && \ git clone https://github.com/gcorso/DiffDock.git /app/DiffDock && \ - git -C /app/DiffDock checkout 85c49b60d3e0b0182a59ee43a34a6d7036981284 \ No newline at end of file + git -C /app/DiffDock checkout 85c49b60d3e0b0182a59ee43a34a6d7036981284 + +# KarmaDock at the pinned commit was written for an older rdkit and crashes on +# the rdkit/torch combo we use today (dimension mismatch in ligand_feature.py +# and the GraphTransformer block). Apply the compatibility patches documented +# in the project README. The script is idempotent. +COPY --chown=appuser:appuser scripts/apply_karmadock_patches.py /tmp/apply_karmadock_patches.py +RUN /app/.venv/bin/python /tmp/apply_karmadock_patches.py /app/KarmaDock && \ + rm -f /tmp/apply_karmadock_patches.py \ No newline at end of file diff --git a/Dockerfile.gnina-bundle b/Dockerfile.gnina-bundle new file mode 100644 index 0000000..08620ef --- /dev/null +++ b/Dockerfile.gnina-bundle @@ -0,0 +1,55 @@ +# syntax=docker/dockerfile:1 +# +# Curated gnina runtime bundle. +# +# Pulls the upstream gnina/gnina image (~9 GB, mostly static archives + MKL +# we don't need at runtime), then extracts the binary plus the ~1.4 GB of +# shared libraries the binary actually loads — boost 1.71, jsoncpp, ICU 66, +# numa, openblas, gfortran, the CUDA 12 runtime libs, and gnina's own +# libtorch / libmolgrid / libopenbabel.so.7 — into ``/export``. +# +# The final stage is ``FROM scratch`` so the published image is just the +# /export tree. Build it once and tag it (locally or in your own registry), +# then point the main Dockerfile's GNINA_BUNDLE_IMAGE build arg at that tag. +# Consumers (see [Dockerfile](Dockerfile) in this repo) do: +# +# docker build -f Dockerfile.gnina-bundle -t gnina-bundle:local . +# docker build --build-arg GNINA_BUNDLE_IMAGE=gnina-bundle:local -t guild . +# +# The main image then does: +# +# ARG GNINA_BUNDLE_IMAGE +# FROM ${GNINA_BUNDLE_IMAGE} AS gnina-source +# ... +# COPY --from=gnina-source /export /opt/gnina +# +# and invokes ``gnina`` with ``LD_LIBRARY_PATH=/opt/gnina/lib`` so the bundle's +# libtorch / libopenbabel don't clash with the main image's CUDA 13 torch +# and system openbabel. +# +# Rebuild this bundle whenever you want to pick up a newer upstream gnina. + +FROM gnina/gnina:latest AS extract +RUN set -eux; \ + mkdir -p /export/bin /export/lib; \ + cp -L /usr/local/bin/gnina /export/bin/; \ + cp -L /usr/local/lib/*.so* /export/lib/ 2>/dev/null || true; \ + cp -L /usr/lib/x86_64-linux-gnu/libboost_*.so.1.71.0 /export/lib/ 2>/dev/null || true; \ + cp -L /usr/lib/x86_64-linux-gnu/libjsoncpp.so* /export/lib/ 2>/dev/null || true; \ + # Ubuntu 20.04 ICU + libnuma + openblas — Debian bookworm ships + # newer/different SONAMEs, so gnina would fail to find these at runtime. + cp -L /usr/lib/x86_64-linux-gnu/libicui18n.so.66* /export/lib/ 2>/dev/null || true; \ + cp -L /usr/lib/x86_64-linux-gnu/libicuuc.so.66* /export/lib/ 2>/dev/null || true; \ + cp -L /usr/lib/x86_64-linux-gnu/libicudata.so.66* /export/lib/ 2>/dev/null || true; \ + cp -L /usr/lib/x86_64-linux-gnu/libnuma.so* /export/lib/ 2>/dev/null || true; \ + cp -L /usr/lib/x86_64-linux-gnu/libopenblas.so* /export/lib/ 2>/dev/null || true; \ + cp -L /usr/lib/x86_64-linux-gnu/libgfortran.so* /export/lib/ 2>/dev/null || true; \ + cp -L /usr/lib/x86_64-linux-gnu/libquadmath.so* /export/lib/ 2>/dev/null || true; \ + if [ -d /usr/local/cuda-12.0/targets/x86_64-linux/lib ]; then \ + cp -RL /usr/local/cuda-12.0/targets/x86_64-linux/lib/*.so* /export/lib/ 2>/dev/null || true; \ + fi + +#################################################################################################### +# Tiny published image: just /export, nothing else. ~1.4 GB on disk. +FROM scratch AS bundle +COPY --from=extract /export /export diff --git a/Makefile b/Makefile index 9806320..dc7a913 100644 --- a/Makefile +++ b/Makefile @@ -1,14 +1,22 @@ -.PHONY : docker-local docker-test dev test run-boltz run-vina run-diffdock run-guild project-init project-setup +.PHONY : docker-local docker-test dev test run-boltz run-vina run-diffdock run-gnina run-plip run-guild project-init project-setup # Run the tests locally test: uv run pytest -v -# Build a local guild image +# Build a local guild image. +# Uses gnina-bundle:local if already present; otherwise builds it first from +# Dockerfile.gnina-bundle (pulls gnina/gnina:latest once, ~9 GB) so the gnina +# docking method works out of the box without needing an external registry. docker-local: + @if ! docker image inspect gnina-bundle:local >/dev/null 2>&1; then \ + echo "gnina-bundle:local not found — building from Dockerfile.gnina-bundle..."; \ + docker build -t gnina-bundle:local -f Dockerfile.gnina-bundle .; \ + fi DOCKER_BUILDKIT=1 \ docker build \ --build-arg APP_NAME=guild \ + --build-arg GNINA_BUNDLE_IMAGE=gnina-bundle:local \ -t guild:latest \ -f Dockerfile \ --target=docker \ @@ -34,8 +42,25 @@ BATCH_SIZE ?= 2 HEAD ?= 0 CLEAN ?= KNOWN_BINDERS ?= +NO_DECOYS ?= +BOX ?= +N_WORKERS ?= PASSWD_FILE ?= /tmp/guild_passwd +# GPU toggle. Default 1 (enabled). Set USE_GPU= (empty) to drop +# `--gpus all --shm-size=8g` from docker run AND forward `--no-gpu` to +# run_guild.py — required on hosts without a usable GPU (gnina then falls back +# to CPU; Boltz is genuinely GPU-bound and shouldn't be combined with USE_GPU=). +USE_GPU ?= 1 +_GPU_FLAGS = $(if $(USE_GPU),--gpus all --shm-size=8g,) +_NO_GPU_FLAG = $(if $(USE_GPU),,--no-gpu) + +# Gnina input mode. Empty (default) → omit the flag (pdbqt default applies). +# Set GNINA_INPUT_MODE=sdf to skip OpenBabel PDBQT prep when gnina is the only +# docking method requested (otherwise BulkRun downgrades to pdbqt with a warning). +GNINA_INPUT_MODE ?= +_GNINA_INPUT_MODE_FLAG = $(if $(GNINA_INPUT_MODE),--gnina-input-mode $(GNINA_INPUT_MODE),) + # Internal docker run flags reused across targets. # Mounts a generated /etc/passwd so pwd.getpwuid() works for the host UID # (required by PyTorch / boltz inside the container). @@ -60,7 +85,10 @@ _HEAD_FLAG = $(if $(filter-out 0,$(HEAD)),--head $(HEAD),) # Collect all optional flags into one variable for DRY target definitions _DECOYS_FLAG = $(if $(DECOYS),--decoys $(DECOYS),) -_OPTIONAL_FLAGS = $(_CLEAN_FLAG) $(_KNOWN_BINDERS_FLAG) $(_HEAD_FLAG) $(_DECOYS_FLAG) +_NO_DECOYS_FLAG = $(if $(NO_DECOYS),--no-decoys,) +_BOX_FLAG = $(if $(BOX),--box $(BOX),) +_N_WORKERS_FLAG = $(if $(N_WORKERS),--n-workers $(N_WORKERS),) +_OPTIONAL_FLAGS = $(_CLEAN_FLAG) $(_KNOWN_BINDERS_FLAG) $(_HEAD_FLAG) $(_DECOYS_FLAG) $(_NO_DECOYS_FLAG) $(_BOX_FLAG) $(_N_WORKERS_FLAG) $(_NO_GPU_FLAG) $(_GNINA_INPUT_MODE_FLAG) # Generate an /etc/passwd that includes the container's original entries plus # the host user. This fixes pwd.getpwuid() failures for LDAP/SSSD users @@ -78,8 +106,7 @@ METHODS ?= boltz run-boltz: _prepare-passwd docker run \ $(DOCKER_COMMON) \ - --gpus all \ - --shm-size=8g \ + $(_GPU_FLAGS) \ -e LD_LIBRARY_PATH=/opt/localcolabfold/.pixi/envs/default/lib:/usr/local/lib:/app/.venv/lib/python3.10/site-packages/nvidia/cu13/lib:/app/.venv/lib/python3.10/site-packages/nvidia/cuda_nvrtc/lib:/app/.venv/lib/python3.10/site-packages/nvidia/cudnn/lib:/app/.venv/lib/python3.10/site-packages/nvidia/cublas/lib \ guild:latest \ python $(MASTER_SCRIPT) \ @@ -117,13 +144,45 @@ run-diffdock: _prepare-passwd --batch-size $(BATCH_SIZE) \ $(_OPTIONAL_FLAGS) -# Generic target — pass METHODS="boltz vina karmadock diffdock" as needed. -# Picks up GPU flags automatically when any GPU method is present. +# Run gnina docking inside the local Docker image. Uses the GPU by default for +# CNN rescoring; pass USE_GPU= (empty) on no-GPU hosts to drop --gpus and run +# gnina CPU-only. Requires: make docker-local first. +GNINA_METHODS ?= gnina +run-gnina: _prepare-passwd + docker run \ + $(DOCKER_COMMON) \ + $(_GPU_FLAGS) \ + guild:latest \ + python $(MASTER_SCRIPT) \ + --project $(PROJECT) \ + --combinations $(COMBINATIONS) \ + --methods $(GNINA_METHODS) \ + --batch-size $(BATCH_SIZE) \ + $(_OPTIONAL_FLAGS) + +# Re-run only the PLIP interactions step over an existing data// tree. +# CPU-safe and skips docking + scoring entirely — useful for regenerating +# plip_interactions.tsv when only the PLIP code changed. Requires the same +# COMBINATIONS / PROJECT used by the original run. +run-plip: _prepare-passwd + docker run \ + $(DOCKER_COMMON) \ + guild:latest \ + python $(MASTER_SCRIPT) \ + --project $(PROJECT) \ + --combinations $(COMBINATIONS) \ + --methods $(METHODS) \ + --batch-size $(BATCH_SIZE) \ + --plip-only \ + $(_OPTIONAL_FLAGS) + +# Generic target — pass METHODS="boltz vina karmadock diffdock gnina" as needed. +# GPU on by default; pass USE_GPU= (empty) for CPU-only hosts (drops --gpus all +# and forwards --no-gpu). Boltz requires a GPU regardless. run-guild: _prepare-passwd docker run \ $(DOCKER_COMMON) \ - --gpus all \ - --shm-size=8g \ + $(_GPU_FLAGS) \ -e LD_LIBRARY_PATH=/opt/localcolabfold/.pixi/envs/default/lib:/usr/local/lib:/app/.venv/lib/python3.10/site-packages/nvidia/cu13/lib:/app/.venv/lib/python3.10/site-packages/nvidia/cuda_nvrtc/lib:/app/.venv/lib/python3.10/site-packages/nvidia/cudnn/lib:/app/.venv/lib/python3.10/site-packages/nvidia/cublas/lib \ guild:latest \ python $(MASTER_SCRIPT) \ diff --git a/README.md b/README.md index d63f5c2..9e9bb7e 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ # Guild -> **Version 1.0.0** — Python ≥3.10, <3.11 +> **Version 1.1.4** — Python ≥3.10, <3.11 Guild is an open-source Protein-Ligand Binding Tools orchestrator that covers the end-to-end pipeline while leveraging multiple docking methods in each step. @@ -13,11 +13,16 @@ Guild is an open-source Protein-Ligand Binding Tools orchestrator that covers th * [Docker (recommended)](#docker) * [Running locally](#how-to-run) * [Installations](#installations) + * [Consuming PLIP interactions output](#consuming-plip-interactions-output) + * [Troubleshooting a failed combination](#troubleshooting-a-failed-combination) * [Usage](#usage) * [Single run](#single-run) * [BulkRun](#bulkrun) * [Methods](#methods) * [Docking](#docking) + * [Vina rescore (automatic with DiffDock and Boltz)](#vina-rescore-automatic-with-diffdock-and-boltz) + * [Custom binding pocket](#custom-binding-pocket) + * [Multi-chain binding pocket](#multi-chain-binding-pocket) * [Post-analysis](#post-analysis) ## Docker @@ -64,12 +69,17 @@ make run-vina \ |---|---|---| | `COMBINATIONS` | *(required)* | Path to the protein–ligand pairs CSV/TSV (use `/workspace/…` paths) | | `PROJECT` | `imagerun` | Output folder name under `data/` (no underscores allowed) | -| `METHODS` | `boltz` | Space-separated list: `boltz`, `vina`, `karmadock`, `diffdock` | +| `METHODS` | `boltz` | Space-separated list: `boltz`, `vina`, `karmadock`, `diffdock`, `gnina` | | `BATCH_SIZE` | `2` | Number of combinations per batch | | `HEAD` | `0` | Take only the first N rows from the combinations table (0 = all) | | `DECOYS` | *(script default)* | Path to the decoys file; omit to use built-in default (`chembl_36_decoys_2.tsv`) | +| `NO_DECOYS` | *(empty)* | Set to `1` to skip decoy expansion entirely (useful for single-protein runs where you only want to score the supplied ligands) | | `CLEAN` | *(empty)* | Set to `1` to delete the project output folder before running | | `KNOWN_BINDERS` | *(empty)* | Set to `1` to enable known-binders expansion | +| `N_WORKERS` | `1` | Vina parallel-worker processes. Vina internally also threads — values >1 may oversubscribe on high-core hosts but are typically fine. | +| `BOX` | *(empty)* | Global fallback Vina box file (`center_{x,y,z}` + `size_{x,y,z}`). Used for combinations whose CSV `box_location` cell is empty; per-row values always take precedence. See [Custom binding pocket](#custom-binding-pocket). | +| `USE_GPU` | `1` | Set empty (`USE_GPU=`) to drop `--gpus all` from `docker run` and forward `--no-gpu` to the python script. Use on no-GPU hosts. gnina falls back to CPU; vina and diffdock are unaffected. **Do not combine with `METHODS=boltz`** — Boltz is genuinely GPU-bound. | +| `GNINA_INPUT_MODE` | *(empty)* | Set to `sdf` to skip OpenBabel PDBQT prep entirely when gnina is the only docking method requested — gnina then reads the RDKit-generated SDF + cleaned PDB directly. Co-requesting Vina or any Vina-rescore (boltz/diffdock auto-add a Vina-rescore) silently falls back to PDBQT with a warning, since OpenBabel still has to run for those methods. | | `MIN_MOL_WT` | `250` | Minimum molecular weight filter for known-binder expansion | | `MAX_MOL_WT` | `450` | Maximum molecular weight filter for known-binder expansion | | `CHEMBL_VERSION` | `chembl_36` | ChEMBL version string used for known-binder lookup | @@ -82,6 +92,8 @@ make run-vina \ | `run-boltz` | Yes | Shortcut for boltz docking | | `run-vina` | No | Shortcut for vina docking (CPU only) | | `run-diffdock` | No | Shortcut for diffdock docking | +| `run-gnina` | Yes* | Shortcut for gnina docking (*GPU used for CNN rescoring; pass `USE_GPU=` for CPU-only) | +| `run-plip` | No | Re-run only the PLIP interactions step over an existing `data//` tree | ### Direct script invocation @@ -106,6 +118,70 @@ python scripts/run_guild.py \ * NVIDIA GPU + [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/) (for GPU methods) +### Consuming PLIP interactions output + +Every `make run-*` invocation runs PLIP after scoring and writes +`data//plip_interactions.tsv` — a tab-separated file with one row +per docked complex that PLIP could analyze. The file is **always written**, +even header-only if no method produced complex PDBs, so downstream code +paths are deterministic. External notebooks should read this file directly +instead of installing `plip` locally (its sdist tries to build openbabel +from source, which fails on most CI / lab hosts): + +```python +import pandas as pd + +plip = pd.read_csv(f"data/{project}/plip_interactions.tsv", sep="\t") +``` + +Schema (see [guild/constants/plip.py](guild/constants/plip.py) for the +canonical list): + +| column | meaning | +|---|---| +| `protein_config_id` | matches the row in `combinations.csv` | +| `smiles` | the docked ligand | +| `n_hbonds` | hydrogen bonds | +| `n_hydrophobic` | hydrophobic contacts | +| `n_pistacking`, `n_pication` | π-stacking and π-cation | +| `n_saltbridges`, `n_halogen`, `n_waterbridges`, `n_metal` | other interaction types | +| `total_interactions` | sum across all categories | +| `n_unique_residues` | unique residue contacts | + +To skip PLIP for a docking run, pass `--no-plip` to `run_guild.py`. To +**re-run only PLIP** over an existing project (no re-docking), use: + +```shell +make run-plip PROJECT=myproject COMBINATIONS=/workspace/path/to/combos.csv METHODS="vina diffdock" +``` + +That iterates the existing `data//batches/*/` tree and regenerates +`plip_interactions.tsv` from whatever complex PDBs are present. + +### Troubleshooting a failed combination + +When a docking method fails for a given combination, the project's +`batch_progress.log` and each batch's `output.log` carry a `FAILED ...` +line that points at a dedicated subprocess transcript: + +| Method | Per-combination log path | +|---|---| +| Boltz | `batches//boltz/.subprocess.log` | +| gnina | `batches//gnina/_.subprocess.log` | +| DiffDock | `batches//diffdock/_batch.subprocess.log` *(batch-level — DiffDock runs once per batch)* | +| Vina | *(uses the Python API — failure trace lands in `output.log`)* | + +The file contains the full argv, exit code, stdout and stderr — written +on every invocation, success or failure. Example FAILED line you'd grep +for: + +``` +FAILED Boltz 6CTA-A-protein_1 (empty manifest after retry) — see /workspace/data/myproject/batches/batch_1/boltz/6CTA-A-protein_1.subprocess.log +``` + +Open that file to see exactly what Boltz / gnina / DiffDock printed before +exiting — no `grep` archaeology in the batch-wide log needed. + --- ## How to run @@ -225,21 +301,27 @@ Running in bulk is necessary to leverage the rank percentile score, as it is emp Your input table should be a pandas DataFrame with the following columns: -|protein_config_id|protein_id|protein_path|protein_chain|original_ligand|original_ligand_chain|ligand_id|smiles|ligand_category|is_pdb| -|-----------------|----------|------------|-------------|---------------|---------------------|---------|------|---------------|------| -|5zk8-A-3C0-A|5zk8|path/to/file.pdb|A|3C0|A|drug_1|CCCC|LOI|1| +|protein_config_id|protein_id|protein_path|protein_chain|original_ligand|original_ligand_chain|ligand_id|smiles|ligand_category|is_pdb|box_location| +|-----------------|----------|------------|-------------|---------------|---------------------|---------|------|---------------|------|------------| +|5zk8-A-3C0-A|5zk8|path/to/file.pdb|A|3C0|A|drug_1|CCCC|LOI|1| *(optional)* | **Column Descriptions:** - `protein_config_id`: Unique identifier for the protein configuration (e.g., `{protein_id}-{chain}-{ligand}-{ligand_chain}`) - `protein_id`: PDB ID or identifier for the protein - `protein_path`: Full path to the protein PDB file -- `protein_chain`: Chain identifier to use for docking +- `protein_chain`: Chain identifier to use for docking. To dock into a pocket + that spans **multiple chains** (e.g. a dimer interface), give a comma-separated + list such as `A,B` (any number of chains). See + [Multi-chain binding pocket](#multi-chain-binding-pocket). - `original_ligand`: Ligand identifier from the PDB file - `original_ligand_chain`: Chain of the original ligand - `ligand_id`: Unique identifier for the ligand - `smiles`: SMILES string of the ligand - `ligand_category`: Category of ligand (e.g., "LOI" for ligand of interest, "known_binder", etc.) - required for plotting - `is_pdb`: Binary indicator (1 if PDB file, 0 otherwise) +- `box_location` *(optional)*: Path to a Vina box file (`center_{x,y,z}` + `size_{x,y,z}`) that + defines the binding pocket for both Vina and Boltz. Supplied per row but conceptually per + protein — see [Custom binding pocket](#custom-binding-pocket). **Basic Example:** @@ -338,12 +420,100 @@ biorxiv: If you use results from any of these tools, please make sure to cite the authors as indicated in the hyperlinks. -### Vina rescore (automatic with DiffDock) +### Vina rescore (automatic with DiffDock and Boltz) + +When `diffdock` or `boltz` is included in the methods list, Guild **automatically adds** a +matching Vina rescore step. The rescore applies Vina's physics-based scoring function to the +predicted pose (score-only, no re-docking), giving a kcal/mol ΔG estimate that's comparable +across methods. + +The two rescore tracks are **independent** — each produces its own column, so a run that uses +both DiffDock and Boltz gets two distinct rescore scores: + +| Upstream method | Auto-enabled rescore | Score column | +|-----------------|------------------------|-------------------------------| +| `boltz` | `vina_rescore_boltz` | `vina_rescore_boltz_score` | +| `diffdock` | `vina_rescore_diffdock`| `vina_rescore_diffdock_score` | + +Both score columns are in kcal/mol (lower = stronger predicted binding). Each is independently +ranked per protein and folded into the `global_rp_score`. + +> **Note:** `boltz_score` itself is the protein-ligand ipTM confidence (range [0, 1], higher = +> more confident structure) — not a binding score. For a binding-strength signal from Boltz, +> use `vina_rescore_boltz_score`. Likewise `gnina`'s `gnina_score` is the Vina-style affinity +> (kcal/mol, lower = better) while `gnina_cnn_score` is a pose-confidence side channel that does +> not participate in guild's rank-percentile aggregation. + +#### Coordinate-frame caveat + +Boltz often recentres its predicted complex into its own internal frame, so a receptor PDBQT +prepared from the *template* PDB will not be in the same physical space as the Boltz-output +ligand. Guild handles this by extracting both the receptor and the ligand from Boltz's complex +PDB on every rescore call. If you call `rescore_boltz_pose` directly outside the bulk pipeline, +do the same — do not reuse the template-frame receptor. + +### Custom binding pocket + +By default, Guild derives the binding pocket from the co-crystal ligand declared in +`original_ligand` / `original_ligand_chain` (for Boltz, residues within 4 Å of that ligand become +`pocket_contacts`; for Vina, a box is built from its centre). When that information isn't +available — apo structures, recombinant assemblies, or pockets predicted by external tools like +fpocket / P2Rank — supply a **Vina box file** instead: + +``` +center_x = -7.470 +center_y = -15.230 +center_z = 5.970 +size_x = 13.830 +size_y = 15.070 +size_z = 15.230 +``` + +There are two ways to wire it in: + +1. **Per-row** via the optional `box_location` column in the combinations CSV — best for + multi-protein runs where each protein has its own pocket file. +2. **Global** via the `BOX=` Makefile flag (or `--box` on the script) — fills the column on rows + where it is empty. Per-row values always win. + +Precedence per combination: `box_location` (explicit) > P2Rank prediction (when +`predict_binding_pocket=True`) > derived from `original_ligand`. + +For Boltz, the box is converted to a residue list at runtime: every residue whose Cα is inside +the axis-aligned box becomes a `contact` constraint in the YAML. + +### Multi-chain binding pocket + +Some pockets sit at the **interface of two (or more) protein chains** — a dimer +interface, a recombinant assembly, an allosteric site between subunits. To dock a +small-molecule ligand into such a pocket, list every chain in the `protein_chain` +column, comma-separated: + +| protein_config_id | protein_id | protein_path | protein_chain | ... | +|-------------------|------------|--------------|---------------|-----| +| 6CTA-A,B--lig1 | 6CTA | /path/6CTA.pdb | `A,B` | ... | + +This works for an arbitrary number of chains (`A`, `A,B`, `A,B,C,D`, …). A single +chain (`A`) behaves exactly as before, so existing tables are unaffected. + +What changes under the hood when more than one chain is listed: + +- **Receptor** (Vina, gnina, KarmaDock, DiffDock): all listed chains are kept in + the prepared receptor, so docking and scoring see the full interface. Residues + are renumbered 1-based *per chain*. +- **Boltz**: emits one `protein` block per chain (each with its own MSA) and a + template/pocket constraint spanning every chain. Per-chain MSA generation and + the larger complex make Boltz runs proportionally more expensive with more + chains. +- **Pocket contacts**: collected from every listed chain, each indexed + independently from 1, so an interface pocket is fully constrained. -When `diffdock` is included in the methods list, Guild **automatically adds** a Vina rescore step. -After DiffDock generates poses, the Vina scoring function is applied to the top-ranked DiffDock pose for each combination (score-only, no re-docking). -This produces an additional `vina_rescore_score` column (kcal/mol, lower = better) alongside the DiffDock confidence score. -Both scores are independently ranked per protein and averaged into the `global_rp_score`. +**Recommended:** supply a `box_location` (see [Custom binding pocket](#custom-binding-pocket)) +that covers the interface. The co-crystal-ligand pocket derivation +(`original_ligand`) still uses the primary (first) chain only; a box covers the +whole interface cleanly. If you encode the chain segment of `protein_config_id`, +use the same comma form (`6CTA-A,B-...`) so DiffDock/complex receptor extraction +keeps both chains. ### Post-analysis diff --git a/UPGRADE_SUMMARY.md b/UPGRADE_SUMMARY.md new file mode 100644 index 0000000..a68ff31 --- /dev/null +++ b/UPGRADE_SUMMARY.md @@ -0,0 +1,146 @@ +# guild_upgrade — sync from guild-internal + +Branch: `guild_upgrade`. **Nothing committed** — all changes are uncommitted in the working tree for your review. + +This brings the public `guild` repo up to date with the latest engineering work in the +internal `guild-internal` repo, **without exposing any internal secrets, infrastructure, +database details, or personal paths**. The two repos have separate git histories (public +was a sanitized import), so this was done by comparing working trees file-by-file, not by +merging git history. + +## Guiding rule + +Internal-only material was **excluded entirely**, not just commented out. Where a shared +file mixed good engineering with internal coupling (e.g. `bulk.py`, the `Dockerfile`, +the `Makefile`), the improvements were ported and the internal parts removed — including +indirect tells like internal function names, the DB API shape, internal CI workflows, +private registry URLs, and example compound names. + +A full token sweep over the result found **zero new internal tokens** introduced by this +upgrade (see "Verification" below). + +## What was brought in (sanitized) + +### New features +- **GNINA docking** — `guild/docking/gnina.py`, `guild/constants/gnina.py`, plus tests + (`tests/single_run/test_gnina_*.py`, `tests/bulk_run/test_gnina_orchestration.py`). + The Python code only references container-internal paths (`/opt/gnina/...`), no secrets. +- **GNINA bundle build** — `Dockerfile.gnina-bundle` builds a curated runtime bundle from + the **public** `gnina/gnina` image. The private `envedancusacr01.azurecr.io` registry + reference was removed; the bundle image is now an overridable build arg + `GNINA_BUNDLE_IMAGE` (default `gnina-bundle:latest`). +- **Multi-chain binding pocket** — comma-separated `protein_chain` (`A,B`) support, threaded + through `run.py`, `bulk.py`, `tools/preparation.py`, `transformers/pdb.py`, and + `scripts/vina_rescore_all.py` (multi-chain port applied onto the public, relative-path + version of that script). +- **Per-subprocess troubleshooting logs** — `guild/tools/subprocess_log.py` + integration. +- **KarmaDock compatibility patches** — `scripts/apply_karmadock_patches.py`, applied at + Docker build time (rdkit/torch ABI fixes). +- **`run_guild.py` CLI** — gained `--methods gnina`, `--box`, `--n-workers`, `--no-gpu`, + `--plip-only`, `--no-plip`, `--gnina-input-mode`. The `--no-database-update` flag and the + `database_update=` wiring were **stripped**. +- **`SHELL_SILENCER` bugfix** — `&>/dev/null` → `> /dev/null 2>&1` (dash-safe). + +### Updated wholesale (verified clean of internal refs) +`guild/run.py`, `guild/docking/{boltz,diffdock,karmadock,vina}.py`, +`guild/tools/{preparation,bulk,utils,scores}.py`, +`guild/transformers/{pdb,msa,converters,chembl}.py`, +`guild/constants/{bulk,guild,p2rank,visualization}.py`, `docs/adfrsuite_class.py`. + +### Config / build (sanitized merges) +- **Dockerfile** — added GNINA bundle stage (genericized) + KarmaDock patches. **Kept** the + public repo's clean `uv sync` (no `GH_TOKEN` secret) and **omitted** the Azure CLI block — + both existed only to serve the excluded `enveda-toolkit`/keyvault/database layer. +- **Makefile** — added `run-gnina`, `run-plip`, a `USE_GPU` toggle, gnina/box/workers flags, + and gnina-bundle auto-build in `docker-local`. **Excluded** all `azurecr.io` push targets, + the `dev`/`run` targets that pulled internal images, the `~/.azure` credential mounts, and + the `github-init`/`init` targets (`gh repo create enveda/guild`, internal remote URL). +- **pyproject.toml** — version bump 1.0.0 → 1.1.4, added one author. **Did not** add + `azure-identity`, `azure-keyvault-secrets`, `enveda-toolkit`, or the `[tool.uv.sources]` + git source. +- **uv.lock** — public lock was already free of internal packages; only the `guild` version + pin was updated to 1.1.4 (no dependency changes were made, so the lock stays consistent). +- **.gitignore** — added safety rules: `/notebooks/database`, `/guild/OLD`, `/temp_data`, and + a `scripts/*` allowlist (prevents committing local experiment scripts that carry + machine-specific paths). Kept the public allowlist for shipped support data + logo. +- **.dockerignore** — un-ignore `scripts/apply_karmadock_patches.py` (now COPY'd by Docker). +- **.pre-commit-config.yaml** — ruff/uv hook version bumps only. + +## What was deliberately EXCLUDED (internal-only) + +- **Database / Azure layer**: `guild/azure/database.py`, `guild/azure/get_secret.py`, + `guild/constants/database.py`, `DATABASE_SCHEMA.md`, `notebooks/database/*`, + `scripts/test_db_real_submission.py`, `scripts/test_db_upload_dummy.py`. The public repo + keeps its existing "no database mode". In `bulk.py` the `database_update` param, the + `retrieve_existing_rp_scores`/`upload_rp_scores` import and all call sites, and the + DB-merge/upload code paths were fully removed (not just disabled). +- **Internal CI**: `.github/workflows/deploy-dev.yml`, `.github/workflows/build-gnina-bundle.yml` + (both use `enveda/reusable-workflows` + ACR + Prefect). The public `test.yml`/`lint.yml` + were kept as-is. +- **Internal tooling/meta**: `CLAUDE.md` (internal dev notes — references ACR, Databricks, + internal compound names), `.cookiecutter.json` (internal template + `it@envedabio.com`), + `.claude/settings.json`, `.gitmodules` (empty), `pyptoject_cpu.toml` (typo-named CPU + variant carrying azure/prefect deps), `environments/installations.sh` (hardcoded + `/home/ec2-user/...` paths). +- **Azure/Prefect constants** in `constants/system.py` (`AZURE_KEYVAULT_URI`, + `AZURE_CLIENT_ID`, `AZURE_ACR_IMAGE_URI`, etc.) — used only by excluded code. +- **New analysis notebooks / plot scripts / result figures** under `notebooks/` — these are + analysis artifacts whose output cells and hardcoded paths (several contained + `/home/antonio.gomes/...`) are a high indirect-leak risk for low value. Existing public + notebooks were left untouched. **Worth a selective, manual revisit if you want any of + these published.** + +## Pre-existing references (NOT introduced by this upgrade — flagged for your call) + +- `.gitignore:55` → `/.databrickscfg` — a protective ignore rule (prevents committing a + Databricks credential file). Left in place; removing it would *reduce* safety. +- `renovate.json:4` → `local>enveda/renovate-config` — was already in the public repo, + unchanged. References an internal renovate config repo. + +## Verification + +- All 15 ported modules import cleanly in the repo venv. +- Full internal-token sweep (`enveda`, `azurecr`, `databricks`, `keyvault`, `enveda_toolkit`, + `enveda_prod`, `cdd-benchling`, the Azure client GUID, `/home/antonio`, `/home/ec2-user`, + `reusable-workflows`, `prefect`, …) over all source/config files → only the two + pre-existing references above; zero new tokens. +- Docker image rebuild + empirical docking results: see the "Empirical results" section below. + +## Empirical results + +Image rebuilt from scratch with `--build-arg GNINA_BUNDLE_IMAGE=gnina-bundle:local` +(`guild:latest`, 38.2 GB). `gnina v1.3.1` runs from the genericized local bundle. + +**Test suite** (`docker run --rm guild-test python -m pytest tests/ -q`): **100 passed**. + +**Empirical docking** — single combination `3pbl-A-ETQ-A / lig1` (D3 receptor + a small +ligand), box-driven pocket, `--no-decoys`, run on an A100 80GB. Every engine produced a +score and completed scoring + PLIP (exit 0): + +| Engine | Mode | Result | +|-----------|------|--------| +| Vina | CPU | `vina_score = -4.892` | +| GNINA | GPU | `gnina_score = -5.09`, `gnina_cnn_score = -0.28` | +| Karmadock | GPU | `karmadock_score = 13.86` (build-time rdkit/torch patches applied OK) | +| DiffDock | GPU | `diffdock_score = -1.7`; auto-rescore `vina_rescore_diffdock_score = -0.029` | +| Boltz | GPU | `boltz_score = 0.327` (ipTM); auto-rescore `vina_rescore_boltz_score = -4.336` | + +Notes: +- The DiffDock→Vina and Boltz→Vina auto-rescore tracks both fired correctly, confirming the + multi-method scoring wiring survived the database-strip of `bulk.py`. +- One benign Boltz PLIP warning (`Binding site LIG:Z:1 not found`) — a known ligand-resname + quirk, not a failure (run exited 0); unrelated to this upgrade. +- A `pandas` `DataFrameGroupBy.apply` FutureWarning appears across runs — pre-existing in + `tools/scores.py`, not introduced here. + +## How to reproduce the build + runs + +```sh +# Build the gnina bundle once (from the public gnina/gnina image), then the app image: +make docker-local # auto-builds gnina-bundle:local if missing, then guild:latest + +# Empirical docking (workspace holds data + combinations.csv at /workspace): +make run-vina PROJECT=testvina COMBINATIONS=/workspace/combinations.csv NO_DECOYS=1 +make run-gnina PROJECT=testgnina COMBINATIONS=/workspace/combinations.csv NO_DECOYS=1 +make run-guild PROJECT=testall COMBINATIONS=/workspace/combinations.csv METHODS="diffdock karmadock boltz" NO_DECOYS=1 +``` diff --git a/docs/adfrsuite_class.py b/docs/adfrsuite_class.py index 9d18027..c5ec1e3 100644 --- a/docs/adfrsuite_class.py +++ b/docs/adfrsuite_class.py @@ -414,7 +414,7 @@ def cleanUpResidues(self, mol, cleanup_list): hohs = hohs + hohs2 if hohs: # remove(hohs) - self.lenHOHS = len(hohs) + lenHOHS = self.lenHOHS = len(hohs) for h in hohs: for b in h.bonds: c = b.atom1 @@ -448,7 +448,7 @@ def cleanUpResidues(self, mol, cleanup_list): if len(mol.chains) > 1: chains_to_delete = [] for c in mol.chains: - len(c.residues) + num_res = len(c.residues) non_std = [] for res in c.residues: if res.type not in self.std_types: @@ -509,6 +509,7 @@ def detectIsPeptide(self): def findNearest(self, atom, bonded_atoms): lenK = len(bonded_atoms) + lenC = 1 nonbonded_atoms = AtomSet([atom]) c = numpy.array(nonbonded_atoms.coords, "f") k = numpy.array(bonded_atoms.coords, "f") @@ -556,13 +557,11 @@ def __init__( cleanup="nphs_lps_waters_nonstdres", outputfilename=None, debug=False, - preserved=None, + preserved={}, delete_single_nonstd_residues=False, dict=None, ): - if preserved is None: - preserved = {} AutoDockMoleculePreparation.__init__( self, molecule, @@ -601,7 +600,7 @@ def summarize(self): msg = "setup " + mol.name + ":\n" if self.newHs != 0: msg = msg + " added %d new hydrogen(s)\n" % self.newHs - if self.chargeType is not None: + if self.chargeType != None: if self.chargeType in ["pdbq", "mol2", "pdbqt", "pdbqs"]: msg = msg + " kept charges from " + self.chargeType + " file\n" else: @@ -610,7 +609,7 @@ def summarize(self): msg = msg + " merged %d non-polar hydrogens\n" % self.lenNPHS if self.lenLPS: msg = msg + " merged %d lone pairs\n" % self.lenLPS - numpy.add.reduce(mol.allAtoms.charge) + totalCh = numpy.add.reduce(mol.allAtoms.charge) if hasattr(self, "chargeError") and abs(self.chargeError) > 0.000005: msg = msg + " total charge error = %6.4f\n" % self.chargeError return msg @@ -657,13 +656,11 @@ def __init__( outputfilename=None, debug=False, version=4, - preserved=None, + preserved={}, delete_single_nonstd_residues=False, dict=None, ): - if preserved is None: - preserved = {} self.dict = dict AutoDockMoleculePreparation.__init__( self, @@ -789,6 +786,7 @@ def write(self, receptor, outputfile): assert len(list(filter(cond, receptor_atoms))) == len_receptor_atoms outptr = open(outputfile, "w") + records = ["ATOM"] # if self.write_CONECT: # records.append('CONECT') # for at in receptor_atoms: @@ -880,7 +878,7 @@ def __init__( dict=None, debug=False, check_for_fragments=False, - bonds_to_inactivate=None, + bonds_to_inactivate=[], inactivate_all_torsions=False, version=3, limit_torsions=False, @@ -889,8 +887,6 @@ def __init__( ): # why aren't repairs, cleanup and allowed_bonds lists?? # to run tests: use allowed_bonds = 'guanidinium' - if bonds_to_inactivate is None: - bonds_to_inactivate = [] if debug: print("LPO: allowed_bonds=", allowed_bonds) self.detect_bonds_between_cycles = detect_bonds_between_cycles @@ -929,7 +925,7 @@ def __init__( # detect aromatic cycle carbons and rename them for AutoDock3 rename = self.version == 3 - self.ACM = AromaticCarbonManager(rename=rename) + ACM = self.ACM = AromaticCarbonManager(rename=rename) self.aromCs = self.ACM.setAromaticCarbons(molecule) # optional output summary filename to write types and number of torsions @@ -949,7 +945,7 @@ def summarize(self): msg = "setup " + mol.name + ":\n" if self.newHs != 0: msg = msg + " added %d new hydrogen(s)\n" % self.newHs - if self.chargeType is not None: + if self.chargeType != None: if self.chargeType in ["pdbq", "mol2"]: msg = msg + " kept charges from " + self.chargeType + " file\n" else: @@ -960,7 +956,7 @@ def summarize(self): msg = msg + " merged %d lone pairs\n" % self.lenLPS if len(self.aromCs): msg = msg + " found %d aromatic carbons\n" % len(self.aromCs) - numpy.add.reduce(mol.allAtoms.charge) + totalCh = numpy.add.reduce(mol.allAtoms.charge) msg = msg + " detected %d rotatable bonds\n" % len(mol.possible_tors_bnds) msg = msg + " set TORSDOF to %d\n" % mol.TORSDOF # if hasattr(self, 'chargeError') and abs(self.chargeError)>0.000005: @@ -1111,6 +1107,7 @@ def toggle_torsion(self, ind1, ind2): self.RBM.toggle_torsion(ind1, ind2) def changePlanarityCriteria(self, cutoff): + oldAromCs = self.aromCs self.aromCs = self.ACM.setAromaticCarbons(self.molecule, cutoff) if self.debug: print("now there are ", len(self.aromCs), " aromaticCs") @@ -1145,7 +1142,7 @@ def __init__( dict=None, debug=False, check_for_fragments=False, - bonds_to_inactivate=None, + bonds_to_inactivate=[], inactivate_all_torsions=False, version=4, typeAtoms=True, @@ -1159,9 +1156,7 @@ def __init__( write_CONECT=False, ): - if bonds_to_inactivate is None: - bonds_to_inactivate = [] - if attach_nonbonded_fragments: + if attach_nonbonded_fragments == True: molecule.attach_nonbonded_fragments(attach_singletons=attach_singletons) # FIX THIS: what if molecule already has autodock_element set??? LigandPreparation.__init__( @@ -1307,6 +1302,7 @@ def write(self, ligand, outputfile): rootnum = rootnum + 1 # before writing, sort rootlist newrootlist = [] + oldrootlist = rootlist r_ctr = 0 for at in ligand.chains.residues.atoms: if hasattr(at, "rnum0"): @@ -1511,6 +1507,7 @@ def writeBreadthFirst(self, outfptr, fromAtom, startAtom): self.outatom_counter += 1 if self.debug: print("self.outatom_counter=", self.outatom_counter) + activeTors = [] # outfptr.write(outstring) for bond in startAtom.bonds: at2 = bond.atom1 @@ -1664,7 +1661,7 @@ class LigandRandomizer: def QuatToAxisAngle(self, q): # (x,y,z,w) assert len(q) == 4 - abs(q[3]) + abs_w = abs(q[3]) if q[3] > 1.0: q[3] = 1.0 if q[3] < -1.0: @@ -1678,7 +1675,7 @@ def QuatToAxisAngle(self, q): quat.append(q[0] * inv_sin_half_angle) quat.append(q[1] * inv_sin_half_angle) quat.append(q[2] * inv_sin_half_angle) - q[3] + w = q[3] # retval.ang = WrpModRad( angle );#line233 qmultiply.cc # define WrpModRad(angle) WrpRad(ModRad(angle)) # define WrpRad(angle) IF (((angle)> PI)? ((angle)-TWOPI) @@ -1743,7 +1740,7 @@ def __init__( ndihe = m.parser.keys.count("BRANCH") else: # do something else here @@ ??FILTER OUT ALL INACTIVE @@#1 - ndihe = len(m.allAtoms.bonds[0].get(lambda x: x.activeTors)) + ndihe = len(m.allAtoms.bonds[0].get(lambda x: x.activeTors == True)) if verbose: print("ndihe =", ndihe) if outputfilename is None: @@ -1781,7 +1778,7 @@ def __init__( orig_d[a] = set(a.bonds.getAtoms()) # setup random torsion angles: dihe = [] - for _ind in range(ndihe): + for ind in range(ndihe): dihe.append(round(random.uniform(-180, 180), 3)) if torsOnly: trans = [0, 0, 0] @@ -1875,6 +1872,7 @@ def __init__( print("quat=", quat) m.state["quat"] = quat # CONVERT TO AXISANGLE in order to use stoc: + retval = [] # to hold x,y,z,ang x, y, z, w = quat # TODO handle big W! Singularities line219 qmultiply.cc # check for identity @@ -1953,8 +1951,8 @@ def __init__( m.allAtoms.updateCoords(newCrds, ind=m.orig_coord_index) # remove all original bonds and set hasBonds to 0 on all levels # @@ REMEMBER WHICH ARE activeTors and not - aBD = {} - RotatableBondManager(m, allowed_bonds="backbone") # 5/2014 + aBD = activeBondDict = {} + rbm = RotatableBondManager(m, allowed_bonds="backbone") # 5/2014 for b in m.allAtoms.bonds[0]: v1 = b.possibleTors v2 = b.activeTors @@ -1973,7 +1971,7 @@ def __init__( a.hasBonds = 0 m.buildBondsByDistance() # do it again here?? - RotatableBondManager(m, allowed_bonds="backbone") # 5/2014 + rbm = RotatableBondManager(m, allowed_bonds="backbone") # 5/2014 newLen = len(m.allAtoms.bonds[0]) if verbose: print( @@ -2006,9 +2004,11 @@ def __init__( # reset bond activity: for k, v in list(aBD.items()): # k = (b.atom1, b.atom2); for b in k[0].bonds: + found = 0 b.possibleTors = 0 b.activeTors = 0 if b.neighborAtom(k[0]) == k[1]: + found = 1 b.possibleTors, b.activeTors = v if verbose: print( @@ -2067,6 +2067,7 @@ def __init__( molecule.has_backbone = allow_backbone_torsions allow_guanidinium_torsions = "guanidinium" in allowed_bond_list molecule.has_guanidinium = allow_guanidinium_torsions + allow_all_torsions = "all" in allowed_bond_list self.molecule = molecule self.debug = debug self.__classifyBonds(molecule.allAtoms, allow_guanidinium_torsions) @@ -2117,7 +2118,7 @@ def __classifyBonds(self, atoms, allow_guanidinium_torsions): # restore appropriate categories: if len(dict["leaf"]): dict["leaf"].leaf = 1 - dict["aromatic"] + aromBnds = dict["aromatic"] if len(dict["cycle"]): dict["cycle"].incycle = 1 if len(dict["rotatable"]): @@ -2565,7 +2566,7 @@ def _setTorsions(self, numTors, type): for k in range(1, numTors + 1): rangeList.append(-k) # turn them all off - self.set_torsions(mol, mol.allAtoms.bonds[0], 0) + ct = self.set_torsions(mol, mol.allAtoms.bonds[0], 0) # print "ct = ", ct, " and mol.torscount=", mol.torscount # torsionMap = mol.torTree.torsionMap # for i in range(tNum): @@ -2622,19 +2623,15 @@ def __init__( ligand, mode="automatic", rigid_filename=None, - residues=None, + residues=[], allowed_bonds="", # deprecated FlexibleDocking: 5/2014: backbone off, amide + guanidinium off - non_rotatable_bonds=None, + non_rotatable_bonds=[], flex_filename=None, reassign=True, debug=False, ): # ?NEED TO TYPE ATOMS? - if non_rotatable_bonds is None: - non_rotatable_bonds = [] - if residues is None: - residues = [] self.molecule = molecule self.debug = debug file_type = os.path.splitext(os.path.basename(molecule.parser.filename))[1] @@ -2849,6 +2846,7 @@ def writeSubtree(self, outfptr, fromAtom, startAtom): """ if startAtom.used == 0: startAtom.used = 1 + charge = startAtom.charge at = startAtom at.newindex = self.outatom_counter at.number = self.outatom_counter @@ -2975,19 +2973,15 @@ def __init__( receptor, mode="automatic", rigid_filename=None, - residues=None, + residues=[], allowed_bonds="backbone", # backbone on; amide + guanidinium off - non_rotatable_bonds=None, + non_rotatable_bonds=[], flexres_filename=None, reassign=True, debug=False, ): # ?NEED TO TYPE ATOMS? - if non_rotatable_bonds is None: - non_rotatable_bonds = [] - if residues is None: - residues = [] self.receptor = receptor self.debug = debug file_type = os.path.splitext(os.path.basename(receptor.parser.filename))[1] @@ -3047,17 +3041,17 @@ def write(self, outputfilenames=None): self.write_flex(self.flex_residues, flexres_filename) self.write_rigid(self.receptor, rigid_filename) self.flex_atoms_to_write = self.flex_residues.atoms.get( - lambda x: not hasattr(x, "used") + lambda x: hasattr(x, "used") == False ) self.rigid_atoms_to_write = self.rigid_residues.atoms.get( - lambda x: not hasattr(x, "used") + lambda x: hasattr(x, "used") == False ) if len(self.rigid_atoms_to_write): if self.debug: print("writing rigid atoms to ", rigid_filename) outptr = open(rigid_filename, "a") for at in self.rigid_atoms_to_write: - if not hasattr(at, "used") or not at.used: + if not hasattr(at, "used") or at.used == False: at.number = self.rigid_outctr self.writer.write_atom(outptr, at) self.rigid_outctr += 1 @@ -3079,7 +3073,7 @@ def write_rigid(self, molecule, rigid_filename): # try to write the non_rotatable_bonds near other atoms in same residue... if not len(resSet) or res not in resSet: # normal output for at in res.atoms: - if not hasattr(at, "used") or not at.used: + if not hasattr(at, "used") or at.used == False: at.number = self.rigid_outctr self.writer.write_atom(outptr, at) self.rigid_outctr += 1 @@ -3101,7 +3095,7 @@ def write_rigid(self, molecule, rigid_filename): ) in ( molecule.allAtoms ): # @@HACKED: write any atoms that haven't been written yet at the end of the rigid file... - if not at.used or hasattr(at, "used") is False: + if at.used == False or hasattr(at, "used") is False: self.writer.write_atom(outptr, at) self.rigid_outctr += 1 at.used = True @@ -3352,7 +3346,7 @@ def writeBreadthFirst(self, outfptr, fromAtom, startAtom): startAtom.used, ) queue = [] - if not startAtom.used: + if startAtom.used == False: startAtom.used = True startAtom.number = startAtom.newindex = self.outatom_counter self.writer.write_atom(outfptr, startAtom) @@ -3362,6 +3356,7 @@ def writeBreadthFirst(self, outfptr, fromAtom, startAtom): self.outatom_counter += 1 if self.debug: print("self.outatom_counter=", self.outatom_counter) + activeTors = [] # outfptr.write(outstring) for bond in startAtom.bonds: if not hasattr(bond, "activeTors"): @@ -3418,7 +3413,7 @@ def writeSubtree(self, outfptr, fromAtom, startAtom): for AutoDock. Specifically, appropriate BRANCH/ENDBRANCH statements are added. """ - if not startAtom.used: + if startAtom.used == False: startAtom.used = True at = startAtom for bond in startAtom.bonds: diff --git a/guild/bulk.py b/guild/bulk.py index aaeeb5d..4e0d31c 100644 --- a/guild/bulk.py +++ b/guild/bulk.py @@ -9,6 +9,7 @@ from concurrent.futures import ( TimeoutError as FuturesTimeoutError, ) +from concurrent.futures.process import BrokenProcessPool import pandas as pd from tqdm import tqdm @@ -34,7 +35,6 @@ PREPROCESSING_TIMEOUT, RP_SCORES_FILE, SMILES_NAMES_DICTIONARY_KEY, - SMILES_TYPE_DICTIONARY_KEY, ) from guild.constants.diffdock import ( DIFFDOCK_COMBINATIONS_FILE, @@ -44,8 +44,11 @@ ALL_AVAILABLE_METHODS, BOLTZ_FOLDER, BOLTZ_PREFIX, + BOX_LOCATION, DIFFDOCK_FOLDER, DIFFDOCK_PREFIX, + GNINA_FOLDER, + GNINA_PREFIX, KARMADOCK_FOLDER, KARMADOCK_PREFIX, LIGAND_CATEGORY, @@ -69,7 +72,8 @@ VINA_FOLDER, # Prefixes VINA_PREFIX, - VINA_RESCORE_PREFIX, + VINA_RESCORE_BOLTZ_PREFIX, + VINA_RESCORE_DIFFDOCK_PREFIX, ) from guild.constants.karmadock import ( KARMADOCK_DATA_FOLDER, @@ -79,6 +83,7 @@ from guild.constants.plip import ( COMPLEX_PDB_SUFFIX, PLIP_COMBINATION_ID, + PLIP_INTERACTION_COUNT_COLUMNS, PLIP_INTERACTIONS_FILE, PLIP_LIGAND_CHAIN, PLIP_LIGAND_RESNAME, @@ -91,32 +96,43 @@ boltz_guild_scoring, deploy_boltz, generate_boltz_yaml, + vina_rescore_boltz_guild_scoring, ) from guild.docking.diffdock import ( deploy_diffdock, diffdock_guild_scoring, + vina_rescore_diffdock_guild_scoring, write_diffdock_combinations_table, ) +from guild.docking.gnina import gnina_guild_scoring from guild.docking.karmadock import deploy_karmadock, karmadock_guild_scoring -from guild.docking.vina import vina_guild_scoring, vina_rescore_guild_scoring +from guild.docking.vina import ( + get_center_and_size_from_box_file, + vina_guild_scoring, +) from guild.run import Guild from guild.tools.bulk import ( available_methods_preparation, extend_all_combinations_table_with_decoys, extend_all_combinations_table_with_known_binders, + generate_boltz_complex_pdbs, generate_diffdock_complex_pdbs, + generate_gnina_complex_pdbs, generate_vina_complex_pdbs, identify_previously_ran_combinations, kickstart_batch_dictionary, ) -from guild.tools.preparation import clean_smiles +from guild.tools.preparation import _normalize_chain_list, clean_smiles from guild.tools.protein_sequence import ( get_original_sequence_dictionary, process_into_fasta_string, ) from guild.tools.scores import compute_rank_percentile_scores from guild.transformers.msa import fetch_protein_msa -from guild.transformers.pdb import get_pocket_contacts_from_ligand +from guild.transformers.pdb import ( + get_pocket_contacts_from_box, + get_pocket_contacts_from_ligand, +) from guild.visualization.plotting import ( bulk_plot_unique_proteins_scorings, bulk_plot_unique_scores, @@ -125,6 +141,22 @@ logger = logging.getLogger(__name__) +def _row_box_location(row): + """ + Extract a usable ``box_location`` path from a combinations-table row. + + Returns ``None`` when the column is absent or the value is missing / empty, + so downstream Vina/Boltz code falls back to existing pocket-derivation logic. + """ + if BOX_LOCATION not in row.index: + return None + value = row[BOX_LOCATION] + if value is None or (isinstance(value, float) and pd.isna(value)): + return None + value = str(value).strip() + return value or None + + class BulkRun: def __init__( self, @@ -141,6 +173,7 @@ def __init__( use_gpu=True, n_workers=None, predict_binding_pocket=False, + gnina_input_mode="pdbqt", ): """ BulkRun class for running multiple docking simulations and performing guild scoring. @@ -157,6 +190,13 @@ def __init__( :param use_gpu: Use GPU. :param n_workers: Number of parallel workers for multiprocessing. If None, uses CPU count. :param predict_binding_pocket: Use P2Rank for binding site prediction instead of original ligand location. + :param gnina_input_mode: ``"pdbqt"`` (default — preserves current + behaviour) or ``"sdf"``. When ``"sdf"`` AND gnina is the sole + docking method requested (no Vina or Vina-rescore co-running), + OpenBabel PDBQT prep is skipped and gnina runs directly on the + RDKit SDF + cleaned PDB. If ``"sdf"`` is requested alongside + Vina/Vina-rescore, this resolves back to ``"pdbqt"`` and a + warning is logged — see the constructor body for the rule. """ # Pathing variables @@ -186,9 +226,50 @@ def __init__( self.methods_to_run = methods_to_run - # Automatically enable Vina rescore when DiffDock is requested - if DIFFDOCK_PREFIX in self.methods_to_run and VINA_RESCORE_PREFIX not in self.methods_to_run: - self.methods_to_run = list(self.methods_to_run) + [VINA_RESCORE_PREFIX] + # Automatically enable Vina rescore for each pose-producing method that + # was requested. The Boltz-rescore and DiffDock-rescore tracks are + # independent — running both methods produces both columns. + # DiffDock gives many ranked poses; Boltz gives a confidence-only score + # (ipTM) which doesn't reflect binding energy — re-scoring the predicted + # complex with Vina's physics-based function provides a comparable ΔG. + if ( + DIFFDOCK_PREFIX in self.methods_to_run + and VINA_RESCORE_DIFFDOCK_PREFIX not in self.methods_to_run + ): + self.methods_to_run = list(self.methods_to_run) + [VINA_RESCORE_DIFFDOCK_PREFIX] + if ( + BOLTZ_PREFIX in self.methods_to_run + and VINA_RESCORE_BOLTZ_PREFIX not in self.methods_to_run + ): + self.methods_to_run = list(self.methods_to_run) + [VINA_RESCORE_BOLTZ_PREFIX] + + # Resolve gnina_input_mode against the (now auto-extended) methods + # list. SDF-mode only applies when gnina is the *sole* PDBQT-relevant + # method requested; if Vina or any of the Vina-rescore tracks are in + # the mix, OpenBabel has to run for them anyway, so we silently fall + # back to PDBQT for gnina (with a warning so the caller knows their + # upstream protonation work will be overwritten by Gasteiger). + _pdbqt_requiring = { + VINA_PREFIX, + VINA_RESCORE_BOLTZ_PREFIX, + VINA_RESCORE_DIFFDOCK_PREFIX, + } + if ( + gnina_input_mode == "sdf" + and GNINA_PREFIX in self.methods_to_run + and not _pdbqt_requiring.intersection(self.methods_to_run) + ): + self.gnina_input_mode = "sdf" + else: + if gnina_input_mode == "sdf": + logger.warning( + "SDF-mode requested for gnina but Vina/Vina-rescore is also " + "in methods_to_run — falling back to PDBQT for gnina. " + "OpenBabel ligand prep will still run, and any upstream " + "ligand protonation will be overwritten by OpenBabel's " + "Gasteiger pathway." + ) + self.gnina_input_mode = "pdbqt" logger.info(f"Methods to run: {self.methods_to_run}") @@ -325,6 +406,23 @@ def _generate_batched_dictionary(self, input_table, batch_size=1000): ) logger.info(f"Prepared available methods for batch {current_batch}") + @staticmethod + def _log_progress(*paths, message): + """ + Append a timestamped line to one or more progress-log files. The + per-batch path is ``{batch}/output.log`` (same file Guild workers + write to via ``logging.basicConfig``, so there's one log per batch + rather than two); the project-level path is + ``{project}/batch_progress.log``. Failures are duplicated to both. + ``None`` entries are skipped so callers can pass an unused slot. + """ + line = f"{message} at {pd.Timestamp.now()}\n" + for path in paths: + if path is None: + continue + with open(path, "a") as f: + f.write(line) + @staticmethod def _run_single_vina_docking(task_params): """ @@ -351,6 +449,8 @@ def _run_single_vina_docking(task_params): use_gpu=task_params["use_gpu"], is_bulk=True, predict_binding_pocket=task_params.get("predict_binding_pocket", False), + box_location=task_params.get("box_location"), + gnina_input_mode=task_params.get("gnina_input_mode", "pdbqt"), ) guild_object.run_autodock_vina() return task_params["ligand_idx"], task_params["protein_idx"] @@ -361,6 +461,42 @@ def _run_single_vina_docking(task_params): ) return None + @staticmethod + def _run_single_gnina_docking(task_params): + """ + Run gnina docking for a single combination (worker function for multiprocessing). + """ + try: + logger.info( + f"Starting gnina docking for SMILES: {task_params['ligand_smile']} " + f"(ligand_idx: {task_params['ligand_idx']}, protein_idx: {task_params['protein_idx']})" + ) + + guild_object = Guild( + ligand_smile=task_params["ligand_smile"], + ligand_idx=task_params["ligand_idx"], + protein_idx=task_params["protein_idx"], + protein_file=task_params["protein_file"], + project_name=task_params["project_name"], + protein_chain=task_params["protein_chain"], + original_ligand=task_params["original_ligand"], + original_ligand_chain=task_params["original_ligand_chain"], + output_log_file=task_params["output_log_file"], + use_gpu=task_params["use_gpu"], + is_bulk=True, + predict_binding_pocket=task_params.get("predict_binding_pocket", False), + box_location=task_params.get("box_location"), + gnina_input_mode=task_params.get("gnina_input_mode", "pdbqt"), + ) + guild_object.run_gnina() + return task_params["ligand_idx"], task_params["protein_idx"] + except Exception as e: + logger.error( + f"Error in gnina docking for SMILES {task_params['ligand_smile']} " + f"(ligand_idx: {task_params['ligand_idx']}): {e}" + ) + return None + @staticmethod def _prepare_single_batch_worker(batch_params): """ @@ -381,6 +517,8 @@ def _prepare_single_batch_worker(batch_params): use_gpu=batch_params["use_gpu"], is_bulk=True, predict_binding_pocket=batch_params.get("predict_binding_pocket", False), + box_location=batch_params.get("box_location"), + gnina_input_mode=batch_params.get("gnina_input_mode", "pdbqt"), ) return batch_params["batch_name"] except Exception as e: @@ -410,27 +548,50 @@ def _prepare_docking(self): # This will download the necessary raw pdb files beforehand - current_subset_row = self.batched_dictionary[current_batch][ + combinations_table = self.batched_dictionary[current_batch][ COMBINATIONS_TABLE_KEY - ].iloc[0] - batch_params_list.append( - { - "batch_name": current_batch, - "ligand_smile": self.batched_dictionary[current_batch][ - SMILES_NAMES_DICTIONARY_KEY - ][current_subset_row[LIGAND_ID]], - "ligand_idx": current_subset_row[LIGAND_ID], - "protein_idx": current_subset_row[PROTEIN_CONF_ID], - "protein_file": current_subset_row[PROTEIN_PATH], - "project_name": f"{self.project_name}/{BATCHES_FOLDER}/{current_batch}", - "protein_chain": current_subset_row[PROTEIN_CHAIN], - "original_ligand": current_subset_row[ORIGINAL_LIGAND], - "original_ligand_chain": current_subset_row[ORIGINAL_LIGAND_CHAIN], - "output_log_file": f"{self.batched_dictionary[current_batch][BATCH_FOLDER]}/{OUTPUT_LOG_FILE}", - "use_gpu": self.use_gpu, - "predict_binding_pocket": self.predict_binding_pocket, - } - ) + ] + # KarmaDock has no per-ligand worker downstream (deploy_karmadock + # runs once per batch on an already-staged data dir), so the prep + # pool must Guild-init every combination. + # Vina, Boltz, gnina and DiffDock have per-ligand or per-batch + # downstream workers that re-prep individual ligands themselves, + # but they all depend on per-protein artifacts that Guild.__init__ + # writes under {batch}/proteins/{protein_idx}_single_chain_clean.pdb. + # Those outputs are keyed on protein_idx only, so we Guild-init one + # row per unique protein in the batch — enough to cover every + # protein's cleaned PDB / PDBQT without wasting work on batches + # with many ligands per protein. Picking only iloc[[0]] previously + # left N-1 proteins unprepped in multi-protein batches, which made + # Boltz silently fall back to the raw template PDB and fail. + if KARMADOCK_PREFIX in self.methods_to_run: + rows_to_prep = combinations_table.iterrows() + else: + rows_to_prep = combinations_table.drop_duplicates( + subset=[PROTEIN_CONF_ID], keep="first" + ).iterrows() + + for _, current_row in rows_to_prep: + batch_params_list.append( + { + "batch_name": current_batch, + "ligand_smile": self.batched_dictionary[current_batch][ + SMILES_NAMES_DICTIONARY_KEY + ][current_row[LIGAND_ID]], + "ligand_idx": current_row[LIGAND_ID], + "protein_idx": current_row[PROTEIN_CONF_ID], + "protein_file": current_row[PROTEIN_PATH], + "project_name": f"{self.project_name}/{BATCHES_FOLDER}/{current_batch}", + "protein_chain": current_row[PROTEIN_CHAIN], + "original_ligand": current_row[ORIGINAL_LIGAND], + "original_ligand_chain": current_row[ORIGINAL_LIGAND_CHAIN], + "output_log_file": f"{self.batched_dictionary[current_batch][BATCH_FOLDER]}/{OUTPUT_LOG_FILE}", + "use_gpu": self.use_gpu, + "predict_binding_pocket": self.predict_binding_pocket, + "box_location": _row_box_location(current_row), + "gnina_input_mode": self.gnina_input_mode, + } + ) with ProcessPoolExecutor(max_workers=self.n_workers) as executor: # Submit all preparation tasks @@ -459,331 +620,796 @@ def _prepare_docking(self): def run_docking(self): """ - Perform docking simulations using AutoDock Vina, KarmaDock, and DiffDock. + Perform docking simulations using the methods listed in + ``self.methods_to_run``. Per-batch method blocks dispatch in the + order the user supplied — e.g. ``methods_to_run=["vina", "boltz"]`` + docks Vina first, ``["boltz", "vina"]`` docks Boltz first. Auto-added + ``vina_rescore_*`` entries are no-ops here (their work runs during + scoring). """ self._prepare_docking() - # Create main progress log file main_progress_log = f"{self.project_folder}/{BATCH_PROGRESS_LOG_FILE}" + # Per-method per-batch runners. Methods not in this map (such as the + # vina_rescore_* auto-additions) are skipped during docking. + method_runners = { + BOLTZ_PREFIX: self._run_boltz_for_batch, + VINA_PREFIX: self._run_vina_for_batch, + GNINA_PREFIX: self._run_gnina_for_batch, + KARMADOCK_PREFIX: self._run_karmadock_for_batch, + DIFFDOCK_PREFIX: self._run_diffdock_for_batch, + } + for current_batch in tqdm(self.batched_dictionary, desc="Running docking"): current_batch_folder = self.batched_dictionary[current_batch][BATCH_FOLDER] + batch_progress_log = f"{current_batch_folder}/{OUTPUT_LOG_FILE}" + + # Mirror the batch-start marker to both logs: the project-level log + # gives a one-line-per-batch overview; the per-batch log opens with + # the same line so it's easy to align timing between the two. + self._log_progress( + main_progress_log, batch_progress_log, + message=f"Starting {current_batch}", + ) - # Log progress to main log file - with open(main_progress_log, "a") as f: - f.write(f"Starting {current_batch} at {pd.Timestamp.now()}\n") - - # Skip batches with no new combinations if len(self.batched_dictionary[current_batch][COMBINATIONS_TO_RUN_KEY]) == 0: logger.info(f"Skipping {current_batch} - no new combinations to dock") - with open(main_progress_log, "a") as f: - f.write( - f"Skipped {current_batch} - no new combinations at {pd.Timestamp.now()}\n" - ) + self._log_progress( + main_progress_log, batch_progress_log, + message=f"Skipped {current_batch} - no new combinations", + ) continue - if BOLTZ_PREFIX in self.methods_to_run: - # Botlz is applied in the whole batch at once, so we need to iterate over the unique protein configuration ids - combinations_table_variable = self.batched_dictionary[current_batch][ - COMBINATIONS_TABLE_KEY + for method in self.methods_to_run: + runner = method_runners.get(method) + if runner is None: + continue + runner(current_batch, current_batch_folder) + + self._log_progress( + main_progress_log, batch_progress_log, + message=f"Completed {current_batch}", + ) + + def _run_boltz_for_batch(self, current_batch, current_batch_folder): + """Run Boltz docking for the batch, then generate complex PDBs.""" + batch_progress_log = f"{current_batch_folder}/{OUTPUT_LOG_FILE}" + main_progress_log = f"{self.project_folder}/{BATCH_PROGRESS_LOG_FILE}" + self._log_progress(batch_progress_log, message="Starting Boltz") + + # Boltz is applied in the whole batch at once, so we need to iterate over the unique protein configuration ids + combinations_table_variable = self.batched_dictionary[current_batch][ + COMBINATIONS_TABLE_KEY + ] + for unique_protein_configuration_id in combinations_table_variable[ + PROTEIN_CONF_ID + ].unique(): + + current_protein_combinations = combinations_table_variable[ + combinations_table_variable[PROTEIN_CONF_ID] + == unique_protein_configuration_id + ] + # ``protein_chain`` may name one chain ("A") or several ("A,B") for a + # pocket that spans multiple chains. Keep only chains actually present + # in the structure, aligned with their sequences and MSAs. + current_protein_chains = _normalize_chain_list( + current_protein_combinations[PROTEIN_CHAIN].values[0] + ) + current_protein_path = current_protein_combinations[PROTEIN_PATH].values[0] + original_sequence_dictionary = get_original_sequence_dictionary( + current_protein_path + ) + current_protein_chains = [ + c for c in current_protein_chains if c in original_sequence_dictionary + ] + current_protein_sequences = [ + process_into_fasta_string(original_sequence_dictionary[c]) + for c in current_protein_chains + ] + current_protein_unique_ligands_ids = list( + current_protein_combinations[LIGAND_ID].unique() + ) + + # Fetch one MSA per chain — cached across all ligands and batches + msa_files = [ + fetch_protein_msa( + sequence=sequence, + protein_id=unique_protein_configuration_id, + protein_chain_id=chain, + output_a3m_dir=self.msa_cache_dir, + ) + for chain, sequence in zip( + current_protein_chains, current_protein_sequences, strict=True + ) + ] + logger.info( + f"Fetched MSA(s) for protein configuration {unique_protein_configuration_id} " + f"(chains {current_protein_chains})" + ) + + # Compute binding-pocket contacts once per protein configuration. + # Precedence: user-supplied box_location > original-ligand derivation. + pocket_contacts = [] + _box_path = None + if BOX_LOCATION in current_protein_combinations.columns: + _box_values = [ + _row_box_location(r) + for _, r in current_protein_combinations.iterrows() ] - for unique_protein_configuration_id in combinations_table_variable[ - PROTEIN_CONF_ID - ].unique(): - - current_protein_combinations = combinations_table_variable[ - combinations_table_variable[PROTEIN_CONF_ID] - == unique_protein_configuration_id - ] - current_protein_chain = current_protein_combinations[PROTEIN_CHAIN].values[0] - current_protein_path = current_protein_combinations[PROTEIN_PATH].values[0] - original_sequence_dictionary = get_original_sequence_dictionary( - current_protein_path + _non_empty_box_values = [b for b in _box_values if b] + _unique_boxes = set(_non_empty_box_values) + if _non_empty_box_values: + # Select deterministically using input order. + _box_path = _non_empty_box_values[0] + else: + _unique_boxes = set() + + if len(_unique_boxes) > 1: + logger.warning( + f"Inconsistent box_location values for protein " + f"{unique_protein_configuration_id}: {sorted(_unique_boxes)}. " + f"Boltz pocket is per-protein; using {_box_path}." + ) + if _box_path: + if os.path.exists(_box_path): + _center, _size = get_center_and_size_from_box_file(_box_path) + pocket_contacts = get_pocket_contacts_from_box( + protein_pdb=current_protein_path, + protein_chain=current_protein_chains, + center=_center, + size=_size, ) - current_protein_sequence = process_into_fasta_string( - original_sequence_dictionary[current_protein_chain] + logger.info( + f"Derived {len(pocket_contacts)} Boltz pocket contacts from " + f"box {_box_path} for protein {unique_protein_configuration_id}" ) - current_protein_unique_ligands_ids = list( - current_protein_combinations[LIGAND_ID].unique() + else: + logger.warning( + f"box_location {_box_path} not found on disk; falling back " + f"to original-ligand pocket derivation." ) - # Fetch MSA once per protein — cached across all ligands and batches - msa_file = fetch_protein_msa( - sequence=current_protein_sequence, - protein_id=unique_protein_configuration_id, - protein_chain_id=current_protein_chain, - output_a3m_dir=self.msa_cache_dir, + if not pocket_contacts: + _orig_lig = current_protein_combinations[ORIGINAL_LIGAND].values[0] + _orig_lig_chain = current_protein_combinations[ + ORIGINAL_LIGAND_CHAIN + ].values[0] + if ( + _orig_lig + and _orig_lig_chain + and pd.notna(_orig_lig) + and pd.notna(_orig_lig_chain) + ): + pocket_contacts = get_pocket_contacts_from_ligand( + protein_pdb=current_protein_path, + protein_chain=current_protein_chains, + original_ligand=_orig_lig, + original_ligand_chain=_orig_lig_chain, + distance_threshold=4.0, ) + + # Run one Boltz job per ligand to avoid O(N²) attention memory growth + for current_ligand_id in current_protein_unique_ligands_ids: + ligand_smiles = self.batched_dictionary[current_batch][ + SMILES_NAMES_DICTIONARY_KEY + ][current_ligand_id] + run_id = f"{unique_protein_configuration_id}_{current_ligand_id}" + + # Saving time by skipping already run Boltz dockings - checking for one of the expected output files. + # Boltz writes outputs under {BOLTZ_FOLDER}/boltz_results_{run_id}_boltz/predictions/{run_id}_boltz/. + boltz_result_fild = ( + f"{BOLTZ_FOLDER}/boltz_results_{run_id}_boltz" + f"/predictions/{run_id}_boltz/confidence_{run_id}_boltz_model_0.json" + ) + if os.path.exists(f"{current_batch_folder}/{boltz_result_fild}"): logger.info( - f"Fetched MSA for protein configuration {unique_protein_configuration_id} (chain {current_protein_chain})" + f"Boltz docking already ran for {current_batch}: " + f"protein {unique_protein_configuration_id}, ligand {current_ligand_id}" ) + continue - # Compute binding-pocket contacts once per protein configuration - _orig_lig = current_protein_combinations[ORIGINAL_LIGAND].values[0] - _orig_lig_chain = current_protein_combinations[ORIGINAL_LIGAND_CHAIN].values[0] - pocket_contacts = [] - if ( - _orig_lig - and _orig_lig_chain - and pd.notna(_orig_lig) - and pd.notna(_orig_lig_chain) - ): - pocket_contacts = get_pocket_contacts_from_ligand( - protein_pdb=current_protein_path, - protein_chain=current_protein_chain, - original_ligand=_orig_lig, - original_ligand_chain=_orig_lig_chain, - distance_threshold=4.0, - ) + os.makedirs(f"{current_batch_folder}/{BOLTZ_FOLDER}", exist_ok=True) - # Run one Boltz job per ligand to avoid O(N²) attention memory growth - for current_ligand_id in current_protein_unique_ligands_ids: - ligand_smiles = self.batched_dictionary[current_batch][ - SMILES_NAMES_DICTIONARY_KEY - ][current_ligand_id] - run_id = f"{unique_protein_configuration_id}_{current_ligand_id}" + # Use single-chain PDB as template to avoid Boltz multi-chain parsing errors + boltz_template_file = ( + f"{current_batch_folder}/{PROTEINS_FOLDER}/" + f"{unique_protein_configuration_id}_single_chain_clean.pdb" + ) + if not os.path.exists(boltz_template_file): + # Fallback to original protein path if single-chain not available + boltz_template_file = current_protein_path + + yaml_file = f"{current_batch_folder}/{BOLTZ_FOLDER}/{run_id}_boltz.yaml" + boltz_out_dir = f"{current_batch_folder}/{BOLTZ_FOLDER}" + + generate_boltz_yaml( + protein_sequence=current_protein_sequences, + protein_chain=current_protein_chains, + ligand_sequences=[ligand_smiles], + ligand_ids=["L"], + output_file=yaml_file, + template_file=boltz_template_file, + pocket_contacts=pocket_contacts if pocket_contacts else None, + msa_file=msa_files, + ) - # Saving time by skipping already run Boltz dockings - checking for one of the expected output files - boltz_result_fild = ( - f"{BOLTZ_FOLDER}/predictions/{run_id}_boltz/confidence_{run_id}_boltz_model_0.json" - ) - if os.path.exists(f"{current_batch_folder}/{boltz_result_fild}"): - logger.info( - f"Boltz docking already ran for {current_batch}: " - f"protein {unique_protein_configuration_id}, ligand {current_ligand_id}" - ) - continue + # Per-combination subprocess log path. Boltz can exit 0 and + # still produce an empty manifest, so we keep the transcript + # regardless of success/failure for downstream troubleshooting. + boltz_subprocess_log = ( + f"{boltz_out_dir}/{run_id}.subprocess.log" + ) - os.makedirs(f"{current_batch_folder}/{BOLTZ_FOLDER}", exist_ok=True) + deploy_boltz( + yaml_file, + out_dir=boltz_out_dir, + use_gpu=self.use_gpu, + subprocess_log_path=boltz_subprocess_log, + ) - # Use single-chain PDB as template to avoid Boltz multi-chain parsing errors - boltz_template_file = ( - f"{current_batch_folder}/{PROTEINS_FOLDER}/" - f"{unique_protein_configuration_id}_single_chain_clean.pdb" + # Check if Boltz produced valid output (manifest with records). + # Template PDB parsing can fail silently in Boltz2, resulting + # in an empty manifest. If that happens, retry without the template. + manifest_path = ( + f"{boltz_out_dir}/boltz_results_{run_id}_boltz/processed/manifest.json" + ) + if os.path.exists(manifest_path): + import json as _json + + with open(manifest_path) as _mf: + _manifest = _json.load(_mf) + if not _manifest.get("records"): + logger.warning( + f"Boltz2 produced empty manifest for {run_id} " + "(likely template parsing failure). Retrying without template..." ) - if not os.path.exists(boltz_template_file): - # Fallback to original protein path if single-chain not available - boltz_template_file = current_protein_path + # Remove the failed output directory + import shutil - yaml_file = f"{current_batch_folder}/{BOLTZ_FOLDER}/{run_id}_boltz.yaml" - boltz_out_dir = f"{current_batch_folder}/{BOLTZ_FOLDER}" + failed_dir = f"{boltz_out_dir}/boltz_results_{run_id}_boltz" + if os.path.isdir(failed_dir): + shutil.rmtree(failed_dir) + # Regenerate YAML without template generate_boltz_yaml( - protein_sequence=current_protein_sequence, - protein_chain=current_protein_chain, + protein_sequence=current_protein_sequences, + protein_chain=current_protein_chains, ligand_sequences=[ligand_smiles], ligand_ids=["L"], output_file=yaml_file, - template_file=boltz_template_file, + template_file=None, pocket_contacts=pocket_contacts if pocket_contacts else None, - msa_file=msa_file, + msa_file=msa_files, ) + # Retry overwrites the same subprocess.log — the + # latter run is the one whose outcome we report on. deploy_boltz( yaml_file, out_dir=boltz_out_dir, use_gpu=self.use_gpu, + subprocess_log_path=boltz_subprocess_log, ) - # Check if Boltz produced valid output (manifest with records). - # Template PDB parsing can fail silently in Boltz2, resulting - # in an empty manifest. If that happens, retry without the template. - manifest_path = ( - f"{boltz_out_dir}/boltz_results_{run_id}_boltz/processed/manifest.json" - ) - if os.path.exists(manifest_path): - import json as _json - - with open(manifest_path) as _mf: - _manifest = _json.load(_mf) - if not _manifest.get("records"): - logger.warning( - f"Boltz2 produced empty manifest for {run_id} " - "(likely template parsing failure). Retrying without template..." - ) - # Remove the failed output directory - import shutil - - failed_dir = f"{boltz_out_dir}/boltz_results_{run_id}_boltz" - if os.path.isdir(failed_dir): - shutil.rmtree(failed_dir) - - # Regenerate YAML without template - generate_boltz_yaml( - protein_sequence=current_protein_sequence, - protein_chain=current_protein_chain, - ligand_sequences=[ligand_smiles], - ligand_ids=["L"], - output_file=yaml_file, - template_file=None, - pocket_contacts=pocket_contacts if pocket_contacts else None, - msa_file=msa_file, - ) - - deploy_boltz( - yaml_file, - out_dir=boltz_out_dir, - use_gpu=self.use_gpu, - ) - - logger.info( - f"Boltz docking completed for {current_batch}: " - f"protein {unique_protein_configuration_id}, ligand {current_ligand_id}" + # After any retry, verify the manifest is non-empty. If not, + # surface this combo as a failure to both progress logs so it's + # visible without diving into Boltz's processed/ tree. + final_failure_reason = None + if not os.path.exists(manifest_path): + final_failure_reason = "no manifest" + else: + import json as _json + + with open(manifest_path) as _mf: + _final_manifest = _json.load(_mf) + if not _final_manifest.get("records"): + final_failure_reason = "empty manifest after retry" + if final_failure_reason is not None: + self._log_progress( + batch_progress_log, + main_progress_log, + message=( + f"FAILED Boltz {run_id} ({final_failure_reason}) " + f"— see {boltz_subprocess_log}" + ), + ) + + logger.info( + f"Boltz docking completed for {current_batch}: " + f"protein {unique_protein_configuration_id}, ligand {current_ligand_id}" + ) + + # Generate complex PDB files — needed both for PLIP and for + # Vina re-scoring of the predicted pose. + logger.info(f"Generating complex PDB files for Boltz results in {current_batch}") + generate_boltz_complex_pdbs(self.batched_dictionary[current_batch]) + + self._log_progress(batch_progress_log, message="Completed Boltz") + + + def _run_vina_for_batch(self, current_batch, current_batch_folder): + """Run Vina docking for the batch, then generate complex PDBs.""" + batch_progress_log = f"{current_batch_folder}/{OUTPUT_LOG_FILE}" + main_progress_log = f"{self.project_folder}/{BATCH_PROGRESS_LOG_FILE}" + self._log_progress(batch_progress_log, message="Starting Vina") + + # Collect all combinations to run + vina_tasks = [] + for _index, current_row in self.batched_dictionary[current_batch][ + COMBINATIONS_TABLE_KEY + ].iterrows(): + protein_path = current_row[PROTEIN_PATH] + current_protein, current_ligand = ( + current_row[PROTEIN_CONF_ID], + current_row[LIGAND_ID], + ) + vina_final_path = f"{self.batched_dictionary[current_batch][BATCH_FOLDER]}/{VINA_FOLDER}/{current_protein}_{current_ligand}.txt" + if not os.path.exists(vina_final_path): + vina_tasks.append( + { + "ligand_smile": self.batched_dictionary[current_batch][ + SMILES_NAMES_DICTIONARY_KEY + ][current_ligand], + "ligand_idx": current_ligand, + "protein_idx": current_protein, + "protein_file": protein_path, + "project_name": f"{self.project_name}/{BATCHES_FOLDER}/{current_batch}", + "protein_chain": current_row[PROTEIN_CHAIN], + "original_ligand": current_row[ORIGINAL_LIGAND], + "original_ligand_chain": current_row[ORIGINAL_LIGAND_CHAIN], + "output_log_file": f"{current_batch_folder}/{OUTPUT_LOG_FILE}", + "use_gpu": self.use_gpu, + "predict_binding_pocket": self.predict_binding_pocket, + "box_location": _row_box_location(current_row), + "gnina_input_mode": self.gnina_input_mode, + } + ) + + if vina_tasks: + logger.info( + f"Running {len(vina_tasks)} Vina docking tasks with {self.n_workers} workers" + ) + # OpenBabel can raise a C-level abort during ligand prep that + # kills the worker process. Once that happens, every still-pending + # future in the same ProcessPoolExecutor raises BrokenProcessPool — + # losing the rest of the batch. Recover by restarting the pool + # with the remaining tasks, skipping the one that crashed. + remaining = list(vina_tasks) + completed = 0 + pool_restarts = 0 + max_pool_restarts = len(vina_tasks) # hard upper bound + + while remaining: + succeeded_indices = set() + pool_broken = False + + with ProcessPoolExecutor(max_workers=self.n_workers) as executor: + future_to_task = {} + for idx, task in enumerate(remaining): + future = executor.submit( + self._run_single_vina_docking, task ) + future_to_task[future] = (idx, task) - if VINA_PREFIX in self.methods_to_run: - # Collect all combinations to run - vina_tasks = [] - for _index, current_row in self.batched_dictionary[current_batch][ - COMBINATIONS_TABLE_KEY - ].iterrows(): - protein_path = current_row[PROTEIN_PATH] - current_protein, current_ligand = ( - current_row[PROTEIN_CONF_ID], - current_row[LIGAND_ID], + desc_suffix = ( + f" (restart {pool_restarts})" if pool_restarts else "" ) - vina_final_path = f"{self.batched_dictionary[current_batch][BATCH_FOLDER]}/{VINA_FOLDER}/{current_protein}_{current_ligand}.txt" - if not os.path.exists(vina_final_path): - vina_tasks.append( - { - "ligand_smile": self.batched_dictionary[current_batch][ - SMILES_NAMES_DICTIONARY_KEY - ][current_ligand], - "ligand_idx": current_ligand, - "protein_idx": current_protein, - "protein_file": protein_path, - "project_name": f"{self.project_name}/{BATCHES_FOLDER}/{current_batch}", - "protein_chain": current_row[PROTEIN_CHAIN], - "original_ligand": current_row[ORIGINAL_LIGAND], - "original_ligand_chain": current_row[ORIGINAL_LIGAND_CHAIN], - "output_log_file": f"{current_batch_folder}/{OUTPUT_LOG_FILE}", - "use_gpu": self.use_gpu, - "predict_binding_pocket": self.predict_binding_pocket, - } - ) + for future in tqdm( + as_completed(future_to_task, timeout=None), + desc=( + f"AutoDock Vina for batch: {current_batch}" + f"{desc_suffix}" + ), + total=len(future_to_task), + ): + idx, task = future_to_task[future] + try: + future.result(timeout=DOCKING_TIMEOUT) + completed += 1 + succeeded_indices.add(idx) + except FuturesTimeoutError: + logger.warning( + f"Task {idx + 1}/{len(remaining)} timed out after " + f"{DOCKING_TIMEOUT}s (SMILES: " + f"{task.get('ligand_smile', 'unknown')})" + ) + succeeded_indices.add(idx) # drop from retry + future.cancel() + self._log_progress( + batch_progress_log, + main_progress_log, + message=f"FAILED Vina {task['protein_idx']}_{task['ligand_idx']} (timeout)", + ) + except BrokenProcessPool: + # Worker died (typically OpenBabel C++ abort). + # The current/queued futures all raise this; + # we identify the actual crasher after the loop. + pool_broken = True + except Exception as e: + logger.error( + f"Task {idx + 1}/{len(remaining)} failed: {e} " + f"(SMILES: {task.get('ligand_smile', 'unknown')})" + ) + succeeded_indices.add(idx) # drop from retry + self._log_progress( + batch_progress_log, + main_progress_log, + message=f"FAILED Vina {task['protein_idx']}_{task['ligand_idx']} ({e})", + ) - if vina_tasks: - logger.info( - f"Running {len(vina_tasks)} Vina docking tasks with {self.n_workers} workers" + if not pool_broken: + break + + not_succeeded = sorted( + set(range(len(remaining))) - succeeded_indices + ) + if not not_succeeded: + break + crasher = remaining[not_succeeded[0]] + self._log_progress( + batch_progress_log, + main_progress_log, + message=( + f"FAILED Vina {crasher['protein_idx']}_{crasher['ligand_idx']} " + "(BrokenProcessPool — presumed OpenBabel C++ abort)" + ), + ) + logger.warning( + f"BrokenProcessPool: skipping presumed crasher " + f"ligand={crasher['ligand_idx']} " + f"SMILES={crasher.get('ligand_smile', 'unknown')}; " + f"restarting pool with {len(not_succeeded) - 1} remaining tasks" + ) + remaining = [remaining[i] for i in not_succeeded[1:]] + pool_restarts += 1 + if pool_restarts > max_pool_restarts: + logger.error( + f"Hit pool-restart cap ({max_pool_restarts}); " + f"bailing with {len(remaining)} tasks unfinished." ) - with ProcessPoolExecutor(max_workers=self.n_workers) as executor: - # Submit all docking tasks - store task info for better error reporting - futures = {} - for i, task in enumerate(vina_tasks): - future = executor.submit(self._run_single_vina_docking, task) - futures[future] = (i, task) # Store index and full task dict - - # Process results as they complete, with per-task timeout - completed = 0 - for future in tqdm( - as_completed(futures, timeout=None), - desc=f"AutoDock Vina for batch: {current_batch}", - total=len(futures), - ): - task_num, task_dict = futures[future] - try: - future.result(timeout=DOCKING_TIMEOUT) - completed += 1 - except FuturesTimeoutError: - logger.warning( - f"Task {task_num+1}/{len(vina_tasks)} timed out after {DOCKING_TIMEOUT} seconds " - f"(SMILES: {task_dict.get('ligand_smile', 'unknown')})" - ) - future.cancel() # Attempt to cancel - except Exception as e: - logger.error( - f"Task {task_num+1}/{len(vina_tasks)} failed: {e} " - f"(SMILES: {task_dict.get('ligand_smile', 'unknown')})" - ) - - logger.info(f"Completed {completed}/{len(vina_tasks)} Vina docking tasks") - else: - logger.info(f"No new Vina combinations to dock for batch {current_batch}") + break + + logger.info( + f"Completed {completed}/{len(vina_tasks)} Vina docking tasks " + f"(pool restarts: {pool_restarts})" + ) + else: + logger.info(f"No new Vina combinations to dock for batch {current_batch}") - # Generate complex PDB files for Vina results - logger.info(f"Generating complex PDB files for Vina results in {current_batch}") - generate_vina_complex_pdbs(self.batched_dictionary[current_batch]) + # Generate complex PDB files for Vina results + logger.info(f"Generating complex PDB files for Vina results in {current_batch}") + generate_vina_complex_pdbs(self.batched_dictionary[current_batch]) - # if BOLTZ_PREFIX in self.methods_to_run: - # # (Boltz docking loop is earlier in this function) - # # Generate complex PDB files for PLIP analysis - # logger.info(f"Generating complex PDB files for Boltz results in {current_batch}") - # generate_boltz_complex_pdbs(self.batched_dictionary[current_batch]) + self._log_progress(batch_progress_log, message="Completed Vina") - # Generate complex PDB files for DiffDock results (for PLIP) - if DIFFDOCK_PREFIX in self.methods_to_run: - diffdock_results_dir = os.path.join( - self.batched_dictionary[current_batch][BATCH_FOLDER], - DIFFDOCK_FOLDER, - DIFFDOCK_RESULTS_FOLDER, + + def _run_gnina_for_batch(self, current_batch, current_batch_folder): + """Run gnina docking for the batch, then generate complex PDBs.""" + batch_progress_log = f"{current_batch_folder}/{OUTPUT_LOG_FILE}" + main_progress_log = f"{self.project_folder}/{BATCH_PROGRESS_LOG_FILE}" + self._log_progress(batch_progress_log, message="Starting gnina") + + gnina_tasks = [] + for _index, current_row in self.batched_dictionary[current_batch][ + COMBINATIONS_TABLE_KEY + ].iterrows(): + protein_path = current_row[PROTEIN_PATH] + current_protein, current_ligand = ( + current_row[PROTEIN_CONF_ID], + current_row[LIGAND_ID], + ) + gnina_final_path = ( + f"{self.batched_dictionary[current_batch][BATCH_FOLDER]}/" + f"{GNINA_FOLDER}/{current_protein}_{current_ligand}.txt" + ) + if not os.path.exists(gnina_final_path): + gnina_tasks.append( + { + "ligand_smile": self.batched_dictionary[current_batch][ + SMILES_NAMES_DICTIONARY_KEY + ][current_ligand], + "ligand_idx": current_ligand, + "protein_idx": current_protein, + "protein_file": protein_path, + "project_name": f"{self.project_name}/{BATCHES_FOLDER}/{current_batch}", + "protein_chain": current_row[PROTEIN_CHAIN], + "original_ligand": current_row[ORIGINAL_LIGAND], + "original_ligand_chain": current_row[ORIGINAL_LIGAND_CHAIN], + "output_log_file": f"{current_batch_folder}/{OUTPUT_LOG_FILE}", + "use_gpu": self.use_gpu, + "predict_binding_pocket": self.predict_binding_pocket, + "box_location": _row_box_location(current_row), + "gnina_input_mode": self.gnina_input_mode, + } ) - if os.path.isdir(diffdock_results_dir): - logger.info( - f"Generating complex PDB files for DiffDock results in {current_batch}" - ) - generate_diffdock_complex_pdbs(self.batched_dictionary[current_batch]) - # Karmadock and DiffDock are triggered by a single run + if gnina_tasks: + logger.info( + f"Running {len(gnina_tasks)} gnina docking tasks with {self.n_workers} workers" + ) + # gnina shares Vina's PDBQT ligand prep, so the same OpenBabel + # C-level abort risk applies; reuse the BrokenProcessPool + # recovery loop with a different worker function. + remaining = list(gnina_tasks) + completed = 0 + pool_restarts = 0 + max_pool_restarts = len(gnina_tasks) + + while remaining: + succeeded_indices = set() + pool_broken = False + + with ProcessPoolExecutor(max_workers=self.n_workers) as executor: + future_to_task = {} + for idx, task in enumerate(remaining): + future = executor.submit( + self._run_single_gnina_docking, task + ) + future_to_task[future] = (idx, task) - if KARMADOCK_PREFIX in self.methods_to_run: - if ( - os.path.exists( - f"{self.batched_dictionary[current_batch][BATCH_FOLDER]}/{KARMADOCK_FOLDER}/{KARMADOCK_RESULTS_FOLDER}/0.csv" - ) - and os.path.exists( - f"{self.batched_dictionary[current_batch][BATCH_FOLDER]}/{KARMADOCK_FOLDER}/{KARMADOCK_RESULTS_FOLDER}/1.csv" - ) - and os.path.exists( - f"{self.batched_dictionary[current_batch][BATCH_FOLDER]}/{KARMADOCK_FOLDER}/{KARMADOCK_RESULTS_FOLDER}/2.csv" - ) - ): - logger.info( - f"Karmadock docking already ran for {current_batch} with {len(self.batched_dictionary[current_batch][COMBINATIONS_TO_RUN_KEY])} combinations" + desc_suffix = ( + f" (restart {pool_restarts})" if pool_restarts else "" ) - else: - deploy_karmadock( - home_path=self.home_path, - karmadock_results_dir=f"{self.batched_dictionary[current_batch][BATCH_FOLDER]}/{KARMADOCK_FOLDER}/{KARMADOCK_RESULTS_FOLDER}", - karmadock_graphs_dir=f"{self.batched_dictionary[current_batch][BATCH_FOLDER]}/{KARMADOCK_FOLDER}/{KARMADOCK_GRAPHS_FOLDER}", - karmadock_data_dir=f"{self.batched_dictionary[current_batch][BATCH_FOLDER]}/{KARMADOCK_FOLDER}/{KARMADOCK_DATA_FOLDER}", - ) - logger.info( - f"Karmadock docking completed for {current_batch} with {len(self.batched_dictionary[current_batch][COMBINATIONS_TO_RUN_KEY])} combinations" + for future in tqdm( + as_completed(future_to_task, timeout=None), + desc=( + f"gnina for batch: {current_batch}" + f"{desc_suffix}" + ), + total=len(future_to_task), + ): + idx, task = future_to_task[future] + try: + future.result(timeout=DOCKING_TIMEOUT) + completed += 1 + succeeded_indices.add(idx) + except FuturesTimeoutError: + logger.warning( + f"gnina task {idx + 1}/{len(remaining)} timed out after " + f"{DOCKING_TIMEOUT}s (SMILES: " + f"{task.get('ligand_smile', 'unknown')})" + ) + succeeded_indices.add(idx) + future.cancel() + gnina_log = ( + f"{current_batch_folder}/{GNINA_FOLDER}/" + f"{task['protein_idx']}_{task['ligand_idx']}.subprocess.log" + ) + self._log_progress( + batch_progress_log, + main_progress_log, + message=( + f"FAILED gnina {task['protein_idx']}_{task['ligand_idx']} " + f"(timeout) — see {gnina_log}" + ), + ) + except BrokenProcessPool: + pool_broken = True + except Exception as e: + logger.error( + f"gnina task {idx + 1}/{len(remaining)} failed: {e} " + f"(SMILES: {task.get('ligand_smile', 'unknown')})" + ) + succeeded_indices.add(idx) + gnina_log = ( + f"{current_batch_folder}/{GNINA_FOLDER}/" + f"{task['protein_idx']}_{task['ligand_idx']}.subprocess.log" + ) + self._log_progress( + batch_progress_log, + main_progress_log, + message=( + f"FAILED gnina {task['protein_idx']}_{task['ligand_idx']} " + f"({e}) — see {gnina_log}" + ), + ) + + if not pool_broken: + break + + not_succeeded = sorted( + set(range(len(remaining))) - succeeded_indices + ) + if not not_succeeded: + break + crasher = remaining[not_succeeded[0]] + logger.warning( + f"BrokenProcessPool (gnina): skipping presumed crasher " + f"ligand={crasher['ligand_idx']} " + f"SMILES={crasher.get('ligand_smile', 'unknown')}; " + f"restarting pool with {len(not_succeeded) - 1} remaining tasks" + ) + gnina_crash_log = ( + f"{current_batch_folder}/{GNINA_FOLDER}/" + f"{crasher['protein_idx']}_{crasher['ligand_idx']}.subprocess.log" + ) + self._log_progress( + batch_progress_log, + main_progress_log, + message=( + f"FAILED gnina {crasher['protein_idx']}_{crasher['ligand_idx']} " + f"(BrokenProcessPool — presumed OpenBabel C++ abort) " + f"— see {gnina_crash_log} (may not exist if the crash happened pre-subprocess)" + ), + ) + remaining = [remaining[i] for i in not_succeeded[1:]] + pool_restarts += 1 + if pool_restarts > max_pool_restarts: + logger.error( + f"Hit gnina pool-restart cap ({max_pool_restarts}); " + f"bailing with {len(remaining)} tasks unfinished." ) + break - if DIFFDOCK_PREFIX in self.methods_to_run: + logger.info( + f"Completed {completed}/{len(gnina_tasks)} gnina docking tasks " + f"(pool restarts: {pool_restarts})" + ) + else: + logger.info(f"No new gnina combinations to dock for batch {current_batch}") + + # Generate complex PDB files for gnina results (shares Vina's PDBQT format) + logger.info(f"Generating complex PDB files for gnina results in {current_batch}") + generate_gnina_complex_pdbs(self.batched_dictionary[current_batch]) + + self._log_progress(batch_progress_log, message="Completed gnina") + + + def _run_karmadock_for_batch(self, current_batch, current_batch_folder): + """Run KarmaDock for the batch (single-shot, batch-level).""" + batch_progress_log = f"{current_batch_folder}/{OUTPUT_LOG_FILE}" + self._log_progress(batch_progress_log, message="Starting KarmaDock") + + if ( + os.path.exists( + f"{self.batched_dictionary[current_batch][BATCH_FOLDER]}/{KARMADOCK_FOLDER}/{KARMADOCK_RESULTS_FOLDER}/0.csv" + ) + and os.path.exists( + f"{self.batched_dictionary[current_batch][BATCH_FOLDER]}/{KARMADOCK_FOLDER}/{KARMADOCK_RESULTS_FOLDER}/1.csv" + ) + and os.path.exists( + f"{self.batched_dictionary[current_batch][BATCH_FOLDER]}/{KARMADOCK_FOLDER}/{KARMADOCK_RESULTS_FOLDER}/2.csv" + ) + ): + logger.info( + f"Karmadock docking already ran for {current_batch} with {len(self.batched_dictionary[current_batch][COMBINATIONS_TO_RUN_KEY])} combinations" + ) + else: + deploy_karmadock( + home_path=self.home_path, + karmadock_results_dir=f"{self.batched_dictionary[current_batch][BATCH_FOLDER]}/{KARMADOCK_FOLDER}/{KARMADOCK_RESULTS_FOLDER}", + karmadock_graphs_dir=f"{self.batched_dictionary[current_batch][BATCH_FOLDER]}/{KARMADOCK_FOLDER}/{KARMADOCK_GRAPHS_FOLDER}", + karmadock_data_dir=f"{self.batched_dictionary[current_batch][BATCH_FOLDER]}/{KARMADOCK_FOLDER}/{KARMADOCK_DATA_FOLDER}", + ) + logger.info( + f"Karmadock docking completed for {current_batch} with {len(self.batched_dictionary[current_batch][COMBINATIONS_TO_RUN_KEY])} combinations" + ) + + self._log_progress(batch_progress_log, message="Completed KarmaDock") + + + def _run_diffdock_for_batch(self, current_batch, current_batch_folder): + """Run DiffDock for the batch. Generates complex PDBs on resume, then docks.""" + batch_progress_log = f"{current_batch_folder}/{OUTPUT_LOG_FILE}" + main_progress_log = f"{self.project_folder}/{BATCH_PROGRESS_LOG_FILE}" + self._log_progress(batch_progress_log, message="Starting DiffDock") + + diffdock_results_dir = os.path.join( + self.batched_dictionary[current_batch][BATCH_FOLDER], + DIFFDOCK_FOLDER, + DIFFDOCK_RESULTS_FOLDER, + ) + if os.path.isdir(diffdock_results_dir): + logger.info( + f"Generating complex PDB files for DiffDock results in {current_batch}" + ) + generate_diffdock_complex_pdbs(self.batched_dictionary[current_batch]) + + batch_csv = f"{current_batch_folder}/{DIFFDOCK_COMBINATIONS_FILE}" + batch_results_dir = ( + f"{current_batch_folder}/{DIFFDOCK_FOLDER}/{DIFFDOCK_RESULTS_FOLDER}" + ) + + # DiffDock writes rank{N}_confidence-{score}.sdf files into + # batch_results_dir/{protein_conf_id}_{ligand_id}/. A batch is + # only "done" when every combination has at least one such SDF. + # Checking len(os.listdir(batch_results_dir)) > 0 is not enough: + # an interrupted run leaves per-combo subfolders that are empty + # (or absent), and the next run would otherwise skip DiffDock + # and produce all-NaN scores downstream. + combinations_df = self.batched_dictionary[current_batch][ + COMBINATIONS_TABLE_KEY + ] - batch_csv = f"{current_batch_folder}/{DIFFDOCK_COMBINATIONS_FILE}" - batch_results_dir = ( - f"{current_batch_folder}/{DIFFDOCK_FOLDER}/{DIFFDOCK_RESULTS_FOLDER}" + missing = [] + for _, row in combinations_df.iterrows(): + combo_dir = os.path.join( + batch_results_dir, + f"{row[PROTEIN_CONF_ID]}_{row[LIGAND_ID]}", + ) + has_pose = os.path.isdir(combo_dir) and any( + "_confidence" in fname and fname.endswith(".sdf") + for fname in os.listdir(combo_dir) + ) + if not has_pose: + missing.append((row[PROTEIN_CONF_ID], row[LIGAND_ID])) + + if not missing: + logger.info( + f"DiffDock batch already ran for {current_batch} " + f"({len(combinations_df)} combinations have poses on disk)" + ) + else: + logger.info( + f"DiffDock for {current_batch}: " + f"{len(missing)}/{len(combinations_df)} combinations " + f"missing docked poses; re-running batch" + ) + # DiffDock runs once per batch (not per combo), so the subprocess + # log is batch-level. All per-combo failures point at this same + # file for context. + diffdock_subprocess_log = ( + f"{current_batch_folder}/{DIFFDOCK_FOLDER}/_batch.subprocess.log" + ) + deploy_diffdock( + home_path=self.home_path, + diffdock_results_dir=batch_results_dir, + input_csv=batch_csv, + subprocess_log_path=diffdock_subprocess_log, + # use_gpu=self.use_gpu, + ) + logger.info( + f"DiffDock batch completed for {current_batch} with " + f"{len(self.batched_dictionary[current_batch][COMBINATIONS_TO_RUN_KEY])} " + f"combinations" + ) + + # Generate per-combination complex PDBs from the SDFs DiffDock + # just wrote. The resume-path call above only fires when the + # results dir already existed at entry, so on a fresh run we'd + # otherwise never get complex PDBs — which then leaves PLIP + # with nothing to analyze for DiffDock. The generator is + # idempotent (skips combos whose *_complex.pdb already exists). + if os.path.isdir(batch_results_dir): + logger.info( + f"Generating complex PDB files for DiffDock results in {current_batch} (post-docking)" ) + generate_diffdock_complex_pdbs(self.batched_dictionary[current_batch]) - # If batch already produced results, skip - if os.path.exists(batch_results_dir) and len(os.listdir(batch_results_dir)) > 0: - logger.info(f"DiffDock batch already ran for {current_batch}") - else: - # Run DiffDock ONCE for the whole batch - deploy_diffdock( - home_path=self.home_path, - diffdock_results_dir=batch_results_dir, - input_csv=batch_csv, - # use_gpu=self.use_gpu, + # Re-check which combos still lack a docked pose after the re-run + # and surface them as explicit failures in both progress logs. + for protein_conf_id, ligand_id in missing: + combo_dir = os.path.join( + batch_results_dir, f"{protein_conf_id}_{ligand_id}" + ) + still_missing = not ( + os.path.isdir(combo_dir) + and any( + "_confidence" in fname and fname.endswith(".sdf") + for fname in os.listdir(combo_dir) ) - logger.info( - f"DiffDock batch completed for {current_batch} with {len(self.batched_dictionary[current_batch][COMBINATIONS_TO_RUN_KEY])} combinations" + ) + if still_missing: + self._log_progress( + batch_progress_log, + main_progress_log, + message=( + f"FAILED DiffDock {protein_conf_id}_{ligand_id} " + f"(no _confidence SDF produced) — see {diffdock_subprocess_log}" + ), ) - # Log batch completion to main progress log - with open(main_progress_log, "a") as f: - f.write(f"Completed {current_batch} at {pd.Timestamp.now()}\n") + self._log_progress(batch_progress_log, message="Completed DiffDock") + def _compute_global_ranks_per_protein(self, all_raw_scores): """ - Compute ranks and rank percentile scores PER PROTEIN across ALL data (all batches + database). + Compute ranks and rank percentile scores PER PROTEIN across ALL data (all batches). This ensures rankings are consistent across batches - all ligands for a given protein are ranked together, not separately per batch. - :param all_raw_scores: DataFrame with raw scores from all batches and database. + :param all_raw_scores: DataFrame with raw scores from all batches. :return: DataFrame with ranks and rank percentile scores added. """ return compute_rank_percentile_scores(all_raw_scores, methods=self.methods_to_run) @@ -796,9 +1422,17 @@ def _process_batch_scoring(self, batch): :param batch: Batch to process. :return: Raw scores table for this batch (COMBINATION_ID, PROTEIN_ID, LIGAND_ID, raw scores). """ + batch_folder = self.batched_dictionary[batch][BATCH_FOLDER] + batch_progress_log = f"{batch_folder}/{OUTPUT_LOG_FILE}" + self._log_progress(batch_progress_log, message="Starting scoring") + logger.info(f"Collecting raw scores for batch {batch}") if len(self.batched_dictionary[batch]["combinations_to_run"]) == 0: logger.info("No new combinations to score") + self._log_progress( + batch_progress_log, + message="Completed scoring (no new combinations)", + ) return None docked_methods = [] @@ -806,19 +1440,33 @@ def _process_batch_scoring(self, batch): if VINA_PREFIX in self.methods_to_run: docked_methods.append(vina_guild_scoring(self.batched_dictionary[batch])) + if GNINA_PREFIX in self.methods_to_run: + docked_methods.append(gnina_guild_scoring(self.batched_dictionary[batch])) + if KARMADOCK_PREFIX in self.methods_to_run: docked_methods.append(karmadock_guild_scoring(self.batched_dictionary[batch])) if DIFFDOCK_PREFIX in self.methods_to_run: docked_methods.append(diffdock_guild_scoring(self.batched_dictionary[batch])) - if VINA_RESCORE_PREFIX in self.methods_to_run: - docked_methods.append(vina_rescore_guild_scoring(self.batched_dictionary[batch])) + if VINA_RESCORE_DIFFDOCK_PREFIX in self.methods_to_run: + docked_methods.append( + vina_rescore_diffdock_guild_scoring(self.batched_dictionary[batch]) + ) + + if VINA_RESCORE_BOLTZ_PREFIX in self.methods_to_run: + docked_methods.append( + vina_rescore_boltz_guild_scoring(self.batched_dictionary[batch]) + ) if BOLTZ_PREFIX in self.methods_to_run: docked_methods.append(boltz_guild_scoring(self.batched_dictionary[batch])) if not docked_methods: + self._log_progress( + batch_progress_log, + message="Completed scoring (no methods produced output)", + ) return None # Merge raw scores from all methods @@ -837,6 +1485,8 @@ def _process_batch_scoring(self, batch): lambda x: self.batched_dictionary[batch][SMILES_NAMES_DICTIONARY_KEY].get(x, "unknown") ) + self._log_progress(batch_progress_log, message="Completed scoring") + # Return RAW scores only - no ranks or rank percentile scores yet return raw_scores_table @@ -845,9 +1495,8 @@ def run_guild_scoring(self, n_processes=None): Scoring function to uniformize multiple docking methods by leveraging the decoy dataset. Steps: 1. Collect RAW scores from all batches (parallel) - 2. Merge with existing database scores (if any) - 3. Compute ranks PER PROTEIN across ALL data (not per batch) - 4. Compute rank percentile scores from global ranks + 2. Compute ranks PER PROTEIN across ALL data (not per batch) + 3. Compute rank percentile scores from global ranks :param n_processes: Number of processes to use for multiprocessing. """ @@ -855,6 +1504,9 @@ def run_guild_scoring(self, n_processes=None): if n_processes is None: n_processes = mp.cpu_count() + main_progress_log = f"{self.project_folder}/{BATCH_PROGRESS_LOG_FILE}" + self._log_progress(main_progress_log, message="Starting scoring") + # Step 1: Collect raw scores from all batches (parallel) if len(self.batched_dictionary) == 0: logger.info("No new batches to process") @@ -882,38 +1534,49 @@ def run_guild_scoring(self, n_processes=None): logger.warning("No scores to process") return - # Ensure ligand category is available before scoring - smiles_type_dictionary = self.unpack_dictionary(what_to_unpack=SMILES_TYPE_DICTIONARY_KEY) + # Ensure ligand category is available before scoring. + # Key the lookup on ligand_id (authoritative per-row category from the + # combinations table), NOT on SMILES. The SMILES-keyed dictionary both + # (a) loses the per-row category when two ligands share a SMILES, and + # (b) silently yields "unknown" whenever an upstream key mishap leaves + # SMILES itself as "unknown". + category_by_ligand = dict( + zip( + self.all_combinations_table[LIGAND_ID], + self.all_combinations_table[LIGAND_CATEGORY], + strict=False, + ) + ) if LIGAND_CATEGORY not in all_raw_scores.columns: - all_raw_scores[LIGAND_CATEGORY] = all_raw_scores[SMILES].apply( - lambda x: smiles_type_dictionary.get(x, "unknown") + all_raw_scores[LIGAND_CATEGORY] = ( + all_raw_scores[LIGAND_ID].map(category_by_ligand).fillna("unknown") ) - # Deduplicate: keep the row that actually has a score (scored rows first, then drop dupes). - score_cols = [SCORES_DICTIONARY[m] for m in self.methods_to_run if SCORES_DICTIONARY[m] in all_raw_scores.columns] - if score_cols: - all_raw_scores = ( - all_raw_scores - .assign(_has_score=all_raw_scores[score_cols].notna().any(axis=1)) - .sort_values("_has_score", ascending=False) - .drop_duplicates(subset=[COMBINATION_ID], keep="first") - .drop(columns=["_has_score"]) - .reset_index(drop=True) - ) + # Collapse any duplicate COMBINATION_ID rows by *coalescing* columns + # (first non-null per column) rather than keeping a single row. Within a + # batch the per-method outer-merge already yields one row per + # combination; coalescing protects against any concatenation that pairs + # rows carrying a complementary set of method-score columns, merging them + # instead of letting one row win and dropping the other's scores. + all_raw_scores = ( + all_raw_scores + .groupby(COMBINATION_ID, as_index=False, sort=False) + .first() + ) - # Step 3: Compute ranks PER PROTEIN across ALL data + # Step 2: Compute ranks PER PROTEIN across ALL data logger.info( f"Computing global ranks per protein for {all_raw_scores.shape[0]} total combinations" ) self.rp_scores_df = self._compute_global_ranks_per_protein(all_raw_scores) - # Step 4: Add ligand category + # Step 3: Add ligand category if LIGAND_CATEGORY not in self.rp_scores_df.columns: - self.rp_scores_df[LIGAND_CATEGORY] = self.rp_scores_df[SMILES].apply( - lambda x: smiles_type_dictionary.get(x, "unknown") + self.rp_scores_df[LIGAND_CATEGORY] = ( + self.rp_scores_df[LIGAND_ID].map(category_by_ligand).fillna("unknown") ) - # Step 5: Save results + # Step 4: Save results self.rp_scores_df.to_csv( self.rp_scores_path, index=False, @@ -921,6 +1584,8 @@ def run_guild_scoring(self, n_processes=None): ) logger.info(f"Saved rank percentile scores to {self.rp_scores_path}") + self._log_progress(main_progress_log, message="Completed scoring") + def unpack_dictionary(self, what_to_unpack="smiles_type_dictionary"): """ Unpack the smiles type dictionary. @@ -1100,8 +1765,59 @@ def run_interactions_analysis(self, methods_to_analyze=None): else: logger.info(f"No DiffDock complex PDB files found for {current_batch}") + if GNINA_PREFIX in methods_to_analyze: + gnina_folder = f"{batch_folder}/{GNINA_FOLDER}" + + complex_metadata = [] + for _, row in combinations_df.iterrows(): + protein_conf_id = row[PROTEIN_CONF_ID] + ligand_id = row[LIGAND_ID] + combination_id = f"{protein_conf_id}_{ligand_id}" + complex_pdb = f"{gnina_folder}/{combination_id}{COMPLEX_PDB_SUFFIX}" + if os.path.exists(complex_pdb): + smiles = batch_dict[SMILES_NAMES_DICTIONARY_KEY][ligand_id] + complex_metadata.append((complex_pdb, protein_conf_id, smiles)) + + if complex_metadata: + logger.info( + f"Analyzing {len(complex_metadata)} gnina complexes in {current_batch}" + ) + batch_interactions = analyze_batch_interactions( + complex_pdb_paths=[x[0] for x in complex_metadata], + combination_ids=[f"{x[1]}_{x[2]}" for x in complex_metadata], + ligand_resname=PLIP_LIGAND_RESNAME, + ligand_chain=PLIP_LIGAND_CHAIN, + ligand_resseq=PLIP_LIGAND_RESSEQ, + ) + if not batch_interactions.empty: + batch_interactions[PROTEIN_CONF_ID] = [x[1] for x in complex_metadata] + batch_interactions[SMILES] = [x[2] for x in complex_metadata] + all_interactions.append(batch_interactions) + else: + logger.info(f"No gnina complex PDB files found for {current_batch}") + + interactions_path = f"{self.project_folder}/{PLIP_INTERACTIONS_FILE}" + if not all_interactions: - logger.warning("No interaction data collected") + # Contract: external consumers (e.g. host-side notebooks that read + # the file with pandas) should always be able to read this path + # after a guild run, even if no method produced complex PDBs. Write + # a header-only TSV with the columns analyze_batch_interactions + # emits so downstream code paths are deterministic. + logger.warning( + "No interaction data collected — writing header-only TSV to " + f"{interactions_path} so the contract 'file always exists' holds." + ) + empty_cols = [ + PROTEIN_CONF_ID, + SMILES, + PLIP_COMBINATION_ID, + *PLIP_INTERACTION_COUNT_COLUMNS, + PLIP_TOTAL_INTERACTIONS, + PLIP_N_UNIQUE_RESIDUES, + ] + self.interactions_df = pd.DataFrame(columns=empty_cols) + self.interactions_df.to_csv(interactions_path, sep="\t", index=False) return None # Combine all batch results @@ -1121,8 +1837,7 @@ def run_interactions_analysis(self, methods_to_analyze=None): ] df_to_save = self.interactions_df[cols_order] - # Save to file - interactions_path = f"{self.project_folder}/{PLIP_INTERACTIONS_FILE}" + # Save to file (interactions_path computed at the top of this method) df_to_save.to_csv(interactions_path, sep="\t", index=False) logger.info(f"Saved interaction analysis to {interactions_path}") diff --git a/guild/constants/bulk.py b/guild/constants/bulk.py index a18e366..fbb6916 100644 --- a/guild/constants/bulk.py +++ b/guild/constants/bulk.py @@ -3,11 +3,15 @@ BOLTZ_SCORE, DIFFDOCK_PREFIX, DIFFDOCK_SCORE, + GNINA_PREFIX, + GNINA_SCORE, KARMADOCK_PREFIX, KARMADOCK_SCORE, VINA_PREFIX, - VINA_RESCORE_PREFIX, - VINA_RESCORE_SCORE, + VINA_RESCORE_BOLTZ_PREFIX, + VINA_RESCORE_BOLTZ_SCORE, + VINA_RESCORE_DIFFDOCK_PREFIX, + VINA_RESCORE_DIFFDOCK_SCORE, VINA_SCORE, ) @@ -86,19 +90,25 @@ KARMADOCK_RP_SCORE = f"rp_{KARMADOCK_SCORE}" DIFFDOCK_RP_SCORE = f"rp_{DIFFDOCK_SCORE}" BOLTZ_RP_SCORE = f"rp_{BOLTZ_SCORE}" -VINA_RESCORE_RP_SCORE = f"rp_{VINA_RESCORE_SCORE}" +GNINA_RP_SCORE = f"rp_{GNINA_SCORE}" +VINA_RESCORE_BOLTZ_RP_SCORE = f"rp_{VINA_RESCORE_BOLTZ_SCORE}" +VINA_RESCORE_DIFFDOCK_RP_SCORE = f"rp_{VINA_RESCORE_DIFFDOCK_SCORE}" RANK_VINA_SCORE = f"rank_{VINA_SCORE}" RANK_KARMADOCK_SCORE = f"rank_{KARMADOCK_SCORE}" RANK_DIFFDOCK_SCORE = f"rank_{DIFFDOCK_SCORE}" RANK_BOLTZ_SCORE = f"rank_{BOLTZ_SCORE}" -RANK_VINA_RESCORE_SCORE = f"rank_{VINA_RESCORE_SCORE}" +RANK_GNINA_SCORE = f"rank_{GNINA_SCORE}" +RANK_VINA_RESCORE_BOLTZ_SCORE = f"rank_{VINA_RESCORE_BOLTZ_SCORE}" +RANK_VINA_RESCORE_DIFFDOCK_SCORE = f"rank_{VINA_RESCORE_DIFFDOCK_SCORE}" SCORES_DIRECTION_DICTIONARY = { VINA_PREFIX: "minimum", KARMADOCK_PREFIX: "maximum", DIFFDOCK_PREFIX: "maximum", BOLTZ_PREFIX: "maximum", - VINA_RESCORE_PREFIX: "minimum", + GNINA_PREFIX: "minimum", + VINA_RESCORE_BOLTZ_PREFIX: "minimum", + VINA_RESCORE_DIFFDOCK_PREFIX: "minimum", } RANKS_DICTIONARY = { @@ -106,7 +116,9 @@ KARMADOCK_PREFIX: RANK_KARMADOCK_SCORE, DIFFDOCK_PREFIX: RANK_DIFFDOCK_SCORE, BOLTZ_PREFIX: RANK_BOLTZ_SCORE, - VINA_RESCORE_PREFIX: RANK_VINA_RESCORE_SCORE, + GNINA_PREFIX: RANK_GNINA_SCORE, + VINA_RESCORE_BOLTZ_PREFIX: RANK_VINA_RESCORE_BOLTZ_SCORE, + VINA_RESCORE_DIFFDOCK_PREFIX: RANK_VINA_RESCORE_DIFFDOCK_SCORE, } RP_SCORES_DICTIONARY = { @@ -114,7 +126,9 @@ KARMADOCK_PREFIX: KARMADOCK_RP_SCORE, DIFFDOCK_PREFIX: DIFFDOCK_RP_SCORE, BOLTZ_PREFIX: BOLTZ_RP_SCORE, - VINA_RESCORE_PREFIX: VINA_RESCORE_RP_SCORE, + GNINA_PREFIX: GNINA_RP_SCORE, + VINA_RESCORE_BOLTZ_PREFIX: VINA_RESCORE_BOLTZ_RP_SCORE, + VINA_RESCORE_DIFFDOCK_PREFIX: VINA_RESCORE_DIFFDOCK_RP_SCORE, } SCORES_TO_USE_DICTIONARY = { @@ -138,9 +152,19 @@ RANK_BOLTZ_SCORE, BOLTZ_RP_SCORE, ], - VINA_RESCORE_PREFIX: [ - VINA_RESCORE_SCORE, - RANK_VINA_RESCORE_SCORE, - VINA_RESCORE_RP_SCORE, + GNINA_PREFIX: [ + GNINA_SCORE, + RANK_GNINA_SCORE, + GNINA_RP_SCORE, + ], + VINA_RESCORE_BOLTZ_PREFIX: [ + VINA_RESCORE_BOLTZ_SCORE, + RANK_VINA_RESCORE_BOLTZ_SCORE, + VINA_RESCORE_BOLTZ_RP_SCORE, + ], + VINA_RESCORE_DIFFDOCK_PREFIX: [ + VINA_RESCORE_DIFFDOCK_SCORE, + RANK_VINA_RESCORE_DIFFDOCK_SCORE, + VINA_RESCORE_DIFFDOCK_RP_SCORE, ], } diff --git a/guild/constants/gnina.py b/guild/constants/gnina.py new file mode 100644 index 0000000..6c79c39 --- /dev/null +++ b/guild/constants/gnina.py @@ -0,0 +1,31 @@ +""" +GNINA constants +""" + +GNINA_BINARY = "/opt/gnina/bin/gnina" + +# gnina ships its own torch/openbabel/boost; we put them under /opt/gnina/lib +# so they only get loaded for the gnina subprocess. Prepend (not replace) to +# the caller's LD_LIBRARY_PATH at invocation time. +GNINA_LIB_PATH = "/opt/gnina/lib" + +# gnina's bundled libopenbabel.so.7 (OB 3.0/3.1) has NO format-plugin tree in +# the image, so its Open Babel can't load the writer for ANY output format — +# every ``--out`` file comes out empty (scores still work because gnina reads +# pdbqt with its own native parser). The only complete plugin set belongs to +# the system Open Babel 3.2 (.so.8) at the paths below. We make gnina load that +# .so.8 under the .so.7 name via a symlink shim (OB 3.1→3.2 is ABI-compatible +# for gnina's calls — verified to leave scores byte-identical) and point its OB +# at the matching plugin/data dirs. See ``_ensure_openbabel_plugin_shim`` in +# guild/docking/gnina.py. +GNINA_OB_SYSTEM_LIB = "/usr/local/lib/libopenbabel.so.8" +GNINA_OB_PLUGIN_DIR = "/usr/local/lib/openbabel/3.2.0" +GNINA_OB_DATA_DIR = "/usr/local/share/openbabel/3.2.0" + +GNINA_DEFAULT_EXHAUSTIVENESS = 8 + +GNINA_DEFAULT_NUMBER_OF_POSES = 9 + +# gnina's default --cnn_scoring is "rescore": Vina-style search produces poses +# and the CNN rescores the top N. Fastest mode that still surfaces a CNN score. +GNINA_DEFAULT_CNN_SCORING = "rescore" diff --git a/guild/constants/guild.py b/guild/constants/guild.py index 590174b..57b8e29 100644 --- a/guild/constants/guild.py +++ b/guild/constants/guild.py @@ -18,6 +18,7 @@ PROTEIN_CHAIN = "protein_chain" ORIGINAL_LIGAND = "original_ligand" ORIGINAL_LIGAND_CHAIN = "original_ligand_chain" +BOX_LOCATION = "box_location" """ @@ -27,12 +28,23 @@ KARMADOCK_PREFIX = "karmadock" DIFFDOCK_PREFIX = "diffdock" BOLTZ_PREFIX = "boltz" -VINA_RESCORE_PREFIX = "vina_rescore" +GNINA_PREFIX = "gnina" +# Vina re-scoring is now split by upstream pose source so a run that produces +# both Boltz and DiffDock complexes gets two distinct re-score columns. +VINA_RESCORE_BOLTZ_PREFIX = "vina_rescore_boltz" +VINA_RESCORE_DIFFDOCK_PREFIX = "vina_rescore_diffdock" VINA_SCORE = f"{VINA_PREFIX}_score" KARMADOCK_SCORE = f"{KARMADOCK_PREFIX}_score" DIFFDOCK_SCORE = f"{DIFFDOCK_PREFIX}_score" BOLTZ_SCORE = f"{BOLTZ_PREFIX}_score" -VINA_RESCORE_SCORE = f"{VINA_RESCORE_PREFIX}_score" +GNINA_SCORE = f"{GNINA_PREFIX}_score" +# CNNscore is a side-channel pose-confidence value emitted by gnina. It is +# saved alongside gnina_score for later analysis but does NOT participate in +# guild's rank-percentile aggregation (no entry in ALL_AVAILABLE_METHODS, +# SCORES_DICTIONARY, RANKS_DICTIONARY, or RP_SCORES_DICTIONARY). +GNINA_CNN_SCORE = f"{GNINA_PREFIX}_cnn_score" +VINA_RESCORE_BOLTZ_SCORE = f"{VINA_RESCORE_BOLTZ_PREFIX}_score" +VINA_RESCORE_DIFFDOCK_SCORE = f"{VINA_RESCORE_DIFFDOCK_PREFIX}_score" SCORES_DICTIONARY = { @@ -40,10 +52,18 @@ KARMADOCK_PREFIX: KARMADOCK_SCORE, DIFFDOCK_PREFIX: DIFFDOCK_SCORE, BOLTZ_PREFIX: BOLTZ_SCORE, - VINA_RESCORE_PREFIX: VINA_RESCORE_SCORE, + GNINA_PREFIX: GNINA_SCORE, + VINA_RESCORE_BOLTZ_PREFIX: VINA_RESCORE_BOLTZ_SCORE, + VINA_RESCORE_DIFFDOCK_PREFIX: VINA_RESCORE_DIFFDOCK_SCORE, } -ALL_AVAILABLE_METHODS = [VINA_PREFIX, KARMADOCK_PREFIX, DIFFDOCK_PREFIX, BOLTZ_PREFIX] +ALL_AVAILABLE_METHODS = [ + VINA_PREFIX, + KARMADOCK_PREFIX, + DIFFDOCK_PREFIX, + BOLTZ_PREFIX, + GNINA_PREFIX, +] RP_SCORES_COLUMNS = [ PROTEIN_CONF_ID, @@ -53,6 +73,7 @@ DIFFDOCK_SCORE, VINA_SCORE, BOLTZ_SCORE, + GNINA_SCORE, ] """ @@ -68,5 +89,7 @@ KARMADOCK_FOLDER = KARMADOCK_PREFIX DIFFDOCK_FOLDER = DIFFDOCK_PREFIX BOLTZ_FOLDER = BOLTZ_PREFIX -VINA_RESCORE_FOLDER = VINA_RESCORE_PREFIX +GNINA_FOLDER = GNINA_PREFIX +VINA_RESCORE_BOLTZ_FOLDER = VINA_RESCORE_BOLTZ_PREFIX +VINA_RESCORE_DIFFDOCK_FOLDER = VINA_RESCORE_DIFFDOCK_PREFIX MSA_FOLDER = "msa" diff --git a/guild/constants/p2rank.py b/guild/constants/p2rank.py index 3be7c44..2c038b7 100644 --- a/guild/constants/p2rank.py +++ b/guild/constants/p2rank.py @@ -4,12 +4,18 @@ from pathlib import Path -from guild.constants.system import WORKING_DIR_PATH +import guild as _guild_pkg """ -P2Rank installation path - assumes p2rank is installed in the workspace +P2Rank installation path — sibling of the guild/ package on disk (so +guild-internal/p2rank_2.4.2 in dev, /app/p2rank_2.4.2 in the container). +Resolving from the guild package location rather than WORKING_DIR_PATH is +deliberate: run_guild.py rebinds WORKING_DIR_PATH to the caller's workspace +for output paths, which would otherwise look for p2rank under the caller's +repo. The guild package's own location is the stable anchor. """ -P2RANK_INSTALL_DIR = Path(WORKING_DIR_PATH) / "p2rank_2.4.2" +_GUILD_PKG_DIR = Path(_guild_pkg.__file__).resolve().parent +P2RANK_INSTALL_DIR = _GUILD_PKG_DIR.parent / "p2rank_2.4.2" P2RANK_EXECUTABLE = P2RANK_INSTALL_DIR / "prank" """ diff --git a/guild/constants/system.py b/guild/constants/system.py index 6f3b065..8c5d990 100644 --- a/guild/constants/system.py +++ b/guild/constants/system.py @@ -14,5 +14,13 @@ """ Shell silencer + +Appended to commands run via ``subprocess.run(..., shell=True)`` to suppress +their stdout/stderr. Must be POSIX-compatible: in dash (Debian's /bin/sh) the +``&>file`` form is parsed as ``&`` (background) followed by ``>file`` (a +no-op redirect), so commands using the bash-only syntax get backgrounded and +the parent returns immediately — silently breaking any code that expects the +subprocess to have finished. ``> /dev/null 2>&1`` is portable across bash +and dash. """ -SHELL_SILENCER = "&>/dev/null" +SHELL_SILENCER = "> /dev/null 2>&1" diff --git a/guild/constants/visualization.py b/guild/constants/visualization.py index b6cba2d..3149621 100644 --- a/guild/constants/visualization.py +++ b/guild/constants/visualization.py @@ -7,7 +7,7 @@ """ Colors """ -GUILD_COLORS = { +ENVEDA_COLORS = { "white": "#FFFFFF", "black": "#000000", "deep-pink": "#F10A84", @@ -35,9 +35,9 @@ } CONTRAST_PALETTE = sns.color_palette(list(CONTRAST_COLORS.values())) -GUILD_PALETTE_RAINBOW = sns.color_palette( - [v for k, v in GUILD_COLORS.items() if k not in ["white", "black"]] +ENVEDA_PALETTE_RAINBOW = sns.color_palette( + [v for k, v in ENVEDA_COLORS.items() if k not in ["white", "black"]] ) -GUILD_PALETTE_RAINBOW_PLOTLY = [ - f"#{int(r*255):02x}{int(g*255):02x}{int(b*255):02x}" for r, g, b in GUILD_PALETTE_RAINBOW +ENVEDA_PALETTE_RAINBOW_PLOTLY = [ + f"#{int(r*255):02x}{int(g*255):02x}{int(b*255):02x}" for r, g, b in ENVEDA_PALETTE_RAINBOW ] diff --git a/guild/docking/boltz.py b/guild/docking/boltz.py index 00350fa..3b2d98f 100644 --- a/guild/docking/boltz.py +++ b/guild/docking/boltz.py @@ -6,7 +6,9 @@ import logging import os import subprocess +import tempfile +import numpy as np import pandas as pd import yaml from tqdm import tqdm @@ -18,12 +20,27 @@ BATCH_FOLDER, COMBINATION_ID, COMBINATIONS_TABLE_KEY, + COMBINATIONS_TO_RUN_KEY, ) +from guild.constants.general import RANDOM_SEED from guild.constants.guild import ( BOLTZ_FOLDER, BOLTZ_SCORE, LIGAND_ID, PROTEIN_CONF_ID, + VINA_RESCORE_BOLTZ_FOLDER, + VINA_RESCORE_BOLTZ_SCORE, +) +from guild.docking.vina import ( + _compute_box_from_pdb_atoms, + _extract_ligand_records, + _extract_protein_from_complex, + vina_score_pose, +) +from guild.transformers.converters import ( + cif_to_pdb, + ligand_pdb_to_pdbqt, + protein_pdb_to_pdbqt, ) logger = logging.getLogger(__name__) @@ -46,8 +63,13 @@ def generate_boltz_yaml( """ Generate the Boltz yaml file. - :param protein_sequence: Protein sequence. - :param protein_chain: Protein chain. + :param protein_sequence: Protein sequence, or a list of sequences for a + multi-chain receptor (one per chain in ``protein_chain``). + :param protein_chain: Protein chain ID, or a list of chain IDs for a + multi-chain receptor (e.g. ``["A", "B"]``). When more + than one chain is given, one ``protein`` block is + emitted per chain and the template/MSA are applied + per chain. :param ligand_sequences: Ligand sequences. :param ligand_ids: Ligand ids. :param output_file: Output file. @@ -57,12 +79,48 @@ def generate_boltz_yaml( :param pocket_contacts: Optional list of ``[chain_id, res_seq]`` protein residue contacts (from ``get_pocket_contacts_from_ligand``) for a Boltz pocket constraint. When provided, a ``constraints`` block is added to - the YAML. + the YAML. May span multiple chains. :param pocket_max_distance: Maximum heavy-atom distance (Å) for the pocket constraint (default 4.0). :param pocket_force: If True, use a potential to *enforce* the pocket constraint (default True). - :param msa_file: Path to a pre-computed a3m MSA file for the protein. When provided, - used instead of ``msa: empty``. + :param msa_file: Path to a pre-computed a3m MSA file for the protein, or a + list of paths (one per chain, same order as + ``protein_chain``). When provided, used instead of + ``msa: empty``. """ + # Normalize protein chains / sequences / MSAs to per-chain lists so single- + # and multi-chain callers share one code path. + protein_chains = [protein_chain] if isinstance(protein_chain, str) else list(protein_chain) + protein_sequences = ( + [protein_sequence] if isinstance(protein_sequence, str) else list(protein_sequence) + ) + if len(protein_sequences) != len(protein_chains): + raise ValueError( + f"generate_boltz_yaml: {len(protein_chains)} chains but " + f"{len(protein_sequences)} sequences — must match." + ) + if msa_file is None or isinstance(msa_file, str): + msa_files = [msa_file] * len(protein_chains) + else: + msa_files = list(msa_file) + if len(msa_files) != len(protein_chains): + raise ValueError( + f"generate_boltz_yaml: {len(protein_chains)} chains but " + f"{len(msa_files)} MSA files — must match." + ) + + protein_dictionary_list = [ + { + "protein": { + "id": chain, + "sequence": sequence, + "msa": msa if msa else "empty", + } + } + for chain, sequence, msa in zip( + protein_chains, protein_sequences, msa_files, strict=True + ) + ] + ligands_dictionary_list = [] for ligand_sequence, ligand_id in zip(ligand_sequences, ligand_ids, strict=True): ligands_dictionary_list.append( @@ -86,13 +144,7 @@ def generate_boltz_yaml( boltz_yaml = { "version": 1, "sequences": [ - { - "protein": { - "id": protein_chain, - "sequence": protein_sequence, - "msa": msa_file if msa_file else "empty", - } - }, + *protein_dictionary_list, *ligands_dictionary_list, ], "properties": affinities_dictionary_list, @@ -101,8 +153,8 @@ def generate_boltz_yaml( if template_file: template_entry = { "pdb": template_file, - "chain_id": [protein_chain], - "template_id": [f"{protein_chain}1"], + "chain_id": list(protein_chains), + "template_id": [f"{chain}1" for chain in protein_chains], } if template_force: template_entry["force"] = True @@ -141,6 +193,7 @@ def deploy_boltz( sampling_steps_affinity=None, diffusion_samples_affinity=None, affinity_mw_correction=True, + subprocess_log_path=None, ): """ Deploy Boltz for docking the ligand to the protein. @@ -223,6 +276,8 @@ def deploy_boltz( ): env[var] = str(cpu_threads) + from guild.tools.subprocess_log import write_subprocess_log + try: result = subprocess.run( boltz_command, @@ -237,11 +292,33 @@ def deploy_boltz( logger.error(f"STDOUT:\n{e.stdout}") if e.stderr: logger.error(f"STDERR:\n{e.stderr}") + if subprocess_log_path is not None: + write_subprocess_log( + subprocess_log_path, + argv=boltz_command, + returncode=-1, + stdout=e.stdout, + stderr=e.stderr, + extra_header=f"timed out after {timeout}s", + ) raise except Exception as e: logger.error(f"Error running Boltz command: {e}") raise + # Always persist the transcript when a log path was supplied — Boltz can + # exit 0 and still produce no records (template parsing failure), so the + # success branch needs the same trace as the failure branch for the + # caller's manifest-emptiness check downstream. + if subprocess_log_path is not None: + write_subprocess_log( + subprocess_log_path, + argv=boltz_command, + returncode=result.returncode, + stdout=result.stdout, + stderr=result.stderr, + ) + if result.returncode != 0: logger.error(f"Boltz failed with exit code {result.returncode}") logger.error(f"STDOUT:\n{result.stdout}") @@ -326,3 +403,207 @@ def boltz_guild_scoring(batch_dictionary): COMBINATION_ID, ], ) + + +# ──────────────────────────────────────────────────────────────────────────── +# Vina score-only re-scoring of Boltz-predicted complexes +# +# Boltz returns a confidence-style score (ipTM) rather than a physics-based +# binding ΔG, so for a comparable energy estimate we rescore the predicted +# complex with Vina's scoring function (score-only, no re-docking). +# +# Critical: both the receptor and the ligand are extracted *from Boltz's +# complex PDB* (per pose) because Boltz typically re-centres the predicted +# complex into its own coordinate frame — a template-frame receptor would not +# be physically near the ligand and the resulting ΔG would be meaningless. +# ──────────────────────────────────────────────────────────────────────────── + + +def _ensure_boltz_complex_pdb(boltz_folder: str, run_id: str) -> str: + """ + Return the path to the relabelled Boltz complex PDB for ``run_id``, + generating it on demand from the source CIF when missing. + """ + complex_pdb = f"{boltz_folder}/{run_id}_complex.pdb" + if os.path.exists(complex_pdb): + return complex_pdb + + cif_file = ( + f"{boltz_folder}/boltz_results_{run_id}_boltz" + f"/predictions/{run_id}_boltz/{run_id}_boltz_model_0.cif" + ) + if not os.path.exists(cif_file): + raise FileNotFoundError( + f"Neither Boltz complex PDB nor source CIF found for {run_id}" + ) + + # Lazy import — guild.transformers.pdb imports from guild.docking.vina so + # routing this through a top-level import would risk a circular load order. + from guild.transformers.pdb import relabel_ligand_chain_in_pdb + + with tempfile.NamedTemporaryFile(suffix=".pdb", delete=False) as tmp: + tmp_pdb = tmp.name + try: + cif_to_pdb(cif_file, tmp_pdb) + relabel_ligand_chain_in_pdb( + input_pdb=tmp_pdb, + output_pdb=complex_pdb, + ligand_chain_id="L", # Boltz YAML always uses ligand chain "L" + ) + finally: + if os.path.exists(tmp_pdb): + os.unlink(tmp_pdb) + return complex_pdb + + +def rescore_boltz_pose( + boltz_folder: str, + protein_conf_id: str, + ligand_id: str, + output_dir: str = None, + box_padding: float = 4.0, + seed: int = RANDOM_SEED, +) -> dict: + """ + Re-score a single Boltz-predicted protein-ligand complex with Vina's + physics-based scoring function (score-only — pose is not re-docked). + + Pipeline (all in Boltz's predicted coordinate frame): + + 1. Ensure the relabelled complex PDB exists (regenerate from CIF if needed). + 2. Extract the receptor (everything that is not the ligand resname) → PDBQT. + 3. Extract the ligand (resname ``LIG``) → PDBQT. + 4. Compute the Vina box from the ligand's bounding box. + 5. Call :func:`guild.docking.vina.vina_score_pose` (score-only). + + :param boltz_folder: Directory containing Boltz outputs for this batch + (``{batch_folder}/boltz``). + :param protein_conf_id: Protein configuration ID. + :param ligand_id: Ligand identifier. + :param output_dir: Where to write intermediate PDB/PDBQT files. + Defaults to ``boltz_folder``. + :param box_padding: Padding around the ligand bounding box (Å). + :param seed: Random seed for Vina. + :return: Dict with ``combination``, ``protein_config_id``, ``ligand_id``, + ``vina_rescore_boltz_score``. + """ + run_id = f"{protein_conf_id}_{ligand_id}" + complex_pdb = _ensure_boltz_complex_pdb(boltz_folder, run_id) + + if output_dir is None: + output_dir = boltz_folder + os.makedirs(output_dir, exist_ok=True) + + # Per-pose receptor in Boltz's frame + receptor_pdb = os.path.join(output_dir, f"{run_id}_receptor_rescore.pdb") + _extract_protein_from_complex(complex_pdb, receptor_pdb, ligand_resname="LIG") + receptor_pdbqt = receptor_pdb.replace(".pdb", ".pdbqt") + protein_pdb_to_pdbqt( + input_pdb=receptor_pdb, + output_pdbqt=receptor_pdbqt, + allow_bad_res=True, + ) + + # Ligand in the same frame + ligand_pdb = os.path.join(output_dir, f"{run_id}_ligand_rescore.pdb") + _extract_ligand_records(complex_pdb, ligand_pdb, resname="LIG") + ligand_pdb_to_pdbqt(pdb=ligand_pdb) + ligand_pdbqt = ligand_pdb.replace(".pdb", ".pdbqt") + + center, size = _compute_box_from_pdb_atoms(ligand_pdb, padding=box_padding) + + score = vina_score_pose( + receptor_pdbqt=receptor_pdbqt, + ligand_pdbqt=ligand_pdbqt, + center=center, + size=size, + seed=seed, + ) + + combination_id = f"{protein_conf_id}_{ligand_id}" + logger.info(f"Vina rescore (Boltz pose): {combination_id} → {score:.3f} kcal/mol") + + return { + COMBINATION_ID: combination_id, + PROTEIN_CONF_ID: protein_conf_id, + LIGAND_ID: ligand_id, + VINA_RESCORE_BOLTZ_SCORE: score, + } + + +def vina_rescore_boltz_batch( + batch_folder: str, + combinations: list, + box_padding: float = 4.0, + seed: int = RANDOM_SEED, +) -> pd.DataFrame: + """ + Batch-rescore every Boltz complex in a batch. + + :param batch_folder: Path to the batch folder. + :param combinations: List of ``(protein_conf_id, ligand_id)`` tuples. + :param box_padding: Padding around the ligand bounding box (Å). + :param seed: Random seed for Vina. + :return: DataFrame with ``combination``, ``protein_config_id``, + ``ligand_id``, ``vina_rescore_boltz_score``. + """ + boltz_folder = os.path.join(batch_folder, BOLTZ_FOLDER) + rescore_output_dir = os.path.join(batch_folder, VINA_RESCORE_BOLTZ_FOLDER) + os.makedirs(rescore_output_dir, exist_ok=True) + + results = [] + for protein_conf_id, ligand_id in combinations: + combination_id = f"{protein_conf_id}_{ligand_id}" + try: + result = rescore_boltz_pose( + boltz_folder=boltz_folder, + protein_conf_id=protein_conf_id, + ligand_id=ligand_id, + output_dir=rescore_output_dir, + box_padding=box_padding, + seed=seed, + ) + except Exception as e: + logger.warning( + f"Vina Boltz rescore failed for {combination_id}: {e}" + ) + result = { + COMBINATION_ID: combination_id, + PROTEIN_CONF_ID: protein_conf_id, + LIGAND_ID: ligand_id, + VINA_RESCORE_BOLTZ_SCORE: np.nan, + } + results.append(result) + + return pd.DataFrame(results) + + +def vina_rescore_boltz_guild_scoring(batch_dictionary): + """ + Vina re-scoring of Boltz-predicted complexes for a batch. + + :param batch_dictionary: Batch information dict (the standard structure + every ``*_guild_scoring`` function receives). + :return: DataFrame with ``COMBINATION_ID``, ``PROTEIN_CONF_ID``, + ``LIGAND_ID``, and ``VINA_RESCORE_BOLTZ_SCORE`` columns. Returns an + empty frame (with those columns) if no Boltz outputs are found. + """ + batch_folder = batch_dictionary[BATCH_FOLDER] + combinations = batch_dictionary[COMBINATIONS_TO_RUN_KEY] + + boltz_root = os.path.join(batch_folder, BOLTZ_FOLDER) + boltz_present = os.path.isdir(boltz_root) and any( + name.startswith("boltz_results_") for name in os.listdir(boltz_root) + ) + if not boltz_present: + logger.warning( + "vina_rescore_boltz: no Boltz outputs found in %s — returning empty score table", + batch_folder, + ) + return pd.DataFrame( + columns=[COMBINATION_ID, PROTEIN_CONF_ID, LIGAND_ID, VINA_RESCORE_BOLTZ_SCORE] + ) + + df = vina_rescore_boltz_batch(batch_folder=batch_folder, combinations=combinations) + keep = [COMBINATION_ID, PROTEIN_CONF_ID, LIGAND_ID, VINA_RESCORE_BOLTZ_SCORE] + return df[[c for c in keep if c in df.columns]] diff --git a/guild/docking/diffdock.py b/guild/docking/diffdock.py index f6fcf3e..dd6e37d 100644 --- a/guild/docking/diffdock.py +++ b/guild/docking/diffdock.py @@ -6,6 +6,7 @@ import os import subprocess +import numpy as np import pandas as pd from tqdm import tqdm @@ -24,6 +25,7 @@ PROTEIN_PATH, PROTEIN_SEQUENCE, ) +from guild.constants.general import RANDOM_SEED from guild.constants.guild import ( DIFFDOCK_FOLDER, DIFFDOCK_SCORE, @@ -31,11 +33,19 @@ PROTEIN_CONF_ID, PROTEINS_FOLDER, SMILES, + VINA_RESCORE_DIFFDOCK_FOLDER, + VINA_RESCORE_DIFFDOCK_SCORE, ) from guild.constants.system import ( PYTHON_EXECUTABLE, SUPPORT_FOLDER, ) +from guild.docking.vina import ( + compute_box_from_sdf, + vina_score_pose, +) +from guild.tools.preparation import _normalize_chain_list +from guild.transformers.converters import sdf_to_pdbqt logger = logging.getLogger(__name__) @@ -121,12 +131,21 @@ def deploy_diffdock_single( return status -def deploy_diffdock(home_path: str, diffdock_results_dir: str, input_csv: str = None): +def deploy_diffdock( + home_path: str, + diffdock_results_dir: str, + input_csv: str = None, + subprocess_log_path: str = None, +): """ Deploy DiffDock for docking the ligand to the protein. Output is saved in the diffdock directory, inside the project folder. :param home_path: Path to the home directory. :param diffdock_results_dir: Path to the diffdock results directory. :param input_csv: Path to the input csv file. + :param subprocess_log_path: Optional path to write the full DiffDock + stdout/stderr transcript. DiffDock runs once per batch (not per + combination), so this is typically + ``batches//diffdock/_batch.subprocess.log``. :return: Status of the deployment. """ os.makedirs(diffdock_results_dir, exist_ok=True) @@ -192,6 +211,16 @@ def deploy_diffdock(home_path: str, diffdock_results_dir: str, input_csv: str = if result.stderr: logger.warning(f"DiffDock STDERR: {result.stderr[-2000:]}") + if subprocess_log_path is not None: + from guild.tools.subprocess_log import write_subprocess_log + write_subprocess_log( + subprocess_log_path, + argv=cmd, + returncode=result.returncode, + stdout=result.stdout, + stderr=result.stderr, + ) + if result.returncode != 0: logger.error(f"DiffDock FAILED (exit code {result.returncode})") return 1 @@ -245,3 +274,255 @@ def diffdock_guild_scoring(batch_dictionary): diffdock_scores_data.append(process_diffdock_files(current_combination, diffdock_folder)) return pd.DataFrame(diffdock_scores_data) + + +# ──────────────────────────────────────────────────────────────────────────── +# Vina score-only re-scoring of DiffDock poses +# +# DiffDock returns a relative confidence score, not a physics-based ΔG. For a +# comparable energy estimate we apply Vina's scoring function to the highest- +# confidence pose (no re-docking). DiffDock writes poses in the coordinate +# frame of the *raw* input PDB, so the rescore receptor must be derived from +# that same raw PDB — the cleaned/centred receptor used by Vina docking is in +# a different frame. +# ──────────────────────────────────────────────────────────────────────────── + + +def _find_best_diffdock_sdf(diffdock_results_dir: str, protein_conf_id: str, ligand_id: str) -> str: + """ + Locate the highest-confidence SDF file for a combination in a DiffDock + results directory. DiffDock names output files like + ``rank1_confidence-1.23.sdf``. + + :return: Absolute path to the best-confidence SDF. + :raises FileNotFoundError: If no SDF files are found for the combination. + """ + folder_name = f"{protein_conf_id}_{ligand_id}" + combo_dir = os.path.join(diffdock_results_dir, folder_name) + + if not os.path.isdir(combo_dir): + raise FileNotFoundError(f"DiffDock results folder not found: {combo_dir}") + + sdf_scores = {} + for fname in os.listdir(combo_dir): + if "_confidence" in fname and fname.endswith(".sdf"): + try: + score = float(fname.split("_confidence")[1].replace(".sdf", "")) + sdf_scores[fname] = score + except ValueError: + continue + + if not sdf_scores: + raise FileNotFoundError(f"No DiffDock SDF files found in {combo_dir}") + + best_fname = max(sdf_scores, key=sdf_scores.get) + return os.path.join(combo_dir, best_fname) + + +def _prepare_receptor_pdbqt_from_raw( + raw_pdb: str, + chain_id, + output_pdbqt: str, +) -> str: + """ + Extract one or more chains (ATOM records only) from a raw PDB and convert + them to PDBQT via OpenBabel, preserving the original crystal coordinates so + the receptor is in the same frame as DiffDock's output SDF files. + + ``chain_id`` may be a single chain ID, a list of IDs, or a comma-separated + string (e.g. ``"A,B"``) so a multi-chain receptor is kept intact. + """ + if not os.path.isfile(raw_pdb): + raise FileNotFoundError(f"Raw PDB not found: {raw_pdb}") + + chain_ids = _normalize_chain_list(chain_id) + chain_pdb = output_pdbqt.replace(".pdbqt", "_chain.pdb") + kept = 0 + with open(raw_pdb) as fin, open(chain_pdb, "w") as fout: + for line in fin: + if line.startswith("ATOM") and len(line) > 21 and line[21] in chain_ids: + fout.write(line) + kept += 1 + fout.write("END\n") + + if kept == 0: + raise ValueError(f"No ATOM records found for chain(s) {chain_ids} in {raw_pdb}") + + result = subprocess.run( + ["obabel", "-ipdb", chain_pdb, "-opdbqt", "-O", output_pdbqt, "-xr"], + capture_output=True, + text=True, + ) + if not os.path.isfile(output_pdbqt) or os.path.getsize(output_pdbqt) == 0: + raise RuntimeError( + f"obabel PDB→PDBQT failed for {chain_pdb}: {result.stderr}" + ) + + logger.info( + f"Prepared receptor PDBQT from raw PDB chain(s) {chain_ids}: " + f"{kept} atoms → {output_pdbqt}" + ) + return output_pdbqt + + +def rescore_diffdock_pose( + receptor_pdbqt: str, + diffdock_results_dir: str, + protein_conf_id: str, + ligand_id: str, + output_dir: str = None, + box_padding: float = 4.0, + seed: int = RANDOM_SEED, +) -> dict: + """ + Re-score a single DiffDock pose with Vina's physics-based scoring function. + + Pipeline: + 1. Find the highest-confidence SDF in ``diffdock_results_dir/{protein}_{ligand}/``. + 2. Convert SDF → PDBQT (preserving 3D coordinates). + 3. Compute a Vina box centred on the SDF pose. + 4. Call :func:`guild.docking.vina.vina_score_pose` (score-only). + + :return: Dict with ``combination``, ``protein_config_id``, ``ligand_id``, + ``vina_rescore_diffdock_score``, ``diffdock_sdf``. + """ + sdf_path = _find_best_diffdock_sdf(diffdock_results_dir, protein_conf_id, ligand_id) + + if output_dir is None: + output_dir = os.path.dirname(sdf_path) + + ligand_pdbqt = os.path.join(output_dir, f"{protein_conf_id}_{ligand_id}_rescore.pdbqt") + sdf_to_pdbqt(sdf_path, pdbqt=ligand_pdbqt) + + center, size = compute_box_from_sdf(sdf_path, padding=box_padding) + + score = vina_score_pose( + receptor_pdbqt=receptor_pdbqt, + ligand_pdbqt=ligand_pdbqt, + center=center, + size=size, + seed=seed, + ) + + combination_id = f"{protein_conf_id}_{ligand_id}" + logger.info(f"Vina rescore (DiffDock pose): {combination_id} → {score:.3f} kcal/mol") + + return { + COMBINATION_ID: combination_id, + PROTEIN_CONF_ID: protein_conf_id, + LIGAND_ID: ligand_id, + VINA_RESCORE_DIFFDOCK_SCORE: score, + "diffdock_sdf": sdf_path, + } + + +def vina_rescore_diffdock_batch( + batch_folder: str, + combinations: list, + receptor_pdbqt_dir: str = None, + box_padding: float = 4.0, + seed: int = RANDOM_SEED, +) -> pd.DataFrame: + """ + Batch re-score DiffDock poses for all combinations in a batch. + + Receptor PDBQT resolution order: + + 1. ``{receptor_pdbqt_dir}/{protein_conf_id}_raw.pdbqt`` — single-chain + PDBQT already in the raw (crystal) coordinate frame. + 2. Auto-generated from ``{protein_conf_id}_raw.pdb`` by extracting the + chain letter encoded in *protein_conf_id* and converting via OpenBabel. + """ + diffdock_results_dir = os.path.join(batch_folder, DIFFDOCK_FOLDER, DIFFDOCK_RESULTS_FOLDER) + if receptor_pdbqt_dir is None: + receptor_pdbqt_dir = os.path.join(batch_folder, PROTEINS_FOLDER) + + rescore_output_dir = os.path.join(batch_folder, VINA_RESCORE_DIFFDOCK_FOLDER) + os.makedirs(rescore_output_dir, exist_ok=True) + + results = [] + for protein_conf_id, ligand_id in combinations: + receptor_pdbqt = os.path.join( + receptor_pdbqt_dir, f"{protein_conf_id}_raw.pdbqt" + ) + if not os.path.exists(receptor_pdbqt): + raw_pdb = os.path.join(receptor_pdbqt_dir, f"{protein_conf_id}_raw.pdb") + if os.path.isfile(raw_pdb): + parts = protein_conf_id.split("-") + chain_id = parts[1] if len(parts) >= 2 else "A" + try: + _prepare_receptor_pdbqt_from_raw(raw_pdb, chain_id, receptor_pdbqt) + except Exception as e: + logger.warning( + f"Could not prepare receptor PDBQT from raw PDB " + f"for {protein_conf_id}: {e}" + ) + receptor_pdbqt = None + else: + logger.warning( + f"Neither {receptor_pdbqt} nor {raw_pdb} found for " + f"{protein_conf_id}, skipping combination" + ) + receptor_pdbqt = None + + if receptor_pdbqt is None or not os.path.exists(receptor_pdbqt): + results.append({ + COMBINATION_ID: f"{protein_conf_id}_{ligand_id}", + PROTEIN_CONF_ID: protein_conf_id, + LIGAND_ID: ligand_id, + VINA_RESCORE_DIFFDOCK_SCORE: np.nan, + "diffdock_sdf": None, + }) + continue + + try: + result = rescore_diffdock_pose( + receptor_pdbqt=receptor_pdbqt, + diffdock_results_dir=diffdock_results_dir, + protein_conf_id=protein_conf_id, + ligand_id=ligand_id, + output_dir=rescore_output_dir, + box_padding=box_padding, + seed=seed, + ) + results.append(result) + except Exception as e: + logger.warning( + f"Vina rescore failed for {protein_conf_id}_{ligand_id}: {e}" + ) + results.append({ + COMBINATION_ID: f"{protein_conf_id}_{ligand_id}", + PROTEIN_CONF_ID: protein_conf_id, + LIGAND_ID: ligand_id, + VINA_RESCORE_DIFFDOCK_SCORE: np.nan, + "diffdock_sdf": None, + }) + + return pd.DataFrame(results) + + +def vina_rescore_diffdock_guild_scoring(batch_dictionary): + """ + Vina re-scoring of DiffDock poses for a batch. + + :return: DataFrame with ``COMBINATION_ID``, ``PROTEIN_CONF_ID``, + ``LIGAND_ID``, and ``VINA_RESCORE_DIFFDOCK_SCORE`` columns. Returns an + empty frame (with those columns) if no DiffDock outputs are found. + """ + batch_folder = batch_dictionary[BATCH_FOLDER] + combinations = batch_dictionary[COMBINATIONS_TO_RUN_KEY] + + diffdock_root = os.path.join(batch_folder, DIFFDOCK_FOLDER, DIFFDOCK_RESULTS_FOLDER) + diffdock_present = os.path.isdir(diffdock_root) and len(os.listdir(diffdock_root)) > 0 + if not diffdock_present: + logger.warning( + "vina_rescore_diffdock: no DiffDock outputs found in %s — returning empty score table", + batch_folder, + ) + return pd.DataFrame( + columns=[COMBINATION_ID, PROTEIN_CONF_ID, LIGAND_ID, VINA_RESCORE_DIFFDOCK_SCORE] + ) + + df = vina_rescore_diffdock_batch(batch_folder=batch_folder, combinations=combinations) + keep = [COMBINATION_ID, PROTEIN_CONF_ID, LIGAND_ID, VINA_RESCORE_DIFFDOCK_SCORE] + return df[[c for c in keep if c in df.columns]] diff --git a/guild/docking/gnina.py b/guild/docking/gnina.py new file mode 100644 index 0000000..b1561ed --- /dev/null +++ b/guild/docking/gnina.py @@ -0,0 +1,352 @@ +""" +GNINA docking tools. + +GNINA (https://github.com/gnina/gnina) is a fork of smina/AutoDock Vina that +rescores Vina-style poses with a CNN. It is invoked here as a standalone +docker — it produces a Vina-style binding affinity (kcal/mol, lower is better) +that drives guild's rank-percentile aggregation, plus a CNN confidence score +(`CNNscore`, [0,1], higher is better) that is carried alongside as a +side-channel column. +""" + +import logging +import os +import re +import subprocess +import tempfile + +import numpy as np +import pandas as pd + +from guild.constants.bulk import ( + BATCH_FOLDER, + COMBINATION_ID, + COMBINATIONS_TO_RUN_KEY, +) +from guild.constants.general import RANDOM_SEED +from guild.constants.gnina import ( + GNINA_BINARY, + GNINA_DEFAULT_CNN_SCORING, + GNINA_DEFAULT_EXHAUSTIVENESS, + GNINA_DEFAULT_NUMBER_OF_POSES, + GNINA_LIB_PATH, + GNINA_OB_DATA_DIR, + GNINA_OB_PLUGIN_DIR, + GNINA_OB_SYSTEM_LIB, +) +from guild.constants.guild import ( + GNINA_CNN_SCORE, + GNINA_FOLDER, + GNINA_SCORE, + LIGAND_ID, + PROTEIN_CONF_ID, +) +from guild.docking.vina import _validate_pdbqt +from guild.tools.subprocess_log import write_subprocess_log + +logger = logging.getLogger(__name__) + +# Bound to avoid hangs from a runaway gnina worker. Matches DOCKING_TIMEOUT in +# guild/constants/bulk.py (kept as a module-local constant to avoid pulling a +# bulk-orchestration dep into the docking module). +GNINA_SUBPROCESS_TIMEOUT = 600 # seconds + +# gnina prints a table to stdout that looks like: +# +# mode | affinity | CNN | CNN +# | (kcal/mol) | pose-score| affinity +# -----+------------+----------+---------- +# 1 -8.345 0.7891 6.234 +# 2 -7.910 0.6512 5.987 +# ... +# +# Each data row starts with the integer mode number; the rest are floats. +_GNINA_POSE_ROW = re.compile( + r"^\s*(\d+)\s+(-?\d+\.\d+)\s+(-?\d+\.\d+)\s+(-?\d+\.\d+)" +) + + +def _ensure_openbabel_plugin_shim() -> tuple[str, str, str] | None: + """ + Make gnina's Open Babel able to write pose files. + + gnina's bundled ``libopenbabel.so.7`` has no format-plugin tree in the + image, so its OB can't write any ``--out`` format and every pose file ends + up empty. The system Open Babel 3.2 (``.so.8``) *does* ship a complete + plugin set. We create a tiny shim dir holding ``libopenbabel.so.7`` -> + system ``.so.8`` symlink; prepending it to ``LD_LIBRARY_PATH`` makes gnina + resolve its OB to the .so.8 (ABI-compatible for gnina's calls), which then + finds the plugins via ``BABEL_LIBDIR``/``BABEL_DATADIR``. + + :return: ``(shim_dir, plugin_dir, data_dir)`` to feed into the gnina + subprocess env, or ``None`` if the system OpenBabel isn't present (in + which case the caller leaves the env untouched — gnina still scores via + its native pdbqt parser, just with empty pose files as before). + """ + if not ( + os.path.exists(GNINA_OB_SYSTEM_LIB) + and os.path.isdir(GNINA_OB_PLUGIN_DIR) + and os.path.isdir(GNINA_OB_DATA_DIR) + ): + return None + + shim_dir = os.path.join(tempfile.gettempdir(), "guild_gnina_ob_shim") + os.makedirs(shim_dir, exist_ok=True) + link = os.path.join(shim_dir, "libopenbabel.so.7") + # Idempotent + race-safe across parallel gnina workers: the target is a + # fixed path, so a pre-existing link is already correct. + try: + if os.path.islink(link): + if os.readlink(link) != GNINA_OB_SYSTEM_LIB: + os.unlink(link) + elif os.path.exists(link): + os.remove(link) + + if not os.path.lexists(link): + os.symlink(GNINA_OB_SYSTEM_LIB, link) + except OSError as e: + logger.warning("gnina: failed to set up OpenBabel shim (%s)", e) + return None + return shim_dir, GNINA_OB_PLUGIN_DIR, GNINA_OB_DATA_DIR + + +def parse_gnina_stdout(stdout: str) -> list[tuple[int, float, float, float]]: + """ + Parse gnina's stdout pose table. + + :param stdout: Captured stdout from a ``gnina`` subprocess call. + :return: List of ``(mode, affinity, cnn_score, cnn_affinity)`` tuples, + one per docked pose, in the order gnina emitted them. + :raises ValueError: when no pose rows can be parsed (treated as a docking + failure by callers). + """ + rows: list[tuple[int, float, float, float]] = [] + for line in stdout.splitlines(): + match = _GNINA_POSE_ROW.match(line) + if not match: + continue + mode = int(match.group(1)) + affinity = float(match.group(2)) + cnn_score = float(match.group(3)) + cnn_affinity = float(match.group(4)) + rows.append((mode, affinity, cnn_score, cnn_affinity)) + if not rows: + raise ValueError("Could not parse any pose rows from gnina stdout") + return rows + + +def deploy_gnina( + receptor: str, + ligand: str, + center: tuple[float, float, float], + size: tuple[float, float, float], + output_pdbqt: str, + output_scores: str, + exhaustiveness: int = GNINA_DEFAULT_EXHAUSTIVENESS, + n_poses: int = GNINA_DEFAULT_NUMBER_OF_POSES, + seed: int = RANDOM_SEED, + cnn_scoring: str = GNINA_DEFAULT_CNN_SCORING, + use_gpu: bool = True, + subprocess_log_path: str | None = None, +) -> dict: + """ + Run docking with the gnina CLI. + + :param receptor: Path to the receptor file. ``.pdbqt`` and ``.pdb`` both + work — gnina sniffs the format from the extension. PDBQT inputs get + an extra structural validation pass; PDB inputs are handed to gnina + verbatim. + :param ligand: Path to the ligand file. ``.pdbqt`` and ``.sdf`` both + work. Same validation rule as above. + :param center: Box center (x, y, z). + :param size: Box size (x, y, z). + :param output_pdbqt: Where gnina should write the multi-pose output PDBQT. + :param output_scores: Path to write the per-pose score file. Format mirrors + Vina's (``mode: affinity``) extended with a tab-separated CNN score: + ``mode: affinity\\tcnn_score``. + :param exhaustiveness: Search exhaustiveness (gnina default 8). + :param n_poses: Number of poses to keep (``--num_modes``). + :param seed: RNG seed. + :param cnn_scoring: One of gnina's CNN modes (``none``/``rescore``/ + ``refinement``/``metrorescore``/``metrorefine``/``all``). Default + ``rescore`` does Vina search then CNN-rescores the top poses. + :param use_gpu: When False, pass ``--no_gpu`` to gnina (CPU-only inference). + :return: ``{"scores": [...affinities], "cnn_scores": [...], "out_pdbqt": + output_pdbqt, "output_scores": output_scores}``. + """ + if receptor.endswith(".pdbqt"): + _validate_pdbqt(receptor, "receptor") + if ligand.endswith(".pdbqt"): + _validate_pdbqt(ligand, "ligand") + + cx, cy, cz = center + sx, sy, sz = size + + argv = [ + GNINA_BINARY, + "--receptor", receptor, + "--ligand", ligand, + "--center_x", str(cx), + "--center_y", str(cy), + "--center_z", str(cz), + "--size_x", str(sx), + "--size_y", str(sy), + "--size_z", str(sz), + "--out", output_pdbqt, + "--num_modes", str(n_poses), + "--exhaustiveness", str(exhaustiveness), + "--seed", str(seed), + "--cnn_scoring", cnn_scoring, + ] + if not use_gpu: + argv.append("--no_gpu") + + # gnina's torch/openbabel/boost live under /opt/gnina/lib, isolated from + # the rest of the image. Prepend that to LD_LIBRARY_PATH for this call + # only — gnina ABI demands its own libtorch + libcudart 12, which would + # clash with our venv-managed CUDA 13 torch if we set it globally. + env = os.environ.copy() + ld_parts = [GNINA_LIB_PATH] + + # Repoint gnina's broken OpenBabel at the system .so.8 + its plugins so the + # ``--out`` pose file isn't written empty. The shim dir goes *ahead* of + # GNINA_LIB_PATH so its libopenbabel.so.7 symlink wins over the bundled + # (plugin-less) one, while gnina's own libtorch/boost still resolve from + # /opt/gnina/lib. No-ops to today's behaviour if the system OB is absent. + ob_shim = _ensure_openbabel_plugin_shim() + if ob_shim is not None: + shim_dir, plugin_dir, data_dir = ob_shim + ld_parts.insert(0, shim_dir) + env["BABEL_LIBDIR"] = plugin_dir + env["BABEL_DATADIR"] = data_dir + + existing_ld = env.get("LD_LIBRARY_PATH", "") + if existing_ld: + ld_parts.append(existing_ld) + env["LD_LIBRARY_PATH"] = ":".join(ld_parts) + + try: + completed = subprocess.run( + argv, + check=True, + capture_output=True, + text=True, + timeout=GNINA_SUBPROCESS_TIMEOUT, + env=env, + ) + except subprocess.CalledProcessError as e: + # Persist the failure transcript before re-raising so the caller can + # point users at it. + if subprocess_log_path is not None: + write_subprocess_log( + subprocess_log_path, + argv=argv, + returncode=e.returncode, + stdout=e.stdout, + stderr=e.stderr, + ) + raise RuntimeError( + f"gnina failed (exit {e.returncode}) for receptor={receptor} " + f"ligand={ligand}: {e.stderr.strip()[:500]}" + ) from e + + if subprocess_log_path is not None: + write_subprocess_log( + subprocess_log_path, + argv=argv, + returncode=completed.returncode, + stdout=completed.stdout, + stderr=completed.stderr, + ) + + poses = parse_gnina_stdout(completed.stdout) + + scores = [affinity for (_mode, affinity, _cnn, _cnn_aff) in poses] + cnn_scores = [cnn for (_mode, _affinity, cnn, _cnn_aff) in poses] + + with open(output_scores, "w") as f: + for (mode, affinity, cnn, _cnn_aff) in poses: + # mode is 1-based in gnina's table; keep that to make the file + # round-trippable to the CLI output. + f.write(f"{mode}: {affinity}\t{cnn}\n") + + logger.info(f"gnina: docking completed. Scores saved to {output_scores}") + + return { + "scores": scores, + "cnn_scores": cnn_scores, + "out_pdbqt": output_pdbqt, + "output_scores": output_scores, + } + + +def process_gnina_output(input_file: str) -> tuple[float, float]: + """ + Read a gnina score file and return the best ``(affinity, cnn_score)``. + + "Best" follows the Vina convention: the row with the **minimum** affinity + (lower kcal/mol = stronger binder). The CNN score returned is the one + reported for that same pose — not the max CNN across all poses. + + :param input_file: Path to the file written by :func:`deploy_gnina`. + :return: ``(affinity, cnn_score)``. If the file is empty/unreadable, + returns ``(nan, nan)`` — matching :func:`process_vina_output`. + """ + df = pd.read_csv(input_file, sep=":", header=None) + if df.empty: + return float("nan"), float("nan") + + # The second column holds "affinity\tcnn_score"; split it back out. + affinity_cnn = df[1].astype(str).str.split("\t", n=1, expand=True) + affinities = affinity_cnn[0].astype(float) + # Older score files might be Vina-format (no CNN column) — guard for that + # so a partial migration doesn't break callers. + if affinity_cnn.shape[1] > 1: + cnn_scores = pd.to_numeric(affinity_cnn[1], errors="coerce") + else: + cnn_scores = pd.Series([np.nan] * len(affinities)) + + best_idx = affinities.idxmin() + return float(affinities.loc[best_idx]), float(cnn_scores.loc[best_idx]) + + +def gnina_guild_scoring(batch_dictionary) -> pd.DataFrame: + """ + Collect per-combination gnina scores into a DataFrame. + + Mirrors :func:`guild.docking.vina.vina_guild_scoring` but emits both the + primary ``gnina_score`` (Vina-style affinity, used in rank-percentile + aggregation) and the side-channel ``gnina_cnn_score`` (CNN confidence, + not registered in ``RANKS_DICTIONARY``/``RP_SCORES_DICTIONARY``). + + :param batch_dictionary: Standard bulk batch dictionary. + :return: DataFrame with columns + ``[COMBINATION_ID, GNINA_SCORE, GNINA_CNN_SCORE, PROTEIN_CONF_ID, LIGAND_ID]``. + """ + combinations_df = pd.DataFrame( + batch_dictionary[COMBINATIONS_TO_RUN_KEY], + columns=[PROTEIN_CONF_ID, LIGAND_ID], + ) + + def score_combination(row): + try: + return process_gnina_output( + f"{batch_dictionary[BATCH_FOLDER]}/{GNINA_FOLDER}/" + f"{row[PROTEIN_CONF_ID]}_{row[LIGAND_ID]}.txt" + ) + except Exception as e: + logger.info( + f"Failed gnina scoring for {(row[PROTEIN_CONF_ID], row[LIGAND_ID])} with error {e}" + ) + return np.nan, np.nan + + score_pairs = combinations_df.apply(score_combination, axis=1) + combinations_df[GNINA_SCORE] = [s[0] for s in score_pairs] + combinations_df[GNINA_CNN_SCORE] = [s[1] for s in score_pairs] + combinations_df[COMBINATION_ID] = ( + combinations_df[PROTEIN_CONF_ID] + "_" + combinations_df[LIGAND_ID] + ) + + return combinations_df[ + [COMBINATION_ID, GNINA_SCORE, GNINA_CNN_SCORE, PROTEIN_CONF_ID, LIGAND_ID] + ] diff --git a/guild/docking/karmadock.py b/guild/docking/karmadock.py index 4b6ab7a..328707b 100644 --- a/guild/docking/karmadock.py +++ b/guild/docking/karmadock.py @@ -7,14 +7,16 @@ import subprocess import time +import numpy as np import pandas as pd +from rdkit import Chem -from guild.constants.bulk import BATCH_FOLDER, COMBINATION_ID +from guild.constants.bulk import BATCH_FOLDER, COMBINATION_ID, COMBINATIONS_TO_RUN_KEY from guild.constants.guild import ( KARMADOCK_FOLDER, KARMADOCK_SCORE, LIGAND_ID, - PROTEIN_ID, + PROTEIN_CONF_ID, ) from guild.constants.karmadock import ( KARMADOCK_COLUMNS, @@ -31,6 +33,125 @@ logger = logging.getLogger(__name__) +# ──────────────────────────────────────────────────────────────────────────── +# Pocket support +# +# KarmaDock does not accept a binding pocket directly — its ``get_pocket()`` +# preprocessing step derives the pocket from residues within 12 Å of the +# *input ligand's* coordinates. By default the input ligand is whatever 3D +# conformer RDKit/OpenBabel happens to embed (typically near the origin), +# which means KarmaDock's pocket lands wherever that conformer falls. To +# control the pocket without modifying KarmaDock's codebase, we translate +# the ligand SDF (and mol2) so its centroid sits at the centre of the user- +# supplied Vina box. KarmaDock's get_pocket() then selects the correct +# residues around that point. +# ──────────────────────────────────────────────────────────────────────────── + + +def _translate_sdf(sdf_path: str, target_center) -> bool: + """Translate every conformer in ``sdf_path`` so its centroid sits at + ``target_center``. Returns True on success.""" + supplier = Chem.SDMolSupplier(sdf_path, removeHs=False, sanitize=False) + mols = [m for m in supplier if m is not None] + if not mols: + return False + + target = np.asarray(target_center, dtype=float) + writer = Chem.SDWriter(sdf_path + ".tmp") + for mol in mols: + conf = mol.GetConformer() + positions = np.array( + [list(conf.GetAtomPosition(i)) for i in range(mol.GetNumAtoms())] + ) + shift = target - positions.mean(axis=0) + for i in range(mol.GetNumAtoms()): + old = conf.GetAtomPosition(i) + conf.SetAtomPosition( + i, (old.x + shift[0], old.y + shift[1], old.z + shift[2]) + ) + writer.write(mol) + writer.close() + os.replace(sdf_path + ".tmp", sdf_path) + return True + + +def _translate_mol2(mol2_path: str, target_center) -> bool: + """Translate ATOM coordinates in a TRIPOS mol2 file in place so the + centroid sits at ``target_center``. Returns True on success.""" + with open(mol2_path) as f: + lines = f.readlines() + + in_atoms = False + atom_indices = [] + coords = [] + for idx, line in enumerate(lines): + stripped = line.strip() + if stripped.startswith("@ATOM"): + in_atoms = True + continue + if stripped.startswith("@") and in_atoms: + in_atoms = False + continue + if in_atoms and stripped: + parts = line.split() + try: + x, y, z = float(parts[2]), float(parts[3]), float(parts[4]) + except (IndexError, ValueError): + continue + atom_indices.append(idx) + coords.append((x, y, z)) + + if not coords: + return False + + arr = np.array(coords) + target = np.asarray(target_center, dtype=float) + shift = target - arr.mean(axis=0) + + for idx, (x, y, z) in zip(atom_indices, coords, strict=True): + nx, ny, nz = x + shift[0], y + shift[1], z + shift[2] + parts = lines[idx].split() + # mol2 ATOM line: id name x y z atom_type [subst_id] [subst_name] [charge] + parts[2] = f"{nx:.4f}" + parts[3] = f"{ny:.4f}" + parts[4] = f"{nz:.4f}" + # Preserve original spacing-ish — single space join is fine for KarmaDock + lines[idx] = " ".join(parts) + "\n" + + with open(mol2_path, "w") as f: + f.writelines(lines) + return True + + +def center_karmadock_ligand_on_box(sdf_path: str, mol2_path: str, box_file: str) -> bool: + """ + Re-centre the karmadock ligand files so their centroid sits at the centre + of the supplied Vina box file. Both the SDF and the mol2 are translated + because KarmaDock's ``get_pocket()`` reads SDF first but falls back to mol2. + + No-op (returns False) if the box file is missing or unreadable. + """ + # Lazy import to avoid a circular load order between karmadock and vina. + from guild.docking.vina import get_center_and_size_from_box_file + + if not os.path.isfile(box_file): + return False + try: + center, _ = get_center_and_size_from_box_file(box_file) + except Exception as e: + logger.warning(f"Could not parse Vina box {box_file}: {e}") + return False + + sdf_ok = _translate_sdf(sdf_path, center) if os.path.isfile(sdf_path) else False + mol2_ok = _translate_mol2(mol2_path, center) if os.path.isfile(mol2_path) else False + if sdf_ok or mol2_ok: + logger.info( + f"Re-centred karmadock ligand on Vina box {box_file} " + f"(SDF={sdf_ok}, mol2={mol2_ok})" + ) + return sdf_ok or mol2_ok + + def process_karmadock_files( input_folder, ligand_name=KARMADOCK_LIGAND, protein_name=KARMADOCK_PROTEIN, mode="single" ): @@ -102,8 +223,22 @@ def deploy_karmadock( :return: Failed steps. """ failed = 0 + # Resolve KarmaDock install dir. run_guild.py patches WORKING_DIR_PATH to + # /workspace (so guild's own output lands in the mounted host folder), but + # KarmaDock itself is baked into the image at /app/KarmaDock. Prefer the + # home_path location if it exists; otherwise fall back to /app/KarmaDock — + # same pattern as deploy_diffdock. + karmadock_root = os.path.join(home_path, "KarmaDock") + if not os.path.isdir(karmadock_root): + karmadock_root = "/app/KarmaDock" + if not os.path.isdir(karmadock_root): + logger.error( + f"KarmaDock install not found at {os.path.join(home_path, 'KarmaDock')} or /app/KarmaDock" + ) + return failed + pre_processing_command = f""" - {PYTHON_EXECUTABLE} -u {home_path}/KarmaDock/utils/pre_processing.py --complex_file_dir {karmadock_data_dir} {SHELL_SILENCER} + {PYTHON_EXECUTABLE} -u {karmadock_root}/utils/pre_processing.py --complex_file_dir {karmadock_data_dir} {SHELL_SILENCER} """ try: @@ -120,7 +255,7 @@ def deploy_karmadock( # Note: failed_steps tracking removed - not returned by this function graph_generation_command = f""" - {PYTHON_EXECUTABLE} -u {home_path}/KarmaDock/utils/generate_graph.py --complex_file_dir {karmadock_data_dir} --graph_file_dir {karmadock_graphs_dir} {SHELL_SILENCER} + {PYTHON_EXECUTABLE} -u {karmadock_root}/utils/generate_graph.py --complex_file_dir {karmadock_data_dir} --graph_file_dir {karmadock_graphs_dir} {SHELL_SILENCER} """ try: @@ -137,9 +272,9 @@ def deploy_karmadock( # Note: failed_steps tracking removed - not returned by this function docking_command = f""" - {PYTHON_EXECUTABLE} -u {home_path}/KarmaDock/utils/ligand_docking.py \ + {PYTHON_EXECUTABLE} -u {karmadock_root}/utils/ligand_docking.py \ --graph_file_dir {karmadock_graphs_dir} \ - --model_file {home_path}/KarmaDock/trained_models/karmadock_screening.pkl \ + --model_file {karmadock_root}/trained_models/karmadock_screening.pkl \ --out_dir {karmadock_results_dir} \ --docking True --scoring True --correct True --batch_size 64 --random_seed 2023 {SHELL_SILENCER} """ @@ -178,8 +313,11 @@ def deploy_karmadock( def karmadock_guild_scoring(batch_dictionary): """ Perform the scoring of the docking results for KarmaDock. - Ensures that output contains: - combination_id, protein_id, ligand_id, karmadock_score + + Returns a DataFrame with columns + ``combination``, ``protein_config_id``, ``ligand_id``, ``karmadock_score`` + matching what the bulk rank/percentile engine expects (groups by + ``protein_config_id``, not ``protein_id``). """ df = ( @@ -191,23 +329,41 @@ def karmadock_guild_scoring(batch_dictionary): .drop(columns=KARMADOCK_COLUMNS_TO_DROP, errors="ignore") ) + # KarmaDock's result CSV only carries the combined ``pdb_id`` string + # (``{protein_config_id}_{ligand_id}``); unlike vina/gnina it has no separate + # protein_config_id / ligand_id fields. Naively splitting on "_" is ambiguous + # whenever EITHER part contains an underscore (e.g. protein_config_id + # "6C97_P0", or ligand_id "pos_3"), which produces wrong keys → the row fails + # the [COMBINATION_ID, PROTEIN_CONF_ID, LIGAND_ID] merge against the other + # methods and gets its scores dropped downstream. Recover the authoritative + # (protein_config_id, ligand_id) by matching the full combined string against + # COMBINATIONS_TO_RUN_KEY — the same tuples vina/gnina use to build their keys. + combo_to_keys = { + f"{conf_id}_{lig_id}": (conf_id, lig_id) + for conf_id, lig_id in batch_dictionary[COMBINATIONS_TO_RUN_KEY] + } + def split_combination(cid: str): + """Resolve ``{protein_config_id}_{ligand_id}`` → (protein_config_id, + ligand_id) via the authoritative combination map, with a best-effort + ``split("_", 1)`` fallback for ids not present in the batch (shouldn't + normally happen).""" + if isinstance(cid, str) and cid in combo_to_keys: + return combo_to_keys[cid] if not isinstance(cid, str) or "_" not in cid: - # fallback case → avoid crash - return ("unknown_protein", "unknown_ligand") - - parts = cid.split("_") - if len(parts) < 2: - return ("unknown_protein", "unknown_ligand") - - protein = parts[0] - ligand = parts[-1] - return (protein, ligand) + return ("unknown_config", "unknown_ligand") + protein_config_id, ligand_id = cid.split("_", 1) + return (protein_config_id, ligand_id) try: - df[PROTEIN_ID], df[LIGAND_ID] = zip( + df[PROTEIN_CONF_ID], df[LIGAND_ID] = zip( *df[COMBINATION_ID].apply(split_combination), strict=True ) - df[[COMBINATION_ID, PROTEIN_ID, LIGAND_ID, KARMADOCK_SCORE]].copy() + # Drop the legacy ``protein_id`` column that ``process_karmadock_files`` + # may have left behind (renamed from pdb_id) — bulk's rank engine only + # needs ``protein_config_id``. + return df[[COMBINATION_ID, PROTEIN_CONF_ID, LIGAND_ID, KARMADOCK_SCORE]].copy() except Exception: - return pd.DataFrame(columns=[COMBINATION_ID, PROTEIN_ID, LIGAND_ID, KARMADOCK_SCORE]) + return pd.DataFrame( + columns=[COMBINATION_ID, PROTEIN_CONF_ID, LIGAND_ID, KARMADOCK_SCORE] + ) diff --git a/guild/docking/vina.py b/guild/docking/vina.py index 1846194..6926429 100644 --- a/guild/docking/vina.py +++ b/guild/docking/vina.py @@ -3,8 +3,6 @@ """ import logging -import os -import subprocess import numpy as np import pandas as pd @@ -16,15 +14,11 @@ COMBINATION_ID, COMBINATIONS_TO_RUN_KEY, ) -from guild.constants.diffdock import DIFFDOCK_RESULTS_FOLDER from guild.constants.general import RANDOM_SEED from guild.constants.guild import ( - DIFFDOCK_FOLDER, LIGAND_ID, PROTEIN_CONF_ID, VINA_FOLDER, - VINA_RESCORE_FOLDER, - VINA_RESCORE_SCORE, VINA_SCORE, ) from guild.constants.vina import ( @@ -35,7 +29,6 @@ radius_of_gyration_from_smiles, vina_box_edge_from_radius_of_gyration, ) -from guild.transformers.converters import sdf_to_pdbqt logger = logging.getLogger(__name__) @@ -241,29 +234,14 @@ def score_combination(row): return combinations_df[[COMBINATION_ID, VINA_SCORE, PROTEIN_CONF_ID, LIGAND_ID]] -def vina_rescore_guild_scoring(batch_dictionary): - """ - Perform Vina re-scoring of DiffDock poses for a batch. - - Thin wrapper around :func:`vina_rescore_diffdock_batch` that accepts a - *batch_dictionary* (the same structure every other ``*_guild_scoring`` - function receives) and returns a DataFrame compatible with the bulk - scoring merge logic. - - :param batch_dictionary: Dictionary containing the batch information. - :return: DataFrame with COMBINATION_ID, PROTEIN_CONF_ID, LIGAND_ID, - and VINA_RESCORE_SCORE columns. - """ - df = vina_rescore_diffdock_batch( - batch_folder=batch_dictionary[BATCH_FOLDER], - combinations=batch_dictionary[COMBINATIONS_TO_RUN_KEY], - ) - # Keep only the columns the merge expects (drop 'diffdock_sdf') - keep = [COMBINATION_ID, PROTEIN_CONF_ID, LIGAND_ID, VINA_RESCORE_SCORE] - return df[[c for c in keep if c in df.columns]] - - # ── Vina score-only re-scoring of pre-docked poses ────────────────────────── +# NOTE: Method-specific orchestration (rescore_boltz_pose, rescore_diffdock_pose, +# vina_rescore_*_batch, vina_rescore_*_guild_scoring) lives in +# guild/docking/boltz.py and guild/docking/diffdock.py — they know the output +# layout of their respective methods. This module keeps only the Vina-grid +# primitives that any pose source can reuse (compute_box_from_sdf, +# _compute_box_from_pdb_atoms, _extract_ligand_records, +# _extract_protein_from_complex, vina_score_pose). def compute_box_from_sdf(sdf_path: str, padding: float = 4.0): @@ -334,248 +312,86 @@ def vina_score_pose( return float(energy[0]) -def _find_best_diffdock_sdf(diffdock_results_dir: str, protein_conf_id: str, ligand_id: str) -> str: - """ - Locate the highest-confidence SDF file for a given combination in a - DiffDock results directory. +# ── Vina score-only re-scoring of Boltz-predicted complexes ───────────────── - DiffDock names output files like ``rank1_confidence-1.23.sdf``. - :return: Absolute path to the best-confidence SDF. - :raises FileNotFoundError: If no SDF files are found for the combination. +def _compute_box_from_pdb_atoms(pdb_path: str, padding: float = 4.0): """ - folder_name = f"{protein_conf_id}_{ligand_id}" - combo_dir = os.path.join(diffdock_results_dir, folder_name) - - if not os.path.isdir(combo_dir): - raise FileNotFoundError(f"DiffDock results folder not found: {combo_dir}") - - sdf_scores = {} - for fname in os.listdir(combo_dir): - if "_confidence" in fname and fname.endswith(".sdf"): + Compute a Vina box (center + size) from the coordinates of all ATOM/HETATM + records in a PDB file. Used to size the Vina scoring grid around a + pre-docked ligand pose extracted from a Boltz complex. + """ + coords = [] + with open(pdb_path) as f: + for line in f: + if not line.startswith(("ATOM", "HETATM")): + continue try: - score = float(fname.split("_confidence")[1].replace(".sdf", "")) - sdf_scores[fname] = score + x = float(line[30:38]) + y = float(line[38:46]) + z = float(line[46:54]) except ValueError: continue - - if not sdf_scores: - raise FileNotFoundError(f"No DiffDock SDF files found in {combo_dir}") - - best_fname = max(sdf_scores, key=sdf_scores.get) - return os.path.join(combo_dir, best_fname) + coords.append((x, y, z)) + if not coords: + raise ValueError(f"No coordinates parsed from {pdb_path}") + arr = np.array(coords) + mins = arr.min(axis=0) + maxs = arr.max(axis=0) + center = tuple(((mins + maxs) / 2.0).tolist()) + span = maxs - mins + size = tuple((span + 2.0 * padding).tolist()) + return center, size -def rescore_diffdock_pose( - receptor_pdbqt: str, - diffdock_results_dir: str, - protein_conf_id: str, - ligand_id: str, - output_dir: str = None, - box_padding: float = 4.0, - seed: int = RANDOM_SEED, -) -> dict: +def _extract_ligand_records(input_pdb: str, output_pdb: str, resname: str = "LIG"): """ - Re-score a DiffDock pose with Vina's physics-based scoring function. - - Pipeline: - 1. Find the highest-confidence SDF in ``diffdock_results_dir/{protein}_{ligand}/`` - 2. Convert SDF → PDBQT (preserving 3D coordinates) - 3. Compute a Vina box centred on the SDF pose - 4. Call ``vina_score_pose()`` (score-only, no re-docking) - - :param receptor_pdbqt: Path to the receptor PDBQT file. - :param diffdock_results_dir: Path to the DiffDock results directory. - :param protein_conf_id: Protein configuration ID (e.g. ``8gut-R-KO8-R``). - :param ligand_id: Ligand identifier. - :param output_dir: Directory for intermediate PDBQT files (defaults to - ``diffdock_results_dir/{protein}_{ligand}/``). - :param box_padding: Padding around the ligand bounding box (Å). - :param seed: Random seed for Vina. - :return: Dict with ``combination``, ``protein_config_id``, ``ligand_id``, - ``vina_rescore_score``, ``diffdock_sdf``. + Write a PDB containing only ATOM/HETATM records whose residue name matches + ``resname``. Used to isolate the ligand of a Boltz complex (whose ligand is + written under residue name ``LIG`` by :func:`relabel_ligand_chain_in_pdb` — + the chain ID is not a reliable marker because ``cif_to_pdb`` may rename it). """ - sdf_path = _find_best_diffdock_sdf(diffdock_results_dir, protein_conf_id, ligand_id) - - if output_dir is None: - output_dir = os.path.dirname(sdf_path) - - # SDF → PDBQT (preserves 3D coordinates) - ligand_pdbqt = os.path.join(output_dir, f"{protein_conf_id}_{ligand_id}_rescore.pdbqt") - sdf_to_pdbqt(sdf_path, pdbqt=ligand_pdbqt) - - # Compute box from the DiffDock pose coordinates - center, size = compute_box_from_sdf(sdf_path, padding=box_padding) - - # Score-only evaluation - score = vina_score_pose( - receptor_pdbqt=receptor_pdbqt, - ligand_pdbqt=ligand_pdbqt, - center=center, - size=size, - seed=seed, - ) - - combination_id = f"{protein_conf_id}_{ligand_id}" - logger.info(f"Vina rescore: {combination_id} → {score:.3f} kcal/mol") - - return { - COMBINATION_ID: combination_id, - PROTEIN_CONF_ID: protein_conf_id, - LIGAND_ID: ligand_id, - VINA_RESCORE_SCORE: score, - "diffdock_sdf": sdf_path, - } + kept = 0 + resname_padded = resname.ljust(3)[:3] + with open(input_pdb) as fin, open(output_pdb, "w") as fout: + for line in fin: + if ( + line.startswith(("ATOM", "HETATM")) + and len(line) > 20 + and line[17:20] == resname_padded + ): + fout.write(line) + kept += 1 + fout.write("END\n") + if kept == 0: + raise ValueError(f"No atoms with resname '{resname}' found in {input_pdb}") + return output_pdb -def _prepare_receptor_pdbqt_from_raw( - raw_pdb: str, - chain_id: str, - output_pdbqt: str, -) -> str: +def _extract_protein_from_complex(complex_pdb: str, output_pdb: str, ligand_resname: str = "LIG"): """ - Extract a single chain (ATOM records only) from a raw PDB and convert to - PDBQT via OpenBabel. This preserves the original crystal coordinates so - the receptor is in the same frame as DiffDock's output SDF files. - - :param raw_pdb: Path to the raw (multi-chain) PDB file. - :param chain_id: Chain letter to extract (e.g. ``"R"``). - :param output_pdbqt: Where to write the receptor PDBQT. - :return: *output_pdbqt* path. - :raises FileNotFoundError: If *raw_pdb* does not exist. - :raises RuntimeError: If ``obabel`` is not available or conversion fails. + Write a PDB containing all ATOM records of the complex EXCEPT those that + are the ligand (resname == ``ligand_resname``). Used to obtain a receptor + PDB in Boltz's predicted coordinate frame — critical because Boltz often + re-centres the whole complex, so the template-frame receptor and the + Boltz-output ligand are not in the same physical space. """ - if not os.path.isfile(raw_pdb): - raise FileNotFoundError(f"Raw PDB not found: {raw_pdb}") - - # Write a single-chain PDB (ATOM records only, no waters/HETATM) - chain_pdb = output_pdbqt.replace(".pdbqt", "_chain.pdb") kept = 0 - with open(raw_pdb) as fin, open(chain_pdb, "w") as fout: + resname_padded = ligand_resname.ljust(3)[:3] + with open(complex_pdb) as fin, open(output_pdb, "w") as fout: for line in fin: - if line.startswith("ATOM") and len(line) > 21 and line[21] == chain_id: + if line.startswith(("ATOM", "HETATM")) and len(line) > 20: + if line[17:20] == resname_padded: + continue fout.write(line) kept += 1 + elif line.startswith(("TER", "END")): + fout.write(line) fout.write("END\n") - if kept == 0: - raise ValueError(f"No ATOM records found for chain {chain_id} in {raw_pdb}") - - # Convert PDB → PDBQT using OpenBabel (receptor mode: -xr) - result = subprocess.run( - ["obabel", "-ipdb", chain_pdb, "-opdbqt", "-O", output_pdbqt, "-xr"], - capture_output=True, - text=True, - ) - if not os.path.isfile(output_pdbqt) or os.path.getsize(output_pdbqt) == 0: - raise RuntimeError( - f"obabel PDB→PDBQT failed for {chain_pdb}: {result.stderr}" + raise ValueError( + f"No protein atoms left in {complex_pdb} after excluding resname '{ligand_resname}'" ) + return output_pdb - logger.info( - f"Prepared receptor PDBQT from raw PDB chain {chain_id}: " - f"{kept} atoms → {output_pdbqt}" - ) - return output_pdbqt - - -def vina_rescore_diffdock_batch( - batch_folder: str, - combinations: list[tuple[str, str]], - receptor_pdbqt_dir: str = None, - box_padding: float = 4.0, - seed: int = RANDOM_SEED, -) -> pd.DataFrame: - """ - Batch re-score DiffDock poses for all combinations in a batch. - - DiffDock produces ligand poses in the coordinate frame of the **raw** - input PDB. The guild pipeline's ``_single_chain_clean`` PDB is - typically re-centred, so it cannot be used directly as the Vina - receptor. This function therefore looks for a receptor PDBQT in the - following order: - - 1. ``{proteins_dir}/{protein_conf_id}_raw.pdbqt`` — single-chain - PDBQT already in the raw (crystal) coordinate frame. - 2. If (1) does not exist, it auto-generates it from - ``{protein_conf_id}_raw.pdb`` by extracting the chain letter - encoded in *protein_conf_id* and converting via OpenBabel. - - :param batch_folder: Path to the batch folder (e.g. ``data/project/batches/batch_1``). - :param combinations: List of (protein_conf_id, ligand_id) tuples. - :param receptor_pdbqt_dir: Directory containing receptor PDBQT files. - Defaults to ``{batch_folder}/proteins/``. - :param box_padding: Padding around the ligand bounding box (Å). - :param seed: Random seed for Vina. - :return: DataFrame with columns: combination, protein_config_id, ligand_id, - vina_rescore_score. - """ - diffdock_results_dir = os.path.join(batch_folder, DIFFDOCK_FOLDER, DIFFDOCK_RESULTS_FOLDER) - if receptor_pdbqt_dir is None: - receptor_pdbqt_dir = os.path.join(batch_folder, "proteins") - - rescore_output_dir = os.path.join(batch_folder, VINA_RESCORE_FOLDER) - os.makedirs(rescore_output_dir, exist_ok=True) - - results = [] - for protein_conf_id, ligand_id in combinations: - # --- Resolve receptor PDBQT in the DiffDock (raw) coordinate frame --- - receptor_pdbqt = os.path.join( - receptor_pdbqt_dir, f"{protein_conf_id}_raw.pdbqt" - ) - if not os.path.exists(receptor_pdbqt): - # Auto-generate from the raw PDB (extract chain, convert) - raw_pdb = os.path.join(receptor_pdbqt_dir, f"{protein_conf_id}_raw.pdb") - if os.path.isfile(raw_pdb): - # Extract chain letter from protein_conf_id (e.g. "8gut-R-KO8-R" → "R") - parts = protein_conf_id.split("-") - chain_id = parts[1] if len(parts) >= 2 else "A" - try: - _prepare_receptor_pdbqt_from_raw(raw_pdb, chain_id, receptor_pdbqt) - except Exception as e: - logger.warning( - f"Could not prepare receptor PDBQT from raw PDB " - f"for {protein_conf_id}: {e}" - ) - receptor_pdbqt = None - else: - logger.warning( - f"Neither {receptor_pdbqt} nor {raw_pdb} found for " - f"{protein_conf_id}, skipping combination" - ) - receptor_pdbqt = None - - if receptor_pdbqt is None or not os.path.exists(receptor_pdbqt): - results.append({ - COMBINATION_ID: f"{protein_conf_id}_{ligand_id}", - PROTEIN_CONF_ID: protein_conf_id, - LIGAND_ID: ligand_id, - VINA_RESCORE_SCORE: np.nan, - "diffdock_sdf": None, - }) - continue - try: - result = rescore_diffdock_pose( - receptor_pdbqt=receptor_pdbqt, - diffdock_results_dir=diffdock_results_dir, - protein_conf_id=protein_conf_id, - ligand_id=ligand_id, - output_dir=rescore_output_dir, - box_padding=box_padding, - seed=seed, - ) - results.append(result) - except Exception as e: - logger.warning( - f"Vina rescore failed for {protein_conf_id}_{ligand_id}: {e}" - ) - results.append({ - COMBINATION_ID: f"{protein_conf_id}_{ligand_id}", - PROTEIN_CONF_ID: protein_conf_id, - LIGAND_ID: ligand_id, - VINA_RESCORE_SCORE: np.nan, - "diffdock_sdf": None, - }) - - return pd.DataFrame(results) diff --git a/guild/run.py b/guild/run.py index 7b6c4e6..1297141 100644 --- a/guild/run.py +++ b/guild/run.py @@ -20,6 +20,8 @@ DATA_FOLDER, DIFFDOCK_FOLDER, DIFFDOCK_PREFIX, + GNINA_FOLDER, + GNINA_PREFIX, KARMADOCK_FOLDER, KARMADOCK_PREFIX, LIGANDS_FOLDER, @@ -45,6 +47,7 @@ from guild.constants.vina import VINA_BOXES_FOLDER from guild.docking.boltz import deploy_boltz, generate_boltz_yaml from guild.docking.diffdock import deploy_diffdock_single +from guild.docking.gnina import deploy_gnina from guild.docking.karmadock import deploy_karmadock from guild.docking.vina import ( deploy_vina, @@ -53,6 +56,7 @@ ) from guild.tools.p2rank import get_binding_site_center_from_p2rank from guild.tools.preparation import ( + _normalize_chain_list, clean_receptor, get_protein_chain, isolate_protein_chain, @@ -70,7 +74,11 @@ smiles_to_sdf_karmadock, ) from guild.transformers.msa import fetch_protein_msa -from guild.transformers.pdb import LigandIdentifier, get_pocket_contacts_from_ligand +from guild.transformers.pdb import ( + LigandIdentifier, + get_pocket_contacts_from_box, + get_pocket_contacts_from_ligand, +) logger = logging.getLogger(__name__) @@ -94,6 +102,8 @@ def __init__( use_gpu: bool = True, is_bulk: bool = False, predict_binding_pocket: bool = False, + box_location: str = None, + gnina_input_mode: str = "pdbqt", ): """ Start the Guild for docking ligands to proteins. @@ -109,6 +119,10 @@ def __init__( :param use_gpu: Use GPU. :param is_bulk: Whether the run is part of a bulk run. :param predict_binding_pocket: Use P2Rank for binding site prediction instead of original ligand location. + :param box_location: Path to a user-supplied Vina box file (center_{x,y,z} + size_{x,y,z}). + When set, takes precedence over P2Rank / original-ligand pocket derivation + for both Vina (used as the Vina box) and Boltz (residues inside the box + become the pocket contact constraint). """ self.original_ligand_smile = ligand_smile self.ligand_idx = ligand_idx @@ -121,11 +135,18 @@ def __init__( ) self.combination_id = f"{self.protein_idx}_{self.ligand_idx}" - # Get first chain of the protein if no target chain is specified - if protein_chain is None: - self.protein_chain = get_protein_chain(self.original_protein_data, self.protein_idx) - else: - self.protein_chain = protein_chain + # Resolve the protein chain(s). ``protein_chain`` may be a single ID + # ("A"), a list, or a comma-separated string ("A,B") for a pocket that + # spans multiple chains (e.g. a dimer interface). ``self.protein_chains`` + # is the canonical list; ``self.protein_chain`` keeps the first chain as + # the "primary" for single-chain code paths (box / original-ligand + # pocket derivation, MSA/file naming). + self.protein_chains = _normalize_chain_list(protein_chain) + if not self.protein_chains: + self.protein_chains = [ + get_protein_chain(self.original_protein_data, self.protein_idx) + ] + self.protein_chain = self.protein_chains[0] self.original_ligand = original_ligand self.original_ligand_chain = original_ligand_chain @@ -135,6 +156,28 @@ def __init__( self.use_gpu = use_gpu self.is_bulk = is_bulk self.predict_binding_pocket = predict_binding_pocket + # ``gnina_input_mode == "sdf"`` skips OpenBabel-driven PDBQT prep for + # both the ligand (in ``_prepare_ligand``) and the receptor (in + # ``run_gnina``). BulkRun resolves the effective mode globally — + # falling back to "pdbqt" when Vina/Vina-rescore is also requested — + # so by the time we get here the value is final. + self.gnina_input_mode = gnina_input_mode + # Validate box_location up-front. An unreadable path is treated as + # "not provided" so the normal P2Rank / original-ligand fallback chain + # still kicks in instead of silently leaving Vina/Boltz without a + # pocket. Downstream code can then trust ``self.box_location is not + # None`` as the single source of truth for "user supplied a box". + if box_location: + if os.path.exists(box_location): + self.box_location = box_location + else: + logger.warning( + f"box_location {box_location} does not exist on disk; " + "falling back to P2Rank / original-ligand for the binding pocket." + ) + self.box_location = None + else: + self.box_location = None self._set_paths(home_path=self.home_path, output_log_file=output_log_file) self._create_directories() @@ -160,8 +203,8 @@ def __init__( logger.error(f"Error in preparing ligand data: {e}") try: - self._prepare_protein(target_chain=self.protein_chain) - logger.info(f"Protein data prepared for chain {self.protein_chain}.") + self._prepare_protein(target_chain=self.protein_chains) + logger.info(f"Protein data prepared for chain(s) {self.protein_chains}.") except Exception as e: logger.error(f"Error in preparing protein data: {e}") @@ -209,6 +252,12 @@ def _set_paths(self, home_path: str = None, output_log_file: str = None): self.vina_output_pdbqt = f"{self.vina_dir}/{self.protein_idx}_{self.ligand_idx}.pdbqt" self.vina_output_scores = f"{self.vina_dir}/{self.protein_idx}_{self.ligand_idx}.txt" + # GNINA directories — reuses the Vina box file (same format) and the + # PDBQT-prepped receptor/ligand produced for Vina. + self.gnina_dir = f"{self.project_dir}/{GNINA_FOLDER}" + self.gnina_output_pdbqt = f"{self.gnina_dir}/{self.protein_idx}_{self.ligand_idx}.pdbqt" + self.gnina_output_scores = f"{self.gnina_dir}/{self.protein_idx}_{self.ligand_idx}.txt" + # KarmaDock directories self.karmadock_root = f"{self.project_dir}/{KARMADOCK_FOLDER}" self.karmadock_graphs_dir = f"{self.karmadock_root}/{KARMADOCK_GRAPHS_FOLDER}" @@ -262,6 +311,7 @@ def _create_directories(self): self.ligand_dir, self.protein_dir, self.vina_dir, + self.gnina_dir, self.karmadock_root, self.karmadock_data_dir, self.karmadock_data_sub_dir, @@ -303,6 +353,16 @@ def _relocate_files(self): if not os.path.exists(self.local_protein): copyfile(self.original_protein_data, self.local_protein) + # User-supplied Vina box short-circuits both _refine_protein_data (which + # would otherwise derive the box from the original co-crystal ligand) and + # P2Rank prediction further down. The path was validated in __init__, so + # self.box_location is guaranteed to be an existing file when non-None. + if self.box_location: + os.makedirs(os.path.dirname(self.vina_box), exist_ok=True) + copyfile(self.box_location, self.vina_box) + logger.info(f"Seeded Vina box from user-supplied {self.box_location}") + return + if ( (self.protein_chain is not None) and (self.original_ligand is not None) @@ -313,6 +373,12 @@ def _relocate_files(self): def _prepare_ligand(self): """ Preprocessing the ligand data for docking. + + When ``self.gnina_input_mode == "sdf"`` the OpenBabel-backed + ``ligand_pdb_to_pdbqt`` step (and its prerequisite ``sdf_to_pdb``) + are skipped: gnina reads the RDKit-generated SDF directly, and no + other method requires the PDBQT in this run (BulkRun has already + downgraded to PDBQT if Vina or a Vina-rescore was co-requested). """ if os.path.exists(self.ligand_sdf): logger.info(f"Ligand {self.ligand_idx} already prepared.") @@ -323,6 +389,13 @@ def _prepare_ligand(self): sdf=self.ligand_sdf, ) + if self.gnina_input_mode == "sdf": + logger.info( + f"Skipping OpenBabel PDBQT prep for ligand {self.ligand_idx} " + "(gnina_input_mode='sdf')." + ) + return + sdf_to_pdb( sdf=self.ligand_sdf, pdb=self.ligand_pdb, @@ -330,9 +403,12 @@ def _prepare_ligand(self): ligand_pdb_to_pdbqt(pdb=self.ligand_pdb) - # Generate ligand-specific box based on this decoy's radius of gyration - # Keep the binding site center, but adjust size for this specific ligand - if os.path.exists(self.vina_box): + # Generate ligand-specific box based on this decoy's radius of gyration. + # Keep the binding site center, but adjust size for this specific ligand. + # A user-supplied box (self.box_location) defines both center AND size + # explicitly, so leave it untouched — only auto-derived boxes (P2Rank or + # original-ligand) get resized to the ligand's radius of gyration. + if os.path.exists(self.vina_box) and not self.box_location: center, _ = get_center_and_size_from_box_file(self.vina_box) generate_vina_box( input_x=center[0], @@ -349,6 +425,13 @@ def _prepare_ligand_for_karmadock(self): """ Prepare ligand specifically for KarmaDock (MOL2 conversion). Only called when KarmaDock is actually being used. + + KarmaDock derives the binding pocket from residues within 12 Å of the + *input ligand's* coordinates (see ``KarmaDock/utils/pre_processing.py`` + ``get_pocket()``). After generating the SDF/mol2 from SMILES we + therefore re-centre them on the Vina box if one is available — that + way KarmaDock picks the user-supplied (or P2Rank- / original-ligand- + derived) pocket instead of wherever obabel's 3D embedding lands. """ if os.path.exists(self.karmadock_mol2): logger.info(f"KarmaDock ligand {self.ligand_idx} already prepared.") @@ -369,14 +452,30 @@ def _prepare_ligand_for_karmadock(self): else: logger.info("Command for converting SMILES to MOL2 succeeded") + # Re-centre on the binding pocket. By this point self.vina_box has been + # populated by box_location (_relocate_files) and/or P2Rank + # (_prepare_protein) and/or _refine_protein_data (original ligand), so + # any available pocket signal drives KarmaDock's get_pocket() correctly. + if os.path.isfile(self.vina_box): + from guild.docking.karmadock import center_karmadock_ligand_on_box + + center_karmadock_ligand_on_box( + sdf_path=self.karmadock_ligand, + mol2_path=self.karmadock_mol2, + box_file=self.vina_box, + ) + def _prepare_protein(self, target_chain=None): """ Preprocessing the protein data for docking. - :param target_chain: The chain to be used for docking. If not provided, the first chain is used. + :param target_chain: Chain ID, list of chain IDs, or comma-separated + string to be used for docking. If not provided, the + first chain is used. Multiple chains are kept + together so a multi-chain pocket survives prep. """ - # Isolate the protein chain + # Isolate the protein chain(s) if not os.path.exists(self.single_chain_protein): isolate_protein_chain( input_file=self.local_protein, @@ -402,14 +501,23 @@ def _prepare_protein(self, target_chain=None): else: logger.info(f"Protein {self.protein_idx} already copied to karmadock data directory.") - # Get the original sequence of the protein + # Get the original sequence of the protein, one per requested chain. + # ``self.original_sequence`` keeps the primary chain's sequence for any + # single-chain consumer; ``self.original_sequences`` is the per-chain + # list (same order as ``self.protein_chains``) used by Boltz. original_sequence_dictionary = get_original_sequence_dictionary(self.cleaned_protein) - self.original_sequence = process_into_fasta_string( - original_sequence_dictionary[self.protein_chain] - ) - - # If P2Rank binding site prediction is enabled, run it after protein is prepared - if self.predict_binding_pocket: + self.sequence_chains = [ + chain for chain in self.protein_chains if chain in original_sequence_dictionary + ] + self.original_sequences = [ + process_into_fasta_string(original_sequence_dictionary[chain]) + for chain in self.sequence_chains + ] + self.original_sequence = self.original_sequences[0] if self.original_sequences else "" + + # If P2Rank binding site prediction is enabled, run it after protein is prepared. + # A user-supplied box overrides P2Rank — it already populated self.vina_box. + if self.predict_binding_pocket and not self.box_location: try: self._run_p2rank_binding_site_prediction() logger.info("P2Rank binding site prediction completed.") @@ -497,6 +605,77 @@ def run_autodock_vina(self): return 0 if failed_steps == 0 else 1 + @timeit() + def run_gnina(self): + """ + Run gnina docking. In the default ``pdbqt`` mode (re)prepares the + receptor PDBQT and feeds gnina the PDBQT receptor + ligand. In + ``sdf`` mode, skips the OpenBabel-backed PDBQT prep entirely and + feeds gnina the PDB receptor + SDF ligand directly (gnina sniffs + the format from the extension). + """ + failed_steps = 0 + + if self.gnina_input_mode == "sdf": + receptor_path = self.cleaned_protein # PDB + ligand_path = self.ligand_sdf # SDF (RDKit-generated upstream) + logger.info("gnina: SDF-mode — skipping PDBQT prep.") + else: + # Receptor PDBQT prep — same as Vina. Run independently of + # run_autodock_vina because gnina can be requested without vina. + try: + protein_pdb_to_pdbqt( + input_pdb=self.cleaned_protein, + output_pdbqt=self.cleaned_protein_pdbqt, + allow_bad_res=True, + ) + logger.info("gnina: receptor prepared.") + except Exception as e: + logger.error(f"gnina: error in preparing receptor: {e}") + failed_steps += 1 + receptor_path = self.cleaned_protein_pdbqt + ligand_path = self.ligand_pdbqt + + # Require a *non-empty* pose file: older runs (before the OpenBabel + # plugin fix) left 0-byte pdbqts behind, and treating those as "done" + # would permanently skip regenerating usable poses on re-run. + try: + output_ok = ( + os.path.isfile(self.gnina_output_pdbqt) + and os.path.getsize(self.gnina_output_pdbqt) > 0 + ) + except OSError: + output_ok = False + + if output_ok: + logger.info("gnina: output file already exists.") + return failed_steps + + gnina_box_center, gnina_box_size = get_center_and_size_from_box_file(self.vina_box) + + # Per-combination subprocess transcript — predictable path that the + # orchestrator points users at when this combo fails. See + # guild/tools/subprocess_log.py for the format. + subprocess_log = ( + f"{self.gnina_dir}/{self.protein_idx}_{self.ligand_idx}.subprocess.log" + ) + try: + deploy_gnina( + receptor=receptor_path, + ligand=ligand_path, + center=gnina_box_center, + size=gnina_box_size, + output_pdbqt=self.gnina_output_pdbqt, + output_scores=self.gnina_output_scores, + use_gpu=self.use_gpu, + subprocess_log_path=subprocess_log, + ) + except Exception as e: + logger.error(f"gnina: error in docking: {e}") + failed_steps += 1 + + return 0 if failed_steps == 0 else 1 + @timeit() def run_karmadock(self): """ @@ -530,18 +709,44 @@ def run_boltz(self): """ Run Boltz for docking the ligand to the protein. Output is saved in the boltz directory, inside the project folder. """ - msa_file = fetch_protein_msa( - sequence=self.original_sequence, - protein_id=self.protein_idx, - protein_chain_id=self.protein_chain, - output_a3m_dir=self.msa_cache_dir, - ) - print(f"MSA file for Boltz: {msa_file}") - - if self.original_ligand is not None and self.original_ligand_chain is not None: + # Fetch one MSA per chain (cache-keyed by chain), aligned with + # self.sequence_chains / self.original_sequences for the Boltz YAML. + msa_files = [ + fetch_protein_msa( + sequence=sequence, + protein_id=self.protein_idx, + protein_chain_id=chain, + output_a3m_dir=self.msa_cache_dir, + ) + for chain, sequence in zip( + self.sequence_chains, self.original_sequences, strict=True + ) + ] + print(f"MSA file(s) for Boltz: {msa_files}") + + # self.box_location is the single source of truth for "user supplied a + # box" (validated in __init__). Don't fall back to checking self.vina_box + # — that file may have been generated from the original co-crystal + # ligand via _refine_protein_data, in which case it isn't actually the + # user-supplied pocket and shouldn't drive the Boltz pocket_contacts. + # Pocket contacts are collected across all docking chains so an + # interface pocket spanning multiple chains is fully constrained. + if self.box_location: + center, size = get_center_and_size_from_box_file(self.vina_box) + pocket_contacts = get_pocket_contacts_from_box( + protein_pdb=self.local_protein, + protein_chain=self.sequence_chains, + center=center, + size=size, + ) + print( + f"Identified {len(pocket_contacts)} pocket contact residues for Boltz " + f"constraints from user-supplied box." + ) + elif self.original_ligand is not None and self.original_ligand_chain is not None: pocket_contacts = get_pocket_contacts_from_ligand( protein_pdb=self.local_protein, - protein_chain=self.protein_chain, + protein_chain=self.sequence_chains, original_ligand=self.original_ligand, original_ligand_chain=self.original_ligand_chain, distance_threshold=4.0, @@ -553,13 +758,13 @@ def run_boltz(self): pocket_contacts = None generate_boltz_yaml( - protein_sequence=self.original_sequence, - protein_chain=self.protein_chain, + protein_sequence=self.original_sequences, + protein_chain=self.sequence_chains, ligand_sequences=[self.original_ligand_smile], ligand_ids=["L"], # Boltz expects ligand Id to be < 5 characters output_file=self.boltz_yaml, template_file=self.local_protein, - msa_file=msa_file, + msa_file=msa_files, pocket_contacts=pocket_contacts, ) deploy_boltz(self.boltz_yaml, out_dir=self.boltz_dir, use_gpu=self.use_gpu) @@ -591,3 +796,11 @@ def dock(self, box_location: str = "", methods: list = ALL_AVAILABLE_METHODS): if BOLTZ_PREFIX in methods: self.run_boltz() logger.info("Boltz: completed.") + + # GNINA + if GNINA_PREFIX in methods: + report_gnina = self.run_gnina() + if report_gnina == 0: + logger.info("gnina: completed.") + else: + logger.error(f"gnina: failed. Failed steps: {report_gnina} / {len(methods)}") diff --git a/guild/tools/bulk.py b/guild/tools/bulk.py index 3b5855b..2011a97 100644 --- a/guild/tools/bulk.py +++ b/guild/tools/bulk.py @@ -16,22 +16,17 @@ ) from guild.constants.bulk import ( BATCH_FOLDER, - BOLTZ_PREFIX, - BOLTZ_RP_SCORE, BULK_TEMPLATE_DICTIONARY, # Batch dictionary keys COMBINATIONS_TABLE_KEY, COMBINATIONS_TO_RUN_KEY, - DIFFDOCK_PREFIX, - DIFFDOCK_RP_SCORE, INPUT_COMBINATIONS_KEY, - KARMADOCK_PREFIX, - KARMADOCK_RP_SCORE, METHODS_TO_SCORE_DICTIONARY_KEY, METHODS_TO_SORT_DICTIONARY_KEY, PRE_EXISTING_RP_SCORES_KEY, PREVIOUS_COMBINATIONS_DF_KEY, PREVIOUS_RP_SCORES_KEY, + RANKS_DICTIONARY, RANKS_LIST_KEY, SCORES_DIRECTION_DICTIONARY, SCORES_TO_USE_DICTIONARY, @@ -39,9 +34,6 @@ SMILES_NAMES_DICTIONARY_KEY, SMILES_TYPE_DICTIONARY_KEY, UNIQUE_PROTEIN_IDS_KEY, - VINA_PREFIX, - VINA_RESCORE_RP_SCORE, - VINA_RP_SCORE, ) from guild.constants.decoys import DECOYS_CATEGORY from guild.constants.diffdock import DIFFDOCK_RESULTS_FOLDER @@ -59,9 +51,9 @@ RP_SCORES_COLUMNS, SMILES, VINA_FOLDER, - VINA_RESCORE_PREFIX, ) from guild.tools.binders import collect_known_binders +from guild.tools.preparation import _normalize_chain_list from guild.transformers.converters import cif_to_pdb, sdf_to_pdb from guild.transformers.pdb import build_complex_pdb, relabel_ligand_chain_in_pdb @@ -156,14 +148,6 @@ def available_methods_preparation(batch_dictionary, methods_to_run): :return: Batch dictionary with the available methods prepared. """ - ranks_dictionary = { - VINA_PREFIX: VINA_RP_SCORE, - KARMADOCK_PREFIX: KARMADOCK_RP_SCORE, - DIFFDOCK_PREFIX: DIFFDOCK_RP_SCORE, - BOLTZ_PREFIX: BOLTZ_RP_SCORE, - VINA_RESCORE_PREFIX: VINA_RESCORE_RP_SCORE, - } - batch_dictionary[METHODS_TO_SCORE_DICTIONARY_KEY] = SCORES_TO_USE_DICTIONARY batch_dictionary[METHODS_TO_SORT_DICTIONARY_KEY] = SCORES_DIRECTION_DICTIONARY @@ -171,7 +155,7 @@ def available_methods_preparation(batch_dictionary, methods_to_run): batch_dictionary[RANKS_LIST_KEY] = [] for current_method in methods_to_run: - batch_dictionary[RANKS_LIST_KEY].append(ranks_dictionary[current_method]) + batch_dictionary[RANKS_LIST_KEY].append(RANKS_DICTIONARY[current_method]) batch_dictionary[SCORES_TO_USE_KEY] += SCORES_TO_USE_DICTIONARY[current_method] return batch_dictionary @@ -286,20 +270,23 @@ def extend_all_combinations_table_with_decoys(input_table, decoys_file_path): return pd.DataFrame(rows_to_add) -def generate_vina_complex_pdbs(batch_dictionary): +def _generate_pdbqt_complex_pdbs(batch_dictionary, output_folder_name: str, method_label: str): """ - Generate complex PDB files for all Vina docking results in a batch. - Merges protein PDB with docked ligand PDBQT into a single complex PDB. + Merge protein PDB with docked-pose PDBQT into a complex PDB, for any + docking method whose output is a multi-pose PDBQT in + ``{BATCH_FOLDER}/{output_folder_name}/{protein}_{ligand}.pdbqt``. - :param batch_dictionary: Dictionary containing batch information including - BATCH_FOLDER, COMBINATIONS_TABLE_KEY, etc. - """ + Used by both Vina and gnina, which share the same PDBQT output format. + :param batch_dictionary: Standard bulk batch dictionary. + :param output_folder_name: Per-method docking-output folder (``VINA_FOLDER``, + ``GNINA_FOLDER``). + :param method_label: Human-readable method name for log messages. + """ batch_folder = batch_dictionary[BATCH_FOLDER] - vina_folder = f"{batch_folder}/{VINA_FOLDER}" + method_folder = f"{batch_folder}/{output_folder_name}" proteins_folder = f"{batch_folder}/proteins" - # Get all combinations for this batch combinations_df = batch_dictionary[COMBINATIONS_TABLE_KEY] complexes_created = 0 @@ -309,19 +296,15 @@ def generate_vina_complex_pdbs(batch_dictionary): protein_conf_id = row[PROTEIN_CONF_ID] ligand_id = row[LIGAND_ID] - # Paths for vina output - ligand_pdbqt = f"{vina_folder}/{protein_conf_id}_{ligand_id}.pdbqt" - complex_pdb = f"{vina_folder}/{protein_conf_id}_{ligand_id}_complex.pdb" + ligand_pdbqt = f"{method_folder}/{protein_conf_id}_{ligand_id}.pdbqt" + complex_pdb = f"{method_folder}/{protein_conf_id}_{ligand_id}_complex.pdb" - # Check if vina output exists if not os.path.exists(ligand_pdbqt): continue - # Skip if complex already exists if os.path.exists(complex_pdb): continue - # Try multiple naming patterns for protein PDB files protein_pdb_candidates = [ f"{proteins_folder}/{protein_conf_id}_single_chain_clean.pdb", f"{proteins_folder}/{protein_conf_id}_clean.pdb", @@ -356,10 +339,32 @@ def generate_vina_complex_pdbs(batch_dictionary): complexes_failed += 1 logger.info( - f"Complex PDB generation complete: {complexes_created} created, {complexes_failed} failed" + f"{method_label} complex PDB generation complete: " + f"{complexes_created} created, {complexes_failed} failed" ) +def generate_vina_complex_pdbs(batch_dictionary): + """ + Generate complex PDB files for all Vina docking results in a batch. + Merges protein PDB with docked ligand PDBQT into a single complex PDB. + + :param batch_dictionary: Dictionary containing batch information including + BATCH_FOLDER, COMBINATIONS_TABLE_KEY, etc. + """ + _generate_pdbqt_complex_pdbs(batch_dictionary, VINA_FOLDER, "Vina") + + +def generate_gnina_complex_pdbs(batch_dictionary): + """ + Generate complex PDB files for all gnina docking results in a batch. + Uses the same PDBQT output format as Vina. + """ + from guild.constants.guild import GNINA_FOLDER + + _generate_pdbqt_complex_pdbs(batch_dictionary, GNINA_FOLDER, "gnina") + + def generate_diffdock_complex_pdbs(batch_dictionary): """ Generate PLIP-ready complex PDB files for DiffDock docking results. @@ -438,22 +443,24 @@ def generate_diffdock_complex_pdbs(batch_dictionary): complexes_failed += 1 continue - # Extract chain from protein_conf_id (e.g. "8gut-R-KO8-R" → "R") + # Extract chain(s) from protein_conf_id. Single chain ("8gut-R-KO8-R" → + # "R") or a comma-joined set for a multi-chain pocket ("8gut-A,B-..." → + # ["A", "B"]); every listed chain is kept so the receptor stays intact. parts = protein_conf_id.split("-") - chain_id = parts[1] if len(parts) >= 2 else "A" + chain_ids = _normalize_chain_list(parts[1]) if len(parts) >= 2 else ["A"] chain_pdb = os.path.join(diffdock_folder, f"{protein_conf_id}_chain.pdb") if not os.path.exists(chain_pdb): kept = 0 with open(raw_pdb) as fin, open(chain_pdb, "w") as fout: for line in fin: - if line.startswith("ATOM") and len(line) > 21 and line[21] == chain_id: + if line.startswith("ATOM") and len(line) > 21 and line[21] in chain_ids: fout.write(line) kept += 1 fout.write("END\n") if kept == 0: logger.warning( - f"No ATOM records for chain {chain_id} in {raw_pdb}" + f"No ATOM records for chain(s) {chain_ids} in {raw_pdb}" ) complexes_failed += 1 continue diff --git a/guild/tools/preparation.py b/guild/tools/preparation.py index d6bcd7d..04bbae9 100644 --- a/guild/tools/preparation.py +++ b/guild/tools/preparation.py @@ -154,6 +154,40 @@ def accept_residue(self, residue): return len(residue) == 1 and (residue.child_list[0].element or "").upper() in METALS_LIST +def _normalize_chain_list(chains): + """ + Normalize a chain specification into a list of chain-ID strings. + + Accepts a single chain ID (``"A"``), a comma-separated string (``"A,B"``), + a list/tuple of IDs, or ``None`` (→ empty list). Whitespace is stripped and + empty tokens dropped, so ``"A, B"`` and ``["A", "B"]`` both yield + ``["A", "B"]``. This is the single place that defines how the + ``protein_chain`` CSV column encodes multiple chains. + """ + if chains is None: + return [] + + # Treat pandas/NumPy missing values (NaN, ) as "no chain specified". + try: + import pandas as pd # type: ignore + + if pd.isna(chains): + return [] + except Exception: + pass + + if isinstance(chains, str): + return [c.strip() for c in chains.split(",") if c.strip()] + + # Iterable of chain IDs (list/tuple/set/etc). If it's a non-iterable scalar, + # treat it as a single token. + try: + return [str(c).strip() for c in chains if str(c).strip()] + except TypeError: + token = str(chains).strip() + return [token] if token else [] + + def get_protein_chain(input_file, input_name): """ Get the chain ID of the first chain in a PDB file. @@ -176,29 +210,37 @@ def get_protein_chain(input_file, input_name): def isolate_protein_chain(input_file, input_name, output_file, target_chain=None): """ - Isolate a protein chain by its chain ID (e.g., 'A'). - Falls back to the first chain if target_chain is None or not found. - Strips HETATM while writing. + Isolate one or more protein chains by their chain IDs (e.g., 'A' or ['A', 'B']). + Falls back to the first chain if ``target_chain`` is None or none of the + requested chains are present. Preserves chain IDs so a multi-chain pocket + (e.g. a dimer interface) stays intact downstream. (HETATM cleanup is done by ``clean_receptor``.) :param input_file: The path to the input PDB file. :param input_name: The name of the input PDB file. :param output_file: The path to the output PDB file. - :param target_chain: The chain to be used for docking. If not provided, the first chain is used. + :param target_chain: Chain ID, list of chain IDs, or comma-separated string + (e.g. ``"A,B"``). If not provided, the first chain is used. """ parser = PDBParser(QUIET=True) structure = parser.get_structure(input_name, input_file) model = structure[0] - # find the requested chain (or default to first) - if target_chain is not None and target_chain in [ch.id for ch in model.get_chains()]: - chain_to_write = model[target_chain] - else: - chain_to_write = next(model.get_chains()) - - # 🔧 rebuild a proper structure hierarchy + available = [ch.id for ch in model.get_chains()] + requested = _normalize_chain_list(target_chain) + chains_to_write = [c for c in requested if c in available] + if not chains_to_write: + # Fall back to the first chain if nothing was requested or matched. + chains_to_write = [next(model.get_chains()).id] + logger.warning( + f"No requested chain {requested} found in {input_file}. " + f"Defaulting to first chain '{chains_to_write[0]}'." + ) + + # 🔧 rebuild a proper structure hierarchy, preserving every requested chain ID structure_builder = StructureBuilder.StructureBuilder() structure_builder.init_structure("filtered") structure_builder.init_model(0) - structure_builder.structure[0].add(chain_to_write.copy()) # preserve chain ID + for chain_id in chains_to_write: + structure_builder.structure[0].add(model[chain_id].copy()) # preserve chain ID io = PDBIO() io.set_structure(structure_builder.get_structure()) @@ -216,9 +258,12 @@ def renumber_pdb_residues(input_pdb, output_pdb, start_index=1): """ parser = PDBParser(QUIET=True) structure = parser.get_structure("renumbered", input_pdb) - i = start_index for model in structure: for chain in model: + # Renumber per chain so each chain is 1-based. Boltz pocket + # contacts index residues per chain, so continuous numbering + # across chains would misalign multi-chain constraints. + i = start_index for residue in chain: res_id = residue.id residue.id = (res_id[0], i, res_id[2]) diff --git a/guild/tools/scores.py b/guild/tools/scores.py index 47166db..a483707 100644 --- a/guild/tools/scores.py +++ b/guild/tools/scores.py @@ -15,8 +15,8 @@ from guild.constants.bulk import ( GLOBAL_RP_SCORE, - RANKS_DICTIONARY, RP_SCORES_DICTIONARY, + RANKS_DICTIONARY, SCORES_DIRECTION_DICTIONARY, ) from guild.constants.guild import PROTEIN_CONF_ID diff --git a/guild/tools/subprocess_log.py b/guild/tools/subprocess_log.py new file mode 100644 index 0000000..badc420 --- /dev/null +++ b/guild/tools/subprocess_log.py @@ -0,0 +1,68 @@ +""" +Helper for writing per-invocation subprocess logs. + +guild's docking helpers (``deploy_boltz``, ``deploy_diffdock``, +``deploy_gnina``) invoke external CLIs via ``subprocess.run``. When one of +them fails (or even succeeds in a degraded way — e.g. Boltz returning an +empty manifest with exit code 0), the orchestrator surfaces a FAILED line +in the progress log but the actual stderr/stdout is left mixed inside the +batch-wide ``output.log``. That makes debugging from outside the container +painful for the user. + +This helper writes a focused, predictable subprocess transcript to +``batches///.subprocess.log`` (or +``batches///_batch.subprocess.log`` for batch-level methods +like DiffDock). Always written, success or failure — the file is small +(typical subprocess stdout under ~100 KB) and "always available" is the +contract external consumers actually want when something looks wrong. +""" + +from __future__ import annotations + +import os +import shlex +from typing import Sequence + + +def write_subprocess_log( + log_path: str | os.PathLike, + argv: Sequence[str], + returncode: int, + stdout: str | bytes | None, + stderr: str | bytes | None, + extra_header: str | None = None, +) -> None: + """ + Write a full subprocess transcript to ``log_path``. + + :param log_path: Destination file. Parent dirs are created if missing. + :param argv: The argv passed to subprocess.run (recorded for reproducibility). + :param returncode: Exit code (``0`` = success). + :param stdout: Captured stdout. ``None`` is treated as empty. + :param stderr: Captured stderr. ``None`` is treated as empty. + :param extra_header: Optional extra context line above the [STDOUT] block — + useful for non-zero-exit-but-logically-failed cases (e.g. Boltz empty + manifest). + """ + os.makedirs(os.path.dirname(log_path) or ".", exist_ok=True) + + status = "SUCCESS" if returncode == 0 else f"FAILED (exit {returncode})" + quoted_argv = " ".join(shlex.quote(str(a)) for a in argv) + + def _decode(blob): + if blob is None: + return "" + if isinstance(blob, bytes): + return blob.decode("utf-8", errors="replace") + return blob + + with open(log_path, "w") as f: + f.write(f"[INVOCATION {status}]\n") + f.write(f"argv: {quoted_argv}\n") + if extra_header: + f.write(f"note: {extra_header}\n") + f.write("\n[STDOUT]\n") + f.write(_decode(stdout)) + f.write("\n\n[STDERR]\n") + f.write(_decode(stderr)) + f.write("\n") diff --git a/guild/tools/utils.py b/guild/tools/utils.py index 7afd509..e0de87f 100644 --- a/guild/tools/utils.py +++ b/guild/tools/utils.py @@ -7,6 +7,8 @@ import time from functools import wraps +import pandas as pd + def read_json_as_dict(input_file): """ diff --git a/guild/transformers/chembl.py b/guild/transformers/chembl.py index 271d886..ddfa089 100644 --- a/guild/transformers/chembl.py +++ b/guild/transformers/chembl.py @@ -3,13 +3,16 @@ import os import random import time +from pathlib import Path from typing import Tuple import chembl_downloader import pandas as pd +from chembl_webresource_client.http_errors import HttpBadRequest from tqdm import tqdm from unipressed import IdMappingClient +import guild as _guild_pkg from guild.constants.chembl import ( ACTIVITY, BINDER_DIR, @@ -24,7 +27,6 @@ UNIPROT_ID, ) from guild.constants.general import RANDOM_SEED -from guild.constants.system import WORKING_DIR_PATH random.seed(RANDOM_SEED) @@ -32,8 +34,14 @@ pd.options.mode.chained_assignment = None # default='warn' -BINDER_DIR_PATH = f"{WORKING_DIR_PATH}/guild/support/{BINDER_DIR}" -DECOYS_DIR_PATH = f"{WORKING_DIR_PATH}/guild/support/{DECOYS_DIR}" +# Anchor the binders/decoys cache directories on the guild package's own +# location, not WORKING_DIR_PATH. This module is imported lazily (e.g. from +# guild.tools.binders) after run_guild.py has rebound WORKING_DIR_PATH to +# the caller's workspace, which has no guild/support tree. The guild package +# location is the stable anchor. +_GUILD_PKG_DIR = Path(_guild_pkg.__file__).resolve().parent +BINDER_DIR_PATH = str(_GUILD_PKG_DIR / "support" / BINDER_DIR) +DECOYS_DIR_PATH = str(_GUILD_PKG_DIR / "support" / DECOYS_DIR) os.makedirs(BINDER_DIR_PATH, exist_ok=True) os.makedirs(DECOYS_DIR_PATH, exist_ok=True) diff --git a/guild/transformers/converters.py b/guild/transformers/converters.py index 13692e9..a31a10a 100644 --- a/guild/transformers/converters.py +++ b/guild/transformers/converters.py @@ -238,6 +238,25 @@ def ligand_pdb_to_pdbqt(pdb: str, timeout: int = 300): if pdbqt_atoms == 0: raise ValueError(f"Generated PDBQT has no ATOM records: {output_pdbqt}") + # OpenBabel sometimes writes non-ligand records such as COMPND/AUTHOR + # before ROOT, which Vina's PDBQT parser rejects with "Unknown or + # inappropriate tag". Strip everything except the tags Vina accepts in a + # ligand PDBQT, including REMARK. + # Match on the leading whitespace-separated token — ``line[:6]`` would + # truncate the 7/9-char tags ENDROOT/ENDBRANCH/TORSDOF, silently + # dropping the very lines that close the torsion tree. Vina tolerates + # the truncated result, but gnina rejects it as malformed. + _VINA_LIGAND_TAGS = { + "REMARK", "ROOT", "ENDROOT", "BRANCH", "ENDBRANCH", + "TORSDOF", "ATOM", "HETATM", "END", + } + filtered = [ + line for line in pdbqt_lines + if line.strip() and line.split(maxsplit=1)[0] in _VINA_LIGAND_TAGS + ] + with open(output_pdbqt, "w") as f: + f.writelines(filtered) + except subprocess.TimeoutExpired as e: raise TimeoutError( f"OpenBabel conversion timed out after {timeout} seconds for {pdb}" diff --git a/guild/transformers/msa.py b/guild/transformers/msa.py index 813c23b..75f81a4 100644 --- a/guild/transformers/msa.py +++ b/guild/transformers/msa.py @@ -16,6 +16,7 @@ import subprocess from pathlib import Path + LOCALCOLABFOLD_DIR_ENV_VAR = "GUILD_LOCALCOLABFOLD_DIR" COLABFOLD_RUN_SCRIPT = "run_colabfoldbatch_sample.sh" @@ -168,18 +169,27 @@ def fetch_protein_msa( ) return None - # Delete temporary files and return the path to the generated a3m, if any - all_files_after_run = list(output_dir_path.glob("*")) - a3m_candidates = [f for f in all_files_after_run if f.suffix == ".a3m"][0] - if not a3m_candidates: + # ColabFold writes the a3m using the FASTA header as the filename stem, which + # matches output_a3m_path exactly. Clean up the other artifacts it produces + # (pickle, png, _env dir, log, config, bibtex) without touching any pre-existing + # a3m files from earlier proteins that share the same output directory. + if not output_a3m_path.exists() or output_a3m_path.stat().st_size == 0: logger.warning( - "fetch_protein_msa: local script completed but no .a3m file was found in %s", - output_dir_path, + "fetch_protein_msa: local script completed but %s was not created. " + "Falling back to empty MSA.", + output_a3m_path, ) return None - # Keep the a3m and clean up other temporary files - files_to_remove = [f for f in all_files_after_run if f != a3m_candidates] - _cleanup_temp_msa_artifacts(files_to_remove) + colabfold_artifacts = [ + output_dir_path / f"{protein_id}_{protein_chain_id}.pickle", + output_dir_path / f"{protein_id}_{protein_chain_id}_coverage.png", + output_dir_path / f"{protein_id}_{protein_chain_id}_env", + output_dir_path / f"{protein_id}_{protein_chain_id}.fasta", + output_dir_path / "cite.bibtex", + output_dir_path / "config.json", + output_dir_path / "log.txt", + ] + _cleanup_temp_msa_artifacts(colabfold_artifacts) logger.info("fetch_protein_msa: MSA saved to %s", output_a3m_path) return str(output_a3m_path) diff --git a/guild/transformers/pdb.py b/guild/transformers/pdb.py index 625d72d..6dc2be5 100644 --- a/guild/transformers/pdb.py +++ b/guild/transformers/pdb.py @@ -13,6 +13,7 @@ from guild.constants.ligands import LIGANDS_TO_IGNORE from guild.docking.vina import generate_vina_box +from guild.tools.preparation import _normalize_chain_list # Suppress PDBConstructionWarning warnings.filterwarnings("ignore", category=PDBConstructionWarning) @@ -330,17 +331,13 @@ def _parse_pdb_atom_fields( def _convert_pdbqt_to_pdb(ligand_pdbqt: str, out_pdb: str) -> None: """ - Convert PDBQT -> PDB, keeping only the **first model** (best-ranked - pose). Vina writes multiple poses as separate MODEL/ENDMDL blocks; - including all of them would inflate PLIP interaction counts. - + Convert PDBQT -> PDB. Prefer OpenBabel CLI (obabel). If not available, fall back to stripping extra PDBQT columns (best-effort). """ obabel = shutil.which("obabel") if obabel: - # -f 1 -l 1: extract only the first molecule (best Vina pose) - cmd = [obabel, ligand_pdbqt, "-O", out_pdb, "-f", "1", "-l", "1"] + cmd = [obabel, ligand_pdbqt, "-O", out_pdb] res = subprocess.run(cmd, capture_output=True, text=True) if res.returncode != 0: raise RuntimeError( @@ -349,25 +346,11 @@ def _convert_pdbqt_to_pdb(ligand_pdbqt: str, out_pdb: str) -> None: return # Fallback: treat PDBQT as PDB and strip after column 66 (keep coords/occ/temp) - # Only keep atoms from the first MODEL block. with ( open(ligand_pdbqt, "r", encoding="utf-8", errors="replace") as fin, open(out_pdb, "w", encoding="utf-8") as fout, ): - seen_model = False for line in fin: - stripped = line.strip() - if stripped.startswith("MODEL"): - if seen_model: - # We've already processed the first model – stop here. - break - seen_model = True - continue - if stripped in ("ENDMDL", "END"): - # End of the first model – stop reading. - if seen_model: - break - continue if _is_atom_line(line): # Keep up to tempFactor (col 66), then try to preserve element if present. base = line[:66] @@ -457,24 +440,9 @@ def _write_complex( ligand_resseq: int, insert_ter_between: bool, ) -> None: - # Read ligand, keep only atom records from the **first MODEL** block. - # Multi-model files (e.g. Vina multi-pose output) should already be - # filtered upstream, but this serves as a safety net so PLIP never - # accidentally analyses overlapping poses. - ligand_atom_lines: list[str] = [] + # Read ligand, keep only atom records with open(ligand_pdb, "r", encoding="utf-8", errors="replace") as f: - seen_first_model = False - for ln in f: - stripped = ln.strip() - if stripped.startswith("MODEL"): - if seen_first_model: - break # second MODEL encountered – stop - seen_first_model = True - continue - if stripped in ("ENDMDL",) and seen_first_model: - break # end of first model - if _is_atom_line(ln): - ligand_atom_lines.append(ln) + ligand_atom_lines = [ln for ln in f if _is_atom_line(ln)] if not ligand_atom_lines: raise ValueError(f"No ATOM/HETATM records found in ligand file: {ligand_pdb}") @@ -540,7 +508,7 @@ def _write_complex( def get_pocket_contacts_from_ligand( protein_pdb: str, - protein_chain: str, + protein_chain, original_ligand: str, original_ligand_chain: str, distance_threshold: float = 6.0, @@ -548,18 +516,22 @@ def get_pocket_contacts_from_ligand( """ Return a list of [chain_id, residue_index] protein residue contacts that lie within ``distance_threshold`` Å of any atom of ``original_ligand`` (by residue - name). ``residue_index`` is 1-based and contiguous along the selected protein - chain (Boltz schema indexing), not the raw PDB residue number. - Suitable for use in a Boltz pocket-constraint ``contacts`` list. + name). ``residue_index`` is 1-based and contiguous *within each chain* (Boltz + schema indexing), not the raw PDB residue number. Suitable for use in a Boltz + pocket-constraint ``contacts`` list. :param protein_pdb: Path to the crystal-structure PDB that contains both protein and the reference ligand. - :param protein_chain: Chain ID of the protein receptor (e.g. ``"A"``). + :param protein_chain: Chain ID, list of chain IDs, or comma-separated string + (e.g. ``"A,B"``) of the protein receptor. Contacts are + collected from every listed chain, each indexed + independently from 1, so a pocket spanning multiple + chains is fully captured. :param original_ligand: Residue name of the reference ligand (e.g. ``"LIG"``). :param original_ligand_chain: Chain ID that contains the reference ligand. :param distance_threshold: Maximum heavy-atom distance (Å) to include a protein residue as a contact (default 6.0). - :return: List of ``[protein_chain, residue_index]`` pairs, or ``[]`` if + :return: List of ``[chain_id, residue_index]`` pairs, or ``[]`` if the ligand residue cannot be found. """ parser = PDBParser(QUIET=True) @@ -582,31 +554,33 @@ def get_pocket_contacts_from_ligand( return [] ligand_coords = np.array([a.get_coord() for a in ligand_atoms]) + available = [c.id for c in model] - # Find protein residues with at least one atom within distance_threshold + # Find protein residues with at least one atom within distance_threshold, + # indexing each chain independently (Boltz expects per-chain 1-based indices). contacts = [] - if protein_chain not in [c.id for c in model]: - logger.warning( - f"get_pocket_contacts_from_ligand: protein chain '{protein_chain}' " - f"not found in {protein_pdb}." - ) - return [] - - seen = set() - sequence_index = 0 - for residue in model[protein_chain]: - if residue.id[0] != " ": # skip heteroatoms / water on the protein chain - continue - sequence_index += 1 - if sequence_index in seen: + for chain_id in _normalize_chain_list(protein_chain): + if chain_id not in available: + logger.warning( + f"get_pocket_contacts_from_ligand: protein chain '{chain_id}' " + f"not found in {protein_pdb}." + ) continue - for atom in residue: - diffs = ligand_coords - atom.get_coord() - dists = np.linalg.norm(diffs, axis=1) - if np.any(dists <= distance_threshold): - contacts.append([protein_chain, sequence_index]) - seen.add(sequence_index) - break # one close atom is enough to include this residue + seen = set() + sequence_index = 0 + for residue in model[chain_id]: + if residue.id[0] != " ": # skip heteroatoms / water on the protein chain + continue + sequence_index += 1 + if sequence_index in seen: + continue + for atom in residue: + diffs = ligand_coords - atom.get_coord() + dists = np.linalg.norm(diffs, axis=1) + if np.any(dists <= distance_threshold): + contacts.append([chain_id, sequence_index]) + seen.add(sequence_index) + break # one close atom is enough to include this residue logger.debug( f"get_pocket_contacts_from_ligand: found {len(contacts)} contact residues " @@ -615,6 +589,65 @@ def get_pocket_contacts_from_ligand( return contacts +def get_pocket_contacts_from_box( + protein_pdb: str, + protein_chain, + center: Tuple[float, float, float], + size: Tuple[float, float, float], +) -> list: + """ + Return ``[[chain_id, residue_index], ...]`` for protein residues whose Cα atom + lies inside the axis-aligned box defined by ``center`` and ``size``. + ``residue_index`` is 1-based and contiguous *within each chain* (Boltz schema + indexing). Same return shape as :func:`get_pocket_contacts_from_ligand`. + + :param protein_pdb: Path to the protein PDB. + :param protein_chain: Chain ID, list of chain IDs, or comma-separated string + (e.g. ``"A,B"``) of the protein receptor. Residues are + collected from every listed chain, each indexed + independently from 1, so an interface pocket spanning + multiple chains is fully captured. + :param center: ``(x, y, z)`` of the box center. + :param size: ``(sx, sy, sz)`` full edge lengths of the box. + :return: List of ``[chain_id, residue_index]`` pairs, or ``[]`` if no + listed chain is found or no residue Cα falls inside the box. + """ + cx, cy, cz = center + sx, sy, sz = size + half = np.array([sx / 2.0, sy / 2.0, sz / 2.0]) + center_arr = np.array([cx, cy, cz]) + + parser = PDBParser(QUIET=True) + structure = parser.get_structure("complex", protein_pdb) + model = structure[0] + available = [c.id for c in model] + + contacts = [] + for chain_id in _normalize_chain_list(protein_chain): + if chain_id not in available: + logger.warning( + f"get_pocket_contacts_from_box: protein chain '{chain_id}' " + f"not found in {protein_pdb}." + ) + continue + sequence_index = 0 + for residue in model[chain_id]: + if residue.id[0] != " ": # skip heteroatoms / water on the protein chain + continue + sequence_index += 1 + if "CA" not in residue: + continue + ca_coord = residue["CA"].get_coord() + if np.all(np.abs(ca_coord - center_arr) <= half): + contacts.append([chain_id, sequence_index]) + + logger.debug( + f"get_pocket_contacts_from_box: found {len(contacts)} contact residues " + f"inside box center={center} size={size} in {protein_pdb}." + ) + return contacts + + def relabel_ligand_chain_in_pdb( input_pdb: str, output_pdb: str, diff --git a/pyproject.toml b/pyproject.toml index 9b9d0e5..4ca53b8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,18 +4,21 @@ build-backend = "hatchling.build" [project] name = "guild" -version = "1.0.0" +version = "1.1.5" description = "A unified framework for multi-method molecular docking and scoring" license = "MIT" authors = [ { name = "António J Preto" }, { name = "Yojana Gadiya" }, + { name = "Tyson Park" }, { name = "Christoph Krettler" }, - { name = "Dani Domingo-Fernandéz" }, { name = "Shweta Bhandare" }, { name = "Keaton Noon" }, - { name = "Tyson Park" } + { name = "Eric Fish" } + { name = "António de La Vega de León" }, + { name = "Dani Domingo-Fernandéz" }, ] + requires-python = ">=3.10,<3.11" readme = "README.md" classifiers = [ diff --git a/scripts/apply_karmadock_patches.py b/scripts/apply_karmadock_patches.py new file mode 100644 index 0000000..76225b4 --- /dev/null +++ b/scripts/apply_karmadock_patches.py @@ -0,0 +1,131 @@ +""" +Apply the KarmaDock compatibility patches documented in the project README. + +KarmaDock is pinned at commit 9a35d0c and was written against an older RDKit +version. With the rdkit/torch versions guild uses today, two adjustments are +required to avoid dimension-mismatch crashes during ligand-feature generation +and inside the GraphTransformer block: + + 1. KarmaDock/dataset/ligand_feature.py — hard-coded ``20`` feature columns is + replaced with the actual ``edge_feature.size(1)``. + 2. KarmaDock/architecture/GraphTransformer_Block.py — the edge encoder is + guarded so an unexpected input width is truncated or padded to match + ``self.edge_encoder.in_features``. + +Idempotent: re-running the script after a patch is already applied is a no-op. + +Usage: + python apply_karmadock_patches.py /app/KarmaDock +""" + +import re +import sys +from pathlib import Path + +LIGAND_FEATURE_ALREADY = "feat_dim = edge_feature.size(1)" +# Indentation-aware: the original snippet appears at multiple indentation +# levels (function body and inside ``if`` blocks). +LIGAND_FEATURE_RE = re.compile( + r"^(?P[ \t]+)edge_feature_new = torch\.zeros\(\(edge_index_new\.size\(1\), 20\)\)\n" + r"(?P=indent)edge_feature_new\[:, \[4, 5, 18\]\] = 1\n", + flags=re.MULTILINE, +) + + +def _ligand_feature_replacement(match: "re.Match[str]") -> str: + indent = match.group("indent") + return ( + f"{indent}feat_dim = edge_feature.size(1)\n" + f"{indent}edge_feature_new = torch.zeros(\n" + f"{indent} (edge_index_new.size(1), feat_dim),\n" + f"{indent} dtype=edge_feature.dtype,\n" + f"{indent} device=edge_feature.device,\n" + f"{indent})\n" + ) + + +GT_BLOCK_TARGET = "edge_feats = self.edge_encoder(edge_s)" +# Lines of the guard block we splice in *before* the target line. Each line +# carries its relative indentation as leading tabs (KarmaDock imports +# ``torch as th`` and indents with tabs throughout). The captured target-line +# indentation is then prepended to every line so the splice matches whatever +# nesting depth ``edge_feats = self.edge_encoder(edge_s)`` happens to sit at. +GT_BLOCK_GUARD_LINES = [ + "if edge_s.size(1) > self.edge_encoder.in_features:", + "\tedge_s = edge_s[:, :self.edge_encoder.in_features]", + "elif edge_s.size(1) < self.edge_encoder.in_features:", + "\tpad = th.zeros(", + "\t\tedge_s.size(0),", + "\t\tself.edge_encoder.in_features - edge_s.size(1),", + "\t\tdevice=edge_s.device,", + "\t\tdtype=edge_s.dtype,", + "\t)", + "\tedge_s = th.cat([edge_s, pad], dim=1)", +] +GT_BLOCK_ALREADY = "edge_s = edge_s[:, :self.edge_encoder.in_features]" + + +def _patch_ligand_feature(path: Path) -> int: + text = path.read_text() + if LIGAND_FEATURE_ALREADY in text: + print(" [skip] ligand_feature.py already patched") + return 0 + new_text, n = LIGAND_FEATURE_RE.subn(_ligand_feature_replacement, text) + if n == 0: + raise RuntimeError( + f"Expected pattern not found in {path}; KarmaDock layout may have " + f"changed and the patch needs updating." + ) + path.write_text(new_text) + print(f" [applied] ligand_feature.py — {n} replacement(s)") + return n + + +def _patch_gt_block(path: Path) -> int: + text = path.read_text() + if GT_BLOCK_ALREADY in text: + print(" [skip] GraphTransformer_Block.py already patched") + return 0 + if GT_BLOCK_TARGET not in text: + raise RuntimeError( + f"Expected line '{GT_BLOCK_TARGET}' not found in {path}; " + f"KarmaDock layout may have changed and the patch needs updating." + ) + # Match the indentation of the target line so we splice cleanly. + line_re = re.compile( + r"^(?P[ \t]*)edge_feats\s*=\s*self\.edge_encoder\(edge_s\)$", + flags=re.MULTILINE, + ) + match = line_re.search(text) + if not match: + raise RuntimeError( + f"Target line found but indentation could not be matched in {path}" + ) + indent = match.group("indent") + guard = "\n".join(indent + line for line in GT_BLOCK_GUARD_LINES) + new_text = ( + text[: match.start()] + + guard + + "\n" + + match.group(0) + + text[match.end():] + ) + path.write_text(new_text) + print(" [applied] GraphTransformer_Block.py") + return 1 + + +def main(karmadock_root: Path): + if not karmadock_root.is_dir(): + sys.exit(f"KarmaDock directory not found: {karmadock_root}") + + print(f"Patching {karmadock_root}") + _patch_ligand_feature(karmadock_root / "dataset" / "ligand_feature.py") + _patch_gt_block(karmadock_root / "architecture" / "GraphTransformer_Block.py") + print("KarmaDock patches applied") + + +if __name__ == "__main__": + if len(sys.argv) != 2: + sys.exit("Usage: apply_karmadock_patches.py ") + main(Path(sys.argv[1])) diff --git a/scripts/run_guild.py b/scripts/run_guild.py index 3b5a32a..b08f271 100644 --- a/scripts/run_guild.py +++ b/scripts/run_guild.py @@ -70,7 +70,7 @@ def parse_args() -> argparse.Namespace: "--methods", "-m", nargs="+", default=["boltz"], - choices=["boltz", "vina", "karmadock", "diffdock"], + choices=["boltz", "vina", "karmadock", "diffdock", "gnina"], help="Docking methods to run (space-separated).", ) parser.add_argument( @@ -96,12 +96,6 @@ def parse_args() -> argparse.Namespace: default="chembl_36", help="ChEMBL version string.", ) - parser.add_argument( - "--no-decoys", - action="store_true", - default=False, - help="Disable decoy generation (useful when decoy file is not available).", - ) parser.add_argument( "--clean", action="store_true", @@ -120,6 +114,77 @@ def parse_args() -> argparse.Namespace: default=False, help="Enable known-binders expansion in BulkRun.", ) + parser.add_argument( + "--no-decoys", + action="store_true", + default=False, + help="Disable decoy generation (useful for single runs).", + ) + parser.add_argument( + "--box", + default=None, + help=( + "Optional path to a Vina box file (center_{x,y,z} + size_{x,y,z}) used as a " + "global fallback for combinations whose 'box_location' column is empty. " + "Per-row values in the CSV always take precedence." + ), + ) + parser.add_argument( + "--n-workers", + type=int, + default=1, + help=( + "Number of parallel worker processes for Vina docking. Vina's internal " + "threading auto-scales, so values >1 oversubscribe slightly but typically " + "still help. Default 1 (safe but serial)." + ), + ) + parser.add_argument( + "--no-gpu", + action="store_true", + default=False, + help=( + "Force CPU-only execution. Passed through to BulkRun(use_gpu=False); " + "gnina respects it via its own --no_gpu flag. Vina and DiffDock are " + "already CPU-only and unaffected; Boltz is genuinely GPU-bound and will " + "still attempt to use a GPU (so don't combine --no-gpu with --methods boltz)." + ), + ) + parser.add_argument( + "--no-plip", + action="store_true", + default=False, + help=( + "Skip the PLIP interactions analysis step. By default, after docking " + "and scoring complete, guild runs PLIP over every method's complex " + "PDBs and writes data//plip_interactions.tsv (always — " + "header-only when nothing was produced). External consumers should " + "read that file rather than installing plip locally; pass this flag " + "only if you don't need the interactions output." + ), + ) + parser.add_argument( + "--plip-only", + action="store_true", + default=False, + help=( + "Skip docking + scoring and run only the PLIP interactions step over " + "an existing data// tree. Useful for re-generating " + "plip_interactions.tsv when only the PLIP code changed." + ), + ) + parser.add_argument( + "--gnina-input-mode", + choices=["pdbqt", "sdf"], + default="pdbqt", + help=( + "Ligand/receptor input format for gnina. 'sdf' skips OpenBabel " + "PDBQT prep entirely (gnina reads RDKit-generated SDF + cleaned " + "PDB natively), but only takes effect when gnina is the sole " + "PDBQT-relevant method requested — co-running with Vina or any " + "Vina-rescore downgrades to PDBQT with a warning." + ), + ) return parser.parse_args() @@ -131,18 +196,39 @@ def main() -> None: args = parse_args() np.random.seed(42) - # scripts/ lives directly under the project root, so one level up is enough. - PROJECT_ROOT = Path(__file__).resolve().parent.parent + # PROJECT_ROOT is the host directory the Makefile bind-mounts at + # ``/workspace`` — the place where ``data//`` outputs land and + # where ``COMBINATIONS``/``BOX`` paths resolve. We read it from the + # ``WORKSPACE_ROOT`` env var (set in the Makefile's DOCKER_COMMON) so the + # script can live anywhere on disk: it does NOT have to sit under + # ``WORKSPACE/scripts/``. The fallback to ``parent.parent`` covers + # invocations that don't go through the Makefile (e.g. someone running + # ``python scripts/run_guild.py`` locally). + workspace_env = os.environ.get("WORKSPACE_ROOT") + if workspace_env: + PROJECT_ROOT = Path(workspace_env) + else: + PROJECT_ROOT = Path(__file__).resolve().parent.parent # Override guild's hardcoded paths so output lands in # PROJECT_ROOT/data/ (the volume-mounted workspace) rather than /app. + import guild as _guild_pkg import guild.bulk as _bulk_mod import guild.constants.system as _sys_const import guild.run as _run_mod + # ``guild/support/`` ships with the guild package itself. Resolving it + # from the imported package's location works both in-repo + # (``./guild/support``) and from an external caller where ``WORKSPACE`` + # is a different repo entirely and the guild source is bind-mounted at + # ``/app/guild`` — there is no ``guild/`` directory at the caller's + # workspace root. + GUILD_PKG_DIR = Path(_guild_pkg.__file__).resolve().parent + GUILD_SUPPORT_DIR = GUILD_PKG_DIR / "support" + _sys_const.WORKING_DIR_PATH = PROJECT_ROOT _sys_const.PROJECTS_FOLDER = str(PROJECT_ROOT / "data") - _sys_const.SUPPORT_FOLDER = str(PROJECT_ROOT / "guild" / "support") + _sys_const.SUPPORT_FOLDER = str(GUILD_SUPPORT_DIR) _sys_const.PYTHON_EXECUTABLE = sys.executable # works both locally & in Docker _bulk_mod.WORKING_DIR_PATH = PROJECT_ROOT _bulk_mod.PROJECTS_FOLDER = str(PROJECT_ROOT / "data") @@ -151,7 +237,7 @@ def main() -> None: # Patch diffdock module's imported constants so they see the overrides. import guild.docking.diffdock as _dd_mod _dd_mod.PYTHON_EXECUTABLE = sys.executable - _dd_mod.SUPPORT_FOLDER = str(PROJECT_ROOT / "guild" / "support") + _dd_mod.SUPPORT_FOLDER = str(GUILD_SUPPORT_DIR) from guild.bulk import BulkRun @@ -160,7 +246,7 @@ def main() -> None: decoys_path = ( args.decoys if args.decoys is not None - else str(PROJECT_ROOT / "guild" / "support" / "decoys" / "chembl_36_decoys_2.tsv") + else str(GUILD_SUPPORT_DIR / "decoys" / "chembl_36_decoys_2.tsv") ) if args.clean: @@ -170,25 +256,48 @@ def main() -> None: # Auto-detect separator (handles .tsv named as .csv, etc.) runs_table = pd.read_csv(combinations_path, sep=None, engine="python") - # Remap protein_path: strip any absolute prefix up to /notebooks/ so paths - # work both locally and inside Docker. - if "protein_path" in runs_table.columns: - runs_table["protein_path"] = runs_table["protein_path"].apply( - lambda p: str(PROJECT_ROOT / re.sub(r"^.*?/notebooks/", "notebooks/", p)) - if pd.notna(p) and "/notebooks/" in str(p) - else p - ) + # Remap protein_path / box_location: strip any absolute host prefix up to + # /notebooks/ or /temp_data/ so paths work both locally and inside Docker + # (where the repo is mounted at /workspace). + _PATH_PREFIX = re.compile(r"^.*?/(notebooks|temp_data)/") + + def _normalize_workspace_path(p): + if pd.isna(p): + return p + s = str(p) + m = _PATH_PREFIX.search(s) + if not m: + return p + return str(PROJECT_ROOT / _PATH_PREFIX.sub(r"\1/", s, count=1)) + + for col in ("protein_path", "box_location"): + if col in runs_table.columns: + runs_table[col] = runs_table[col].apply(_normalize_workspace_path) # Optionally take only the first N rows if args.head > 0: runs_table = runs_table.head(args.head).reset_index(drop=True) + # Global box fallback: fill empty 'box_location' rows with --box, when supplied. + # Per-row CSV values win. + if args.box: + if "box_location" not in runs_table.columns: + runs_table["box_location"] = args.box + else: + mask = runs_table["box_location"].isna() | ( + runs_table["box_location"].astype(str).str.strip() == "" + ) + runs_table.loc[mask, "box_location"] = args.box + print(f"Box (global): {args.box}") + print(f"Project: {project_name}") print(f"Combinations: {combinations_path} ({len(runs_table)} rows)") print(f"Decoys: {decoys_path}") print(f"Methods: {args.methods}") print(f"Batch size: {args.batch_size}") print(f"Known bndrs: {args.use_known_binders}") + print(f"GPU: {not args.no_gpu}") + bulk = BulkRun( runs_table, project_name, @@ -200,16 +309,31 @@ def main() -> None: decoys=decoys_path, use_decoys=not args.no_decoys, use_known_binders=args.use_known_binders, - n_workers=1, + n_workers=args.n_workers, + use_gpu=not args.no_gpu, + gnina_input_mode=args.gnina_input_mode, ) - t0 = time.time() - bulk.run_docking() - print(f"Docking time: {time.time() - t0:.1f}s") + if not args.plip_only: + t0 = time.time() + bulk.run_docking() + print(f"Docking time: {time.time() - t0:.1f}s") + + t0 = time.time() + bulk.run_guild_scoring() + print(f"Scoring time: {time.time() - t0:.1f}s") + else: + print("Skipping docking + scoring (--plip-only).") - t0 = time.time() - bulk.run_guild_scoring() - print(f"Scoring time: {time.time() - t0:.1f}s") + # PLIP interaction analysis: always run by default. The contract is that + # data//plip_interactions.tsv exists after every run (header-only + # if no method produced complex PDBs), so external notebooks can read it + # without installing plip locally. --no-plip skips this step; --plip-only + # runs ONLY this step. + if args.plip_only or not args.no_plip: + t0 = time.time() + bulk.run_interactions_analysis() + print(f"PLIP time: {time.time() - t0:.1f}s") # ── Print final scores summary ────────────────────────────────────── if hasattr(bulk, "rp_scores_df") and bulk.rp_scores_df is not None and not bulk.rp_scores_df.empty: diff --git a/scripts/vina_rescore_all.py b/scripts/vina_rescore_all.py index b98d748..47065d0 100644 --- a/scripts/vina_rescore_all.py +++ b/scripts/vina_rescore_all.py @@ -55,25 +55,28 @@ def prepare_receptor_pdbqt(protein_conf_id: str) -> str | None: parts = protein_conf_id.split("-") pdb_id = parts[0] - chain_id = parts[1] if len(parts) >= 2 else "A" + # Single chain or a comma-joined set ("A,B") for a multi-chain receptor. + chain_ids = ( + [c.strip() for c in parts[1].split(",") if c.strip()] if len(parts) >= 2 else ["A"] + ) raw_pdb = os.path.join(PDB_DIR, f"{pdb_id}.pdb") if not os.path.isfile(raw_pdb): logger.warning(f"Raw PDB not found: {raw_pdb}") return None - # Extract single chain + # Extract chain(s) — a comma-joined set ("A,B") keeps a multi-chain receptor chain_pdb = output_pdbqt.replace(".pdbqt", "_chain.pdb") kept = 0 with open(raw_pdb) as fin, open(chain_pdb, "w") as fout: for line in fin: - if line.startswith("ATOM") and len(line) > 21 and line[21] == chain_id: + if line.startswith("ATOM") and len(line) > 21 and line[21] in chain_ids: fout.write(line) kept += 1 fout.write("END\n") if kept == 0: - logger.warning(f"No ATOM records for chain {chain_id} in {raw_pdb}") + logger.warning(f"No ATOM records for chain(s) {chain_ids} in {raw_pdb}") return None result = subprocess.run( diff --git a/tests/bulk_run/bulk_run_imports.py b/tests/bulk_run/bulk_run_imports.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/bulk_run/test_diffdock_resume.py b/tests/bulk_run/test_diffdock_resume.py new file mode 100644 index 0000000..d9cac16 --- /dev/null +++ b/tests/bulk_run/test_diffdock_resume.py @@ -0,0 +1,124 @@ +""" +Tests for the hardened DiffDock resume-skip check in +``BulkRun.run_docking()``. The previous check skipped DiffDock when the +top-level results directory was non-empty, which silently masked +interrupted runs whose per-combination subfolders were empty. The check now +requires every combination to have at least one ``*_confidence*.sdf`` on +disk before declaring the batch done. +""" + +import shutil +from pathlib import Path +from unittest.mock import patch + +import pandas as pd +import pytest + +from guild.bulk import BulkRun +from guild.constants.bulk import BATCH_FOLDER +from guild.constants.diffdock import DIFFDOCK_RESULTS_FOLDER +from guild.constants.guild import ( + DIFFDOCK_FOLDER, + DIFFDOCK_PREFIX, + LIGAND_ID, + PROTEIN_CONF_ID, +) + +TEST_DIR = Path(__file__).parent.parent +TEST_DATA_DIR = TEST_DIR / "test_data" + + +@pytest.fixture +def test_input_table(): + df = pd.read_csv(TEST_DATA_DIR / "bulk_dummy.csv") + df["protein_path"] = str(TEST_DATA_DIR / df["protein_path"].iloc[0]) + return df + + +@pytest.fixture +def cleanup(): + yield + test_project = Path.cwd() / "data" / "test-diffdock-resume" + if test_project.exists(): + shutil.rmtree(test_project, ignore_errors=True) + + +def _build_bulk(test_input_table): + """Construct a minimal BulkRun configured for DiffDock only.""" + return BulkRun( + input_table=test_input_table, + project_name="test-diffdock-resume", + methods_to_run=[DIFFDOCK_PREFIX], + use_decoys=False, + use_known_binders=False, + use_gpu=False, + n_workers=1, + ) + + +def _results_dir(bulk): + """Resolve the DiffDock results dir for the first (and only) batch.""" + batch_name = next(iter(bulk.batched_dictionary)) + batch_folder = bulk.batched_dictionary[batch_name][BATCH_FOLDER] + return Path(batch_folder) / DIFFDOCK_FOLDER / DIFFDOCK_RESULTS_FOLDER, batch_name + + +def _populate_full_results(bulk): + """Drop one valid *_confidence*.sdf into every per-combo subfolder.""" + results_dir, batch_name = _results_dir(bulk) + combinations_df = bulk.batched_dictionary[batch_name]["combinations_table"] + for _, row in combinations_df.iterrows(): + combo_dir = results_dir / f"{row[PROTEIN_CONF_ID]}_{row[LIGAND_ID]}" + combo_dir.mkdir(parents=True, exist_ok=True) + (combo_dir / "rank1_confidence-1.23.sdf").write_text("MOCK\n") + + +def _populate_empty_subfolders(bulk): + """Create per-combo subfolders but leave them empty (the bug scenario).""" + results_dir, batch_name = _results_dir(bulk) + combinations_df = bulk.batched_dictionary[batch_name]["combinations_table"] + for _, row in combinations_df.iterrows(): + combo_dir = results_dir / f"{row[PROTEIN_CONF_ID]}_{row[LIGAND_ID]}" + combo_dir.mkdir(parents=True, exist_ok=True) + + +def test_skip_when_all_combos_have_confidence_sdf(test_input_table, cleanup): + bulk = _build_bulk(test_input_table) + _populate_full_results(bulk) + + with ( + patch.object(BulkRun, "_prepare_docking", lambda self: None), + patch("guild.bulk.deploy_diffdock") as mock_deploy, + patch("guild.bulk.generate_diffdock_complex_pdbs"), + ): + bulk.run_docking() + + mock_deploy.assert_not_called() + + +def test_rerun_when_combo_folder_empty(test_input_table, cleanup): + bulk = _build_bulk(test_input_table) + _populate_empty_subfolders(bulk) + + with ( + patch.object(BulkRun, "_prepare_docking", lambda self: None), + patch("guild.bulk.deploy_diffdock") as mock_deploy, + patch("guild.bulk.generate_diffdock_complex_pdbs"), + ): + bulk.run_docking() + + mock_deploy.assert_called_once() + + +def test_rerun_when_combo_folder_missing(test_input_table, cleanup): + bulk = _build_bulk(test_input_table) + # Intentionally do NOT create any per-combo subfolders. + + with ( + patch.object(BulkRun, "_prepare_docking", lambda self: None), + patch("guild.bulk.deploy_diffdock") as mock_deploy, + patch("guild.bulk.generate_diffdock_complex_pdbs"), + ): + bulk.run_docking() + + mock_deploy.assert_called_once() diff --git a/tests/bulk_run/test_gnina_orchestration.py b/tests/bulk_run/test_gnina_orchestration.py new file mode 100644 index 0000000..aeda2be --- /dev/null +++ b/tests/bulk_run/test_gnina_orchestration.py @@ -0,0 +1,88 @@ +""" +Tests for the gnina integration in the bulk pipeline: + +- gnina is registered as a standalone method (no auto-enable rescore rules). +- `gnina_guild_scoring` returns the expected columns including the + side-channel CNN score, and `vina_guild_scoring` does not leak gnina columns. +""" + +import shutil +from pathlib import Path + +import pandas as pd +import pytest + +from guild.bulk import BulkRun +from guild.constants.guild import ( + ALL_AVAILABLE_METHODS, + GNINA_CNN_SCORE, + GNINA_PREFIX, + GNINA_SCORE, + SCORES_DICTIONARY, + VINA_RESCORE_BOLTZ_PREFIX, + VINA_RESCORE_DIFFDOCK_PREFIX, + VINA_SCORE, +) + +TEST_DIR = Path(__file__).parent.parent +TEST_DATA_DIR = TEST_DIR / "test_data" + + +@pytest.fixture +def test_input_table(): + df = pd.read_csv(TEST_DATA_DIR / "bulk_dummy.csv") + df["protein_path"] = str(TEST_DATA_DIR / df["protein_path"].iloc[0]) + return df + + +@pytest.fixture +def cleanup(): + yield + test_project = Path.cwd() / "data" / "test-gnina" + if test_project.exists(): + shutil.rmtree(test_project, ignore_errors=True) + + +def test_gnina_registered_as_available_method(): + """gnina is in the list of methods picked up when none are specified.""" + assert GNINA_PREFIX in ALL_AVAILABLE_METHODS + + +def test_gnina_score_dictionary_entry(): + assert SCORES_DICTIONARY[GNINA_PREFIX] == GNINA_SCORE + + +def test_gnina_alone_does_not_enable_any_rescore(test_input_table, cleanup): + """gnina is a standalone docker; selecting it doesn't add any vina_rescore_* tracks.""" + bulk = BulkRun( + input_table=test_input_table, + project_name="test-gnina", + methods_to_run=[GNINA_PREFIX], + use_decoys=False, + use_known_binders=False, + use_gpu=False, + n_workers=1, + ) + assert GNINA_PREFIX in bulk.methods_to_run + assert VINA_RESCORE_BOLTZ_PREFIX not in bulk.methods_to_run + assert VINA_RESCORE_DIFFDOCK_PREFIX not in bulk.methods_to_run + + +def test_gnina_guild_scoring_emits_score_and_cnn_columns(): + """gnina_guild_scoring returns both gnina_score and gnina_cnn_score; vina_guild_scoring doesn't.""" + from guild.constants.bulk import BATCH_FOLDER, COMBINATIONS_TO_RUN_KEY + from guild.docking.gnina import gnina_guild_scoring + from guild.docking.vina import vina_guild_scoring + + empty_batch = {BATCH_FOLDER: "/nonexistent", COMBINATIONS_TO_RUN_KEY: []} + + gnina_df = gnina_guild_scoring(empty_batch) + vina_df = vina_guild_scoring(empty_batch) + + assert GNINA_SCORE in gnina_df.columns + assert GNINA_CNN_SCORE in gnina_df.columns + assert VINA_SCORE not in gnina_df.columns + + assert VINA_SCORE in vina_df.columns + assert GNINA_SCORE not in vina_df.columns + assert GNINA_CNN_SCORE not in vina_df.columns diff --git a/tests/bulk_run/test_prepare_docking_multi_protein.py b/tests/bulk_run/test_prepare_docking_multi_protein.py new file mode 100644 index 0000000..c02095a --- /dev/null +++ b/tests/bulk_run/test_prepare_docking_multi_protein.py @@ -0,0 +1,160 @@ +""" +Tests for ``BulkRun._prepare_docking`` row-selection behavior. + +Multi-protein Boltz runs were failing because the prep pool only Guild-inited +``combinations_table.iloc[[0]]`` for non-KarmaDock methods. Protein prep +writes per-protein artifacts (e.g. ``{batch}/proteins/_single_chain_clean.pdb``), +so for batches whose combinations span multiple distinct ``protein_config_id`` +values, only the first protein got its cleaned PDB. Boltz then silently fell +back to the raw input PDB and produced empty manifests. + +The fix Guild-inits one row per unique ``protein_config_id`` for non-KarmaDock +methods; KarmaDock keeps preparing every row because its downstream worker +needs the per-ligand data tree. +""" + +import shutil +from pathlib import Path +from unittest.mock import patch + +import pandas as pd +import pytest + +from guild.bulk import BulkRun +from guild.constants.guild import ( + BOLTZ_PREFIX, + KARMADOCK_PREFIX, + PROTEIN_CONF_ID, +) + +TEST_DIR = Path(__file__).parent.parent +TEST_DATA_DIR = TEST_DIR / "test_data" + + +@pytest.fixture +def multi_protein_table(): + """Four rows across two distinct ``protein_config_id`` values.""" + base = pd.read_csv(TEST_DATA_DIR / "bulk_dummy.csv") + protein_path = str(TEST_DATA_DIR / base["protein_path"].iloc[0]) + + template = base.iloc[0].to_dict() + rows = [] + # Two proteins × two ligands each. + for protein_conf_id in ("3pbl-A-ETQ-A", "3pbl-B-ETQ-A"): + for ligand_id, smiles in ( + ("lig1", "CC(=COC=O)CCC1=C(C)CCCC1(C)C"), + ("lig2", "CCO"), + ): + row = template.copy() + row["protein_config_id"] = protein_conf_id + row["ligand_id"] = ligand_id + row["smiles"] = smiles + row["protein_path"] = protein_path + rows.append(row) + return pd.DataFrame(rows) + + +@pytest.fixture +def cleanup(): + yield + test_project = Path.cwd() / "data" / "test-prepare-multi-protein" + if test_project.exists(): + shutil.rmtree(test_project, ignore_errors=True) + + +class _SyncExecutor: + """Drop-in for ``ProcessPoolExecutor`` that runs `submit` jobs inline, + avoiding the pickling boundary so tests can use local closures.""" + + def __init__(self, *args, **kwargs): + pass + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + def submit(self, fn, *args, **kwargs): + future = _CompletedFuture() + try: + future._result = fn(*args, **kwargs) + except Exception as e: # pragma: no cover — surface in test failure + future._exc = e + return future + + +class _CompletedFuture: + _result = None + _exc = None + + def result(self, timeout=None): + if self._exc is not None: + raise self._exc + return self._result + + def cancel(self): + pass + + +def _captured_prep_params(table, methods): + """ + Build a BulkRun, replace the prep pool's executor with a synchronous + stand-in, and capture the params handed to ``_prepare_single_batch_worker``. + """ + bulk = BulkRun( + input_table=table, + project_name="test-prepare-multi-protein", + methods_to_run=methods, + use_decoys=False, + use_known_binders=False, + use_gpu=False, + n_workers=1, + ) + + captured = [] + + def _capture(params): + captured.append(params) + return params.get("batch_name") + + # as_completed is just `iter()` for our sync fake. + with ( + patch.object(BulkRun, "_prepare_single_batch_worker", staticmethod(_capture)), + patch("guild.bulk.ProcessPoolExecutor", _SyncExecutor), + patch("guild.bulk.as_completed", lambda fs, **kw: list(fs)), + ): + bulk._prepare_docking() + + return captured + + +def test_non_karmadock_preps_one_row_per_unique_protein(multi_protein_table, cleanup): + """Boltz/Vina/DiffDock/gnina: prep tasks = unique proteins, not 1 and not N rows.""" + captured = _captured_prep_params(multi_protein_table, [BOLTZ_PREFIX]) + + prepped_protein_ids = {p["protein_idx"] for p in captured} + expected = set(multi_protein_table[PROTEIN_CONF_ID].unique()) + + assert len(captured) == len(expected), ( + f"Expected one prep task per unique protein ({len(expected)}), " + f"got {len(captured)}: {[p['protein_idx'] for p in captured]}" + ) + assert prepped_protein_ids == expected + + +def test_karmadock_still_preps_every_row(multi_protein_table, cleanup): + """KarmaDock: must prep every (protein, ligand) row because its per-ligand + data tree is staged here, not by a downstream worker.""" + captured = _captured_prep_params(multi_protein_table, [KARMADOCK_PREFIX]) + + assert len(captured) == len(multi_protein_table) + captured_pairs = {(p["protein_idx"], p["ligand_idx"]) for p in captured} + expected_pairs = set( + zip( + multi_protein_table[PROTEIN_CONF_ID], + multi_protein_table["ligand_id"], + strict=True, + ) + ) + assert captured_pairs == expected_pairs diff --git a/tests/bulk_run/test_progress_logging.py b/tests/bulk_run/test_progress_logging.py new file mode 100644 index 0000000..4037bae --- /dev/null +++ b/tests/bulk_run/test_progress_logging.py @@ -0,0 +1,133 @@ +""" +Tests for the per-batch and project-level progress logs written by +``BulkRun.run_docking`` / ``run_guild_scoring``. + +Each log is a simple append-only file of timestamped milestones: + + Starting Vina at 2026-05-25 12:34:56.789012 + Completed Vina at 2026-05-25 12:34:57.123456 + FAILED Vina 3pbl-A-ETQ-A_lig1 (timeout) at ... + +The per-batch milestones are appended to ``batches//output.log`` — +the same file Guild workers write to via ``logging.basicConfig``, so there's +a single log per batch rather than two. The project-level copy at +``{project}/batch_progress.log`` carries batch-level milestones plus a +copy of every FAILED line so a single grep surfaces every failure across +the run. +""" + +import shutil +from pathlib import Path +from unittest.mock import patch + +import pandas as pd +import pytest + +from guild.bulk import BulkRun +from guild.constants.bulk import BATCH_FOLDER, BATCH_PROGRESS_LOG_FILE, OUTPUT_LOG_FILE +from guild.constants.guild import BOLTZ_PREFIX, VINA_PREFIX + +TEST_DIR = Path(__file__).parent.parent +TEST_DATA_DIR = TEST_DIR / "test_data" + + +@pytest.fixture +def test_input_table(): + df = pd.read_csv(TEST_DATA_DIR / "bulk_dummy.csv") + df["protein_path"] = str(TEST_DATA_DIR / df["protein_path"].iloc[0]) + return df + + +@pytest.fixture +def cleanup(): + yield + test_project = Path.cwd() / "data" / "test-progress-log" + if test_project.exists(): + shutil.rmtree(test_project, ignore_errors=True) + + +def test_per_batch_log_records_method_start_and_end(test_input_table, cleanup): + """Each per-method helper writes Starting/Completed lines into the per-batch log.""" + bulk = BulkRun( + input_table=test_input_table, + project_name="test-progress-log", + methods_to_run=[VINA_PREFIX, BOLTZ_PREFIX], + use_decoys=False, + use_known_binders=False, + use_gpu=False, + n_workers=1, + ) + + # Mock the actual docking work — we only care about the log writes that + # the helpers do at their start/end. + def make_noop(name): + # Each helper's body would otherwise hit real subprocess / I/O work. + # Patch them to a no-op so the test stays fast; the log writes happen + # in the helpers themselves, which we leave intact via wraps=... . + def _noop(self, current_batch, current_batch_folder): + # Reproduce the helpers' own progress lines so the test still + # exercises the format the user will see in real runs. + BulkRun._log_progress( + f"{current_batch_folder}/{OUTPUT_LOG_FILE}", + message=f"Starting {name}", + ) + BulkRun._log_progress( + f"{current_batch_folder}/{OUTPUT_LOG_FILE}", + message=f"Completed {name}", + ) + return _noop + + with ( + patch.object(BulkRun, "_prepare_docking", lambda self: None), + patch.object(BulkRun, "_run_vina_for_batch", make_noop("Vina")), + patch.object(BulkRun, "_run_boltz_for_batch", make_noop("Boltz")), + ): + bulk.run_docking() + + batch_name = next(iter(bulk.batched_dictionary)) + batch_folder = bulk.batched_dictionary[batch_name][BATCH_FOLDER] + batch_log_text = Path(batch_folder, OUTPUT_LOG_FILE).read_text() + + # All method milestones should be present in the per-batch log. + for marker in ("Starting Vina", "Completed Vina", "Starting Boltz", "Completed Boltz"): + assert marker in batch_log_text, f"missing {marker!r} in:\n{batch_log_text}" + + # The batch's start/completion bookends should also be there. + assert "Starting batch_1" in batch_log_text + assert "Completed batch_1" in batch_log_text + + +def test_failure_line_appears_in_both_logs(test_input_table, cleanup): + """``FAILED`` entries land in both the per-batch and the project log.""" + bulk = BulkRun( + input_table=test_input_table, + project_name="test-progress-log", + methods_to_run=[VINA_PREFIX], + use_decoys=False, + use_known_binders=False, + use_gpu=False, + n_workers=1, + ) + + batch_name = next(iter(bulk.batched_dictionary)) + batch_folder = bulk.batched_dictionary[batch_name][BATCH_FOLDER] + batch_log = f"{batch_folder}/{OUTPUT_LOG_FILE}" + main_log = f"{bulk.project_folder}/{BATCH_PROGRESS_LOG_FILE}" + + # Directly exercise _log_progress as the helpers do when a combo fails. + BulkRun._log_progress( + batch_log, main_log, + message="FAILED Vina 3pbl-A-ETQ-A_lig1 (timeout)", + ) + + batch_text = Path(batch_log).read_text() + main_text = Path(main_log).read_text() + assert "FAILED Vina 3pbl-A-ETQ-A_lig1 (timeout)" in batch_text + assert "FAILED Vina 3pbl-A-ETQ-A_lig1 (timeout)" in main_text + + +def test_log_progress_skips_none_paths(tmp_path: Path): + """``None`` entries in the *paths argpack are silently ignored.""" + log = tmp_path / "progress.log" + BulkRun._log_progress(str(log), None, message="Sentinel") + assert "Sentinel at" in log.read_text() diff --git a/tests/bulk_run/test_run_docking_order.py b/tests/bulk_run/test_run_docking_order.py new file mode 100644 index 0000000..5b20e21 --- /dev/null +++ b/tests/bulk_run/test_run_docking_order.py @@ -0,0 +1,120 @@ +""" +Tests that ``BulkRun.run_docking`` dispatches each batch's method blocks in +the order the user supplied via ``methods_to_run`` (rather than a hardcoded +order). +""" + +import shutil +from pathlib import Path +from unittest.mock import patch + +import pandas as pd +import pytest + +from guild.bulk import BulkRun +from guild.constants.guild import ( + BOLTZ_PREFIX, + DIFFDOCK_PREFIX, + GNINA_PREFIX, + KARMADOCK_PREFIX, + VINA_PREFIX, +) + +TEST_DIR = Path(__file__).parent.parent +TEST_DATA_DIR = TEST_DIR / "test_data" + + +@pytest.fixture +def test_input_table(): + df = pd.read_csv(TEST_DATA_DIR / "bulk_dummy.csv") + df["protein_path"] = str(TEST_DATA_DIR / df["protein_path"].iloc[0]) + return df + + +@pytest.fixture +def cleanup(): + yield + test_project = Path.cwd() / "data" / "test-run-docking-order" + if test_project.exists(): + shutil.rmtree(test_project, ignore_errors=True) + + +def _run_and_capture(test_input_table, methods): + """ + Build a BulkRun, mock out every per-method helper plus ``_prepare_docking``, + and return the order in which the helpers were called for the (single) + batch. + """ + bulk = BulkRun( + input_table=test_input_table, + project_name="test-run-docking-order", + methods_to_run=methods, + use_decoys=False, + use_known_binders=False, + use_gpu=False, + n_workers=1, + ) + + call_order = [] + + def make_recorder(name): + def _record(self, current_batch, current_batch_folder): + call_order.append(name) + return _record + + with ( + patch.object(BulkRun, "_prepare_docking", lambda self: None), + patch.object(BulkRun, "_run_boltz_for_batch", make_recorder(BOLTZ_PREFIX)), + patch.object(BulkRun, "_run_vina_for_batch", make_recorder(VINA_PREFIX)), + patch.object(BulkRun, "_run_gnina_for_batch", make_recorder(GNINA_PREFIX)), + patch.object(BulkRun, "_run_karmadock_for_batch", make_recorder(KARMADOCK_PREFIX)), + patch.object(BulkRun, "_run_diffdock_for_batch", make_recorder(DIFFDOCK_PREFIX)), + ): + bulk.run_docking() + + return call_order, bulk.methods_to_run + + +def test_user_order_vina_then_boltz(test_input_table, cleanup): + """``[vina, boltz]`` dispatches Vina before Boltz.""" + order, expanded = _run_and_capture(test_input_table, [VINA_PREFIX, BOLTZ_PREFIX]) + + # methods_to_run is auto-extended with vina_rescore_boltz when boltz is + # requested, but that prefix has no docking-time runner, so it doesn't + # show up in the call sequence. + assert order == [VINA_PREFIX, BOLTZ_PREFIX] + # Sanity: the auto-added rescore prefix is in the expanded list. + assert any("vina_rescore_boltz" == m for m in expanded) + + +def test_user_order_boltz_then_vina(test_input_table, cleanup): + """``[boltz, vina]`` dispatches Boltz before Vina.""" + order, _ = _run_and_capture(test_input_table, [BOLTZ_PREFIX, VINA_PREFIX]) + # First two entries (skipping the auto-added rescore which is a no-op). + docking_order = [m for m in order if m in {BOLTZ_PREFIX, VINA_PREFIX}] + assert docking_order == [BOLTZ_PREFIX, VINA_PREFIX] + + +def test_five_methods_in_explicit_order(test_input_table, cleanup): + """A custom interleaved order is respected end-to-end.""" + desired = [ + GNINA_PREFIX, + DIFFDOCK_PREFIX, + VINA_PREFIX, + KARMADOCK_PREFIX, + BOLTZ_PREFIX, + ] + order, _ = _run_and_capture(test_input_table, desired) + + # Filter out any non-docking prefixes (vina_rescore_* auto-additions). + docking_order = [m for m in order if m in set(desired)] + assert docking_order == desired + + +def test_rescore_prefixes_do_not_dispatch_runners(test_input_table, cleanup): + """vina_rescore_* prefixes are skipped during the docking phase.""" + order, expanded = _run_and_capture(test_input_table, [DIFFDOCK_PREFIX]) + # DiffDock alone auto-adds vina_rescore_diffdock; it must not produce a + # docking-helper call (the rescore happens in run_guild_scoring). + assert "vina_rescore_diffdock" in expanded + assert order == [DIFFDOCK_PREFIX] diff --git a/tests/bulk_run/test_subprocess_log.py b/tests/bulk_run/test_subprocess_log.py new file mode 100644 index 0000000..1e7485b --- /dev/null +++ b/tests/bulk_run/test_subprocess_log.py @@ -0,0 +1,79 @@ +""" +Tests for ``guild.tools.subprocess_log.write_subprocess_log``. + +That helper backs the per-combination ``*.subprocess.log`` files our docking +helpers emit so external consumers can read a focused, predictable +transcript when a method fails — rather than grepping the batch-wide +``output.log``. The helper's contract is: + +- Always writes a complete transcript (success or failure). +- Creates parent dirs if missing. +- Decodes bytes stdout/stderr; tolerates None. +- Quotes argv for reproducibility. +""" + +from pathlib import Path + +from guild.tools.subprocess_log import write_subprocess_log + + +def test_success_transcript(tmp_path: Path): + log = tmp_path / "method" / "combo.subprocess.log" + write_subprocess_log( + log, + argv=["/usr/bin/echo", "hello world"], + returncode=0, + stdout="hello world\n", + stderr="", + ) + + text = log.read_text() + assert "[INVOCATION SUCCESS]" in text + assert "argv: /usr/bin/echo 'hello world'" in text + assert "[STDOUT]\nhello world" in text + assert "[STDERR]\n" in text + + +def test_failure_transcript_with_header(tmp_path: Path): + log = tmp_path / "boltz.subprocess.log" + write_subprocess_log( + log, + argv=["boltz", "predict", "in.yaml"], + returncode=1, + stdout="some progress...\n", + stderr="ImportError: foo\n", + extra_header="exited 1 — Boltz crashed before writing manifest", + ) + + text = log.read_text() + assert "[INVOCATION FAILED (exit 1)]" in text + assert "note: exited 1" in text + assert "ImportError: foo" in text + + +def test_bytes_decoded(tmp_path: Path): + log = tmp_path / "b.log" + write_subprocess_log( + log, + argv=["true"], + returncode=0, + stdout=b"binary out\n", + stderr=b"\xc3\xa9rror\n", # 'érror' UTF-8 + ) + text = log.read_text() + assert "binary out" in text + assert "érror" in text + + +def test_none_blobs_tolerated(tmp_path: Path): + log = tmp_path / "n.log" + write_subprocess_log(log, argv=["x"], returncode=0, stdout=None, stderr=None) + # Should not raise; file should still exist with the header block. + assert "[INVOCATION SUCCESS]" in log.read_text() + + +def test_parent_dirs_created(tmp_path: Path): + log = tmp_path / "deep" / "nested" / "path" / "combo.log" + assert not log.parent.exists() + write_subprocess_log(log, argv=["x"], returncode=0, stdout="", stderr="") + assert log.exists() diff --git a/tests/bulk_run/test_vina_rescore_split.py b/tests/bulk_run/test_vina_rescore_split.py new file mode 100644 index 0000000..eafab25 --- /dev/null +++ b/tests/bulk_run/test_vina_rescore_split.py @@ -0,0 +1,130 @@ +""" +Tests for the split Vina re-scoring tracks: ``vina_rescore_boltz`` and +``vina_rescore_diffdock`` are independent and must coexist when both upstream +methods are present in ``methods_to_run``. +""" + +import shutil +from pathlib import Path + +import pandas as pd +import pytest + +from guild.bulk import BulkRun +from guild.constants.guild import ( + BOLTZ_PREFIX, + DIFFDOCK_PREFIX, + SCORES_DICTIONARY, + VINA_PREFIX, + VINA_RESCORE_BOLTZ_PREFIX, + VINA_RESCORE_BOLTZ_SCORE, + VINA_RESCORE_DIFFDOCK_PREFIX, + VINA_RESCORE_DIFFDOCK_SCORE, +) + +TEST_DIR = Path(__file__).parent.parent +TEST_DATA_DIR = TEST_DIR / "test_data" + + +@pytest.fixture +def test_input_table(): + df = pd.read_csv(TEST_DATA_DIR / "bulk_dummy.csv") + df["protein_path"] = str(TEST_DATA_DIR / df["protein_path"].iloc[0]) + return df + + +@pytest.fixture +def cleanup(): + yield + test_project = Path.cwd() / "data" / "test-rescore-split" + if test_project.exists(): + shutil.rmtree(test_project, ignore_errors=True) + + +def test_scores_dictionary_distinguishes_rescore_methods(): + """Both rescore prefixes have their own distinct score columns.""" + assert SCORES_DICTIONARY[VINA_RESCORE_BOLTZ_PREFIX] == VINA_RESCORE_BOLTZ_SCORE + assert SCORES_DICTIONARY[VINA_RESCORE_DIFFDOCK_PREFIX] == VINA_RESCORE_DIFFDOCK_SCORE + assert VINA_RESCORE_BOLTZ_SCORE != VINA_RESCORE_DIFFDOCK_SCORE + + +def test_boltz_auto_enables_only_boltz_rescore(test_input_table, cleanup): + """Requesting boltz auto-adds vina_rescore_boltz but NOT vina_rescore_diffdock.""" + bulk = BulkRun( + input_table=test_input_table, + project_name="test-rescore-split", + methods_to_run=[BOLTZ_PREFIX], + use_decoys=False, + use_known_binders=False, + use_gpu=False, + n_workers=1, + ) + assert BOLTZ_PREFIX in bulk.methods_to_run + assert VINA_RESCORE_BOLTZ_PREFIX in bulk.methods_to_run + assert VINA_RESCORE_DIFFDOCK_PREFIX not in bulk.methods_to_run + + +def test_diffdock_auto_enables_only_diffdock_rescore(test_input_table, cleanup): + """Requesting diffdock auto-adds vina_rescore_diffdock but NOT vina_rescore_boltz.""" + bulk = BulkRun( + input_table=test_input_table, + project_name="test-rescore-split", + methods_to_run=[DIFFDOCK_PREFIX], + use_decoys=False, + use_known_binders=False, + use_gpu=False, + n_workers=1, + ) + assert DIFFDOCK_PREFIX in bulk.methods_to_run + assert VINA_RESCORE_DIFFDOCK_PREFIX in bulk.methods_to_run + assert VINA_RESCORE_BOLTZ_PREFIX not in bulk.methods_to_run + + +def test_both_methods_enable_both_rescores(test_input_table, cleanup): + """Requesting both Boltz and DiffDock auto-enables both rescore tracks.""" + bulk = BulkRun( + input_table=test_input_table, + project_name="test-rescore-split", + methods_to_run=[BOLTZ_PREFIX, DIFFDOCK_PREFIX], + use_decoys=False, + use_known_binders=False, + use_gpu=False, + n_workers=1, + ) + assert VINA_RESCORE_BOLTZ_PREFIX in bulk.methods_to_run + assert VINA_RESCORE_DIFFDOCK_PREFIX in bulk.methods_to_run + + +def test_vina_alone_does_not_enable_any_rescore(test_input_table, cleanup): + """Plain Vina docking does not auto-enable either rescore.""" + bulk = BulkRun( + input_table=test_input_table, + project_name="test-rescore-split", + methods_to_run=[VINA_PREFIX], + use_decoys=False, + use_known_binders=False, + use_gpu=False, + n_workers=1, + ) + assert VINA_PREFIX in bulk.methods_to_run + assert VINA_RESCORE_BOLTZ_PREFIX not in bulk.methods_to_run + assert VINA_RESCORE_DIFFDOCK_PREFIX not in bulk.methods_to_run + + +def test_rescore_scoring_functions_emit_distinct_columns(): + """The two guild_scoring entry points return distinct score columns.""" + from guild.constants.bulk import BATCH_FOLDER, COMBINATIONS_TO_RUN_KEY + from guild.docking.boltz import vina_rescore_boltz_guild_scoring + from guild.docking.diffdock import vina_rescore_diffdock_guild_scoring + + # Empty batch (no outputs on disk) → both functions return frames with + # their respective score columns, no overlap. + empty_batch = {BATCH_FOLDER: "/nonexistent", COMBINATIONS_TO_RUN_KEY: []} + + boltz_df = vina_rescore_boltz_guild_scoring(empty_batch) + diffdock_df = vina_rescore_diffdock_guild_scoring(empty_batch) + + assert VINA_RESCORE_BOLTZ_SCORE in boltz_df.columns + assert VINA_RESCORE_DIFFDOCK_SCORE not in boltz_df.columns + assert VINA_RESCORE_DIFFDOCK_SCORE in diffdock_df.columns + assert VINA_RESCORE_BOLTZ_SCORE not in diffdock_df.columns diff --git a/tests/scores/test_multimethod_underscore_combination.py b/tests/scores/test_multimethod_underscore_combination.py new file mode 100644 index 0000000..0f4ba2a --- /dev/null +++ b/tests/scores/test_multimethod_underscore_combination.py @@ -0,0 +1,107 @@ +""" +Regression test for the multi-method scoring bug where a protein_config_id +containing an underscore (e.g. ``6C97_P0``) caused: + - KarmaDock rows to be keyed wrong (split("_") mis-parse) → they failed the + cross-method merge and got dropped by dedup (one method per row), and + - smiles / ligand_category to come out as "unknown". + +The fix makes ``karmadock_guild_scoring`` recover the authoritative +(protein_config_id, ligand_id) from ``COMBINATIONS_TO_RUN_KEY`` instead of +string-splitting, plus a coalescing dedup and ligand_id-keyed category lookup. +""" + +import pandas as pd + +from guild.constants.bulk import ( + BATCH_FOLDER, + COMBINATION_ID, + COMBINATIONS_TO_RUN_KEY, +) +from guild.constants.guild import ( + KARMADOCK_FOLDER, + KARMADOCK_SCORE, + LIGAND_ID, + PROTEIN_CONF_ID, +) +from guild.constants.karmadock import KARMADOCK_RESULTS_FOLDER +from guild.docking.karmadock import karmadock_guild_scoring + +# protein_config_id with an underscore — the case that broke split("_") +PROTEIN = "6C97_P0" +LIGANDS = ["lig1", "drug_2", "pos_3"] # ligand ids, one also containing "_" + + +def _write_karmadock_results(results_dir): + """Emit a KarmaDock-style results CSV (pdb_id == combined combination id).""" + results_dir.mkdir(parents=True, exist_ok=True) + rows = [] + for i, lig in enumerate(LIGANDS): + rows.append( + { + "pdb_id": f"{PROTEIN}_{lig}", + "score": 10.0 + i, + "RMSD": 1.0, + "FF_RMSD": 1.0, + "Aligned_RMSD": 1.0, + } + ) + pd.DataFrame(rows).to_csv(results_dir / "0.csv", index=False) + + +def _batch_dictionary(batch_folder): + return { + BATCH_FOLDER: str(batch_folder), + COMBINATIONS_TO_RUN_KEY: [(PROTEIN, lig) for lig in LIGANDS], + } + + +def test_karmadock_recovers_authoritative_keys(tmp_path): + """KarmaDock must recover protein_config_id='6C97_P0' and the full ligand_id, + not the split('_') mis-parse ('6C97', 'P0_lig1').""" + results_dir = tmp_path / KARMADOCK_FOLDER / KARMADOCK_RESULTS_FOLDER + _write_karmadock_results(results_dir) + + df = karmadock_guild_scoring(_batch_dictionary(tmp_path)) + + assert set(df[PROTEIN_CONF_ID].unique()) == {PROTEIN}, ( + f"expected all rows under protein_config_id={PROTEIN!r}, " + f"got {df[PROTEIN_CONF_ID].unique().tolist()}" + ) + assert set(df[LIGAND_ID]) == set(LIGANDS) + # combination id is preserved intact + assert set(df[COMBINATION_ID]) == {f"{PROTEIN}_{lig}" for lig in LIGANDS} + # the underscore-containing ligand id survived + row = df[df[LIGAND_ID] == "drug_2"] + assert len(row) == 1 + assert row.iloc[0][PROTEIN_CONF_ID] == PROTEIN + + +def test_karmadock_keys_merge_with_vina_keys(tmp_path): + """KarmaDock's recovered keys must match how vina/gnina build theirs + (protein_config_id + '_' + ligand_id), so an outer merge unites them into + one row per combination instead of doubling.""" + results_dir = tmp_path / KARMADOCK_FOLDER / KARMADOCK_RESULTS_FOLDER + _write_karmadock_results(results_dir) + karma = karmadock_guild_scoring(_batch_dictionary(tmp_path)) + + # vina-style frame, built the way vina_guild_scoring does + vina = pd.DataFrame( + { + PROTEIN_CONF_ID: [PROTEIN] * len(LIGANDS), + LIGAND_ID: LIGANDS, + "vina_score": [-5.0, -6.0, -7.0], + } + ) + vina[COMBINATION_ID] = vina[PROTEIN_CONF_ID] + "_" + vina[LIGAND_ID] + + merged = pd.merge( + vina, + karma, + on=[COMBINATION_ID, PROTEIN_CONF_ID, LIGAND_ID], + how="outer", + ) + + # One row per combination, both method columns populated (no NaN). + assert len(merged) == len(LIGANDS) + assert merged["vina_score"].notna().all() + assert merged[KARMADOCK_SCORE].notna().all() diff --git a/tests/single_run/single_run_imports.py b/tests/single_run/single_run_imports.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/single_run/test_gnina_input_mode.py b/tests/single_run/test_gnina_input_mode.py new file mode 100644 index 0000000..f816d8e --- /dev/null +++ b/tests/single_run/test_gnina_input_mode.py @@ -0,0 +1,200 @@ +""" +Tests for the ``gnina_input_mode`` knob. + +- ``deploy_gnina`` should accept SDF + PDB inputs and pass them through + to gnina's CLI without forcing PDBQT validation. +- The existing PDBQT call path should be unchanged. +- ``BulkRun`` should silently downgrade ``"sdf"`` to ``"pdbqt"`` when + Vina or a Vina-rescore is co-requested, logging a warning. +""" + +import logging +import shutil +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pandas as pd +import pytest + +from guild.bulk import BulkRun +from guild.constants.guild import ( + BOLTZ_PREFIX, + GNINA_PREFIX, + VINA_PREFIX, +) +from guild.docking import gnina as gnina_module +from guild.docking.gnina import deploy_gnina + +TEST_DIR = Path(__file__).parent.parent +TEST_DATA_DIR = TEST_DIR / "test_data" + + +@pytest.fixture +def fake_inputs(tmp_path: Path): + """Create minimal placeholder receptor + ligand files for both formats.""" + pdb = tmp_path / "receptor.pdb" + pdb.write_text("ATOM 1 N ALA A 1 0.000 0.000 0.000 1.00 0.00 N\nEND\n") + + sdf = tmp_path / "ligand.sdf" + sdf.write_text("dummy\n\n\n 0 0 0 0 0 0 0 0 0 0999 V2000\nM END\n$$$$\n") + + pdbqt_r = tmp_path / "receptor.pdbqt" + pdbqt_r.write_text( + "ROOT\n" + "ATOM 1 N ALA A 1 0.000 0.000 0.000 1.00 0.00 +0.000 N\n" + "ENDROOT\nTORSDOF 0\n" + ) + + pdbqt_l = tmp_path / "ligand.pdbqt" + pdbqt_l.write_text( + "ROOT\n" + "ATOM 1 C UNL 1 0.000 0.000 0.000 1.00 0.00 +0.000 C\n" + "ENDROOT\nTORSDOF 0\n" + ) + + return { + "pdb": str(pdb), "sdf": str(sdf), + "pdbqt_receptor": str(pdbqt_r), "pdbqt_ligand": str(pdbqt_l), + "out_pdbqt": str(tmp_path / "out.pdbqt"), + "out_scores": str(tmp_path / "scores.txt"), + } + + +def _fake_completed(stdout=""): + m = MagicMock() + m.stdout = stdout + m.stderr = "" + m.returncode = 0 + return m + + +_FAKE_GNINA_STDOUT = ( + "mode | affinity | CNN | CNN\n" + " | (kcal/mol) | pose-score| affinity\n" + "-----+------------+----------+----------\n" + " 1 -8.345 0.789 6.234\n" +) + + +def test_sdf_mode_emits_sdf_and_pdb_argv(fake_inputs, tmp_path): + """SDF + PDB inputs land in argv unmodified; PDBQT validation is skipped.""" + captured_argv = {} + + def _capture(argv, **kwargs): + captured_argv["argv"] = list(argv) + return _fake_completed(_FAKE_GNINA_STDOUT) + + with patch.object(gnina_module.subprocess, "run", side_effect=_capture): + deploy_gnina( + receptor=fake_inputs["pdb"], + ligand=fake_inputs["sdf"], + center=(0.0, 0.0, 0.0), + size=(20.0, 20.0, 20.0), + output_pdbqt=fake_inputs["out_pdbqt"], + output_scores=fake_inputs["out_scores"], + use_gpu=False, + ) + + argv = captured_argv["argv"] + assert fake_inputs["pdb"] in argv, f"pdb receptor missing from argv: {argv}" + assert fake_inputs["sdf"] in argv, f"sdf ligand missing from argv: {argv}" + # gnina's CLI passes the path right after --receptor / --ligand + assert argv[argv.index("--receptor") + 1] == fake_inputs["pdb"] + assert argv[argv.index("--ligand") + 1] == fake_inputs["sdf"] + + +def test_pdbqt_mode_unchanged(fake_inputs): + """PDBQT path still emits the PDBQT argv and validation runs (no crash).""" + captured_argv = {} + + def _capture(argv, **kwargs): + captured_argv["argv"] = list(argv) + return _fake_completed(_FAKE_GNINA_STDOUT) + + with patch.object(gnina_module.subprocess, "run", side_effect=_capture): + deploy_gnina( + receptor=fake_inputs["pdbqt_receptor"], + ligand=fake_inputs["pdbqt_ligand"], + center=(0.0, 0.0, 0.0), + size=(20.0, 20.0, 20.0), + output_pdbqt=fake_inputs["out_pdbqt"], + output_scores=fake_inputs["out_scores"], + use_gpu=False, + ) + + argv = captured_argv["argv"] + assert argv[argv.index("--receptor") + 1] == fake_inputs["pdbqt_receptor"] + assert argv[argv.index("--ligand") + 1] == fake_inputs["pdbqt_ligand"] + + +@pytest.fixture +def bulk_table(): + df = pd.read_csv(TEST_DATA_DIR / "bulk_dummy.csv") + df["protein_path"] = str(TEST_DATA_DIR / df["protein_path"].iloc[0]) + return df + + +@pytest.fixture +def cleanup_bulk(): + yield + test_project = Path.cwd() / "data" / "test-gnina-input-mode" + if test_project.exists(): + shutil.rmtree(test_project, ignore_errors=True) + + +def test_bulkrun_downgrades_sdf_to_pdbqt_when_vina_also_requested( + bulk_table, cleanup_bulk, caplog +): + """Co-requesting Vina forces gnina back to PDBQT mode with a warning.""" + caplog.set_level(logging.WARNING, logger="guild.bulk") + + bulk = BulkRun( + input_table=bulk_table, + project_name="test-gnina-input-mode", + methods_to_run=[VINA_PREFIX, GNINA_PREFIX], + gnina_input_mode="sdf", + use_decoys=False, + use_known_binders=False, + use_gpu=False, + n_workers=1, + ) + + assert bulk.gnina_input_mode == "pdbqt" + assert any( + "SDF-mode requested for gnina but Vina" in rec.message + for rec in caplog.records + ), f"expected downgrade warning, got: {[r.message for r in caplog.records]}" + + +def test_bulkrun_keeps_sdf_when_gnina_alone(bulk_table, cleanup_bulk): + """Gnina-only run with sdf preserves the requested mode.""" + bulk = BulkRun( + input_table=bulk_table, + project_name="test-gnina-input-mode", + methods_to_run=[GNINA_PREFIX], + gnina_input_mode="sdf", + use_decoys=False, + use_known_binders=False, + use_gpu=False, + n_workers=1, + ) + assert bulk.gnina_input_mode == "sdf" + + +def test_bulkrun_keeps_sdf_when_gnina_with_boltz(bulk_table, cleanup_bulk): + """Boltz doesn't require PDBQT, so gnina+boltz still honours sdf-mode.""" + bulk = BulkRun( + input_table=bulk_table, + project_name="test-gnina-input-mode", + methods_to_run=[GNINA_PREFIX, BOLTZ_PREFIX], + gnina_input_mode="sdf", + use_decoys=False, + use_known_binders=False, + use_gpu=False, + n_workers=1, + ) + # BulkRun auto-adds vina_rescore_boltz when boltz is requested; that DOES + # require PDBQT, so gnina downgrades. Surfaces an important user-facing + # consequence of the auto-rescore: combining gnina-sdf with boltz also + # falls back to PDBQT (because of the implicit vina_rescore_boltz). + assert bulk.gnina_input_mode == "pdbqt" diff --git a/tests/single_run/test_gnina_scoring.py b/tests/single_run/test_gnina_scoring.py new file mode 100644 index 0000000..1e596da --- /dev/null +++ b/tests/single_run/test_gnina_scoring.py @@ -0,0 +1,65 @@ +""" +Tests for gnina stdout parsing and the per-combination score file produced by +:func:`guild.docking.gnina.deploy_gnina` / :func:`process_gnina_output`. +""" + +import math +from pathlib import Path + +import pytest + +from guild.docking.gnina import parse_gnina_stdout, process_gnina_output + +# Example gnina stdout (truncated to the relevant section). Columns are +# mode | affinity (kcal/mol) | CNNscore | CNNaffinity. +GNINA_STDOUT_SAMPLE = """\ +gnina v1.3 Linux64 + _______ _____ _________ _____ _____ + | ____ | | | | | | + ... + +mode | affinity | CNN | CNN + | (kcal/mol) | pose-score| affinity +-----+------------+----------+---------- + 1 -8.345 0.7891 6.234 + 2 -7.910 0.6512 5.987 + 3 -7.500 0.5800 5.612 +""" + + +def test_parse_gnina_stdout_extracts_all_poses(): + poses = parse_gnina_stdout(GNINA_STDOUT_SAMPLE) + assert len(poses) == 3 + assert poses[0] == (1, -8.345, 0.7891, 6.234) + assert poses[2] == (3, -7.500, 0.58, 5.612) + + +def test_parse_gnina_stdout_raises_on_empty(): + with pytest.raises(ValueError): + parse_gnina_stdout("nothing useful here\nno table rows\n") + + +def test_process_gnina_output_picks_best_affinity(tmp_path: Path): + score_file = tmp_path / "scores.txt" + score_file.write_text( + "1: -8.345\t0.7891\n" + "2: -7.910\t0.6512\n" + "3: -7.500\t0.5800\n" + ) + + affinity, cnn = process_gnina_output(str(score_file)) + # "Best" = minimum affinity, matching the Vina convention. + assert affinity == pytest.approx(-8.345) + # CNN score returned is the one from the same pose, not max across poses. + assert cnn == pytest.approx(0.7891) + + +def test_process_gnina_output_handles_missing_cnn_column(tmp_path: Path): + # Vina-formatted score file (no CNN column) — older score files must still + # parse without crashing. + score_file = tmp_path / "vina_format.txt" + score_file.write_text("1: -8.0\n2: -7.5\n") + + affinity, cnn = process_gnina_output(str(score_file)) + assert affinity == pytest.approx(-8.0) + assert math.isnan(cnn) diff --git a/tests/single_run/test_multi_chain_support.py b/tests/single_run/test_multi_chain_support.py new file mode 100644 index 0000000..0c83154 --- /dev/null +++ b/tests/single_run/test_multi_chain_support.py @@ -0,0 +1,111 @@ +""" +Tests for multi-chain receptor support: chain isolation keeps every requested +chain, residue renumbering restarts per chain, and the Boltz YAML emits one +``protein`` block per chain. See README "Multi-chain binding pocket". +""" + +from pathlib import Path + +import yaml +from Bio.PDB import PDBParser + +from guild.docking.boltz import generate_boltz_yaml +from guild.tools.preparation import _normalize_chain_list, isolate_protein_chain + +TEST_DATA_DIR = Path(__file__).parent.parent / "test_data" +PDB = str(TEST_DATA_DIR / "3pbl.pdb") # has chains A and B + + +def _chains_and_resnums(pdb_path): + model = PDBParser(QUIET=True).get_structure("x", pdb_path)[0] + return {ch.id: [res.id[1] for res in ch] for ch in model} + + +def test_normalize_chain_list_forms(): + assert _normalize_chain_list("A") == ["A"] + assert _normalize_chain_list("A,B") == ["A", "B"] + assert _normalize_chain_list("A, B ") == ["A", "B"] + assert _normalize_chain_list(["A", "B"]) == ["A", "B"] + assert _normalize_chain_list(None) == [] + + +def test_isolate_keeps_all_requested_chains(tmp_path): + out = tmp_path / "multi.pdb" + isolate_protein_chain(PDB, "3pbl", str(out), target_chain="A,B") + chains = _chains_and_resnums(str(out)) + assert set(chains) == {"A", "B"}, f"expected both chains, got {set(chains)}" + + +def test_renumber_restarts_per_chain(tmp_path): + """Each isolated chain must be 1-based independently — required so Boltz + pocket contacts ([chain, index]) line up with each chain's sequence.""" + out = tmp_path / "multi.pdb" + isolate_protein_chain(PDB, "3pbl", str(out), target_chain="A,B") + chains = _chains_and_resnums(str(out)) + for chain_id, resnums in chains.items(): + assert resnums[0] == 1, f"chain {chain_id} should start at residue 1" + # Residues are contiguous within the chain. + assert resnums == list(range(1, len(resnums) + 1)) + + +def test_single_chain_still_works(tmp_path): + out = tmp_path / "single.pdb" + isolate_protein_chain(PDB, "3pbl", str(out), target_chain="A") + chains = _chains_and_resnums(str(out)) + assert set(chains) == {"A"} + + +def test_boltz_yaml_single_chain_unchanged(tmp_path): + out = tmp_path / "single.yaml" + generate_boltz_yaml( + protein_sequence="MKT", + protein_chain="A", + ligand_sequences=["CCO"], + ligand_ids=["L"], + output_file=str(out), + template_file=None, + ) + doc = yaml.safe_load(out.read_text()) + proteins = [s for s in doc["sequences"] if "protein" in s] + assert len(proteins) == 1 + assert proteins[0]["protein"]["id"] == "A" + assert proteins[0]["protein"]["sequence"] == "MKT" + + +def test_boltz_yaml_multi_chain(tmp_path): + out = tmp_path / "multi.yaml" + generate_boltz_yaml( + protein_sequence=["MKT", "GGS"], + protein_chain=["A", "B"], + ligand_sequences=["CCO"], + ligand_ids=["L"], + output_file=str(out), + template_file=PDB, + msa_file=["a.a3m", "b.a3m"], + pocket_contacts=[["A", 1], ["B", 2]], + ) + doc = yaml.safe_load(out.read_text()) + proteins = [s["protein"] for s in doc["sequences"] if "protein" in s] + assert [p["id"] for p in proteins] == ["A", "B"] + assert [p["sequence"] for p in proteins] == ["MKT", "GGS"] + assert [p["msa"] for p in proteins] == ["a.a3m", "b.a3m"] + # Template constrains both chains. + template = doc["templates"][0] + assert template["chain_id"] == ["A", "B"] + assert template["template_id"] == ["A1", "B1"] + # Pocket constraint spans both chains. + pocket = doc["constraints"][0]["pocket"] + assert pocket["contacts"] == [["A", 1], ["B", 2]] + + +def test_boltz_yaml_chain_sequence_mismatch_raises(tmp_path): + import pytest + + with pytest.raises(ValueError): + generate_boltz_yaml( + protein_sequence=["MKT"], + protein_chain=["A", "B"], + ligand_sequences=["CCO"], + ligand_ids=["L"], + output_file=str(tmp_path / "bad.yaml"), + ) diff --git a/tests/single_run/test_pocket_contacts_from_box.py b/tests/single_run/test_pocket_contacts_from_box.py new file mode 100644 index 0000000..60c3f68 --- /dev/null +++ b/tests/single_run/test_pocket_contacts_from_box.py @@ -0,0 +1,95 @@ +""" +Tests for ``guild.transformers.pdb.get_pocket_contacts_from_box``. +""" + +from pathlib import Path + +from guild.docking.vina import get_center_and_size_from_box_file +from guild.transformers.pdb import get_pocket_contacts_from_box + +TEST_DATA_DIR = Path(__file__).parent.parent / "test_data" +PDB = str(TEST_DATA_DIR / "3pbl.pdb") +BOX = str(TEST_DATA_DIR / "3pbl_box.txt") + + +def test_returns_residues_inside_box(): + center, size = get_center_and_size_from_box_file(BOX) + contacts = get_pocket_contacts_from_box( + protein_pdb=PDB, + protein_chain="A", + center=center, + size=size, + ) + assert contacts, "expected at least one residue Cα inside the box" + # Schema: list of [chain_id, 1-based contiguous sequence index] + for chain, idx in contacts: + assert chain == "A" + assert isinstance(idx, int) and idx >= 1 + + +def test_empty_when_box_is_far_away(): + contacts = get_pocket_contacts_from_box( + protein_pdb=PDB, + protein_chain="A", + center=(10000.0, 10000.0, 10000.0), + size=(1.0, 1.0, 1.0), + ) + assert contacts == [] + + +def test_empty_when_chain_missing(): + contacts = get_pocket_contacts_from_box( + protein_pdb=PDB, + protein_chain="Z", + center=(0.0, 0.0, 0.0), + size=(20.0, 20.0, 20.0), + ) + assert contacts == [] + + +def test_multi_chain_indexes_each_chain_from_one(): + """A pocket spanning two chains must return contacts from both, with each + chain's residue index restarting at 1 (Boltz per-chain indexing).""" + # A large box centred on the structure should catch residues on both chains. + center, _ = get_center_and_size_from_box_file(BOX) + contacts = get_pocket_contacts_from_box( + protein_pdb=PDB, + protein_chain="A,B", + center=center, + size=(200.0, 200.0, 200.0), + ) + chains = {c for c, _ in contacts} + assert chains == {"A", "B"}, f"expected contacts on both chains, got {chains}" + # Each chain is indexed independently from 1. + for chain in ("A", "B"): + idxs = [idx for c, idx in contacts if c == chain] + assert min(idxs) == 1, f"chain {chain} should be 1-based, got min {min(idxs)}" + + +def test_multi_chain_is_union_of_single_chains(): + """contacts('A,B') == contacts('A') + contacts('B'), so the list form and + the comma-string form are equivalent and chains compose cleanly.""" + center, _ = get_center_and_size_from_box_file(BOX) + size = (200.0, 200.0, 200.0) + only_a = get_pocket_contacts_from_box(PDB, "A", center, size) + only_b = get_pocket_contacts_from_box(PDB, "B", center, size) + both_str = get_pocket_contacts_from_box(PDB, "A,B", center, size) + both_list = get_pocket_contacts_from_box(PDB, ["A", "B"], center, size) + assert both_str == both_list + assert both_str == only_a + only_b + + +def test_cuboid_axis_independence(): + """An asymmetric box must include residues outside the smaller axis if they + sit along a larger one — cuboid AABB, not min-axis sphere.""" + center, _ = get_center_and_size_from_box_file(BOX) + narrow_cubic = get_pocket_contacts_from_box( + protein_pdb=PDB, protein_chain="A", center=center, size=(2.0, 2.0, 2.0) + ) + elongated_x = get_pocket_contacts_from_box( + protein_pdb=PDB, protein_chain="A", center=center, size=(40.0, 2.0, 2.0) + ) + assert len(elongated_x) >= len(narrow_cubic), ( + "elongating the X axis must not shrink the contact set " + f"(narrow={len(narrow_cubic)}, elongated={len(elongated_x)})" + ) diff --git a/uv.lock b/uv.lock index d23e505..f54ee93 100644 --- a/uv.lock +++ b/uv.lock @@ -963,7 +963,7 @@ wheels = [ [[package]] name = "guild" -version = "1.0.0" +version = "1.1.4" source = { editable = "." } dependencies = [ { name = "boltz", extra = ["cuda"] }, From 0836e28aad97c959dad128a2d076ea971210a733 Mon Sep 17 00:00:00 2001 From: AJPreto Date: Fri, 12 Jun 2026 11:15:31 +0000 Subject: [PATCH 2/2] Ruff formatting fixes --- guild/bulk.py | 1 - guild/tools/scores.py | 2 +- guild/tools/utils.py | 2 -- guild/transformers/chembl.py | 1 - guild/transformers/msa.py | 1 - pyproject.toml | 2 +- uv.lock | 2 +- 7 files changed, 3 insertions(+), 8 deletions(-) diff --git a/guild/bulk.py b/guild/bulk.py index 4e0d31c..8b705ae 100644 --- a/guild/bulk.py +++ b/guild/bulk.py @@ -66,7 +66,6 @@ # Scores lists RP_SCORES_COLUMNS, # Dictionaries - SCORES_DICTIONARY, SMILES, # Folders VINA_FOLDER, diff --git a/guild/tools/scores.py b/guild/tools/scores.py index a483707..47166db 100644 --- a/guild/tools/scores.py +++ b/guild/tools/scores.py @@ -15,8 +15,8 @@ from guild.constants.bulk import ( GLOBAL_RP_SCORE, - RP_SCORES_DICTIONARY, RANKS_DICTIONARY, + RP_SCORES_DICTIONARY, SCORES_DIRECTION_DICTIONARY, ) from guild.constants.guild import PROTEIN_CONF_ID diff --git a/guild/tools/utils.py b/guild/tools/utils.py index e0de87f..7afd509 100644 --- a/guild/tools/utils.py +++ b/guild/tools/utils.py @@ -7,8 +7,6 @@ import time from functools import wraps -import pandas as pd - def read_json_as_dict(input_file): """ diff --git a/guild/transformers/chembl.py b/guild/transformers/chembl.py index ddfa089..ca76365 100644 --- a/guild/transformers/chembl.py +++ b/guild/transformers/chembl.py @@ -8,7 +8,6 @@ import chembl_downloader import pandas as pd -from chembl_webresource_client.http_errors import HttpBadRequest from tqdm import tqdm from unipressed import IdMappingClient diff --git a/guild/transformers/msa.py b/guild/transformers/msa.py index 75f81a4..288b3cb 100644 --- a/guild/transformers/msa.py +++ b/guild/transformers/msa.py @@ -16,7 +16,6 @@ import subprocess from pathlib import Path - LOCALCOLABFOLD_DIR_ENV_VAR = "GUILD_LOCALCOLABFOLD_DIR" COLABFOLD_RUN_SCRIPT = "run_colabfoldbatch_sample.sh" diff --git a/pyproject.toml b/pyproject.toml index 4ca53b8..cd63151 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,7 +14,7 @@ authors = [ { name = "Christoph Krettler" }, { name = "Shweta Bhandare" }, { name = "Keaton Noon" }, - { name = "Eric Fish" } + { name = "Eric Fish" }, { name = "António de La Vega de León" }, { name = "Dani Domingo-Fernandéz" }, ] diff --git a/uv.lock b/uv.lock index f54ee93..deaf934 100644 --- a/uv.lock +++ b/uv.lock @@ -963,7 +963,7 @@ wheels = [ [[package]] name = "guild" -version = "1.1.4" +version = "1.1.5" source = { editable = "." } dependencies = [ { name = "boltz", extra = ["cuda"] },