import sys from pathlib import Path DIR = Path(__file__).resolve().parent.parent sys.path.insert(0, str(DIR)) import argparse import copy import csv import os from pathlib import Path from types import SimpleNamespace import numpy as np import torch import yaml from rdkit.Chem import RemoveAllHs from torch_geometric.loader import DataLoader from onescience.datapipes.diffdock.process_mols import write_mol_with_coords from onescience.utils.diffdock.diffusion_utils import get_t_schedule from onescience.utils.diffdock.inference_utils import InferenceDataset, set_nones from onescience.utils.diffdock.logging_utils import configure_logger, get_logger from onescience.utils.diffdock.sampling import randomize_position, sampling from onescience.utils.diffdock.validation import validate_sampling_entrypoint try: from models.score_wrapper import load_model_args, load_score_model, model_uses_lm_embeddings except ImportError: from models.score_wrapper import load_model_args, load_score_model, model_uses_lm_embeddings def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("--config", required=True, help="Path to the sampling YAML config.") return parser.parse_args() def _resolve_env_vars(obj): if isinstance(obj, str): return os.path.expandvars(obj) if isinstance(obj, dict): return {k: _resolve_env_vars(v) for k, v in obj.items()} if isinstance(obj, list): return [_resolve_env_vars(v) for v in obj] return obj # def load_config(config_path): # with open(config_path, "r", encoding="utf-8") as handle: # return yaml.safe_load(handle) or {} def load_config(config_path): with open(config_path, "r", encoding="utf-8") as handle: return _resolve_env_vars(yaml.safe_load(handle) or {}) def flatten_config(config): flat = {} for key, value in config.items(): if isinstance(value, dict): flat.update(value) else: flat[key] = value return flat def to_namespace(config): return SimpleNamespace(**config) def resolve_device(device_name): if device_name in {None, "auto"}: return torch.device("cuda" if torch.cuda.is_available() else "cpu") return torch.device(device_name) def load_inputs(args): if args.protein_ligand_csv is not None: with open(args.protein_ligand_csv, "r", encoding="utf-8", newline="") as handle: rows = list(csv.DictReader(handle)) complex_names = set_nones([row.get("complex_name") for row in rows]) protein_paths = set_nones([row.get("protein_path") for row in rows]) protein_sequences = set_nones([row.get("protein_sequence") for row in rows]) ligand_descriptions = set_nones([row.get("ligand_description") for row in rows]) else: complex_names = [args.complex_name or "complex_0"] protein_paths = [args.protein_path] protein_sequences = [args.protein_sequence] ligand_descriptions = [args.ligand_description] complex_names = [name if name is not None else f"complex_{idx}" for idx, name in enumerate(complex_names)] return complex_names, protein_paths, protein_sequences, ligand_descriptions def ensure_output_dirs(out_dir, complex_names): os.makedirs(out_dir, exist_ok=True) for name in complex_names: os.makedirs(os.path.join(out_dir, name), exist_ok=True) def get_ligand_mol(complex_graph): mol = complex_graph.mol return mol[0] if isinstance(mol, (list, tuple)) else mol def resolve_lm_embeddings_flag(args, model_args): lm_embeddings = getattr(args, "lm_embeddings", None) if lm_embeddings is None: return model_uses_lm_embeddings(model_args) return lm_embeddings def build_inference_dataset( args, model_args, *, complex_names, protein_paths, protein_sequences, ligand_descriptions, lm_embeddings, ): return InferenceDataset( out_dir=args.out_dir, complex_names=complex_names, protein_files=protein_paths, ligand_descriptions=ligand_descriptions, protein_sequences=protein_sequences, lm_embeddings=lm_embeddings, receptor_radius=model_args.receptor_radius, remove_hs=model_args.remove_hs, c_alpha_max_neighbors=model_args.c_alpha_max_neighbors, all_atoms=model_args.all_atoms, atom_radius=model_args.atom_radius, atom_max_neighbors=model_args.atom_max_neighbors, knn_only_graph=not getattr(model_args, "not_knn_only_graph", False), ) def inference_graph_signature(model_args, lm_embeddings): return ( getattr(model_args, "receptor_radius", None), getattr(model_args, "remove_hs", None), getattr(model_args, "c_alpha_max_neighbors", None), getattr(model_args, "all_atoms", None), getattr(model_args, "atom_radius", None), getattr(model_args, "atom_max_neighbors", None), not getattr(model_args, "not_knn_only_graph", False), bool(lm_embeddings), ) def extract_confidence_scores(confidence, confidence_model_args): confidence_scores = confidence if isinstance(getattr(confidence_model_args, "rmsd_classification_cutoff", None), list): confidence_scores = confidence_scores[:, 0] return np.asarray(confidence_scores.detach().cpu().numpy()).reshape(-1) def load_optional_confidence_model(args, device, logger): confidence_model_dir = getattr(args, "confidence_model_dir", None) if confidence_model_dir is None: return None, None confidence_model_dir = Path(confidence_model_dir) confidence_checkpoint_path = confidence_model_dir / args.confidence_ckpt if not confidence_checkpoint_path.exists(): raise FileNotFoundError( f"Confidence checkpoint not found: {confidence_checkpoint_path}. " "This example does not auto-download models." ) confidence_model_args_preview = load_model_args(confidence_model_dir) validate_sampling_entrypoint( confidence_model_args_preview, context="DiffDock sampling confidence-model checkpoint", include_confidence=True, confidence_mode=True, ) confidence_model, confidence_model_args, _ = load_score_model( model_dir=confidence_model_dir, ckpt=args.confidence_ckpt, device=device, no_parallel=True, confidence_mode=True, old=getattr(args, "old_confidence_model", False), ) if hasattr(args, "crop_beyond"): confidence_model_args.crop_beyond = args.crop_beyond logger.info("Loaded optional confidence model from %s", confidence_model_dir) return confidence_model, confidence_model_args def main(): parsed = parse_args() raw_config = load_config(parsed.config) args = to_namespace(flatten_config(raw_config)) device = resolve_device(getattr(args, "device", "auto")) validate_sampling_entrypoint( args, include_confidence=( getattr(args, "confidence_model_dir", None) is not None or getattr(args, "old_confidence_model", False) ), ) configure_logger(getattr(args, "loglevel", "INFO")) logger = get_logger() model_dir = Path(args.model_dir) checkpoint_path = model_dir / args.ckpt if not checkpoint_path.exists(): raise FileNotFoundError( f"Checkpoint not found: {checkpoint_path}. This example does not auto-download models." ) score_model_args_preview = load_model_args(model_dir) validate_sampling_entrypoint( score_model_args_preview, context="DiffDock sampling score-model checkpoint", ) model, score_model_args, t_to_sigma = load_score_model( model_dir=model_dir, ckpt=args.ckpt, device=device, no_parallel=True, old=getattr(args, "old_score_model", False), ) if hasattr(args, "crop_beyond"): score_model_args.crop_beyond = args.crop_beyond logger.info("DiffDock sampling will run on %s", device) confidence_model, confidence_model_args = load_optional_confidence_model(args, device, logger) complex_names, protein_paths, protein_sequences, ligand_descriptions = load_inputs(args) ensure_output_dirs(args.out_dir, complex_names) score_lm_embeddings = resolve_lm_embeddings_flag(args, score_model_args) test_dataset = build_inference_dataset( args, score_model_args, complex_names=complex_names, protein_paths=protein_paths, protein_sequences=protein_sequences, ligand_descriptions=ligand_descriptions, lm_embeddings=score_lm_embeddings, ) test_loader = DataLoader(dataset=test_dataset, batch_size=1, shuffle=False) confidence_loader = None if confidence_model is not None: confidence_lm_embeddings = resolve_lm_embeddings_flag(args, confidence_model_args) confidence_needs_independent_graph = ( inference_graph_signature(score_model_args, score_lm_embeddings) != inference_graph_signature(confidence_model_args, confidence_lm_embeddings) ) if confidence_needs_independent_graph: logger.info( "Confidence rerank requires independent inference graphs; building a separate confidence dataset." ) confidence_dataset = build_inference_dataset( args, confidence_model_args, complex_names=complex_names, protein_paths=protein_paths, protein_sequences=protein_sequences, ligand_descriptions=ligand_descriptions, lm_embeddings=confidence_lm_embeddings, ) confidence_loader = DataLoader(dataset=confidence_dataset, batch_size=1, shuffle=False) else: logger.info("Confidence rerank will reuse the score-model inference graphs.") tr_schedule = get_t_schedule( sigma_schedule=args.sigma_schedule, inference_steps=args.inference_steps, inf_sched_alpha=args.inf_sched_alpha, inf_sched_beta=args.inf_sched_beta, ) failures = 0 skipped = 0 num_samples = args.samples_per_complex test_ds_size = len(test_dataset) logger.info("Size of test dataset: %s", test_ds_size) if confidence_loader is None: loader_iter = ((orig_complex_graph, None) for orig_complex_graph in test_loader) else: loader_iter = zip(test_loader, confidence_loader) for idx, (orig_complex_graph, confidence_orig_complex_graph) in enumerate(loader_iter): if not orig_complex_graph.success[0]: skipped += 1 logger.warning( "Skipping %s because preprocessing failed.", test_dataset.complex_names[idx], ) continue if confidence_orig_complex_graph is not None and not confidence_orig_complex_graph.success[0]: skipped += 1 logger.warning( "Skipping %s because confidence preprocessing failed.", test_dataset.complex_names[idx], ) continue try: data_list = [copy.deepcopy(orig_complex_graph) for _ in range(num_samples)] confidence_data_list = None if confidence_orig_complex_graph is not None: confidence_data_list = [copy.deepcopy(confidence_orig_complex_graph) for _ in range(num_samples)] randomize_position( data_list, score_model_args.no_torsion, args.no_random, score_model_args.tr_sigma_max, initial_noise_std_proportion=args.initial_noise_std_proportion, choose_residue=args.choose_residue, ) ligand = get_ligand_mol(orig_complex_graph) data_list, confidence = sampling( data_list=data_list, model=model, inference_steps=args.actual_steps if args.actual_steps is not None else args.inference_steps, tr_schedule=tr_schedule, rot_schedule=tr_schedule, tor_schedule=tr_schedule, device=device, t_to_sigma=t_to_sigma, model_args=score_model_args, no_random=args.no_random, ode=args.ode, confidence_model=confidence_model, confidence_data_list=confidence_data_list, confidence_model_args=confidence_model_args, batch_size=args.batch_size, no_final_step_noise=args.no_final_step_noise, temp_sampling=[ args.temp_sampling_tr, args.temp_sampling_rot, args.temp_sampling_tor, ], temp_psi=[ args.temp_psi_tr, args.temp_psi_rot, args.temp_psi_tor, ], temp_sigma_data=[ args.temp_sigma_data_tr, args.temp_sigma_data_rot, args.temp_sigma_data_tor, ], ) ligand_positions = np.asarray( [ complex_graph["ligand"].pos.cpu().numpy() + orig_complex_graph.original_center.cpu().numpy() for complex_graph in data_list ] ) rerank_order = np.arange(len(ligand_positions)) confidence_scores = None if confidence is not None: confidence_scores = extract_confidence_scores(confidence, confidence_model_args) confidence_scores = np.nan_to_num(confidence_scores, nan=-1e-6) rerank_order = np.argsort(confidence_scores)[::-1] logger.info( "Applied confidence rerank for %s. Ranked confidences: %s", complex_names[idx], np.array2string(confidence_scores[rerank_order], precision=4), ) write_dir = os.path.join(args.out_dir, complex_names[idx]) for rank, sample_idx in enumerate(rerank_order, start=1): pos = ligand_positions[sample_idx] mol_pred = copy.deepcopy(ligand) if score_model_args.remove_hs: mol_pred = RemoveAllHs(mol_pred) filename = f"rank{rank}.sdf" if confidence_scores is not None: filename = f"rank{rank}_conf{confidence_scores[sample_idx]:.4f}.sdf" write_mol_with_coords(mol_pred, pos, os.path.join(write_dir, filename)) except Exception as exc: logger.warning("Failed on %s with error: %s", complex_names[idx], exc) failures += 1 logger.info("Failed for %s / %s complexes.", failures, test_ds_size) logger.info("Skipped %s / %s complexes.", skipped, test_ds_size) logger.info("Results saved in %s", args.out_dir) if __name__ == "__main__": main()