From a1f98464037a4bfc1f1c505a686ce40685fc0e68 Mon Sep 17 00:00:00 2001 From: akshat998 Date: Mon, 6 Jul 2026 01:12:51 -0700 Subject: [PATCH] Add optional ML classifier for molecule prioritization within collections Adds an FNN classifier (Morgan fingerprints -> 512 -> 256 -> 128 -> 1, trained with Adam/BCE and early stopping) that plugs into the existing tranche/collection-based screening workflow: train it once on a prescreen run's docking-score summaries (tools/train_ml_classifier.py), then each queue (Slurm) or AWS Batch subjob filters its own collection(s) through it before docking, keeping only molecules predicted worth the cost of a full docking run. Implemented consistently across both execution paths: - Slurm/HPC: one-queue.sh invokes a small CLI bridge (ml_classifier_predict.py) once per collection, since the rest of the pipeline is Bash and cannot run PyTorch directly. - AWS Batch: vf_aws_run.py calls the shared ml_classifier module in-process, batched once per subjob. Activated via all.ctrl (use_ml_classifier); disabled by default with no change to existing behavior. Also unifies the two execution paths' docking-score summary file schema (both now include a SMILES column, and both use the best/most-negative score for the "maximum-score" column, matching what the column name implies), and updates vf_report.sh's column extraction to match the corrected schema. --- README.md | 14 + docs/vfvs/features.md | 4 + docs/vfvs/installation/prerequisites.md | 4 + .../vfvs/using-vfvs/preparing-the-workflow.md | 20 + input-files/ml_classifier/README.md | 4 + tools/docker/Dockerfile | 6 + tools/templates/all.ctrl | 19 + tools/templates/ml_classifier.py | 555 ++++++++++++++++++ tools/templates/ml_classifier_predict.py | 139 +++++ tools/templates/one-queue.sh | 94 +++ tools/templates/vf_aws_run.py | 108 +++- tools/train_ml_classifier.py | 184 ++++++ tools/vf_report.sh | 8 +- 13 files changed, 1153 insertions(+), 6 deletions(-) create mode 100644 input-files/ml_classifier/README.md create mode 100644 tools/templates/ml_classifier.py create mode 100644 tools/templates/ml_classifier_predict.py create mode 100644 tools/train_ml_classifier.py diff --git a/README.md b/README.md index abfe707..fe577f6 100644 --- a/README.md +++ b/README.md @@ -14,6 +14,20 @@ Currently, there exist two versions of VirtualFlow, which are tailored to differ They use the same core technology regarding the workflow management and parallelization, and they can be used individually or in concert with each other. Additional versions are expected to arrive in the future. Pre-built ready-to-dock ligand libraries for VFVS are available for free (in the download section). +### Optional: Machine Learning Classifier for Molecule Prioritization + +For ultra-large screens, VFVS supports an optional machine learning classifier that predicts which molecules within a collection are likely to be high-affinity binders, so only those molecules are docked. It is a fully-connected feedforward neural network (Morgan fingerprints in, binding probability out) trained on the docking scores from a small representative "prescreen" run. It is disabled by default and works identically on both the Slurm/HPC and AWS Batch execution paths. + +1. **Prescreen.** Run the existing, unmodified VFVS workflow on a small representative set of collections, to get real docking scores in the usual per-collection summary files. +2. **Train.** From the `tools` directory, run: + ``` + python3 train_ml_classifier.py --scenario + ``` + This trains the classifier on the prescreen's summary files and saves it to the path given by `ml_classifier_model_path` in `all.ctrl` (default `ml_classifier/model.pt`, relative to `input-files/`). +3. **Primary screen.** Set `use_ml_classifier=true` in `all.ctrl` (see the "Machine Learning Classifier" section of `tools/templates/all.ctrl` for all options) and submit the primary screen as usual. + +Requires PyTorch and RDKit (CPU-only PyTorch is sufficient); only needed if `use_ml_classifier=true`. See [Preparing the Workflow](docs/vfvs/using-vfvs/preparing-the-workflow.md#machine-learning-classifier-for-molecule-prioritization) for full details. + ### Overview of Resources diff --git a/docs/vfvs/features.md b/docs/vfvs/features.md index 8ead416..7041976 100644 --- a/docs/vfvs/features.md +++ b/docs/vfvs/features.md @@ -23,3 +23,7 @@ Due to the many options regarding the dockings scenarios, VFVS is suitable for c ## Multiple Docking Scenarios VFVS allows to carry out multiple docking scenarios per ligand. A docking scenario in VFVS is defined by the receptor structure, by the docking parameters (such as exhaustiveness), rigid or flexible receptor docking, the choice of flexible receptor side chains or the docking program. This allows also for ensemble dockings. + +## Machine Learning Classifier for Molecule Prioritization + +VFVS supports an optional machine learning classifier that prioritizes which molecules within a collection are worth docking, based on a small representative prescreen run. Molecules predicted unlikely to be high-affinity binders are skipped before docking, reducing computational cost for large screens while maintaining hit enrichment. This is disabled by default and available on both the Slurm/HPC and AWS Batch execution paths; see [Preparing the Workflow](using-vfvs/preparing-the-workflow.md#machine-learning-classifier-for-molecule-prioritization). diff --git a/docs/vfvs/installation/prerequisites.md b/docs/vfvs/installation/prerequisites.md index 351b7a7..db1d761 100644 --- a/docs/vfvs/installation/prerequisites.md +++ b/docs/vfvs/installation/prerequisites.md @@ -34,3 +34,7 @@ SPORES can be obtained here: The SPORES binary has to be placed in the folder tools/bin/ and has to be named `spores`. +## Machine Learning Classifier (Optional) + +If the optional machine learning classifier for molecule prioritization is enabled (`use_ml_classifier=true` in `all.ctrl`, see [Preparing the Workflow](../using-vfvs/preparing-the-workflow.md#machine-learning-classifier-for-molecule-prioritization)), PyTorch and RDKit need to be installed and importable by the `python3` used to run `tools/train_ml_classifier.py` and, on the Slurm/HPC path, `tools/templates/ml_classifier_predict.py`. A CPU-only build of PyTorch is sufficient. This is not required when `use_ml_classifier` is left at its default of `false`. For AWS Batch runs, both packages are already installed in the provided Docker image. + diff --git a/docs/vfvs/using-vfvs/preparing-the-workflow.md b/docs/vfvs/using-vfvs/preparing-the-workflow.md index 9ee58fd..33f1e62 100644 --- a/docs/vfvs/using-vfvs/preparing-the-workflow.md +++ b/docs/vfvs/using-vfvs/preparing-the-workflow.md @@ -65,6 +65,26 @@ The required ligand input database can be prepared either from scratch with VFLP There, specific subsets of the database can be downloaded, as well as the required collection lengths file which is needed to set up the `tools` folder as described earlier [here](../../using-virtualflow/preparing-the-workflow.md#central-task-list) (i.e. this file can be used for the central task `todo.all` list which is used by VirtualFlow). +## Machine Learning Classifier for Molecule Prioritization + +For ultra-large screens, VFVS supports an optional machine learning classifier that predicts which molecules within a collection are likely to be high-affinity binders, so that only those molecules are docked. It is a fully-connected feedforward neural network trained on Morgan fingerprints and the docking scores from a small representative prescreen run. It is disabled by default; enabling it does not change how collections/tranches are selected for screening, only which molecules within an already-selected collection get docked. + +The workflow has three steps: + +1. **Prescreen.** Run the existing, unmodified VFVS workflow (`all.ctrl` + `vf_prepare_folders.sh` + `vf_start_jobline.sh`, or the AWS Batch equivalents) on a small representative set of collections, to produce real docking scores in the usual per-collection summary files under `output-files/{complete,incomplete}//summaries/`. No configuration changes are needed for this step. +2. **Train.** From the `tools` directory, run the training script once (a single interactive invocation, not a batch job): + + ``` + python3 train_ml_classifier.py --scenario + ``` + + This collects every summary file for the given docking scenario, extracts (SMILES, docking score) pairs, and trains+saves the classifier to the path given by `ml_classifier_model_path` in `workflow/control/all.ctrl` (default `ml_classifier/model.pt`, relative to `input-files/`). +3. **Primary screen.** Set `use_ml_classifier=true` in `all.ctrl` (see the "Machine Learning Classifier" section of the control file for all available options), then submit the primary screen as usual. Each queue (Slurm) or AWS Batch array child loads the trained classifier once per collection/subjob and filters out molecules scoring at or below `ml_classifier_probability_cutoff`, before they reach the docking programs. + +A classifier trained once for a given docking scenario can be reused across later primary screens without retraining, as long as `ml_classifier_model_path` still points at the saved file. Filtered-out molecules are recorded in the same per-collection ligand-list status files used for other pre-docking checks (such as the element and duplicate-coordinate checks), so every input molecule remains accounted for. + +Requires PyTorch and RDKit to be installed (see [Prerequisites](../installation/prerequisites.md)); these are only needed if `use_ml_classifier=true`. + diff --git a/input-files/ml_classifier/README.md b/input-files/ml_classifier/README.md new file mode 100644 index 0000000..dec5d8c --- /dev/null +++ b/input-files/ml_classifier/README.md @@ -0,0 +1,4 @@ +This folder holds the trained ML tranche-prioritization classifier (`model.pt` by default). +It is populated by running `tools/train_ml_classifier.py` once against a prescreen run's summary files. +Its contents are bundled into `vf_input.tar.gz` automatically for AWS Batch runs, the same way as the rest of `input-files/`. +See the "Machine Learning Classifier" section of `all.ctrl` and the main VFVS README for usage. diff --git a/tools/docker/Dockerfile b/tools/docker/Dockerfile index 8f1f6c9..300cf4a 100644 --- a/tools/docker/Dockerfile +++ b/tools/docker/Dockerfile @@ -20,6 +20,12 @@ FROM amazonlinux:2 RUN yum update -y && yum -y install python3 RUN pip3 install boto3 +# Dependencies for the optional ML tranche-prioritization classifier (tools/templates/ml_classifier.py). +# CPU-only torch wheel: this Dockerfile provisions no GPU AWS Batch compute environments, and the +# default PyPI torch wheel would pull in unneeded CUDA runtime bloat. +RUN pip3 install torch --index-url https://download.pytorch.org/whl/cpu +RUN pip3 install rdkit + ADD . /opt/vf/tools ENV USER ec2-user diff --git a/tools/templates/all.ctrl b/tools/templates/all.ctrl index 5f99f52..44c002f 100644 --- a/tools/templates/all.ctrl +++ b/tools/templates/all.ctrl @@ -268,6 +268,25 @@ energy_max=10000 # Maximum allowed energy value. Recommended: 10000 # Possible values: Positive integer +***************************************************************************************************************************************************************** +************************************************************* Machine Learning Classifier ***************************************************************** +***************************************************************************************************************************************************************** + +use_ml_classifier=false +# Optional per-collection molecule-prioritization filter, applied once per collection before docking (not per docking scenario). When enabled, every ligand in a collection is scored by a pretrained FNN classifier (Morgan fingerprints -> binding-probability) and ligands scoring at or below ml_classifier_probability_cutoff are skipped, exactly like the existing ligand_elements/ligand_coordinates pre-docking checks. Train a model first with tools/train_ml_classifier.py. +# Possible values: +# * true +# * false +# Settable via range control files: Yes + +ml_classifier_model_path=ml_classifier/model.pt +# Path to the trained classifier (.pt file), relative to the input-files/ folder (NOT to tools/, unlike collection_folder). Only used if use_ml_classifier=true. +# Settable via range control files: Yes + +ml_classifier_probability_cutoff=0.5 +# Minimum predicted binding probability (exclusive) required to keep a ligand for docking. Per the manuscript's ML classifier spec, 0.5 is the recommended default. Only used if use_ml_classifier=true. +# Settable via range control files: Yes + ***************************************************************************************************************************************************************** ******************************************************************* Terminating Variables ***************************************************************** ***************************************************************************************************************************************************************** diff --git a/tools/templates/ml_classifier.py b/tools/templates/ml_classifier.py new file mode 100644 index 0000000..f2ca843 --- /dev/null +++ b/tools/templates/ml_classifier.py @@ -0,0 +1,555 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Optional machine learning classifier for prioritizing candidate compounds before docking. + +Implements the "Machine Learning Classifier for Tranche Prioritization" described in the +AdaptiveFlow manuscript (Methods section): a fully-connected feedforward neural network (FNN) +trained on Morgan fingerprints and prescreen docking scores, used to filter a large pool of +undocked candidate compounds down to those predicted to be high-confidence binders before they +are passed into the docking pipeline. + +Shared between VFVS's two execution paths: imported directly by vf_aws_run.py (AWS Batch), and +invoked via the small CLI wrapper ml_classifier_predict.py from one-queue.sh (Slurm/HPC, which +cannot run PyTorch directly since it is a Bash script). Filtering operates on whatever ligands a +collection already contains -- this module has no knowledge of tranches, collections, or joblines. + +@author: akshat +""" +import os +import random +import uuid + +import numpy as np +import torch +import torch.nn as nn +from torch.utils.data import TensorDataset, DataLoader + +from rdkit import Chem +from rdkit.Chem import AllChem + + +def get_device(device=None): + """ + Resolves the torch device to use, preferring CUDA when available. + + Args: + device (str or torch.device or None): explicit device to use. If None, 'cuda' is used + when available, else 'cpu'. + + Returns: + torch.device + """ + if device is not None: + return torch.device(device) + return torch.device('cuda' if torch.cuda.is_available() else 'cpu') + + +def set_seed(seed): + """ + Seeds python's random module, numpy, and torch (including CUDA, if available) for + reproducibility. + + Args: + seed (int): random seed to use. + + Returns: + None + """ + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + +class FNNClassifier(nn.Module): + """ + Fully-connected feedforward neural network for binary binding-likelihood classification. + + Architecture (matches the manuscript exactly): an input layer of 1,024 nodes (matching the + Morgan fingerprint length), followed by hidden layers of 512, 256, and 128 nodes, each using + ReLU activation, followed by a single sigmoid-activated output node representing the + predicted binding probability. + """ + + def __init__(self, input_dim=1024): + super(FNNClassifier, self).__init__() + self.net = nn.Sequential( + nn.Linear(input_dim, 512), + nn.ReLU(), + nn.Linear(512, 256), + nn.ReLU(), + nn.Linear(256, 128), + nn.ReLU(), + nn.Linear(128, 1), + nn.Sigmoid(), + ) + + def forward(self, x): + return self.net(x) + + +def smi_to_fingerprint(smi, radius=2, n_bits=1024): + """ + Converts a SMILES string into a Morgan fingerprint (radius 2, 1024 bits by default, as + specified in the manuscript). + + Args: + smi (str): SMILES string of the compound. + radius (int): Morgan fingerprint radius. + n_bits (int): Morgan fingerprint bit-vector length. + + Returns: + np.ndarray of shape (n_bits,), dtype float32, or None if the SMILES could not be parsed + (including an empty string, which VFVS uses to represent a ligand whose SMILES could not + be extracted from its structure file). + """ + if not smi: + return None + mol = Chem.MolFromSmiles(smi) + if mol is None: + print('WARNING: Invalid SMILES provided, skipping: {}'.format(smi)) + return None + fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius, nBits=n_bits) + arr = np.zeros((n_bits,), dtype=np.float32) + for bit in fp.GetOnBits(): + arr[bit] = 1.0 + return arr + + +def _fingerprint_all(smiles_list, radius=2, n_bits=1024): + """ + Fingerprints a list of SMILES, dropping any that fail to parse. + + Returns: + (kept_indices, fingerprint_matrix) where kept_indices maps rows of the matrix back to + indices in the original smiles_list. + """ + kept_indices = [] + fps = [] + for i, smi in enumerate(smiles_list): + fp = smi_to_fingerprint(smi, radius=radius, n_bits=n_bits) + if fp is not None: + kept_indices.append(i) + fps.append(fp) + if len(fps) == 0: + return [], np.zeros((0, n_bits), dtype=np.float32) + return kept_indices, np.stack(fps, axis=0) + + +def label_by_top_percent(scores, top_percent=25): + """ + Converts docking scores into binary labels via percentile thresholding. + + More negative docking scores indicate better (higher-affinity) binding. The top + `top_percent` % of compounds with the most negative scores are labeled 1 ("high-confidence + binder"); the rest are labeled 0. This is implemented directly as a percentile cutoff on the + raw (ascending) scores to avoid any ambiguity about which "direction" a percentile is taken + from. + + Args: + scores (np.ndarray): docking scores (more negative = better). + top_percent (float): percentage (0-100) of best (most negative) scores to label positive. + + Returns: + (labels_ndarray, threshold_value): labels is an int array (0/1) aligned with `scores`, + threshold_value is the score cutoff at/below which a compound is labeled positive. + """ + threshold = np.percentile(scores, top_percent) + labels = (scores <= threshold).astype(np.int64) + return labels, float(threshold) + + +def undersample_majority(X, y, seed=42): + """ + Randomly undersamples the majority class so that both classes are equally represented. + + Args: + X (np.ndarray): feature matrix, shape (n_samples, n_features). + y (np.ndarray): binary labels (0/1), shape (n_samples,). + seed (int): random seed for reproducible undersampling. + + Returns: + (X_balanced, y_balanced) + """ + rng = np.random.RandomState(seed) + pos_idx = np.where(y == 1)[0] + neg_idx = np.where(y == 0)[0] + + if len(pos_idx) == 0 or len(neg_idx) == 0: + return X, y + + if len(pos_idx) < len(neg_idx): + minority_idx, majority_idx = pos_idx, neg_idx + else: + minority_idx, majority_idx = neg_idx, pos_idx + + sampled_majority_idx = rng.choice(majority_idx, size=len(minority_idx), replace=False) + balanced_idx = np.concatenate([minority_idx, sampled_majority_idx]) + rng.shuffle(balanced_idx) + return X[balanced_idx], y[balanced_idx] + + +def stratified_val_split(X, y, val_fraction=0.10, seed=42): + """ + Splits a (class-balanced) dataset into train/validation partitions, preserving the class + balance of each partition. + + Args: + X (np.ndarray): feature matrix. + y (np.ndarray): binary labels. + val_fraction (float): fraction of samples to hold out for validation. + seed (int): random seed. + + Returns: + (X_train, y_train, X_val, y_val) + """ + rng = np.random.RandomState(seed) + train_idx_parts = [] + val_idx_parts = [] + for cls in np.unique(y): + cls_idx = np.where(y == cls)[0] + rng.shuffle(cls_idx) + n_val = max(1, int(round(len(cls_idx) * val_fraction))) if len(cls_idx) > 1 else 0 + val_idx_parts.append(cls_idx[:n_val]) + train_idx_parts.append(cls_idx[n_val:]) + + train_idx = np.concatenate(train_idx_parts) + val_idx = np.concatenate(val_idx_parts) + rng.shuffle(train_idx) + rng.shuffle(val_idx) + return X[train_idx], y[train_idx], X[val_idx], y[val_idx] + + +def train_classifier(smiles_list, scores, model_out_path, top_percent=25, val_fraction=0.10, + batch_size=256, max_epochs=50, patience=5, learning_rate=0.001, + seed=42, receptor=None, fp_radius=2, fp_nbits=1024, device=None, + verbose=True): + """ + Trains the FNN tranche-prioritization classifier on prescreen docking results. + + Args: + smiles_list (list of str): SMILES strings of the prescreen compounds. + scores (list or np.ndarray of float): docking scores aligned with smiles_list (more + negative = better binder), as produced by a prescreen run (see + train_ml_classifier.py, which parses these out of VFVS's summary files). + model_out_path (str): path to save the trained model (.pt file). + top_percent (float): percentage of most-negative-scoring prescreen compounds labeled as + positive ("high-confidence binder"). Default 25, per the manuscript. + val_fraction (float): validation split fraction. Default 0.10, per the manuscript. + batch_size (int): training batch size. Default 256, per the manuscript. + max_epochs (int): maximum number of training epochs. Default 50, per the manuscript. + patience (int): early-stopping patience (epochs without validation-loss improvement). + Default 5, per the manuscript. + learning_rate (float): Adam learning rate. Default 0.001, per the manuscript. + seed (int): random seed for undersampling, splitting, and model initialization. + receptor (str or None): optional receptor identifier/path stored as metadata, so a + saved model can later be associated with the target it was trained for. + fp_radius (int): Morgan fingerprint radius. + fp_nbits (int): Morgan fingerprint length. + device (str or None): torch device override. Auto-detected (CUDA if available) if None. + verbose (bool): whether to print per-epoch training/validation loss. + + Returns: + (model, history): the trained FNNClassifier (in eval mode, on `device`), and a dict with + keys 'train_loss', 'val_loss' (lists, one entry per epoch actually run), and + 'best_epoch'. + + Raises: + Exception: if fewer than 2 classes are present after labeling (cannot train a + classifier), or if no SMILES in the prescreen data could be parsed. + """ + set_seed(seed) + device = get_device(device) + + scores = np.asarray(scores, dtype=np.float32) + kept_indices, X_all = _fingerprint_all(smiles_list, radius=fp_radius, n_bits=fp_nbits) + if X_all.shape[0] == 0: + raise Exception('No valid SMILES found in prescreen data.') + scores_kept = scores[kept_indices] + + labels, threshold = label_by_top_percent(scores_kept, top_percent=top_percent) + if len(np.unique(labels)) < 2: + raise Exception('Only one class present after labeling (top_percent={}); ' + 'cannot train a binary classifier. Provide more diverse prescreen ' + 'docking scores.'.format(top_percent)) + + X_bal, y_bal = undersample_majority(X_all, labels, seed=seed) + X_train, y_train, X_val, y_val = stratified_val_split(X_bal, y_bal, val_fraction=val_fraction, + seed=seed) + + train_ds = TensorDataset(torch.from_numpy(X_train), torch.from_numpy(y_train.astype(np.float32))) + train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True) + + X_val_t = torch.from_numpy(X_val).to(device) + y_val_t = torch.from_numpy(y_val.astype(np.float32)).to(device) + + model = FNNClassifier(input_dim=fp_nbits).to(device) + optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) + loss_fn = nn.BCELoss() + + history = {'train_loss': [], 'val_loss': []} + best_val_loss = float('inf') + best_state_dict = None + best_epoch = 0 + epochs_without_improvement = 0 + + for epoch in range(max_epochs): + model.train() + epoch_train_losses = [] + for xb, yb in train_loader: + xb, yb = xb.to(device), yb.to(device) + optimizer.zero_grad() + preds = model(xb).squeeze(1) + loss = loss_fn(preds, yb) + loss.backward() + optimizer.step() + epoch_train_losses.append(loss.item()) + train_loss = float(np.mean(epoch_train_losses)) + + model.eval() + with torch.no_grad(): + val_preds = model(X_val_t).squeeze(1) + val_loss = float(loss_fn(val_preds, y_val_t).item()) + + history['train_loss'].append(train_loss) + history['val_loss'].append(val_loss) + if verbose: + print('Epoch {}/{} - train_loss: {:.4f} - val_loss: {:.4f}'.format( + epoch + 1, max_epochs, train_loss, val_loss)) + + if val_loss < best_val_loss: + best_val_loss = val_loss + best_state_dict = {k: v.clone() for k, v in model.state_dict().items()} + best_epoch = epoch + epochs_without_improvement = 0 + else: + epochs_without_improvement += 1 + if epochs_without_improvement >= patience: + if verbose: + print('Early stopping at epoch {} (patience={})'.format(epoch + 1, patience)) + break + + model.load_state_dict(best_state_dict) + model.eval() + + metadata = { + 'input_dim': fp_nbits, + 'fp_radius': fp_radius, + 'fp_nbits': fp_nbits, + 'top_percent': top_percent, + 'train_threshold': threshold, + 'seed': seed, + 'receptor': receptor, + 'best_epoch': best_epoch, + 'best_val_loss': best_val_loss, + } + save_classifier(model, model_out_path, metadata) + history['best_epoch'] = best_epoch + + return model, history + + +def save_classifier(model, out_path, metadata): + """ + Saves a trained FNNClassifier's weights and associated metadata to disk. + + Args: + model (FNNClassifier): the trained model. + out_path (str): destination path (.pt file). Parent directories are created if needed. + metadata (dict): metadata to store alongside the weights (fingerprint params, threshold, + seed, receptor, etc.), retrievable via load_classifier(). + + Returns: + None + + Note: + Written atomically (temp file + os.replace) so that concurrent readers -- e.g. many + independent Slurm queue processes or AWS Batch array children calling load_classifier() + on this same path -- never observe a partially-written checkpoint, even if this function + is called again (retraining) while a previous screen using the model is still running. + """ + out_dir = os.path.dirname(out_path) + if out_dir != '' and not os.path.exists(out_dir): + os.makedirs(out_dir, exist_ok=True) + checkpoint = {'state_dict': model.state_dict()} + checkpoint.update(metadata) + + tmp_path = '{}.tmp-{}'.format(out_path, uuid.uuid4().hex) + torch.save(checkpoint, tmp_path) + os.replace(tmp_path, out_path) + + +def load_classifier(model_path, device=None): + """ + Loads a previously trained FNNClassifier from disk. + + Args: + model_path (str): path to a .pt file created by save_classifier()/train_classifier(). + device (str or None): torch device override. Auto-detected (CUDA if available) if None. + + Returns: + (model, metadata): the FNNClassifier in eval mode on `device`, and the metadata dict + saved alongside the weights. + + Raises: + Exception: if model_path does not exist or is not a file. + """ + if not os.path.isfile(model_path): + raise Exception('Model path {} not found.'.format(model_path)) + device = get_device(device) + checkpoint = torch.load(model_path, map_location=device) + input_dim = checkpoint.get('input_dim', 1024) + model = FNNClassifier(input_dim=input_dim).to(device) + model.load_state_dict(checkpoint['state_dict']) + model.eval() + metadata = {k: v for k, v in checkpoint.items() if k != 'state_dict'} + return model, metadata + + +def predict_proba(model, smiles_list, device=None, fp_radius=2, fp_nbits=1024): + """ + Predicts binding probability for a list of SMILES using a trained classifier. + + Args: + model (FNNClassifier): a trained (eval-mode) model, as returned by load_classifier() or + train_classifier(). + smiles_list (list of str): SMILES strings to score. + device (str or None): torch device override. Auto-detected (CUDA if available) if None. + fp_radius (int): Morgan fingerprint radius (must match training). + fp_nbits (int): Morgan fingerprint length (must match training). + + Returns: + list of (smi, probability_or_None), aligned with `smiles_list` in order. Invalid or empty + SMILES yield a probability of None. + """ + device = get_device(device) + model = model.to(device) + model.eval() + + kept_indices, X = _fingerprint_all(smiles_list, radius=fp_radius, n_bits=fp_nbits) + results = [None] * len(smiles_list) + + if X.shape[0] > 0: + with torch.no_grad(): + X_t = torch.from_numpy(X).to(device) + probs = model(X_t).squeeze(1).cpu().numpy() + for local_i, orig_i in enumerate(kept_indices): + results[orig_i] = float(probs[local_i]) + + return [(smi, results[i]) for i, smi in enumerate(smiles_list)] + + +def filter_smiles_mask(model, smiles_list, probability_cutoff=0.5, device=None, + fp_radius=2, fp_nbits=1024): + """ + Scores a list of SMILES and returns a keep/discard mask, for filtering an in-memory batch of + candidate ligands before they are submitted for docking. + + Ligands whose SMILES could not be extracted/parsed (probability is None) are always treated + as discards (fail-closed) -- if the ML classifier is enabled, an unscoreable ligand is an + unexpected state and it is safer to not spend docking budget on it than to silently bypass + the filter the user opted into. + + Args: + model (FNNClassifier): a trained (eval-mode) model. + smiles_list (list of str): SMILES strings to score. + probability_cutoff (float): minimum predicted probability (exclusive) to retain a + compound for docking. Default 0.5, per the manuscript. + device (str or None): torch device override. Auto-detected (CUDA if available) if None. + fp_radius (int): Morgan fingerprint radius (must match training). + fp_nbits (int): Morgan fingerprint length (must match training). + + Returns: + (keep_mask, probabilities): keep_mask is a list of bool aligned with smiles_list (True = + retain for docking); probabilities is the aligned list of predicted probabilities + (None for compounds with invalid/empty SMILES, which are always discarded). + """ + scored = predict_proba(model, smiles_list, device=device, fp_radius=fp_radius, fp_nbits=fp_nbits) + keep_mask = [prob is not None and prob > probability_cutoff for _, prob in scored] + probabilities = [prob for _, prob in scored] + return keep_mask, probabilities + + +if __name__ == '__main__': + # Synthetic self-test: exercises train -> save -> load -> filter end-to-end without + # requiring any real docking runs. Fabricated docking scores are correlated with a simple + # RDKit descriptor (heavy-atom count) plus noise, so there is real learnable signal. + from rdkit.Chem import Descriptors + + SCRATCH_DIR = './ml_classifier_selftest' + os.makedirs(SCRATCH_DIR, exist_ok=True) + + demo_smiles = [ + 'CCO', 'CCCO', 'CCCCO', 'CCCCCO', 'CCN', 'CCCN', 'CCCCN', 'c1ccccc1', 'c1ccccc1C', + 'c1ccccc1CC', 'c1ccccc1O', 'c1ccccc1N', 'CC(=O)O', 'CC(=O)N', 'CC(=O)OC', 'CC(=O)OCC', + 'CCOC(=O)C', 'c1ccncc1', 'c1ccoc1', 'c1ccsc1', 'C1CCCCC1', 'C1CCCCC1O', 'C1CCCCC1N', + 'CC(C)C', 'CC(C)CC', 'CC(C)(C)C', 'C1=CC(=CC=C1CSCC2C(C(C(O2)N3C=NC4=C(N=CN=C43)N)O)O)Cl', + 'CCCCCCCCCC', 'c1ccc2ccccc2c1', 'c1ccc2[nH]ccc2c1', 'CC(N)C(=O)O', 'NC(Cc1ccccc1)C(=O)O', + 'CC1CCC(C)CC1', 'OC1CCCCC1', 'NC1CCCCC1', 'FC1=CC=CC=C1', 'ClC1=CC=CC=C1', + 'BrC1=CC=CC=C1', 'CCOCC', 'CCSCC', 'CCC(=O)CC', 'CCC(=O)NCC', 'c1ccc(cc1)C(=O)O', + 'c1ccc(cc1)C(=O)N', 'CC(=O)Nc1ccccc1', 'COc1ccccc1', 'CSc1ccccc1', 'CNc1ccccc1', + 'CC1=CC=CC=C1C', 'CC1=CC=CC=C1CC', 'NCCc1ccccc1', 'OCCc1ccccc1', 'CC(C)Oc1ccccc1', + ] + + rng = np.random.RandomState(42) + + def fabricate_score(smi): + mol = Chem.MolFromSmiles(smi) + heavy_atoms = Descriptors.HeavyAtomCount(mol) + return -1.0 * heavy_atoms + rng.normal(0, 1.5) + + all_scores = [fabricate_score(s) for s in demo_smiles] + + prescreen_smiles = demo_smiles[:30] + prescreen_scores = all_scores[:30] + candidate_smiles = demo_smiles[30:] + ['not_a_smiles', 'also_invalid???', ''] + + model_path = os.path.join(SCRATCH_DIR, 'classifier.pt') + + print('--- Step 1: training classifier on synthetic prescreen data ---') + model, history = train_classifier(prescreen_smiles, prescreen_scores, model_path, + receptor='receptor.pdbqt', seed=42, verbose=True) + assert os.path.exists(model_path), 'Model file was not written.' + assert history['best_epoch'] <= 49, 'Training ran past max_epochs=50.' + + labels_check, _ = label_by_top_percent(np.array(prescreen_scores, dtype=np.float32), top_percent=25) + frac_positive = labels_check.mean() + print('Fraction of prescreen labeled positive: {:.2f} (expected ~0.25)'.format(frac_positive)) + assert 0.10 < frac_positive < 0.40, 'Positive label fraction far from expected ~25%.' + + print('--- Step 2: loading saved classifier and checking architecture ---') + loaded_model, metadata = load_classifier(model_path) + shapes = [tuple(p.shape) for p in loaded_model.net.parameters() if p.dim() == 2] + expected_shapes = [(512, 1024), (256, 512), (128, 256), (1, 128)] + assert shapes == expected_shapes, 'Unexpected layer shapes: {}'.format(shapes) + print('Layer shapes match spec: {}'.format(shapes)) + + print('--- Step 3: filtering an in-memory candidate batch ---') + keep_mask, probabilities = filter_smiles_mask(loaded_model, candidate_smiles, probability_cutoff=0.5) + n_kept = sum(keep_mask) + print('Filter stats: kept {}/{} candidates.'.format(n_kept, len(candidate_smiles))) + assert n_kept < len(candidate_smiles), 'Expected some candidates to be filtered out.' + for smi, keep, prob in zip(candidate_smiles, keep_mask, probabilities): + if keep: + assert prob is not None and prob > 0.5, '{} kept but prob={}'.format(smi, prob) + else: + assert prob is None or prob <= 0.5, '{} dropped but prob={}'.format(smi, prob) + assert keep_mask[candidate_smiles.index('')] is False, 'Empty SMILES must always be discarded (fail-closed).' + print('Threshold correctness verified for all candidates (kept iff prob > 0.5); empty SMILES fail-closed.') + + print('--- Step 4: determinism check (retrain with identical seed) ---') + model_path_2 = os.path.join(SCRATCH_DIR, 'classifier_run2.pt') + model_2, _ = train_classifier(prescreen_smiles, prescreen_scores, model_path_2, + receptor='receptor.pdbqt', seed=42, verbose=False) + for (n1, p1), (n2, p2) in zip(model.named_parameters(), model_2.named_parameters()): + assert torch.equal(p1, p2), 'Parameter {} differs between identical-seed runs.'.format(n1) + print('Determinism verified: identical seeds produce bit-identical weights.') + + print('--- Step 5: invalid SMILES handling ---') + invalid_present = any(prob is None for _, prob in predict_proba(loaded_model, ['not_a_smiles', 'CCO'])) + assert invalid_present, 'Invalid SMILES should yield probability None.' + print('Invalid SMILES correctly skipped with a warning rather than raising.') + + print('\nAll self-tests passed.') diff --git a/tools/templates/ml_classifier_predict.py b/tools/templates/ml_classifier_predict.py new file mode 100644 index 0000000..ad67070 --- /dev/null +++ b/tools/templates/ml_classifier_predict.py @@ -0,0 +1,139 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Slurm-side inference bridge for the optional ML tranche-prioritization classifier. + +one-queue.sh (a Bash script) cannot run PyTorch directly, so it shells out to this small CLI once +per collection (not once per ligand -- loading a model and importing torch/rdkit per ligand would +be a severe performance antipattern across a collection's hundreds-to-thousands of ligands), +passing every ligand's name and SMILES in a single input file and reading back one decision per +ligand. + +Usage: + python3 templates/ml_classifier_predict.py \\ + --model-path \\ + --probability-cutoff 0.5 \\ + --input-tsv /_ml_classifier_input.tsv \\ + --output-tsv /_ml_classifier_decisions.tsv + +Input TSV format: one row per ligand, `ligand_namesmiles` (smiles may be empty, meaning the +SMILES could not be extracted from that ligand's structure file). + +Output TSV format: exactly one row per input row (rows are never dropped), in input order: +`ligand_nameKEEP|DISCARDprobability_or_NA`. Ligands with an empty/unparsable SMILES are +always written as DISCARD/NA (fail-closed). + +Exit codes: 0 only if the output file was fully written with one row per input row. Non-zero on +any hard failure (model missing/corrupt, input file missing/unreadable), with an actionable +message on stderr. + +@author: akshat +""" +import os +import sys +import csv +import argparse + +from ml_classifier import load_classifier, filter_smiles_mask + + +def read_input_tsv(input_tsv_path): + """ + Reads the per-collection ligand-name/SMILES input file. + + Malformed rows (wrong field count) are skipped with a stderr warning rather than aborting the + whole run -- this file is machine-generated by one-queue.sh in the same job, so malformed rows + should be rare/impossible, but defensive parsing is cheap insurance. + + Returns: + list of (ligand_name, smiles) tuples, in file order. + """ + rows = [] + with open(input_tsv_path, 'r', newline='') as f: + reader = csv.reader(f, delimiter='\t') + for line_number, fields in enumerate(reader, start=1): + if len(fields) != 2: + print('WARNING: Skipping malformed input row {}:{}: {}'.format( + input_tsv_path, line_number, fields), file=sys.stderr) + continue + rows.append((fields[0], fields[1])) + return rows + + +def write_output_tsv_atomic(output_tsv_path, decisions): + """ + Writes the decisions file atomically (temp file + os.replace), so a Bash-side reader can never + observe a partially-written output file. + + Args: + output_tsv_path (str): destination path. + decisions (list of (ligand_name, decision, probability_str)): rows to write, in order. + """ + out_dir = os.path.dirname(output_tsv_path) + if out_dir != '' and not os.path.exists(out_dir): + os.makedirs(out_dir, exist_ok=True) + + tmp_path = '{}.tmp-{}'.format(output_tsv_path, os.getpid()) + with open(tmp_path, 'w', newline='') as f: + # lineterminator explicitly set to '\n': csv.writer defaults to '\r\n', which would leave + # a stray '\r' at the end of each row's last field -- harmless in Python, but a Bash-side + # reader doing `grep|awk` on this file would pick up that '\r' as part of the probability + # field. + writer = csv.writer(f, delimiter='\t', lineterminator='\n') + for row in decisions: + writer.writerow(row) + os.replace(tmp_path, output_tsv_path) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--model-path', required=True, help='Path to the trained classifier (.pt file).') + parser.add_argument('--probability-cutoff', type=float, default=0.5, + help='Minimum predicted probability (exclusive) to KEEP a ligand.') + parser.add_argument('--input-tsv', required=True, help='Path to the ligand_namesmiles input file.') + parser.add_argument('--output-tsv', required=True, help='Path to write ligand_namedecisionprobability to.') + args = parser.parse_args() + + try: + model, metadata = load_classifier(args.model_path) + except Exception as e: + print('ERROR: Could not load ML classifier from {} ({}). ' + 'Train one first with: python3 train_ml_classifier.py'.format(args.model_path, e), file=sys.stderr) + return 1 + + if not os.path.isfile(args.input_tsv): + print('ERROR: Input TSV not found: {}'.format(args.input_tsv), file=sys.stderr) + return 1 + + try: + rows = read_input_tsv(args.input_tsv) + except Exception as e: + print('ERROR: Could not read input TSV {} ({})'.format(args.input_tsv, e), file=sys.stderr) + return 1 + + ligand_names = [name for name, _ in rows] + smiles_list = [smi for _, smi in rows] + + fp_radius = metadata.get('fp_radius', 2) + fp_nbits = metadata.get('fp_nbits', 1024) + keep_mask, probabilities = filter_smiles_mask( + model, smiles_list, probability_cutoff=args.probability_cutoff, + fp_radius=fp_radius, fp_nbits=fp_nbits) + + decisions = [] + for ligand_name, keep, probability in zip(ligand_names, keep_mask, probabilities): + decision = 'KEEP' if keep else 'DISCARD' + probability_str = '{:.6f}'.format(probability) if probability is not None else 'NA' + decisions.append((ligand_name, decision, probability_str)) + + try: + write_output_tsv_atomic(args.output_tsv, decisions) + except Exception as e: + print('ERROR: Could not write output TSV {} ({})'.format(args.output_tsv, e), file=sys.stderr) + return 1 + + return 0 + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/tools/templates/one-queue.sh b/tools/templates/one-queue.sh index 2162e2b..7094e58 100644 --- a/tools/templates/one-queue.sh +++ b/tools/templates/one-queue.sh @@ -151,6 +151,24 @@ error_response_ligand_coordinates() { echo "Ligand ${next_ligand} ${ligand_list_entry} on $(date)." } +error_response_ml_classifier_filtered() { + + # Variables + probability=$1 + ligand_list_entry="filtered(ml_classifier:${probability})" + + # Printing some information + echo | tee -a /dev/stderr + echo "The ligand was filtered out by the ML tranche-prioritization classifier (probability=${probability})." | tee -a /dev/stderr + echo "Skipping this ligand and continuing with next one." | tee -a /dev/stderr + + # Updating the ligand list file + echo "${next_ligand} ${ligand_list_entry}" >> ${VF_TMPDIR}/${USER}/VFVS/${VF_JOBLETTER}/${VF_QUEUE_NO_12}/${VF_QUEUE_NO}/workflow/ligand-collections/ligand-lists/${next_ligand_collection_metatranch}/${next_ligand_collection_tranch}/${next_ligand_collection_ID}.status + + # Printing some information + echo "Ligand ${next_ligand} ${ligand_list_entry} on $(date)." +} + obabel_check_energy() { # Checking format @@ -428,6 +446,50 @@ prepare_collection_files_tmp() { # Extracting all the ligands at the same time (faster than individual for each ligand separately) tar -xf ${VF_TMPDIR}/${USER}/VFVS/${VF_JOBLETTER}/${VF_QUEUE_NO_12}/${VF_QUEUE_NO}/input-files/ligands/${next_ligand_collection_metatranch}/${next_ligand_collection_tranch}/${next_ligand_collection_ID}.tar -C ${VF_TMPDIR}/${USER}/VFVS/${VF_JOBLETTER}/${VF_QUEUE_NO_12}/${VF_QUEUE_NO}/input-files/ligands/${next_ligand_collection_metatranch}/${next_ligand_collection_tranch} + # Optional ML tranche-prioritization classifier: score every ligand of this collection once + # (not per ligand later, which would mean re-loading a PyTorch model per ligand -- a severe + # antipattern over a collection's many ligands), and record a keep/discard decision for each. + # Disabled by default; behavior above/below is unchanged unless use_ml_classifier=true. + use_ml_classifier="$(grep -m 1 "^use_ml_classifier=" ${VF_CONTROLFILE_TEMP} | tr -d '[[:space:]]' | awk -F '[=#]' '{print $2}')" + if [[ "${use_ml_classifier}" == "true" ]]; then + + ml_classifier_model_path="$(grep -m 1 "^ml_classifier_model_path=" ${VF_CONTROLFILE_TEMP} | tr -d '[[:space:]]' | awk -F '[=#]' '{print $2}')" + ml_classifier_probability_cutoff="$(grep -m 1 "^ml_classifier_probability_cutoff=" ${VF_CONTROLFILE_TEMP} | tr -d '[[:space:]]' | awk -F '[=#]' '{print $2}')" + # Falling back to the same defaults vf_aws_run.py's process_config() uses, in case a + # range/override controlfile omits these keys -- keeps the two execution paths consistent + # instead of the Slurm path hard-failing on a key the AWS path would silently default. + if [[ -z "${ml_classifier_model_path}" ]]; then + ml_classifier_model_path="ml_classifier/model.pt" + fi + if [[ -z "${ml_classifier_probability_cutoff}" ]]; then + ml_classifier_probability_cutoff="0.5" + fi + ml_classifier_ligand_dir="${VF_TMPDIR}/${USER}/VFVS/${VF_JOBLETTER}/${VF_QUEUE_NO_12}/${VF_QUEUE_NO}/input-files/ligands/${next_ligand_collection_metatranch}/${next_ligand_collection_tranch}/${next_ligand_collection_ID}" + + # Building the ligand-name/SMILES input file for the whole collection. Reuses the exact + # same "grep SMILES | last field" extraction rule already used later for the + # summary file (see the "Getting the ligand SMILES" step below), just applied once per + # ligand in this collection rather than inline in the main per-ligand loop. + > ${ml_classifier_ligand_dir}/_ml_classifier_input.tsv + for ml_classifier_ligand_file in $(ls ${ml_classifier_ligand_dir}/*.${ligand_library_format} 2>/dev/null); do + ml_classifier_ligand_name="$(basename "${ml_classifier_ligand_file}" .${ligand_library_format})" + ml_classifier_ligand_smiles="$(grep SMILES "${ml_classifier_ligand_file}" | awk '{print $NF}')" + printf "%s\t%s\n" "${ml_classifier_ligand_name}" "${ml_classifier_ligand_smiles}" >> ${ml_classifier_ligand_dir}/_ml_classifier_input.tsv + done + + # Running the ML classifier once for the whole collection + python3 templates/ml_classifier_predict.py --model-path ../input-files/${ml_classifier_model_path} --probability-cutoff ${ml_classifier_probability_cutoff} --input-tsv ${ml_classifier_ligand_dir}/_ml_classifier_input.tsv --output-tsv ${ml_classifier_ligand_dir}/_ml_classifier_decisions.tsv + ml_classifier_exit_code=$? + + ml_classifier_input_line_count=$(wc -l < ${ml_classifier_ligand_dir}/_ml_classifier_input.tsv) + ml_classifier_output_line_count=$(wc -l < ${ml_classifier_ligand_dir}/_ml_classifier_decisions.tsv 2>/dev/null || echo 0) + + if [ "${ml_classifier_exit_code}" -ne "0" ] || [ ! -f ${ml_classifier_ligand_dir}/_ml_classifier_decisions.tsv ] || [ "${ml_classifier_input_line_count}" -ne "${ml_classifier_output_line_count}" ]; then + echo " * Error: The ML tranche-prioritization classifier could not be applied to collection ${next_ligand_collection_tranch}_${next_ligand_collection_ID}." | tee -a /dev/stderr + error_response_std $LINENO + fi + fi + # Copying the required old output files if continuing old collection for docking_scenario_name in "${docking_scenario_names[@]}"; do if [ "${new_collection}" = "false" ]; then @@ -997,9 +1059,41 @@ while true; do # Getting the ligand SMILES if [[ ${ligand_library_format} == "pdb" || ${ligand_library_format} == "pdbqt" || ${ligand_library_format} == "mol2" ]]; then next_ligand_smiles="$(grep SMILES ${VF_TMPDIR}/${USER}/VFVS/${VF_JOBLETTER}/${VF_QUEUE_NO_12}/${VF_QUEUE_NO}/input-files/ligands/${next_ligand_collection_metatranch}/${next_ligand_collection_tranch}/${next_ligand_collection_ID}/${next_ligand}.${ligand_library_format} | awk '{print $NF}')" + # If no SMILES remark line was found, fall back to "NA" rather than an empty string, so + # the summary file always has the same number of whitespace-separated fields per row + # (an empty field here would silently shift every later column by one for that row, + # corrupting train_ml_classifier.py's whitespace-based parsing). + if [[ -z "${next_ligand_smiles}" ]]; then + next_ligand_smiles="NA" + fi else next_ligand_smiles="NA" fi + + # Checking the ML tranche-prioritization classifier's decision for this ligand, if enabled. + # This is a lookup against the per-collection decisions file computed once in + # prepare_collection_files_tmp() (not a fresh Python/model invocation per ligand). + if [[ "${use_ml_classifier}" == "true" ]]; then + # Note: the tab is passed as a real ANSI-C-quoted tab character ($'\t'), not the two-char + # sequence "\t" -- GNU grep's BRE does not reliably expand "\t" to a tab across all + # locales/versions, and -P (which would) is unavailable in some locales ("-P supports + # only unibyte and UTF-8 locales"). Concatenating a double-quoted string with an + # ANSI-C-quoted one is standard bash and produces a single argument to grep. + ml_classifier_decision_line=$(grep "^${next_ligand}"$'\t' ${ml_classifier_ligand_dir}/_ml_classifier_decisions.tsv 2>/dev/null | head -n 1 || true) + if [[ -z "${ml_classifier_decision_line}" ]]; then + # No decision found for this ligand (unexpected state) -- fail closed, same as an + # unscoreable SMILES. + error_response_ml_classifier_filtered "missing_decision" + continue + fi + ml_classifier_decision=$(awk -F'\t' '{print $2}' <<< "${ml_classifier_decision_line}") + if [[ "${ml_classifier_decision}" == "DISCARD" ]]; then + ml_classifier_probability=$(awk -F'\t' '{print $3}' <<< "${ml_classifier_decision_line}") + error_response_ml_classifier_filtered "${ml_classifier_probability}" + continue + fi + fi + # Loop for each docking type for docking_scenario_index in $(seq ${docking_scenario_index_start} ${docking_scenario_index_end}); do diff --git a/tools/templates/vf_aws_run.py b/tools/templates/vf_aws_run.py index 5220d50..f209ee1 100644 --- a/tools/templates/vf_aws_run.py +++ b/tools/templates/vf_aws_run.py @@ -44,6 +44,8 @@ import time from pathlib import Path +from ml_classifier import load_classifier, filter_smiles_mask + # Given a config file, parse out all of the configuration options @@ -98,6 +100,17 @@ def process_config(ctx): 'replicas': int(new_config['docking_scenario_replicas'][index]) } + # Optional ML tranche-prioritization classifier (disabled by default). Resolved to an + # absolute path the same way docking_scenarios[...]['config'] is above, since + # ml_classifier_model_path (like docking_scenario_inputfolders) is relative to input-files/. + new_config['use_ml_classifier'] = ctx['config.temp'].get('use_ml_classifier', 'false') + new_config['ml_classifier_model_path'] = os.path.join( + ctx['temp_dir'], "vf_input", "input-files", + ctx['config.temp'].get('ml_classifier_model_path', 'ml_classifier/model.pt') + ) + new_config['ml_classifier_probability_cutoff'] = float( + ctx['config.temp'].get('ml_classifier_probability_cutoff', 0.5)) + return new_config # Retrieve the config file (eventually can be non-S3) @@ -383,7 +396,7 @@ def create_summary_file(ctx, scenario, collection, scenario_result): with gzip.open(f"{collection['number']}.txt.gz", "wt") as summmary_fp: summmary_fp.write( - "Tranch Compound average-score maximum-score number-of-dockings ") + "Tranch Compound SMILES average-score maximum-score number-of-dockings ") for replica_index in range(scenario['replicas']): replica_str = f"score-replica-{replica_index}" @@ -396,11 +409,16 @@ def create_summary_file(ctx, scenario, collection, scenario_result): if(len(ligand['scores']) > 0): - max_score = max(ligand['scores']) + # Despite the column being named "maximum-score" (kept for schema parity with the + # Slurm/HPC summary format), this is the BEST (most negative) docking score across + # replicas, matching one-queue.sh's update_summary() -- NOT the literal maximum + # (worst score), which this used to compute prior to the ML classifier work. + min_score = min(ligand['scores']) avg_score = sum(ligand['scores']) / len(ligand['scores']) + smiles = collection['ligands'].get(ligand_key, {}).get('smiles') or "NA" summmary_fp.write( - f"{collection['key']} {ligand_key} {avg_score:3.1f} {max_score:3.1f} {len(ligand['scores']):5d} ") + f"{collection['key']} {ligand_key} {smiles} {avg_score:3.1f} {min_score:3.1f} {len(ligand['scores']):5d} ") for replica_index in range(scenario['replicas']): summmary_fp.write( f"{ligand['scores'][replica_index]:3.1f} ") @@ -473,6 +491,70 @@ def scenario_collection_output_directory_txt_gz(ctx, scenario, collection, resul return os.path.join(*scenario_collection_output(ctx, scenario, collection, result_type, skip_num=skip_num, tmp_prefix=tmp_prefix, append=".txt.gz")) +def apply_ml_classifier_filter(collections, model_path, probability_cutoff): + """ + Filters ligands across all given collections using a pretrained ML tranche-prioritization + classifier, mutating `collections` in place. + + Ligands not predicted worth docking (probability <= probability_cutoff), or whose SMILES + could not be extracted (ligand['smiles'] is None -- fail-closed), are popped from + collection['ligands'] and logged to collection['log']/['log_json'], exactly like the existing + B/Si/Sn/duplicate-coordinate skip checks in process(). + + Deliberately takes only plain data (no ctx/boto3), so it can be unit-tested in isolation with + a fabricated `collections` dict, without mocking any AWS/S3 state. + + Args: + collections (dict): collection_key -> collection dict, each with a 'ligands' dict of + ligand_key -> {'path': ..., 'smiles': ...}, plus 'log'/'log_json' lists. + model_path (str): absolute path to the trained classifier (.pt file). + probability_cutoff (float): minimum predicted probability (exclusive) to keep a ligand. + + Returns: + None (collections is mutated in place). + + Raises: + Exception: if the model cannot be loaded. This is a whole-subjob configuration problem + (missing/corrupt model file), not a per-ligand transient failure, so it is + deliberately allowed to propagate and fail the whole container task -- unlike this + file's per-ligand docking-exception handling elsewhere in process_ligand(). + """ + model, metadata = load_classifier(model_path) + fp_radius = metadata.get('fp_radius', 2) + fp_nbits = metadata.get('fp_nbits', 1024) + + ligand_refs = [] + smiles_list = [] + for collection_key in collections: + collection = collections[collection_key] + for ligand_key in collection['ligands']: + ligand = collection['ligands'][ligand_key] + ligand_refs.append((collection_key, ligand_key)) + smiles_list.append(ligand.get('smiles') or '') + + if len(ligand_refs) == 0: + return + + keep_mask, probabilities = filter_smiles_mask( + model, smiles_list, probability_cutoff=probability_cutoff, + fp_radius=fp_radius, fp_nbits=fp_nbits) + + ligands_to_skip = {} + for (collection_key, ligand_key), keep, probability in zip(ligand_refs, keep_mask, probabilities): + if not keep: + probability_str = f"{probability:.4f}" if probability is not None else "invalid_smiles" + skip_reason = f"filtered(ml_classifier:{probability_str})" + skip_reason_json = f"ml classifier probability: {probability_str}" + collections[collection_key]['log'].append(f"{ligand_key} {skip_reason}") + collections[collection_key]['log_json'].append( + {'ligand': ligand_key, 'status': 'filtered', 'info': skip_reason_json}) + ligands_to_skip.setdefault(collection_key, []).append(ligand_key) + + for collection_key, ligand_keys in ligands_to_skip.items(): + for ligand_key in ligand_keys: + collections[collection_key]['ligands'].pop(ligand_key, None) + + def process(ctx): # Figure out who I am... @@ -541,6 +623,7 @@ def process(ctx): skip_ligand = 0 skip_reason = "" skip_reason_json = "" + ligand['smiles'] = None # Check to see if ligand contains B, Si, Sn or has duplicate coordinates with open(ligand['path'], "r") as read_file: @@ -570,6 +653,15 @@ def process(ctx): break coords[coord_str] = 1 + # Getting the ligand SMILES (same "line containing SMILES, take the last + # whitespace token" rule used on the Slurm/HPC side, for consistency and so + # summary files carry the same schema regardless of execution path). Does not + # break the loop -- unlike the checks above, finding this doesn't disqualify + # the ligand, and the B/Si/Sn/coordinate scan should continue regardless. + match = re.search(r'SMILES', line) + if(match): + ligand['smiles'] = line.split()[-1] + if skip_ligand: collection['log'].append(f"{ligand_key} {skip_reason}") collection['log_json'].append( @@ -579,6 +671,16 @@ def process(ctx): for ligand_key in ligands_to_skip: collection['ligands'].pop(ligand_key, None) + # Optional ML tranche-prioritization classifier: filter every remaining ligand across every + # collection in this subjob, once (a single batched inference call), before task-list + # construction below. + if(ctx['config'].get('use_ml_classifier', 'false') == 'true'): + apply_ml_classifier_filter( + collections, + ctx['config']['ml_classifier_model_path'], + ctx['config']['ml_classifier_probability_cutoff'] + ) + # Create the task list based on the scenarios and replicas required tasklist = [] for scenario_key in ctx['config']['docking_scenarios']: diff --git a/tools/train_ml_classifier.py b/tools/train_ml_classifier.py new file mode 100644 index 0000000..9b16bf9 --- /dev/null +++ b/tools/train_ml_classifier.py @@ -0,0 +1,184 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Trains the optional ML tranche-prioritization classifier (see tools/templates/ml_classifier.py) on +the results of a prescreen run of VFVS. + +Usage: run VFVS's existing, unmodified workflow (all.ctrl + submit/AWS Batch + one-queue.sh or +vf_aws_run.py) on a small representative "prescreen" set of collections first, to produce real +docking scores. This generates the usual per-collection summary files under +output-files/{complete,incomplete}//summaries/. Then, run this script once (a single +interactive invocation from the tools/ directory -- no Slurm/AWS Batch job needed, since training +the classifier only takes seconds): + + python3 train_ml_classifier.py --scenario + +This reads ml_classifier_model_path from workflow/control/all.ctrl, collects every summary file +for the given docking scenario, extracts (SMILES, docking score) pairs, and trains+saves the +classifier. Once saved, set use_ml_classifier=true in all.ctrl and run the primary screen as usual +-- both the Slurm and AWS Batch paths will load this classifier and filter collections before +docking. + +Both VFVS execution paths write the same summary schema (Tranch Compound SMILES average-score +maximum-score number-of-dockings score-replica-*), so this script works identically regardless of +whether the prescreen ran on Slurm or AWS Batch. + +@author: akshat +""" +import os +import sys +import glob +import gzip +import argparse + +sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), 'templates')) +# ml_classifier.py lives in templates/ (not alongside this script) so that it is automatically +# bundled into the AWS Batch Docker image (which ADDs the whole tools/ tree) and importable via a +# relative path from one-queue.sh's cwd. This script is not part of that runtime path, so it needs +# to add templates/ to sys.path explicitly. +from ml_classifier import train_classifier, load_classifier # noqa: E402 + + +def read_config_file(filename): + """ + Minimal, standalone re-implementation of VFVS's `key=value` control-file convention (see + all.ctrl). Deliberately not shared with one-queue.sh/vf_aws_run.py's own config-reading code + (bash grep-idiom / parse_config()), since neither is meant to be invoked from a plain + standalone Python script like this one; a few lines of duplication here is simpler and safer + than trying to reuse either. + """ + params = {} + with open(filename, 'r') as f: + for line in f: + line = line.strip() + if not line or line.startswith('#'): + continue + if '=' not in line: + continue + key, value = line.split('=', 1) + key = key.strip() + value = value.split('#', 1)[0].strip() + if key: + params[key] = value + return params + + +def parse_summary_files(summaries_glob): + """ + Parses VFVS summary files (gzip'd, whitespace-columnar) into (smiles, docking_score) pairs. + + Expected columns (both Slurm and AWS Batch write this same schema): Tranch, Compound, SMILES, + average-score, maximum-score, number-of-dockings, score-replica-0, .... The "maximum-score" + column is (despite its name) the best/most-negative docking score across replicas -- used + here as the training label, per the manuscript's "docking scores obtained during the + prescreen" definition. The header line (starting with "Tranch") is skipped. Malformed lines + (wrong column count, non-numeric score) are skipped with a warning rather than aborting the + whole run. + + Args: + summaries_glob (str): glob pattern matching gzip'd summary files, e.g. + '../output-files/complete//summaries/**/*.txt.gz'. + + Returns: + (smiles_list, scores_list) + + Raises: + Exception: if no matching summary files are found, or no valid rows could be parsed. + """ + summary_files = sorted(glob.glob(summaries_glob, recursive=True)) + if len(summary_files) == 0: + raise Exception('No summary files found matching: {}'.format(summaries_glob)) + + smiles_list = [] + scores_list = [] + n_malformed = 0 + for summary_file in summary_files: + with gzip.open(summary_file, 'rt') as f: + for line_number, line in enumerate(f, start=1): + line = line.strip() + if not line or line.startswith('Tranch'): + continue + fields = line.split() + if len(fields) < 5: + n_malformed += 1 + print('WARNING: Skipping malformed line {}:{}: {}'.format( + summary_file, line_number, line)) + continue + smi = fields[2] + try: + score = float(fields[4]) + except ValueError: + n_malformed += 1 + print('WARNING: Skipping malformed line {}:{}: {}'.format( + summary_file, line_number, line)) + continue + smiles_list.append(smi) + scores_list.append(score) + + print('Parsed {} summary files: {} valid docking scores, {} malformed lines skipped.'.format( + len(summary_files), len(smiles_list), n_malformed)) + + if len(smiles_list) == 0: + raise Exception('No valid docking scores found in summary files matching: {}'.format(summaries_glob)) + + return smiles_list, scores_list + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--scenario', required=True, + help='Docking scenario name whose prescreen summaries to train on.') + parser.add_argument('--ctrl-file', default='../workflow/control/all.ctrl', + help='Path to all.ctrl (for ml_classifier_model_path).') + parser.add_argument('--summaries-glob', default=None, + help='Glob pattern for prescreen summary files. Defaults to ' + '../output-files/{complete,incomplete}//summaries/**/*.txt.gz') + parser.add_argument('--top-percent', type=float, default=25, + help='Percentage of most-negative-scoring prescreen compounds labeled positive.') + parser.add_argument('--seed', type=int, default=42, help='Random seed.') + args = parser.parse_args() + + config_params = read_config_file(args.ctrl_file) + model_relative_path = config_params.get('ml_classifier_model_path', 'ml_classifier/model.pt') + model_out_path = os.path.join('..', 'input-files', model_relative_path) + + if args.summaries_glob is not None: + summaries_globs = [args.summaries_glob] + else: + summaries_globs = [ + '../output-files/complete/{}/summaries/**/*.txt.gz'.format(args.scenario), + '../output-files/incomplete/{}/summaries/**/*.txt.gz'.format(args.scenario), + ] + + smiles_list = [] + scores_list = [] + for summaries_glob in summaries_globs: + try: + smi_part, score_part = parse_summary_files(summaries_glob) + smiles_list.extend(smi_part) + scores_list.extend(score_part) + except Exception as e: + print('NOTE: {}'.format(e)) + + if len(smiles_list) == 0: + raise Exception('No valid docking scores found across any of: {}'.format(summaries_globs)) + + if os.path.isfile(model_out_path): + try: + _, existing_metadata = load_classifier(model_out_path) + existing_receptor = existing_metadata.get('receptor') + if existing_receptor is not None and existing_receptor != args.scenario: + print('WARNING: {} already exists and was trained for scenario "{}". ' + 'Overwriting with a model trained for scenario "{}".'.format( + model_out_path, existing_receptor, args.scenario)) + except Exception: + pass # existing file is unreadable/corrupt; fine to overwrite + + model, history = train_classifier(smiles_list, scores_list, model_out_path, + top_percent=args.top_percent, receptor=args.scenario, seed=args.seed) + + print('Classifier trained on {} prescreen compounds (best epoch {}, best val loss {:.4f}).'.format( + len(smiles_list), history['best_epoch'] + 1, min(history['val_loss']))) + print('Saved to: {}'.format(model_out_path)) + print('To use it for the primary screen, set use_ml_classifier=true and ml_classifier_model_path={} ' + 'in {}.'.format(model_relative_path, args.ctrl_file)) diff --git a/tools/vf_report.sh b/tools/vf_report.sh index 1c60b96..8639e18 100755 --- a/tools/vf_report.sh +++ b/tools/vf_report.sh @@ -471,7 +471,9 @@ if [[ "${category}" = "vs" ]]; then folder=$(basename ${folder}) for file in $(ls ${tempdir}/output-files/${docking_scenario_name}/summaries/${metatranch}/${folder} 2>/dev/null || true); do file=$(basename ${file} || true) - zcat ${tempdir}/output-files/${docking_scenario_name}/summaries/${metatranch}/${folder}/${file} 2>/dev/null | awk '{print $1, $2, $4}' >> ${tempdir}/summaries.all || true + # Summary columns: Tranch Compound SMILES average-score maximum-score number-of-dockings ... + # Field 5 (maximum-score) is the best/most-negative docking score; field 4 is the average. + zcat ${tempdir}/output-files/${docking_scenario_name}/summaries/${metatranch}/${folder}/${file} 2>/dev/null | awk '{print $1, $2, $5}' >> ${tempdir}/summaries.all || true done rm -r ${tempdir}/output-files/${docking_scenario_name}/summaries/${metatranch}/${folder} done @@ -482,7 +484,7 @@ if [[ "${category}" = "vs" ]]; then for metatranch in $(ls -A ${folder}/summaries/); do for tranch in $(ls -A ${folder}/summaries/${metatranch}); do for file in $(ls -A ${folder}/summaries/${metatranch}/${tranch}); do - zcat ${folder}/summaries/${metatranch}/${tranch}/${file} 2>/dev/null | awk '{print $1, $2, $4}' >> ${tempdir}/summaries.all 2>/dev/null || true + zcat ${folder}/summaries/${metatranch}/${tranch}/${file} 2>/dev/null | awk '{print $1, $2, $5}' >> ${tempdir}/summaries.all 2>/dev/null || true summary_flag="true" done done @@ -498,7 +500,7 @@ if [[ "${category}" = "vs" ]]; then for metatranch in $(ls -A ${folder}/summaries/); do for tranch in $(ls -A ${folder}/summaries/${metatranch}); do for file in $(ls -A ${folder}/summaries/${metatranch}/${tranch}); do - zcat ${folder}/summaries/${metatranch}/${tranch}/${file} 2>/dev/null | awk '{print $1, $2, $4}' >> ${tempdir}/summaries.all 2>/dev/null || true + zcat ${folder}/summaries/${metatranch}/${tranch}/${file} 2>/dev/null | awk '{print $1, $2, $5}' >> ${tempdir}/summaries.all 2>/dev/null || true summary_flag="true" done done