Add optional ML classifier for molecule prioritization within collections

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 <docking_scenario_name>
+   ```
+   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
 

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).

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.
+

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}/<scenario>/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 <docking_scenario_name>
+   ```
+
+   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`.
+
 
 
 

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.

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

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    *****************************************************************
 *****************************************************************************************************************************************************************