OCScore Pipeline Examples¶
These examples demonstrate how to train and use OCScore models for consensus scoring.
Complete OCScore Pipeline¶
Complete end-to-end pipeline to obtain OCScore results from scratch:
Complete OCScore pipeline¶
#!/usr/bin/env python3
"""
Example: Complete OCScore Pipeline
This example demonstrates the complete pipeline to obtain OCScore results from scratch:
1. Receptor and ligand preparation
2. Docking with multiple engines (Vina, PLANTS, GNINA)
3. Pose clustering to find representative poses
4. Rescoring with multiple scoring functions (ODDT, PLANTS, Vina, SMINA, GNINA)
5. Feature extraction (receptor and ligand descriptors)
6. Model inference using trained OCScore model
The rescoring results are automatically mapped to database column names for consistency.
Usage:
# Update paths in the script to match your system
python examples/11_python_api_complete_ocscore_pipeline.py
"""
###############################################################################
# USER CONFIGURATION - Update these variables to match your system
###############################################################################
# OCDocker configuration file path
# Set this to the absolute path of your OCDocker.cfg file
# If None, will use OCDOCKER_CONFIG environment variable or search for OCDocker.cfg
OCDOCKER_CONFIG_FILE = "/path/to/OCDocker.cfg" # Update this path
# Receptor configuration
RECEPTOR_PATH = "/path/to/receptor_directory/receptor.pdb"
RECEPTOR_NAME = "Receptor"
PREPARED_RECEPTOR_PDBQT = "/path/to/receptor_directory/prepared_receptor.pdbqt"
PREPARED_RECEPTOR_MOL2 = "/path/to/receptor_directory/prepared_receptor.mol2"
BOX_CENTER = (0.0, 0.0, 0.0) # (x, y, z) coordinates of the box center
# Ligand configuration - List of ligand paths (base directories)
# Each path should contain: {ligand_name}.smi, boxes/box0.pdb, and subdirectories for outputs
# Ligand names are automatically extracted from the last folder name in each path
LIGAND_PATHS = [
]
# Add more ligand paths here (you can use glob to get all ligand folders):
# "/path/to/ligand2",
# "/path/to/ligand3",
# Model configuration
MODEL_NAME = "OCScore" # Name of your trained model (without extension)
MODELS_DIR = "OCScore_models" # Directory containing models
PCA_MODEL_PATH = None # Path to PCA model if used, e.g., f"{MODELS_DIR}/{MODEL_NAME}_pca.pkl"
# Preprocessing configuration (should match training settings)
SCALER = "standard" # "standard" or "minmax"
INVERT_CONDITIONALLY = True
NORMALIZE = True
USE_MASK = True
SCORE_COLUMNS_LIST = ["SMINA", "VINA", "ODDT", "PLANTS"]
###############################################################################
# !!! CRITICAL WARNING: SCORING FUNCTION COLUMN ORDER !!!
###############################################################################
#
# THE FOLLOWING ORDER MUST BE STRICTLY RESPECTED WHEN APPLYING MASKS:
#
# 1. SMINA_VINA
# 2. SMINA_SCORING_DKOES
# 3. SMINA_VINARDO
# 4. SMINA_OLD_SCORING_DKOES
# 5. SMINA_FAST_DKOES
# 6. SMINA_SCORING_AD4
# 7. VINA_VINA
# 8. VINA_VINARDO
# 9. PLANTS_CHEMPLP
# 10. PLANTS_PLP
# 11. PLANTS_PLP95
# 12. ODDT_RFSCORE_V1
# 13. ODDT_RFSCORE_V2
# 14. ODDT_RFSCORE_V3
# 15. ODDT_PLECRF_P5_L1_S65536
# 16. ODDT_NNSCORE
#
# !!! WARNING: If you change the order of scoring function columns in the output,
# the mask will be applied incorrectly, leading to wrong predictions!
#
# The mask is a 16-element array where each position corresponds to one of the
# scoring functions above in the exact order listed. Position 0 = SMINA_VINA,
# position 1 = SMINA_SCORING_DKOES, etc.
#
# DO NOT MODIFY THE ORDER OF SCORING FUNCTION COLUMNS WITHOUT UPDATING THE MASK!
#
###############################################################################
# GPU configuration
USE_GPU = True # Set to False to force CPU usage (useful if CUDA is not available or to avoid GPU memory issues)
# Output configuration
OUTPUT_FILE = "ocscore_results.csv" # CSV file to save results (None to skip saving)
SAVE_TO_FILE = True # Set to False to only store results in memory
# Multiprocessing configuration
N_JOBS = 4 # Number of parallel jobs (cores) to use. Set to -1 for all available cores
USE_MULTIPROCESSING = True # Set to False to process ligands sequentially
# Pipeline stage toggles — enable each engine independently
RUN_VINA_DOCKING = True
RUN_PLANTS_DOCKING = True
RUN_GNINA_DOCKING = True
RUN_ODDT_RESCORING = True
RUN_PLANTS_RESCORING = True
RUN_VINA_RESCORING = True
RUN_SMINA_RESCORING = True
RUN_GNINA_RESCORING = True
# When True, partial reruns keep scores from disabled engines (read from disk/cache).
# Use this to add GNINA (or any engine) to an already-completed protocol without losing other scores.
MERGE_EXISTING_RESCORING = True
# Example — GNINA rescoring only on an already-docked/scored ligand folder:
# RUN_VINA_DOCKING = RUN_PLANTS_DOCKING = RUN_GNINA_DOCKING = False
# RUN_ODDT_RESCORING = RUN_PLANTS_RESCORING = RUN_VINA_RESCORING = RUN_SMINA_RESCORING = False
# RUN_GNINA_RESCORING = True
# MERGE_EXISTING_RESCORING = True
# Existing vina/smina/plants/oddt outputs are left untouched; only gninaFiles/rescoring is written.
# Checkpoint/resume configuration
# When enabled, each ligand writes a per-ligand checkpoint and can resume from the
# last completed stage. This is backward-compatible with runs done before this code:
# if no checkpoint file exists, completion is inferred from existing output artifacts.
ENABLE_LIGAND_CHECKPOINT = True
LIGAND_CHECKPOINT_FILE = ".ocscore_pipeline_checkpoint.json"
LIGAND_FEATURES_CACHE_FILE = ".ocscore_pipeline_features.json"
###############################################################################
# END USER CONFIGURATION
###############################################################################
import argparse
import json
import os
import time
import warnings
import numpy as np
import pandas as pd
from glob import glob
from typing import Any, Dict, Optional, Tuple
# Configure sklearn/joblib to use threading backend for parallel execution
# This allows sklearn models to use multiple threads while main process uses multiprocessing
# The threading backend avoids the "Loky-backed parallel loops cannot be called in multiprocessing" issue
warnings.filterwarnings('ignore', message='.*Loky-backed parallel loops cannot be called in a multiprocessing.*')
# Note: We keep the default multiprocessing start method ('fork' on Linux)
# which is faster and works well with proper tmp directory isolation
try:
import joblib
from joblib import parallel_backend, Parallel, delayed
# Set default backend to threading so sklearn can parallelize within multiprocessing workers
# Threading backend works inside multiprocessing contexts (unlike Loky)
joblib.parallel.DEFAULT_BACKEND = 'threading'
JOBLIB_AVAILABLE = True
except (ImportError, AttributeError):
# If joblib not available, try setting environment variable
os.environ['JOBLIB_BACKEND'] = 'threading'
JOBLIB_AVAILABLE = False
USE_MULTIPROCESSING = False
print("Warning: joblib not available. Multiprocessing disabled.")
# Explicitly bootstrap OCDocker with the specified config file BEFORE other imports
# This ensures the config is loaded correctly regardless of working directory
# Set OCDOCKER_NO_AUTO_BOOTSTRAP to prevent auto-bootstrap from running first
os.environ['OCDOCKER_NO_AUTO_BOOTSTRAP'] = '1'
if OCDOCKER_CONFIG_FILE and os.path.isfile(OCDOCKER_CONFIG_FILE):
import OCDocker.Error as ocerror
import OCDocker.Initialise as init
print(f"Loading OCDocker configuration from: {OCDOCKER_CONFIG_FILE}")
bootstrap_ns = argparse.Namespace(
multiprocess=USE_MULTIPROCESSING,
update=False,
config_file=OCDOCKER_CONFIG_FILE,
output_level=ocerror.ReportLevel.WARNING,
overwrite=False
)
init.bootstrap(bootstrap_ns)
print("OCDocker configuration loaded successfully.\n")
else:
# Fall back to auto-bootstrap if config file not specified or not found
if OCDOCKER_CONFIG_FILE:
print(f"Warning: Config file not found at {OCDOCKER_CONFIG_FILE}, using auto-bootstrap...")
# Re-enable auto-bootstrap
os.environ.pop('OCDOCKER_NO_AUTO_BOOTSTRAP', None)
# Now import other OCDocker modules (they won't trigger auto-bootstrap since we already bootstrapped)
import OCDocker.Docking.Gnina as ocgnina
import OCDocker.Docking.PLANTS as ocplants
import OCDocker.Docking.Smina as ocsmina
import OCDocker.Docking.Vina as ocvina
import OCDocker.Ligand as ocl
import OCDocker.OCScore.Scoring as ocscoring
import OCDocker.OCScore.Utils.Data as ocscoredata
import OCDocker.OCScore.Utils.IO as ocscoreio
import OCDocker.Processing.Preprocessing.RMSDClustering as ocrmsdclust
import OCDocker.Receptor as ocr
import OCDocker.Rescoring.ODDT as ocoddt
import OCDocker.Toolbox.Conversion as occonversion
import OCDocker.Toolbox.MoleculeProcessing as ocmolproc
import OCDocker.Toolbox.Security as ocsec
# Allow unsafe deserialization
ocsec.allow_unsafe_runtime(deserialization=True, script_exec=False)
def main():
'''Main function to process all ligands.'''
_validate_pipeline_config()
# OCDocker auto-bootstraps on import, so configuration is already loaded
# If you need to verify bootstrap or use custom settings, you can:
# 1. Set OCDOCKER_NO_AUTO_BOOTSTRAP=1 environment variable
# 2. Import OCDocker.Initialise and call bootstrap() explicitly
# Automatically derive ligand names from paths (last folder name)
ligand_names = [os.path.basename(os.path.normpath(path)) for path in LIGAND_PATHS]
# Create receptor object
receptor = ocr.Receptor(RECEPTOR_PATH, name=RECEPTOR_NAME)
# Prepare list of (ligand_path, ligand_name) tuples
ligand_tasks = list(zip(LIGAND_PATHS, ligand_names))
print(f"\n{'='*60}")
print(f"OCSCORE PIPELINE")
print(f"{'='*60}")
print(f"Receptor: {RECEPTOR_NAME}")
print(f"Number of ligands: {len(ligand_tasks)}")
print(f"Use mask: {USE_MASK}")
print(f"Merge existing rescoring: {MERGE_EXISTING_RESCORING}")
print(f"Docking: vina={RUN_VINA_DOCKING}, plants={RUN_PLANTS_DOCKING}, gnina={RUN_GNINA_DOCKING}")
print(
"Rescoring: "
f"oddt={RUN_ODDT_RESCORING}, plants={RUN_PLANTS_RESCORING}, "
f"vina={RUN_VINA_RESCORING}, smina={RUN_SMINA_RESCORING}, gnina={RUN_GNINA_RESCORING}"
)
print(f"Multiprocessing: {USE_MULTIPROCESSING}")
if USE_MULTIPROCESSING:
print(f"Number of jobs: {N_JOBS}")
print(f"{'='*60}\n")
# Process ligands
if USE_MULTIPROCESSING and JOBLIB_AVAILABLE and len(ligand_tasks) > 1:
# Use joblib for parallel processing
print(f"Processing {len(ligand_tasks)} ligands in parallel using {N_JOBS} cores...")
results = Parallel(n_jobs=N_JOBS)(
delayed(process_single_ligand)(ligand_path, ligand_name, receptor)
for ligand_path, ligand_name in ligand_tasks
)
else:
# Process sequentially
print(f"Processing {len(ligand_tasks)} ligands sequentially...")
results = []
for ligand_path, ligand_name in ligand_tasks:
print(f"Processing ligand: {ligand_name}")
result = process_single_ligand(ligand_path, ligand_name, receptor)
results.append(result)
# Filter out None results (failed processing)
results = [r for r in results if r is not None]
if not results:
print("No ligands were successfully processed.")
return
# Batch all ligands together for model inference
# This ensures proper normalization (scaler fit on all data, not single rows)
print(f"\n{'='*60}")
print(f"BATCH MODEL INFERENCE")
print(f"{'='*60}")
# Convert all results to a single DataFrame
if results:
# Get all unique keys from all dictionaries
all_keys = set()
for result in results:
if result is not None:
all_keys.update(result.keys())
# Ensure all dictionaries have all keys (fill missing with None)
normalized_results = []
for result in results:
if result is not None:
normalized_result = {key: result.get(key, None) for key in all_keys}
normalized_results.append(normalized_result)
# Create DataFrame from normalized dictionaries
feature_df = pd.DataFrame(normalized_results)
else:
feature_df = pd.DataFrame()
if feature_df.empty:
print("No features to process for model inference.")
return
# Path to your trained model
model_path = f"{MODELS_DIR}/{MODEL_NAME}.pt"
mask_path = f"{MODELS_DIR}/{MODEL_NAME}_mask.pkl"
scaler_path = f"{MODELS_DIR}/{MODEL_NAME}_scaler.pkl" # Path to saved scaler
# Load the mask if it exists
mask = None
if os.path.isfile(mask_path) and USE_MASK:
try:
mask = ocscoreio.load_mask(MODEL_NAME, models_dir=MODELS_DIR)
except Exception as e:
print(f"Warning: Could not load mask: {e}")
mask = None
# Check if scaler exists (required for proper normalization)
if NORMALIZE and not os.path.isfile(scaler_path):
print(f"WARNING: Scaler file not found at {scaler_path}")
print(" This means normalization will use a NEW scaler fitted on prediction data,")
print(" which is INCORRECT. The scaler should be saved during training.")
print(" Predictions may be inaccurate!")
scaler_path = None # Will create new scaler (incorrect but won't crash)
elif NORMALIZE:
print(f"Using saved scaler from: {scaler_path}")
# Get OCScore predictions for all ligands at once
try:
print(f"Running model inference on {len(feature_df)} ligands...")
print(f"Feature DataFrame shape: {feature_df.shape}")
ocscore_predictions = ocscoring.get_score(
model_path=model_path,
data=feature_df,
pca_model=PCA_MODEL_PATH,
mask=mask,
score_columns_list=SCORE_COLUMNS_LIST,
scaler=SCALER,
scaler_path=scaler_path if NORMALIZE else None, # Use saved scaler if normalization is enabled
invert_conditionally=INVERT_CONDITIONALLY,
normalize=NORMALIZE,
serialization_method="auto", # Auto-detect model format
use_gpu=USE_GPU # Use GPU if available and USE_GPU=True
)
if isinstance(ocscore_predictions, pd.DataFrame):
print(f"Prediction DataFrame shape: {ocscore_predictions.shape}")
print(f"Prediction DataFrame columns: {list(ocscore_predictions.columns)}")
if 'predicted_score' in ocscore_predictions.columns:
print(f"All predicted_score values: {ocscore_predictions['predicted_score'].tolist()}")
print(f"Unique predicted_score values: {ocscore_predictions['predicted_score'].nunique()}")
elif isinstance(ocscore_predictions, (pd.Series, np.ndarray)):
predictions_array = np.asarray(ocscore_predictions)
print(f"Prediction array shape: {predictions_array.shape}")
print(f"All prediction values: {predictions_array.tolist()}")
print(f"Unique prediction values: {len(np.unique(predictions_array))}")
# Add OCScore predictions to results
if isinstance(ocscore_predictions, pd.DataFrame):
if 'predicted_score' in ocscore_predictions.columns:
# Map predictions back to results by index
for idx, result in enumerate(results):
if result is not None and idx < len(ocscore_predictions):
result['OCSCORE'] = ocscore_predictions['predicted_score'].iloc[idx]
print(f" Mapped prediction {idx} to ligand {result.get('ligand', 'unknown')}: {result['OCSCORE']}")
elif len(ocscore_predictions.columns) == 1:
# Single prediction column
for idx, result in enumerate(results):
if result is not None and idx < len(ocscore_predictions):
result['OCSCORE'] = ocscore_predictions.iloc[idx, 0]
else:
# Try to find a numeric column
numeric_cols = ocscore_predictions.select_dtypes(include=[np.number]).columns
if len(numeric_cols) > 0:
for idx, result in enumerate(results):
if result is not None and idx < len(ocscore_predictions):
result['OCSCORE'] = ocscore_predictions[numeric_cols[0]].iloc[idx]
elif isinstance(ocscore_predictions, (pd.Series, np.ndarray)):
# Array/Series of predictions
predictions_array = np.asarray(ocscore_predictions)
for idx, result in enumerate(results):
if result is not None and idx < len(predictions_array):
result['OCSCORE'] = float(predictions_array[idx])
print(f"Model inference completed for {len(results)} ligands.")
except FileNotFoundError as e:
print(f"Warning: Model file not found: {e}")
for result in results:
if result is not None:
result['OCSCORE'] = None
except Exception as e:
print(f"Error during model inference: {e}")
import traceback
traceback.print_exc()
for result in results:
if result is not None:
result['OCSCORE'] = None
# Convert results to DataFrame
# Use orient='index' and transpose to preserve order, then convert properly
# First, ensure all dictionaries have the same keys (fill missing with None)
if results:
# Get all unique keys from all dictionaries
all_keys = set()
for result in results:
if result is not None:
all_keys.update(result.keys())
# Ensure all dictionaries have all keys (fill missing with None)
normalized_results = []
for result in results:
if result is not None:
normalized_result = {key: result.get(key, None) for key in all_keys}
normalized_results.append(normalized_result)
else:
normalized_results.append({key: None for key in all_keys})
# Create DataFrame from normalized dictionaries
results_df = pd.DataFrame(normalized_results)
else:
results_df = pd.DataFrame()
# Reorder columns to match the data source order (from training data file)
# !!! CRITICAL: This ensures all columns (especially SFs) are in the exact same order
# as the training data, which is essential for proper mask application and model inference!
if not results_df.empty:
# Get the column order from config (no file path needed)
# Uses reference_column_order from OCDocker.cfg
source_order = ocscoredata.get_column_order() # Uses config by default
# Use the reorder function to match the config column order
# This handles OCSCORE insertion and extra columns automatically
results_df = ocscoredata.reorder_columns_to_match_data_order(
df=results_df,
data_source=None, # Uses config.reference_column_order by default
keep_extra_columns=True, # Keep OCSCORE and any other extra columns
fill_missing_columns=False # Don't add missing columns as NaN
)
# Manually insert OCSCORE right after 'ligand' if it exists
if 'OCSCORE' in results_df.columns:
cols = list(results_df.columns)
if 'ligand' in cols:
# Remove OCSCORE from current position
cols.remove('OCSCORE')
# Insert after 'ligand'
ligand_idx = cols.index('ligand')
cols.insert(ligand_idx + 1, 'OCSCORE')
results_df = results_df[cols]
# Print summary
print(f"\n{'='*60}")
print(f"PROCESSING COMPLETE")
print(f"{'='*60}")
print(f"Successfully processed: {len(results)}/{len(ligand_tasks)} ligands")
if 'OCSCORE' in results_df.columns:
print(f"OCScore predictions: {results_df['OCSCORE'].notna().sum()}/{len(results_df)}")
if results_df['OCSCORE'].notna().any():
print(f"OCScore range: {results_df['OCSCORE'].min():.4f} - {results_df['OCSCORE'].max():.4f}")
print(f"{'='*60}\n")
# Save to file if requested
if SAVE_TO_FILE and OUTPUT_FILE:
output_path = OUTPUT_FILE
results_df.to_csv(output_path, index=False)
print(f"Results saved to: {output_path}")
return results_df
# Mapping function to convert rescoring keys to database column names
def map_rescoring_key_to_db_column(key: str, engine: Optional[str] = None) -> str:
'''Map rescoring result keys to database column names.
Parameters
----------
key : str
The key from rescoring results (e.g., 'vina_vina_rescoring', 'smina_vinardo_rescoring', etc.)
engine : str
The engine name (e.g., 'vina', 'smina')
If None, the engine will be inferred from the key.
If provided, the engine will be used to determine the database column name.
If not provided, the engine will be inferred from the key.
Returns
-------
str
The database column name (e.g., 'VINA_VINA', 'SMINA_VINARDO', etc.)
'''
key_lower = key.lower()
# Mapping dictionary for rescoring keys to database column names
# !!! CRITICAL: These must match SCORING_FUNCTION_ORDER for mask application !!!
mapping = {
# VINA mappings
'vina_vina_rescoring': 'VINA_VINA',
'vina_vinardo_rescoring': 'VINA_VINARDO',
# SMINA mappings - MUST match SCORING_FUNCTION_ORDER
'smina_vina_rescoring': 'SMINA_VINA',
'smina_vinardo_rescoring': 'SMINA_VINARDO',
'smina_dkoes_scoring_rescoring': 'SMINA_SCORING_DKOES',
'smina_scoring_dkoes_rescoring': 'SMINA_SCORING_DKOES', # Alternative format
'smina_old_scoring_dkoes_rescoring': 'SMINA_OLD_SCORING_DKOES',
'smina_fast_dkoes_rescoring': 'SMINA_FAST_DKOES',
'smina_ad4_scoring_rescoring': 'SMINA_SCORING_AD4',
# Handle alternative SMINA key formats (from actual rescoring output)
'smina_dkoes_fast': 'SMINA_FAST_DKOES',
'smina_dkoes_scoring_old': 'SMINA_OLD_SCORING_DKOES',
'smina_dkoes_fast_rescoring': 'SMINA_FAST_DKOES',
'smina_dkoes_scoring_old_rescoring': 'SMINA_OLD_SCORING_DKOES',
# PLANTS mappings
'plants_chemplp': 'PLANTS_CHEMPLP',
'plants_plp': 'PLANTS_PLP',
'plants_plp95': 'PLANTS_PLP95',
# GNINA mappings
'gnina_vina_rescoring': 'GNINA_VINA',
'gnina_vinardo_rescoring': 'GNINA_VINARDO',
'gnina_dkoes_scoring_rescoring': 'GNINA_SCORING_DKOES',
'gnina_scoring_dkoes_rescoring': 'GNINA_SCORING_DKOES', # Alternative format
'gnina_old_scoring_dkoes_rescoring': 'GNINA_OLD_SCORING_DKOES',
'gnina_fast_dkoes_rescoring': 'GNINA_FAST_DKOES',
'gnina_ad4_scoring_rescoring': 'GNINA_SCORING_AD4',
'gnina_dkoes_fast': 'GNINA_FAST_DKOES',
'gnina_dkoes_scoring_old': 'GNINA_OLD_SCORING_DKOES',
'gnina_dkoes_fast_rescoring': 'GNINA_FAST_DKOES',
'gnina_dkoes_scoring_old_rescoring': 'GNINA_OLD_SCORING_DKOES',
# ODDT mappings (these come from the dataframe columns, already prefixed with oddt_)
'oddt_rfscore_v1': 'ODDT_RFSCORE_V1',
'oddt_rfscore_v2': 'ODDT_RFSCORE_V2',
'oddt_rfscore_v3': 'ODDT_RFSCORE_V3',
'oddt_plecrf_p5_l1_s65536': 'ODDT_PLECRF_P5_L1_S65536',
'oddt_plec_p5_l1_s65536': 'ODDT_PLECRF_P5_L1_S65536', # Alternative naming
'oddt_nnscore': 'ODDT_NNSCORE',
}
# Check if exact match exists
if key_lower in mapping:
return mapping[key_lower]
# Handle ODDT keys that are already prefixed with oddt_
if key_lower.startswith('oddt_'):
# Remove oddt_ prefix
inner_key = key_lower[5:] # Remove 'oddt_'
# Try to match with known ODDT patterns
if inner_key.startswith('rfscore_v'):
# Extract version number (handle both 'rfscore_v1' and 'rfscorev1' formats)
version = inner_key.replace('rfscore_v', '').replace('rfscorev', '')
return f'ODDT_RFSCORE_V{version.upper()}'
elif 'plec' in inner_key.lower():
# Handle PLEC variations (case-insensitive)
inner_key_lower = inner_key.lower()
if 'p5_l1_s65536' in inner_key_lower or 'p5l1s65536' in inner_key_lower:
return 'ODDT_PLECRF_P5_L1_S65536'
elif inner_key.lower() == 'nnscore':
return 'ODDT_NNSCORE'
# Handle new format: rescoring_{scoring_function}_{pose_number} or rescoring_{pose_number}
if key_lower.startswith('rescoring_'):
# Extract scoring function and pose number
# Format: rescoring_{scoring_function}_{pose_number} or rescoring_{pose_number}
parts = key_lower.split('_')
if len(parts) >= 2:
# Check if last part is a number (pose number)
if parts[-1].isdigit():
pose_number = parts[-1]
# Remove 'rescoring' and pose number, rest is scoring function
scoring_function_parts = parts[1:-1]
else:
# No pose number, just scoring function after 'rescoring'
scoring_function_parts = parts[1:]
if scoring_function_parts:
# Reconstruct scoring function name
scoring_function = '_'.join(scoring_function_parts)
# Use provided engine if available, otherwise try to detect
engines_to_try = [engine] if engine else ['vina', 'smina', 'gnina']
# Try to match with known formats
for eng in engines_to_try:
test_key_old = f'{eng}_{scoring_function}_rescoring'
test_key_new = f'{eng}_{scoring_function}'
if test_key_old in mapping:
return mapping[test_key_old]
if test_key_new in mapping:
return mapping[test_key_new]
# If not found in mapping, construct based on engine or scoring function
if engine:
# Use provided engine
if engine == 'vina':
return f'VINA_{scoring_function.upper()}'
elif engine == 'gnina':
sf_mapping = {
'dkoes_scoring': 'SCORING_DKOES',
'scoring_dkoes': 'SCORING_DKOES',
'old_scoring_dkoes': 'OLD_SCORING_DKOES',
'dkoes_scoring_old': 'OLD_SCORING_DKOES',
'fast_dkoes': 'FAST_DKOES',
'ad4_scoring': 'SCORING_AD4',
}
if scoring_function in sf_mapping:
return f'GNINA_{sf_mapping[scoring_function]}'
return f'GNINA_{scoring_function.upper()}'
elif engine == 'smina':
sf_mapping = {
'dkoes_scoring': 'SCORING_DKOES',
'scoring_dkoes': 'SCORING_DKOES',
'old_scoring_dkoes': 'OLD_SCORING_DKOES',
'dkoes_scoring_old': 'OLD_SCORING_DKOES',
'fast_dkoes': 'FAST_DKOES',
'ad4_scoring': 'SCORING_AD4',
}
if scoring_function in sf_mapping:
return f'SMINA_{sf_mapping[scoring_function]}'
return f'SMINA_{scoring_function.upper()}'
else:
# Try to detect engine from scoring function
if scoring_function in ['vina', 'vinardo']:
return f'VINA_{scoring_function.upper()}'
else:
# Assume smina for other scoring functions
sf_mapping = {
'dkoes_scoring': 'SCORING_DKOES',
'scoring_dkoes': 'SCORING_DKOES',
'old_scoring_dkoes': 'OLD_SCORING_DKOES',
'dkoes_scoring_old': 'OLD_SCORING_DKOES',
'fast_dkoes': 'FAST_DKOES',
'ad4_scoring': 'SCORING_AD4',
}
if scoring_function in sf_mapping:
return f'SMINA_{sf_mapping[scoring_function]}'
return f'SMINA_{scoring_function.upper()}'
else:
# Just 'rescoring_{pose_number}' - no scoring function specified
# Use engine if provided, otherwise default to vina_vina
if engine == 'smina':
return 'SMINA_VINARDO' # Default SMINA scoring function
if engine == 'gnina':
return 'GNINA_VINARDO' # Default GNINA scoring function
return 'VINA_VINA' # Default VINA scoring function
# Handle old format: VINA/SMINA rescoring keys (remove _rescoring suffix if present)
if key_lower.endswith('_rescoring'):
key_without_suffix = key_lower[:-10] # Remove '_rescoring'
if key_without_suffix in mapping:
return mapping[key_without_suffix]
# Try to construct the mapping
if key_without_suffix.startswith('vina_'):
sf = key_without_suffix.replace('vina_', '')
return f'VINA_{sf.upper()}'
elif key_without_suffix.startswith('gnina_'):
sf = key_without_suffix.replace('gnina_', '')
sf_mapping = {
'dkoes_scoring': 'SCORING_DKOES',
'scoring_dkoes': 'SCORING_DKOES',
'old_scoring_dkoes': 'OLD_SCORING_DKOES',
'dkoes_scoring_old': 'OLD_SCORING_DKOES',
'fast_dkoes': 'FAST_DKOES',
'dkoes_fast': 'FAST_DKOES',
'ad4_scoring': 'SCORING_AD4',
'vina': 'VINA',
'vinardo': 'VINARDO',
}
if sf in sf_mapping:
return f'GNINA_{sf_mapping[sf]}'
return f'GNINA_{sf.upper()}'
elif key_without_suffix.startswith('smina_'):
sf = key_without_suffix.replace('smina_', '')
# Handle special SMINA scoring function names
# !!! CRITICAL: These must match SCORING_FUNCTION_ORDER !!!
sf_mapping = {
'dkoes_scoring': 'SCORING_DKOES',
'scoring_dkoes': 'SCORING_DKOES', # Alternative format
'old_scoring_dkoes': 'OLD_SCORING_DKOES',
'dkoes_scoring_old': 'OLD_SCORING_DKOES', # Alternative format
'fast_dkoes': 'FAST_DKOES',
'dkoes_fast': 'FAST_DKOES', # Alternative format
'ad4_scoring': 'SCORING_AD4',
'vina': 'VINA',
'vinardo': 'VINARDO',
}
if sf in sf_mapping:
return f'SMINA_{sf_mapping[sf]}'
else:
return f'SMINA_{sf.upper()}'
# If no mapping found, return uppercase version of the key
return key.upper()
CHECKPOINT_STAGE_DOCKING = "docking"
CHECKPOINT_STAGE_RESCORING = "rescoring"
CHECKPOINT_STAGE_FEATURES = "features"
def _enabled_docking_engines() -> Tuple[str, ...]:
engines = []
if RUN_VINA_DOCKING:
engines.append("vina")
if RUN_PLANTS_DOCKING:
engines.append("plants")
if RUN_GNINA_DOCKING:
engines.append("gnina")
return tuple(engines)
def _enabled_rescoring_engines() -> Tuple[str, ...]:
engines = []
if RUN_ODDT_RESCORING:
engines.append("oddt")
if RUN_PLANTS_RESCORING:
engines.append("plants")
if RUN_VINA_RESCORING:
engines.append("vina")
if RUN_SMINA_RESCORING:
engines.append("smina")
if RUN_GNINA_RESCORING:
engines.append("gnina")
return tuple(engines)
def _validate_pipeline_config() -> None:
if not _enabled_rescoring_engines():
raise ValueError(
"At least one rescoring engine must be enabled (ODDT, PLANTS, Vina, SMINA, or GNINA)."
)
if RUN_ODDT_RESCORING:
has_prepared_receptor = bool(PREPARED_RECEPTOR_PDBQT and os.path.isfile(PREPARED_RECEPTOR_PDBQT))
has_pdbqt_preparer = any([
RUN_VINA_DOCKING,
RUN_VINA_RESCORING,
RUN_SMINA_RESCORING,
RUN_GNINA_DOCKING,
RUN_GNINA_RESCORING,
])
if not has_prepared_receptor and not has_pdbqt_preparer:
raise ValueError(
"ODDT rescoring requires PREPARED_RECEPTOR_PDBQT or a PDBQT-capable engine "
"(Vina, SMINA, or GNINA docking/rescoring)."
)
def _docking_outputs_ready(vina_poses: list, plants_poses: list, gnina_poses: list) -> bool:
if not _enabled_docking_engines():
return bool(vina_poses or plants_poses or gnina_poses)
ready = True
if RUN_VINA_DOCKING:
ready = ready and bool(vina_poses)
if RUN_PLANTS_DOCKING:
ready = ready and bool(plants_poses)
if RUN_GNINA_DOCKING:
ready = ready and bool(gnina_poses)
return ready
def _checkpoint_medoid(checkpoint: Dict[str, Any]) -> Optional[str]:
stage_data = checkpoint.get("stage_data", {})
if not isinstance(stage_data, dict):
return None
rescoring_data = stage_data.get(CHECKPOINT_STAGE_RESCORING, {})
if not isinstance(rescoring_data, dict):
return None
medoid = rescoring_data.get("medoid")
if isinstance(medoid, str) and medoid and os.path.isfile(medoid):
return medoid
return None
def _rescoring_engine_outputs_present(
ligand_path: str,
ligand_name: str,
engine: str,
vina_ligand: Optional[ocvina.Vina] = None,
plants_ligand: Optional[ocplants.PLANTS] = None,
smina_ligand: Optional[ocsmina.Smina] = None,
gnina_ligand: Optional[ocgnina.Gnina] = None,
) -> bool:
if engine == "oddt":
return os.path.isfile(os.path.join(ligand_path, "oddt", f"{ligand_name}.csv"))
if engine == "plants":
if plants_ligand is None:
return False
try:
return bool(plants_ligand.read_rescore_logs(f"{ligand_path}/plantsFiles"))
except Exception:
return False
if engine == "vina":
if vina_ligand is None:
return False
try:
return bool(vina_ligand.read_rescore_logs(f"{ligand_path}/vinaFiles/rescoring"))
except Exception:
return False
if engine == "smina":
if smina_ligand is None:
return False
try:
return bool(smina_ligand.read_rescore_logs(f"{ligand_path}/sminaFiles/rescoring"))
except Exception:
return False
if engine == "gnina":
if gnina_ligand is None:
return False
try:
return bool(gnina_ligand.read_rescore_logs(f"{ligand_path}/gninaFiles/rescoring"))
except Exception:
return False
return False
def _enabled_rescoring_outputs_present(
ligand_path: str,
ligand_name: str,
vina_ligand: Optional[ocvina.Vina] = None,
plants_ligand: Optional[ocplants.PLANTS] = None,
smina_ligand: Optional[ocsmina.Smina] = None,
gnina_ligand: Optional[ocgnina.Gnina] = None,
) -> bool:
enabled = _enabled_rescoring_engines()
if not enabled:
return True
return all(
_rescoring_engine_outputs_present(
ligand_path,
ligand_name,
engine,
vina_ligand,
plants_ligand,
smina_ligand,
gnina_ligand,
)
for engine in enabled
)
def _read_preserved_rescoring_outputs(
ligand_path: str,
ligand_name: str,
vina_ligand: Optional[ocvina.Vina] = None,
plants_ligand: Optional[ocplants.PLANTS] = None,
smina_ligand: Optional[ocsmina.Smina] = None,
gnina_ligand: Optional[ocgnina.Gnina] = None,
) -> Dict[str, Any]:
"""Read on-disk rescoring for engines that are disabled this run."""
if not MERGE_EXISTING_RESCORING:
return {}
preserved_engines = tuple(
engine for engine in ("oddt", "plants", "vina", "smina", "gnina")
if engine not in _enabled_rescoring_engines()
)
if not preserved_engines:
return {}
preserved: Dict[str, Any] = {}
for engine in preserved_engines:
if not _rescoring_engine_outputs_present(
ligand_path,
ligand_name,
engine,
vina_ligand,
plants_ligand,
smina_ligand,
gnina_ligand,
):
continue
partial = _read_single_engine_rescoring(
ligand_path,
ligand_name,
engine,
vina_ligand,
plants_ligand,
smina_ligand,
gnina_ligand,
)
if partial:
preserved.update(partial)
return preserved
def _read_single_engine_rescoring(
ligand_path: str,
ligand_name: str,
engine: str,
vina_ligand: Optional[ocvina.Vina] = None,
plants_ligand: Optional[ocplants.PLANTS] = None,
smina_ligand: Optional[ocsmina.Smina] = None,
gnina_ligand: Optional[ocgnina.Gnina] = None,
) -> Dict[str, Any]:
result: Dict[str, Any] = {}
if engine == "oddt":
oddt_csv = os.path.join(ligand_path, "oddt", f"{ligand_name}.csv")
if not os.path.isfile(oddt_csv):
return result
try:
oddt_df = pd.read_csv(oddt_csv)
oddt_dict = ocoddt.df_to_dict(oddt_df)
if not oddt_dict:
return result
first_key = list(oddt_dict.keys())[0]
oddt_scores = oddt_dict[first_key]
if isinstance(oddt_scores, dict):
for key, value in oddt_scores.items():
result[map_rescoring_key_to_db_column(f"oddt_{key}")] = value
except Exception:
return result
elif engine == "plants" and plants_ligand is not None:
try:
plants_rescoring = plants_ligand.read_rescore_logs(f"{ligand_path}/plantsFiles")
if plants_rescoring:
_add_plants_scores_to_rescoring_result(plants_rescoring, result)
except Exception:
return result
elif engine == "vina" and vina_ligand is not None:
try:
vina_rescoring = vina_ligand.read_rescore_logs(f"{ligand_path}/vinaFiles/rescoring")
for key, value in (vina_rescoring or {}).items():
result[map_rescoring_key_to_db_column(key, engine="vina")] = _value_from_rescore_entry(value)
except Exception:
return result
elif engine == "smina" and smina_ligand is not None:
try:
smina_rescoring = smina_ligand.read_rescore_logs(f"{ligand_path}/sminaFiles/rescoring")
for key, value in (smina_rescoring or {}).items():
result[map_rescoring_key_to_db_column(key, engine="smina")] = _value_from_rescore_entry(value)
except Exception:
return result
elif engine == "gnina" and gnina_ligand is not None:
try:
gnina_rescoring = gnina_ligand.read_rescore_logs(f"{ligand_path}/gninaFiles/rescoring")
for key, value in (gnina_rescoring or {}).items():
result[map_rescoring_key_to_db_column(key, engine="gnina")] = _value_from_rescore_entry(value)
except Exception:
return result
return _normalize_sf_names(result)
def _merge_rescoring_results(*parts: Optional[Dict[str, Any]]) -> Dict[str, Any]:
merged: Dict[str, Any] = {}
for part in parts:
if isinstance(part, dict):
merged.update(part)
return merged
def _engine_output_dir_exists(ligand_path: str, engine: str) -> bool:
dirs = {
"vina": f"{ligand_path}/vinaFiles",
"plants": f"{ligand_path}/plantsFiles",
"gnina": f"{ligand_path}/gninaFiles",
"smina": f"{ligand_path}/sminaFiles",
}
path = dirs.get(engine, "")
return bool(path and os.path.isdir(path))
def _prepare_receptor_pdbqt_for_oddt(
vina_ligand: Optional[ocvina.Vina],
smina_ligand: Optional[ocsmina.Smina],
gnina_ligand: Optional[ocgnina.Gnina],
) -> str:
if PREPARED_RECEPTOR_PDBQT and os.path.isfile(PREPARED_RECEPTOR_PDBQT):
return PREPARED_RECEPTOR_PDBQT
for ligand_obj in (vina_ligand, smina_ligand, gnina_ligand):
if ligand_obj is None:
continue
ligand_obj.run_prepare_receptor(overwrite=False)
prepared = getattr(ligand_obj, "prepared_receptor", "")
if prepared and os.path.isfile(prepared):
return prepared
raise ValueError(
"ODDT rescoring requires a prepared receptor (PDBQT). "
"Set PREPARED_RECEPTOR_PDBQT or enable Vina, SMINA, or GNINA rescoring."
)
def _checkpoint_path(ligand_path: str) -> str:
return os.path.join(ligand_path, LIGAND_CHECKPOINT_FILE)
def _features_cache_path(ligand_path: str) -> str:
return os.path.join(ligand_path, LIGAND_FEATURES_CACHE_FILE)
def _write_json_atomic(path: str, payload: Dict[str, Any]) -> None:
tmp_path = f"{path}.tmp"
with open(tmp_path, "w", encoding="utf-8") as handle:
json.dump(payload, handle, indent=2, sort_keys=True, default=_json_default)
os.replace(tmp_path, path)
def _json_default(value: Any) -> Any:
if isinstance(value, np.generic):
return value.item()
if isinstance(value, (set, tuple)):
return list(value)
return str(value)
def _default_checkpoint(ligand_name: str) -> Dict[str, Any]:
return {
"version": 1,
"ligand": ligand_name,
"status": "new",
"completed_stages": [],
"stage_data": {},
"failed_stage": "",
"last_error": "",
"updated_at": time.time(),
}
def _load_checkpoint(ligand_path: str, ligand_name: str) -> Dict[str, Any]:
checkpoint = _default_checkpoint(ligand_name)
if not ENABLE_LIGAND_CHECKPOINT:
return checkpoint
path = _checkpoint_path(ligand_path)
if not os.path.isfile(path):
return checkpoint
try:
with open(path, "r", encoding="utf-8") as handle:
loaded = json.load(handle)
except Exception:
return checkpoint
if not isinstance(loaded, dict):
return checkpoint
checkpoint["status"] = str(loaded.get("status", checkpoint["status"]))
checkpoint["failed_stage"] = str(loaded.get("failed_stage", ""))
checkpoint["last_error"] = str(loaded.get("last_error", ""))
checkpoint["updated_at"] = loaded.get("updated_at", checkpoint["updated_at"])
completed = loaded.get("completed_stages", [])
if isinstance(completed, list):
checkpoint["completed_stages"] = sorted({str(stage) for stage in completed})
stage_data = loaded.get("stage_data", {})
if isinstance(stage_data, dict):
checkpoint["stage_data"] = stage_data
return checkpoint
def _save_checkpoint(ligand_path: str, checkpoint: Dict[str, Any]) -> None:
if not ENABLE_LIGAND_CHECKPOINT:
return
checkpoint["updated_at"] = time.time()
_write_json_atomic(_checkpoint_path(ligand_path), checkpoint)
def _is_stage_complete(checkpoint: Dict[str, Any], stage: str) -> bool:
completed = checkpoint.get("completed_stages", [])
return stage in completed if isinstance(completed, list) else False
def _clear_stage(checkpoint: Dict[str, Any], stage: str) -> None:
completed = checkpoint.get("completed_stages", [])
if isinstance(completed, list):
checkpoint["completed_stages"] = sorted([st for st in completed if st != stage])
stage_data = checkpoint.get("stage_data", {})
if isinstance(stage_data, dict):
stage_data.pop(stage, None)
checkpoint["status"] = "in_progress"
def _mark_stage_complete(checkpoint: Dict[str, Any], stage: str, stage_data: Optional[Dict[str, Any]] = None) -> None:
completed = checkpoint.get("completed_stages", [])
completed_set = set(completed if isinstance(completed, list) else [])
completed_set.add(stage)
checkpoint["completed_stages"] = sorted(completed_set)
checkpoint["status"] = "in_progress" if stage != CHECKPOINT_STAGE_FEATURES else "completed"
checkpoint["failed_stage"] = ""
checkpoint["last_error"] = ""
if stage_data:
checkpoint.setdefault("stage_data", {})[stage] = stage_data
def _mark_stage_failed(checkpoint: Dict[str, Any], stage: str, message: str) -> None:
checkpoint["status"] = "failed"
checkpoint["failed_stage"] = stage
checkpoint["last_error"] = message
def _load_cached_features(ligand_path: str) -> Optional[Dict[str, Any]]:
path = _features_cache_path(ligand_path)
if not os.path.isfile(path):
return None
try:
with open(path, "r", encoding="utf-8") as handle:
payload = json.load(handle)
if isinstance(payload, dict):
return payload
except Exception:
return None
return None
def _save_cached_features(ligand_path: str, features: Dict[str, Any]) -> None:
if not ENABLE_LIGAND_CHECKPOINT:
return
_write_json_atomic(_features_cache_path(ligand_path), features)
def _normalize_sf_names(rescoring_result: Dict[str, Any]) -> Dict[str, Any]:
sf_name_corrections = {
"SMINA_DKOES_FAST": "SMINA_FAST_DKOES",
"SMINA_DKOES_SCORING_OLD": "SMINA_OLD_SCORING_DKOES",
"SMINA_SCORING_DKOES_OLD": "SMINA_OLD_SCORING_DKOES",
"SMINA_FAST_DKOES_RESCORING": "SMINA_FAST_DKOES",
"SMINA_OLD_SCORING_DKOES_RESCORING": "SMINA_OLD_SCORING_DKOES",
}
for old_name, correct_name in sf_name_corrections.items():
if old_name in rescoring_result and correct_name not in rescoring_result:
rescoring_result[correct_name] = rescoring_result.pop(old_name)
return rescoring_result
def _value_from_rescore_entry(value: Any) -> Any:
if isinstance(value, list) and len(value) > 0:
return value[0]
return value
def _add_plants_scores_to_rescoring_result(plants_rescoring: Dict[str, Any], rescoring_result: Dict[str, Any]) -> None:
for outer_key, inner_dict in plants_rescoring.items():
db_column_name = map_rescoring_key_to_db_column(outer_key)
if isinstance(inner_dict, dict):
if "PLANTS_TOTAL_SCORE" in inner_dict:
total_score = inner_dict["PLANTS_TOTAL_SCORE"]
rescoring_result[db_column_name] = _value_from_rescore_entry(total_score)
else:
print(
f"Warning: PLANTS_TOTAL_SCORE not found in inner dict for {outer_key}. "
f"Available keys: {list(inner_dict.keys())}"
)
else:
rescoring_result[db_column_name] = _value_from_rescore_entry(inner_dict)
def _read_rescoring_outputs(
ligand_path: str,
ligand_name: str,
vina_ligand: Optional[ocvina.Vina] = None,
plants_ligand: Optional[ocplants.PLANTS] = None,
smina_ligand: Optional[ocsmina.Smina] = None,
gnina_ligand: Optional[ocgnina.Gnina] = None,
engines: Optional[Tuple[str, ...]] = None,
) -> Optional[Dict[str, Any]]:
enabled = engines or _enabled_rescoring_engines()
if not enabled:
return {}
rescoring_result: Dict[str, Any] = {}
for engine in enabled:
partial = _read_single_engine_rescoring(
ligand_path,
ligand_name,
engine,
vina_ligand,
plants_ligand,
smina_ligand,
gnina_ligand,
)
if not partial:
return None
rescoring_result.update(partial)
if not rescoring_result:
return None
return _normalize_sf_names(rescoring_result)
def _run_vina_docking(ligand_path: str, vina_ligand: ocvina.Vina) -> list:
vina_ligand.run_prepare_receptor(overwrite=True)
vina_ligand.run_prepare_ligand(overwrite=True)
vina_ligand.run_docking(overwrite=True)
vina_ligand.split_poses()
vina_poses_dir = (
os.path.dirname(vina_ligand.output_vina)
if hasattr(vina_ligand, "output_vina")
else f"{ligand_path}/vinaFiles"
)
vina_pattern = f"{vina_poses_dir}/*_split_*.pdbqt"
pose_files_found = False
for _ in range(50):
found_files = glob(vina_pattern)
if found_files and wait_for_files_ready(found_files, max_wait=2.0):
pose_files_found = True
break
time.sleep(0.2)
if not pose_files_found:
print("Warning: No stable pose files found for Vina after waiting, proceeding anyway...")
time.sleep(0.5)
vina_poses = vina_ligand.get_docked_poses()
if vina_poses and not wait_for_files_ready(vina_poses, max_wait=5.0):
print("Warning: Some Vina pose files may not be fully ready, but proceeding...")
return vina_poses or []
def _run_plants_docking(ligand_path: str, plants_ligand: ocplants.PLANTS) -> list:
plants_ligand.run_prepare_receptor(overwrite=True)
plants_ligand.run_prepare_ligand(overwrite=True)
plants_ligand.run_docking(overwrite=True)
plants_output_dir = (
plants_ligand.output_plants
if hasattr(plants_ligand, "output_plants")
else f"{ligand_path}/plantsFiles"
)
plants_run_dir = os.path.join(plants_output_dir, "run")
plants_pattern = f"{plants_run_dir}/*.mol2"
plants_files_found = False
for _ in range(100):
found_files = [
f for f in glob(plants_pattern)
if not f.endswith("_protein.mol2") and not f.endswith("_fixed.mol2")
]
if found_files and wait_for_files_ready(found_files, max_wait=2.0):
plants_files_found = True
break
time.sleep(0.2)
if not plants_files_found:
print("Warning: No stable PLANTS output files found after waiting, proceeding anyway...")
time.sleep(0.5)
plants_poses = plants_ligand.get_docked_poses()
if plants_poses and not wait_for_files_ready(plants_poses, max_wait=3.0):
print("Warning: Some PLANTS pose files may not be fully ready, but proceeding...")
return plants_poses or []
def _run_gnina_docking(ligand_path: str, gnina_ligand: ocgnina.Gnina) -> list:
gnina_ligand.run_prepare_receptor(overwrite=True)
gnina_ligand.run_prepare_ligand(overwrite=True)
gnina_ligand.run_docking(overwrite=True)
gnina_ligand.split_poses()
gnina_poses_dir = (
os.path.dirname(gnina_ligand.output_gnina)
if hasattr(gnina_ligand, "output_gnina")
else f"{ligand_path}/gninaFiles"
)
gnina_pattern = f"{gnina_poses_dir}/*_split_*.pdbqt"
pose_files_found = False
for _ in range(50):
found_files = glob(gnina_pattern)
if found_files and wait_for_files_ready(found_files, max_wait=2.0):
pose_files_found = True
break
time.sleep(0.2)
if not pose_files_found:
print("Warning: No stable pose files found for GNINA after waiting, proceeding anyway...")
time.sleep(0.5)
gnina_poses = gnina_ligand.get_docked_poses()
if gnina_poses and not wait_for_files_ready(gnina_poses, max_wait=5.0):
print("Warning: Some GNINA pose files may not be fully ready, but proceeding...")
return gnina_poses or []
def _run_docking_stage(
ligand_path: str,
vina_ligand: Optional[ocvina.Vina] = None,
plants_ligand: Optional[ocplants.PLANTS] = None,
gnina_ligand: Optional[ocgnina.Gnina] = None,
) -> Tuple[list, list, list]:
vina_poses: list = []
plants_poses: list = []
gnina_poses: list = []
if RUN_VINA_DOCKING:
if vina_ligand is None:
raise ValueError("Vina docking is enabled but no Vina ligand object was provided.")
vina_poses = _run_vina_docking(ligand_path, vina_ligand)
if not vina_poses:
raise ValueError("Vina docking did not generate valid pose files.")
if RUN_PLANTS_DOCKING:
if plants_ligand is None:
raise ValueError("PLANTS docking is enabled but no PLANTS ligand object was provided.")
plants_poses = _run_plants_docking(ligand_path, plants_ligand)
if not plants_poses:
raise ValueError("PLANTS docking did not generate valid pose files.")
if RUN_GNINA_DOCKING:
if gnina_ligand is None:
raise ValueError("GNINA docking is enabled but no GNINA ligand object was provided.")
gnina_poses = _run_gnina_docking(ligand_path, gnina_ligand)
if not gnina_poses:
raise ValueError("GNINA docking did not generate valid pose files.")
if not _docking_outputs_ready(vina_poses, plants_poses, gnina_poses):
raise ValueError(
f"Docking did not generate valid pose files for enabled engines: {', '.join(_enabled_docking_engines())}."
)
return vina_poses, plants_poses, gnina_poses
def _load_existing_docking_poses(
vina_ligand: Optional[ocvina.Vina] = None,
plants_ligand: Optional[ocplants.PLANTS] = None,
gnina_ligand: Optional[ocgnina.Gnina] = None,
) -> Tuple[list, list, list]:
vina_poses: list = []
plants_poses: list = []
gnina_poses: list = []
if vina_ligand is not None:
vina_poses = vina_ligand.get_docked_poses() or []
if vina_poses:
_ = wait_for_files_ready(vina_poses, max_wait=2.0)
if plants_ligand is not None:
plants_poses = plants_ligand.get_docked_poses() or []
if plants_poses:
_ = wait_for_files_ready(plants_poses, max_wait=2.0)
if gnina_ligand is not None:
gnina_poses = gnina_ligand.get_docked_poses() or []
if gnina_poses:
_ = wait_for_files_ready(gnina_poses, max_wait=2.0)
return vina_poses, plants_poses, gnina_poses
def _select_representative_medoid(
ligand_path: str,
ligand_name: str,
vina_poses: list,
plants_poses: list,
vina_ligand: ocvina.Vina,
plants_ligand: ocplants.PLANTS,
gnina_poses: Optional[list] = None,
gnina_ligand: Optional[ocgnina.Gnina] = None,
) -> str:
gnina_poses = gnina_poses or []
poses_list = vina_poses + plants_poses + gnina_poses
valid_poses = []
max_retries = 5
retry_delay = 0.3
for pose_file in poses_list:
validated = False
for attempt in range(max_retries):
if validate_molecule_file(pose_file):
valid_poses.append(pose_file)
validated = True
break
if attempt < max_retries - 1:
time.sleep(retry_delay * (attempt + 1))
if not validated:
print(f"Warning: Could not validate pose file {pose_file} after {max_retries} attempts, skipping.")
if not valid_poses:
raise ValueError(f"No valid pose files found for ligand {ligand_name} after validation")
if len(valid_poses) == 1:
return valid_poses[0]
if len(valid_poses) < 2:
raise ValueError(f"Need at least 2 valid poses for RMSD calculation, found {len(valid_poses)} for ligand {ligand_name}")
mol2_poses_dir = f"{ligand_path}/poses_mol2"
os.makedirs(mol2_poses_dir, exist_ok=True)
mol2_poses = []
mol2_to_original_map = {}
for pose_file in valid_poses:
pose_ext = os.path.splitext(pose_file)[1].lower()
pose_basename = os.path.basename(pose_file)
if pose_ext == ".mol2":
mol2_poses.append(pose_file)
mol2_to_original_map[pose_file] = pose_file
else:
mol2_path = os.path.join(mol2_poses_dir, f"{os.path.splitext(pose_basename)[0]}.mol2")
occonversion.convert_mols(pose_file, mol2_path, overwrite=True)
if wait_for_file_stable(mol2_path, max_wait=3.0):
mol2_poses.append(mol2_path)
mol2_to_original_map[mol2_path] = pose_file
else:
print(f"Warning: Could not convert {pose_file} to MOL2 format, skipping for RMSD calculation")
if len(mol2_poses) == 1:
return mol2_to_original_map.get(mol2_poses[0], mol2_poses[0])
if len(mol2_poses) < 2:
raise ValueError(f"Need at least 2 valid MOL2 poses for RMSD calculation, found {len(mol2_poses)} for ligand {ligand_name}")
if not wait_for_files_ready(mol2_poses, max_wait=5.0):
print("Warning: Some MOL2 pose files may not be fully ready for RMSD calculation, but proceeding...")
time.sleep(0.5)
rmsd_matrix = ocmolproc.get_rmsd_matrix(mol2_poses)
pose_engine_map = {}
for mol2_pose in mol2_poses:
original_pose = mol2_to_original_map.get(mol2_pose, mol2_pose)
if original_pose in vina_poses:
pose_engine_map[mol2_pose] = "vina"
elif original_pose in plants_poses:
pose_engine_map[mol2_pose] = "plants"
elif original_pose in gnina_poses:
pose_engine_map[mol2_pose] = "gnina"
clusters = ocrmsdclust.cluster_rmsd(
rmsd_matrix,
algorithm="agglomerativeClustering",
outputPlot=f"{ligand_path}/medoids.png",
pose_engine_map=pose_engine_map,
)
medoids_mol2 = ocrmsdclust.get_medoids(rmsd_matrix, clusters, onlyBiggest=True)
medoids = [mol2_to_original_map.get(medoid_mol2, medoid_mol2) for medoid_mol2 in medoids_mol2]
medoids_dict = {}
for medoid in medoids:
if medoid in vina_poses:
medoids_dict[medoid] = vina_ligand.read_log(onlyBest=False)[ocvina.get_pose_index_from_file_path(medoid)]
elif medoid in plants_poses:
medoids_dict[medoid] = plants_ligand.read_log(onlyBest=False)[ocplants.get_pose_index_from_file_path(medoid)]
elif medoid in gnina_poses and gnina_ligand is not None:
medoids_dict[medoid] = gnina_ligand.read_log(onlyBest=False)[ocgnina.get_pose_index_from_file_path(medoid)]
if not medoids_dict:
raise ValueError(f"Could not determine medoid for ligand {ligand_name}")
return list(medoids_dict.keys())[0]
def _run_rescoring_stage(
ligand_path: str,
ligand_name: str,
ligand: ocl.Ligand,
medoid: str,
vina_ligand: Optional[ocvina.Vina] = None,
plants_ligand: Optional[ocplants.PLANTS] = None,
smina_ligand: Optional[ocsmina.Smina] = None,
gnina_ligand: Optional[ocgnina.Gnina] = None,
) -> Dict[str, Any]:
if not _enabled_rescoring_engines():
return {}
if RUN_ODDT_RESCORING:
prepared_receptor = _prepare_receptor_pdbqt_for_oddt(vina_ligand, smina_ligand, gnina_ligand)
oddt_outdir = os.path.join(ligand_path, "oddt")
try:
from joblib import parallel_backend
with parallel_backend("threading"):
_ = ocoddt.run_oddt(
prepared_receptor,
medoid,
ligand.name,
oddt_outdir,
overwrite=True,
)
except ImportError:
_ = ocoddt.run_oddt(
prepared_receptor,
medoid,
ligand.name,
oddt_outdir,
overwrite=True,
)
medoid_extension = os.path.splitext(medoid)[1].lower()
if RUN_PLANTS_RESCORING:
if plants_ligand is None:
raise ValueError("PLANTS rescoring is enabled but no PLANTS ligand object was provided.")
if medoid_extension != ".mol2":
plants_input = medoid.replace(medoid_extension, ".mol2")
occonversion.convert_mols(medoid, plants_input, overwrite=True)
else:
plants_input = medoid
plants_pose_list = f"{ligand_path}/plantsFiles/plants_pose_list.txt"
ocplants.write_pose_list(plants_input, plants_pose_list)
plants_ligand.run_rescore(plants_pose_list, logFile="", overwrite=True)
time.sleep(0.3)
pdbqt_engines = RUN_VINA_RESCORING or RUN_SMINA_RESCORING or RUN_GNINA_RESCORING
if pdbqt_engines:
if medoid_extension != ".pdbqt":
vina_smina_input = medoid.replace(medoid_extension, ".pdbqt")
occonversion.convert_mols(medoid, vina_smina_input, overwrite=True)
else:
vina_smina_input = medoid
if RUN_VINA_RESCORING:
if vina_ligand is None:
raise ValueError("Vina rescoring is enabled but no Vina ligand object was provided.")
vina_ligand.run_rescore(
f"{ligand_path}/vinaFiles/rescoring",
vina_smina_input,
overwrite=True,
splitLigand=False,
)
if RUN_SMINA_RESCORING:
if smina_ligand is None:
raise ValueError("SMINA rescoring is enabled but no SMINA ligand object was provided.")
smina_ligand.run_rescore(
f"{ligand_path}/sminaFiles/rescoring",
vina_smina_input,
overwrite=True,
splitLigand=False,
)
if RUN_GNINA_RESCORING:
if gnina_ligand is None:
raise ValueError("GNINA rescoring is enabled but no GNINA ligand object was provided.")
gnina_ligand.run_rescore(
f"{ligand_path}/gninaFiles/rescoring",
vina_smina_input,
overwrite=True,
splitLigand=False,
)
rescoring_result = _read_rescoring_outputs(
ligand_path,
ligand_name,
vina_ligand,
plants_ligand,
smina_ligand,
gnina_ligand,
)
if rescoring_result is None:
raise ValueError(
f"Rescoring outputs are incomplete or invalid for ligand {ligand_name} "
f"(enabled: {', '.join(_enabled_rescoring_engines())})."
)
return rescoring_result
def _infer_completed_stages(
ligand_path: str,
ligand_name: str,
checkpoint: Dict[str, Any],
vina_ligand: Optional[ocvina.Vina] = None,
plants_ligand: Optional[ocplants.PLANTS] = None,
smina_ligand: Optional[ocsmina.Smina] = None,
gnina_ligand: Optional[ocgnina.Gnina] = None,
) -> None:
if _is_stage_complete(checkpoint, CHECKPOINT_STAGE_FEATURES) and _is_stage_complete(checkpoint, CHECKPOINT_STAGE_RESCORING):
return
if (
isinstance(_load_cached_features(ligand_path), dict)
and _enabled_rescoring_outputs_present(
ligand_path,
ligand_name,
vina_ligand,
plants_ligand,
smina_ligand,
gnina_ligand,
)
and not _is_stage_complete(checkpoint, CHECKPOINT_STAGE_FEATURES)
):
_mark_stage_complete(checkpoint, CHECKPOINT_STAGE_FEATURES)
if _enabled_rescoring_outputs_present(
ligand_path,
ligand_name,
vina_ligand,
plants_ligand,
smina_ligand,
gnina_ligand,
):
if not _is_stage_complete(checkpoint, CHECKPOINT_STAGE_RESCORING):
_mark_stage_complete(checkpoint, CHECKPOINT_STAGE_RESCORING)
if not _is_stage_complete(checkpoint, CHECKPOINT_STAGE_DOCKING):
_mark_stage_complete(checkpoint, CHECKPOINT_STAGE_DOCKING)
return
vina_poses, plants_poses, gnina_poses = _load_existing_docking_poses(vina_ligand, plants_ligand, gnina_ligand)
if _docking_outputs_ready(vina_poses, plants_poses, gnina_poses) and not _is_stage_complete(checkpoint, CHECKPOINT_STAGE_DOCKING):
_mark_stage_complete(checkpoint, CHECKPOINT_STAGE_DOCKING)
def process_single_ligand(ligand_path: str, ligand_name: str, receptor: ocr.Receptor) -> Optional[dict]:
''' Process a single ligand through the complete OCScore pipeline with checkpoint/resume support.
Parameters
----------
ligand_path : str
Base path to the ligand directory
ligand_name : str
Name of the ligand (without extension)
receptor : ocr.Receptor
Receptor object
Returns
-------
dict | None
Dictionary containing all features and OCScore prediction. None if processing fails.
'''
checkpoint: Optional[Dict[str, Any]] = None
current_stage = "initialization"
try:
ligand = ocl.Ligand(f"{ligand_path}/{ligand_name}.smi", name=ligand_name)
ligand.create_box(centroid=BOX_CENTER, save_path=f"{ligand_path}/boxes/")
vina_ligand: Optional[ocvina.Vina] = None
if (
RUN_VINA_DOCKING
or RUN_VINA_RESCORING
or (MERGE_EXISTING_RESCORING and _engine_output_dir_exists(ligand_path, "vina"))
):
os.makedirs(f"{ligand_path}/vinaFiles", exist_ok=True)
vina_ligand = ocvina.Vina(
f"{ligand_path}/vinaFiles/conf_vina.txt",
f"{ligand_path}/boxes/box0.pdb",
receptor, PREPARED_RECEPTOR_PDBQT,
ligand, f"{ligand_path}/prepared_ligand.pdbqt",
f"{ligand_path}/vinaFiles/vina.log", f"{ligand_path}/vinaFiles/vina.pdbqt",
name=f"Vina {receptor.name}-{ligand_name}",
)
plants_ligand: Optional[ocplants.PLANTS] = None
if (
RUN_PLANTS_DOCKING
or RUN_PLANTS_RESCORING
or (MERGE_EXISTING_RESCORING and _engine_output_dir_exists(ligand_path, "plants"))
):
os.makedirs(f"{ligand_path}/plantsFiles", exist_ok=True)
plants_ligand = ocplants.PLANTS(
f"{ligand_path}/plantsFiles/conf_plants.txt",
f"{ligand_path}/boxes/box0.pdb",
receptor, PREPARED_RECEPTOR_MOL2,
ligand, f"{ligand_path}/prepared_ligand.mol2",
f"{ligand_path}/plantsFiles/plants.log", f"{ligand_path}/plantsFiles",
name=f"Plants {receptor.name}-{ligand_name}",
)
smina_ligand: Optional[ocsmina.Smina] = None
if RUN_SMINA_RESCORING or (
MERGE_EXISTING_RESCORING and _engine_output_dir_exists(ligand_path, "smina")
):
os.makedirs(f"{ligand_path}/sminaFiles", exist_ok=True)
smina_ligand = ocsmina.Smina(
f"{ligand_path}/sminaFiles/conf_smina.txt",
f"{ligand_path}/boxes/box0.pdb",
receptor, PREPARED_RECEPTOR_PDBQT,
ligand, f"{ligand_path}/prepared_ligand.pdbqt",
f"{ligand_path}/sminaFiles/smina.log", f"{ligand_path}/sminaFiles/smina.pdbqt",
name=f"Smina {receptor.name}-{ligand_name}",
)
gnina_ligand: Optional[ocgnina.Gnina] = None
if RUN_GNINA_DOCKING or RUN_GNINA_RESCORING:
os.makedirs(f"{ligand_path}/gninaFiles", exist_ok=True)
gnina_ligand = ocgnina.Gnina(
f"{ligand_path}/gninaFiles/conf_gnina.txt",
f"{ligand_path}/boxes/box0.pdb",
receptor, PREPARED_RECEPTOR_PDBQT,
ligand, f"{ligand_path}/prepared_ligand.pdbqt",
f"{ligand_path}/gninaFiles/gnina.log", f"{ligand_path}/gninaFiles/gnina.pdbqt",
name=f"Gnina {receptor.name}-{ligand_name}",
)
# ODDT-only: create a minimal PDBQT preparer when no prepared receptor path is configured
if (
RUN_ODDT_RESCORING
and not (PREPARED_RECEPTOR_PDBQT and os.path.isfile(PREPARED_RECEPTOR_PDBQT))
and vina_ligand is None
and smina_ligand is None
and gnina_ligand is None
):
os.makedirs(f"{ligand_path}/vinaFiles", exist_ok=True)
vina_ligand = ocvina.Vina(
f"{ligand_path}/vinaFiles/conf_vina.txt",
f"{ligand_path}/boxes/box0.pdb",
receptor, PREPARED_RECEPTOR_PDBQT,
ligand, f"{ligand_path}/prepared_ligand.pdbqt",
f"{ligand_path}/vinaFiles/vina.log", f"{ligand_path}/vinaFiles/vina.pdbqt",
name=f"Vina {receptor.name}-{ligand_name}",
)
current_stage = "checkpoint"
checkpoint = _load_checkpoint(ligand_path, ligand_name)
resume_enabled = ENABLE_LIGAND_CHECKPOINT
if resume_enabled:
_infer_completed_stages(
ligand_path,
ligand_name,
checkpoint,
vina_ligand,
plants_ligand,
smina_ligand,
gnina_ligand,
)
_save_checkpoint(ligand_path, checkpoint)
if resume_enabled and _is_stage_complete(checkpoint, CHECKPOINT_STAGE_FEATURES):
cached_features = _load_cached_features(ligand_path)
if (
isinstance(cached_features, dict)
and _enabled_rescoring_outputs_present(
ligand_path,
ligand_name,
vina_ligand,
plants_ligand,
smina_ligand,
gnina_ligand,
)
):
print(f"[{ligand_name}] Resume: using cached features.")
return cached_features
print(f"[{ligand_name}] Cached features are stale or missing newly enabled rescoring outputs. Refreshing.")
_clear_stage(checkpoint, CHECKPOINT_STAGE_FEATURES)
_save_checkpoint(ligand_path, checkpoint)
rescoring_result: Optional[Dict[str, Any]] = None
if resume_enabled and _is_stage_complete(checkpoint, CHECKPOINT_STAGE_RESCORING):
if _enabled_rescoring_outputs_present(
ligand_path,
ligand_name,
vina_ligand,
plants_ligand,
smina_ligand,
gnina_ligand,
):
rescoring_result = _read_rescoring_outputs(
ligand_path,
ligand_name,
vina_ligand,
plants_ligand,
smina_ligand,
gnina_ligand,
)
if rescoring_result is None:
print(f"[{ligand_name}] Checkpoint says rescoring complete, but outputs are missing/corrupted. Recomputing.")
_clear_stage(checkpoint, CHECKPOINT_STAGE_RESCORING)
_clear_stage(checkpoint, CHECKPOINT_STAGE_FEATURES)
_save_checkpoint(ligand_path, checkpoint)
else:
print(f"[{ligand_name}] Resume: using existing rescoring outputs for enabled engines.")
else:
print(f"[{ligand_name}] Newly enabled rescoring engine(s) missing outputs. Running rescoring.")
_clear_stage(checkpoint, CHECKPOINT_STAGE_RESCORING)
_clear_stage(checkpoint, CHECKPOINT_STAGE_FEATURES)
_save_checkpoint(ligand_path, checkpoint)
if rescoring_result is None:
vina_poses: list = []
plants_poses: list = []
gnina_poses: list = []
medoid = _checkpoint_medoid(checkpoint)
if medoid is None:
if resume_enabled and _is_stage_complete(checkpoint, CHECKPOINT_STAGE_DOCKING):
vina_poses, plants_poses, gnina_poses = _load_existing_docking_poses(
vina_ligand,
plants_ligand,
gnina_ligand,
)
if not _docking_outputs_ready(vina_poses, plants_poses, gnina_poses):
print(f"[{ligand_name}] Checkpoint says docking complete, but poses are missing. Re-running docking.")
_clear_stage(checkpoint, CHECKPOINT_STAGE_DOCKING)
_save_checkpoint(ligand_path, checkpoint)
else:
print(f"[{ligand_name}] Resume: using existing docking outputs.")
if not _docking_outputs_ready(vina_poses, plants_poses, gnina_poses):
if not _enabled_docking_engines():
vina_poses, plants_poses, gnina_poses = _load_existing_docking_poses(
vina_ligand,
plants_ligand,
gnina_ligand,
)
if not _docking_outputs_ready(vina_poses, plants_poses, gnina_poses):
if not _enabled_docking_engines():
raise ValueError(
f"No docking poses found for {ligand_name}. "
"Enable a docking engine or run docking first."
)
current_stage = CHECKPOINT_STAGE_DOCKING
vina_poses, plants_poses, gnina_poses = _run_docking_stage(
ligand_path,
vina_ligand,
plants_ligand,
gnina_ligand,
)
_mark_stage_complete(checkpoint, CHECKPOINT_STAGE_DOCKING)
_save_checkpoint(ligand_path, checkpoint)
current_stage = CHECKPOINT_STAGE_RESCORING
medoid = _select_representative_medoid(
ligand_path,
ligand_name,
vina_poses,
plants_poses,
vina_ligand,
plants_ligand,
gnina_poses=gnina_poses,
gnina_ligand=gnina_ligand,
)
else:
current_stage = CHECKPOINT_STAGE_RESCORING
print(f"[{ligand_name}] Resume: using checkpoint medoid {medoid}.")
rescoring_result = _run_rescoring_stage(
ligand_path,
ligand_name,
ligand,
medoid,
vina_ligand,
plants_ligand,
smina_ligand,
gnina_ligand,
)
_mark_stage_complete(checkpoint, CHECKPOINT_STAGE_RESCORING, stage_data={"medoid": medoid})
_save_checkpoint(ligand_path, checkpoint)
preserved_rescoring = _read_preserved_rescoring_outputs(
ligand_path,
ligand_name,
vina_ligand,
plants_ligand,
smina_ligand,
gnina_ligand,
)
rescoring_result = _merge_rescoring_results(preserved_rescoring, rescoring_result)
current_stage = CHECKPOINT_STAGE_FEATURES
cached_features = _load_cached_features(ligand_path) if MERGE_EXISTING_RESCORING else None
receptor_descriptors = receptor.get_descriptors()
ligand_descriptors = ligand.get_descriptors()
all_features: Dict[str, Any] = {}
if isinstance(cached_features, dict):
all_features.update(cached_features)
all_features.update(rescoring_result if rescoring_result is not None else {})
all_features.update(receptor_descriptors)
all_features.update(ligand_descriptors)
all_features["name"] = f"{receptor.name}_{ligand.name}"
all_features["receptor"] = receptor.name
all_features["ligand"] = ligand.name
_save_cached_features(ligand_path, all_features)
_mark_stage_complete(checkpoint, CHECKPOINT_STAGE_FEATURES, stage_data={"feature_count": len(all_features)})
_save_checkpoint(ligand_path, checkpoint)
return all_features
except Exception as e:
if checkpoint is not None:
try:
_mark_stage_failed(checkpoint, current_stage, str(e))
_save_checkpoint(ligand_path, checkpoint)
except Exception:
pass
print(f"Error processing ligand {ligand_name}: {e}")
import traceback
traceback.print_exc()
return None
def validate_molecule_file(file_path: str) -> bool:
'''Validate that a molecule file can be loaded and is complete.
Parameters
----------
file_path : str
Path to the molecule file
Returns
-------
bool
True if file is valid and can be loaded
'''
from OCDocker.Toolbox import Validation as ocvalidation
# First check if file is stable (not being written)
if not wait_for_file_stable(file_path, max_wait=2.0):
return False
# Then validate the molecule structure
try:
return ocvalidation.is_molecule_valid(file_path)
except Exception:
return False
def wait_for_file_stable(file_path: str, max_wait: float = 5.0, check_interval: float = 0.1) -> bool:
'''Wait for a file to stabilize (size stops changing).
Parameters
----------
file_path : str
Path to the file to check
max_wait : float
Maximum time to wait in seconds
check_interval : float
Time between checks in seconds
Returns
-------
bool
True if file stabilized, False if timeout
'''
if not os.path.isfile(file_path):
return False
start_time = time.time()
last_size = -1
stable_count = 0
required_stable_checks = 3 # File must be stable for 3 consecutive checks
while time.time() - start_time < max_wait:
try:
current_size = os.path.getsize(file_path)
if current_size == last_size:
stable_count += 1
if stable_count >= required_stable_checks:
return True
else:
stable_count = 0
last_size = current_size
time.sleep(check_interval)
except (OSError, IOError):
# File might be locked or deleted
time.sleep(check_interval)
continue
return False
def wait_for_files_ready(file_paths: list, max_wait: float = 8.0, check_interval: float = 0.2) -> bool:
'''Wait for all files in a list to exist, be stable, and be readable.
Parameters
----------
file_paths : list
List of file paths to wait for
max_wait : float
Maximum time to wait in seconds
check_interval : float
Time between checks in seconds
Returns
-------
bool
True if all files are ready, False if timeout
'''
if not file_paths:
return True
start_time = time.time()
ready_files = set()
while time.time() - start_time < max_wait:
all_ready = True
for file_path in file_paths:
if file_path in ready_files:
continue
if wait_for_file_stable(file_path, max_wait=check_interval * 2, check_interval=check_interval / 2):
ready_files.add(file_path)
else:
all_ready = False
if all_ready and len(ready_files) == len(file_paths):
return True
time.sleep(check_interval)
return len(ready_files) == len(file_paths)
if __name__ == "__main__":
results = main()
This script demonstrates:
Receptor and ligand preparation
Multi-engine docking (Vina, PLANTS)
Pose clustering to find representative poses
Rescoring with multiple scoring functions (ODDT, PLANTS, Vina, SMINA)
Feature extraction (receptor and ligand descriptors)
Model inference using trained OCScore model
Automatic mapping of rescoring results to database column names
Multiprocessing support for processing multiple ligands
Inference from CSV¶
Example of OCScore inference loading features directly from a CSV file:
OCScore inference from CSV¶
#!/usr/bin/env python3
"""
Example: OCScore Inference from CSV
This example demonstrates how to run OCScore model inference using feature data
loaded from a CSV file. The script:
1. Loads OCDocker configuration to enforce `reference_column_order`
2. Loads input CSV data and keeps the original table for output
3. Loads model artifacts (model, optional mask, optional scaler)
4. Runs `ocscoring.get_score(...)` for inference
5. Exports an output CSV preserving original rows/columns plus `OCSCORE`
Usage:
# Recommended: pass config path explicitly
python examples/13_python_api_inference_from_csv.py \
--csv-path /path/to/features.csv \
--model-name OCScore \
--config-path /path/to/OCDocker.cfg \
--output-csv /path/to/scored.csv
# Alternative: use OCDOCKER_CONFIG environment variable
export OCDOCKER_CONFIG=/path/to/OCDocker.cfg
python examples/13_python_api_inference_from_csv.py \
--csv-path /path/to/features.csv \
--model-name OCScore
"""
import argparse
import os
import sys
import numpy as np
import pandas as pd
# Allow running the example directly from source checkout without installation.
_script_dir = os.path.dirname(os.path.abspath(__file__))
_parent_dir = os.path.dirname(_script_dir)
if _parent_dir not in sys.path:
sys.path.insert(0, _parent_dir)
import OCDocker.Config as occonfig
import OCDocker.OCScore.Scoring as ocscoring
import OCDocker.OCScore.Utils.IO as ocscoreio
def _resolve_model_path(model_name: str, models_dir: str, model_path: str | None) -> str:
"""Resolve the model path from an explicit path or common model-name patterns."""
if model_path:
if not os.path.isfile(model_path):
raise FileNotFoundError(f"Model file not found: {model_path}")
return model_path
candidates = [
os.path.join(models_dir, f"{model_name}.pt"),
os.path.join(models_dir, f"{model_name}.pth"),
os.path.join(models_dir, f"{model_name}.pkl"),
]
for candidate in candidates:
if os.path.isfile(candidate):
return candidate
raise FileNotFoundError(
f"Could not find model for '{model_name}' in {models_dir}. "
"Tried .pt, .pth, and .pkl."
)
def _resolve_config_path(config_path: str | None) -> str | None:
"""Resolve config path from argument/env/common local locations."""
if config_path:
candidate = os.path.abspath(config_path)
if not os.path.isfile(candidate):
raise FileNotFoundError(f"Config file not found: {candidate}")
return candidate
env_cfg = os.getenv("OCDOCKER_CONFIG", "").strip()
if env_cfg:
env_candidate = os.path.abspath(env_cfg)
if os.path.isfile(env_candidate):
return env_candidate
repo_cfg = os.path.join(_parent_dir, "OCDocker.cfg")
if os.path.isfile(repo_cfg):
return os.path.abspath(repo_cfg)
cwd_cfg = os.path.abspath("OCDocker.cfg")
if os.path.isfile(cwd_cfg):
return cwd_cfg
return None
def _ensure_reference_order_config(config_path: str | None) -> str | None:
"""Load config into OCDocker singleton so get_score can enforce column order."""
resolved = _resolve_config_path(config_path)
if resolved is None:
return None
os.environ["OCDOCKER_CONFIG"] = resolved
loaded = occonfig.OCDockerConfig.from_config_file(resolved)
occonfig.set_config(loaded)
return resolved
def main() -> None:
'''Run OCScore inference on a feature CSV and print the resulting scores.'''
parser = argparse.ArgumentParser(description="Run OCScore inference from a CSV file.")
parser.add_argument(
"--csv-path",
required=True,
help="Path to input CSV file containing features for inference.",
)
parser.add_argument(
"--model-name",
default="OCScore",
help="Model name used to locate model/mask/scaler files (default: OCScore).",
)
parser.add_argument(
"--models-dir",
default=None,
help="Directory containing model artifacts. Defaults to OCScore_models.",
)
parser.add_argument(
"--model-path",
default=None,
help="Explicit model path. If omitted, inferred from --model-name in --models-dir.",
)
parser.add_argument(
"--config-path",
default=None,
help="Path to OCDocker config file (.cfg/.yml). Needed to enforce reference_column_order.",
)
parser.add_argument(
"--scaler-path",
default=None,
help="Optional scaler path. If omitted, tries <model_name>_scaler.pkl in models dir.",
)
parser.add_argument(
"--no-mask",
action="store_true",
help="Disable mask loading and run inference without a mask.",
)
parser.add_argument(
"--use-gpu",
action="store_true",
help="Use GPU when available (PyTorch models).",
)
parser.add_argument(
"--output-csv",
default="predictions_from_csv.csv",
help="Output CSV file for predictions (default: predictions_from_csv.csv).",
)
parser.add_argument(
"--score-column",
default="OCSCORE",
help="Column name to store predictions in the output CSV (default: OCSCORE).",
)
parser.add_argument(
"--disable-reference-order",
action="store_true",
help="Disable enforcement of reference_column_order (not recommended).",
)
parser.add_argument(
"--invert-conditionally",
dest="invert_conditionally",
action="store_true",
default=False,
help="Invert VINA/SMINA/PLANTS-like columns before inference (default: enabled).",
)
args = parser.parse_args()
if not os.path.isfile(args.csv_path):
raise FileNotFoundError(f"Input CSV file not found: {args.csv_path}")
enforce_reference_order = not args.disable_reference_order
loaded_config_path = _ensure_reference_order_config(args.config_path)
if enforce_reference_order:
cfg = occonfig.get_config()
if not cfg.paths.reference_column_order:
raise ValueError(
"reference_column_order is not set in the active config. "
"Pass --config-path /path/to/OCDocker.cfg (or set OCDOCKER_CONFIG)."
)
print(f"Using OCDocker config: {loaded_config_path if loaded_config_path else 'active runtime config'}")
# Keep a full copy for output (same input rows/columns + OCScore column).
input_df = pd.read_csv(args.csv_path)
# Use OCScore loader for inference input (drops rows with invalid NaNs).
scoring_df = ocscoreio.load_data(args.csv_path)
if scoring_df.empty:
raise ValueError("No valid rows found for inference after CSV preprocessing.")
models_dir = os.path.abspath(args.models_dir) if args.models_dir else ocscoreio.get_models_dir()
model_path = _resolve_model_path(args.model_name, models_dir, args.model_path)
mask = None
if not args.no_mask:
try:
mask = ocscoreio.load_mask(args.model_name, models_dir=models_dir)
print(f"Loaded mask for model '{args.model_name}'.")
except FileNotFoundError:
print(f"Mask file not found for model '{args.model_name}'. Continuing without mask.")
scaler_path = args.scaler_path
if scaler_path is None:
candidate_scaler = os.path.join(models_dir, f"{args.model_name}_scaler.pkl")
if os.path.isfile(candidate_scaler):
scaler_path = candidate_scaler
if scaler_path:
print(f"Using scaler: {scaler_path}")
else:
print("No scaler provided/found. Inference will fit a scaler on the prediction data.")
predictions = ocscoring.get_score(
model_path=model_path,
data=scoring_df,
mask=mask,
scaler_path=scaler_path,
use_gpu=args.use_gpu,
enforce_reference_column_order=enforce_reference_order,
invert_conditionally=args.invert_conditionally,
)
if "predicted_score" not in predictions.columns:
raise ValueError("Prediction output does not contain 'predicted_score' column.")
# Build output with the same rows/columns as input plus an OCScore column.
# Rows removed during preprocessing keep NaN in the score column.
output_df = input_df.copy()
output_df[args.score_column] = np.nan
output_df.loc[scoring_df.index, args.score_column] = predictions["predicted_score"].to_numpy()
print(f"Input shape: {input_df.shape}")
print(f"Output shape: {output_df.shape}")
print(f"Predictions assigned to {predictions['predicted_score'].notna().sum()} rows.")
print(output_df[[args.score_column]].head())
if args.output_csv:
output_df.to_csv(args.output_csv, index=False)
print(f"Saved output CSV to: {args.output_csv}")
if __name__ == "__main__":
main()
This script demonstrates:
Loading input features from a
.csvfileLoading OCDocker config so
reference_column_orderis enforcedResolving model artifacts from
OCScore_models(or a custom directory)Optional mask and scaler loading
Running model inference and exporting an output CSV with original rows/columns plus
OCSCORE
Example command:
python examples/13_python_api_inference_from_csv.py \
--csv-path /path/to/features.csv \
--model-name OCScore \
--config-path /path/to/OCDocker.cfg \
--output-csv /path/to/scored.csv
Staged pipeline CLI¶
Examples 14–16 are also available as ocdocker ocscore subcommands (requires pip install "ocdocker[ml]"):
ocdocker ocscore reduce --pdbbind-archive ... --dudez-archive ... --output-dir ...
ocdocker ocscore train --protocol development --raw-input-dir ... --output-dir ...
ocdocker ocscore score --export-dir ... --raw-archive ... --output-csv ...
See Usage for the full command list.
Reusable SHAP Plots¶
The SHAP plotting utilities are reusable and policy-agnostic. Use
OCDocker.OCScore.Analysis.SHAP.Plots.save_shap_plot_suite with in-memory
arrays, or save_shap_plot_suite_from_paths with explicit SHAP, feature,
metadata, and label paths.
The default feature-family suggestion is available as an editable example:
Default suggested SHAP feature families¶
families:
PMI:
- ligand_PMI*
AUTOCORR2D:
- ligand_AUTOCORR2D_*
VSA/EState:
- ligand_EState_VSA*
- ligand_PEOE_VSA*
- ligand_SMR_VSA*
- ligand_SlogP_VSA*
- ligand_VSA_EState*
shape/size/topology:
- ligand_NPR*
- ligand_Asphericity
- ligand_Eccentricity
- ligand_InertialShapeFactor
- ligand_RadiusOfGyration
- ligand_SpherocityIndex
- ligand_BertzCT
- ligand_MolWt
- ligand_ExactMolWt
- ligand_HeavyAtomMolWt
- ligand_HeavyAtomCount
- ligand_TPSA
- ligand_MolLogP
- ligand_MolMR
- ligand_RingCount
- ligand_Num*Ring*
- ligand_Num*Cycle*
- ligand_NumHAcceptors
- ligand_NumHDonors
- ligand_NumHeteroatoms
- ligand_NumValenceElectrons
- ligand_FractionCSP3
fragments:
- ligand_fr_*
BCUT2D:
- ligand_BCUT2D_*
Chi/Kappa/topological indices:
- ligand_Chi*
- ligand_Kappa*
- ligand_BalabanJ
- ligand_HallKierAlpha
partial-charge / scalar EState:
- ligand_*PartialCharge
- ligand_MaxEStateIndex
- ligand_MinEStateIndex
- ligand_MaxAbsEStateIndex
- ligand_MinAbsEStateIndex
PLANTS:
- plants_*
Vina/Smina:
- vina_*
- smina_*
GNINA:
- gnina_*
ODDT:
- oddt_*
receptor:
- receptor_*
other ligand:
- ligand_*
other:
- "*"
Example CLI usage:
ocdocker ocscore shap \
--export-dir /path/to/best_model \
--reduction-archive /path/to/reduction.tar.gz \
--output-dir /path/to/shap \
--policy full \
--eval-split validation \
--family-spec examples/ocscore_shap_feature_families.yml \
--dependence-features ligand_PMI1 plants_plp
Configuration¶
The complete pipeline script includes a configuration section at the top where you can customize:
Receptor and ligand paths
Model paths and names
Preprocessing settings
Output file paths
Multiprocessing settings
Example configuration:
# Receptor configuration
RECEPTOR_PATH = "/path/to/receptor.pdb"
RECEPTOR_NAME = "Receptor"
# Ligand configuration
LIGAND_PATHS = [
"/path/to/ligand1",
"/path/to/ligand2",
]
# Model configuration
MODEL_NAME = "OCScore"
MODELS_DIR = "OCScore_models"
# Multiprocessing
N_JOBS = 4 # Number of parallel jobs
USE_MULTIPROCESSING = True