OneScience-Group/DiffDock

1from argparse import Namespace2from functools import partial3from pathlib import Path4 5import torch6import yaml7from torch_geometric.nn.data_parallel import DataParallel8 9from onescience.utils.diffdock.diffusion_utils import get_timestep_embedding, t_to_sigma as t_to_sigma_compl10from onescience.utils.diffdock.utils import ExponentialMovingAverage11 12from .aa_model import AAModel13from .cg_model import CGModel14from .old_aa_model import AAOldModel15 16 17_LM_EMBEDDING_KEYS = (18    "moad_esm_embeddings_path",19    "pdbbind_esm_embeddings_path",20    "pdbsidechain_esm_embeddings_path",21    "esm_embeddings_path",22    "esm_embeddings_model",23)24 25 26def load_model_args(model_dir):27    model_dir = Path(model_dir)28    config_path = model_dir / "model_parameters.yml"29    with config_path.open("r", encoding="utf-8") as handle:30        config = yaml.full_load(handle) or {}31    return Namespace(**config)32 33 34def model_uses_lm_embeddings(model_args):35    return any(getattr(model_args, key, None) is not None for key in _LM_EMBEDDING_KEYS)36 37 38def _has_arg(args, name):39    try:40        return name in args41    except TypeError:42        return hasattr(args, name)43 44 45def _get_arg(args, name, default=None):46    if isinstance(args, dict):47        return args.get(name, default)48    if _has_arg(args, name):49        return getattr(args, name)50    return default51 52 53def get_model(args, device, t_to_sigma, no_parallel=False, confidence_mode=False, old=False):54    timestep_emb_func = get_timestep_embedding(55        embedding_type=_get_arg(args, "embedding_type", "sinusoidal"),56        embedding_dim=args.sigma_embed_dim,57        embedding_scale=_get_arg(args, "embedding_scale", 10000),58    )59 60    all_atoms = _get_arg(args, "all_atoms", False)61    if old and not all_atoms:62        raise NotImplementedError(63            "The old coarse-grained DiffDock model path is not migrated yet. "64            "Use old=True only with all_atoms=True, or migrate old_cg_model.py first."65        )66 67    lm_embedding_type = None68    if (69        _get_arg(args, "moad_esm_embeddings_path") is not None70        or _get_arg(args, "pdbbind_esm_embeddings_path") is not None71        or _get_arg(args, "pdbsidechain_esm_embeddings_path") is not None72        or _get_arg(args, "esm_embeddings_path") is not None73    ):74        lm_embedding_type = "precomputed"75    if _get_arg(args, "esm_embeddings_model") is not None:76        lm_embedding_type = args.esm_embeddings_model77 78    if old:79        model_class = AAOldModel80    elif all_atoms:81        model_class = AAModel82    else:83        model_class = CGModel84 85    model_kwargs = dict(86        t_to_sigma=t_to_sigma,87        device=device,88        no_torsion=args.no_torsion,89        timestep_emb_func=timestep_emb_func,90        num_conv_layers=args.num_conv_layers,91        lig_max_radius=args.max_radius,92        scale_by_sigma=args.scale_by_sigma,93        sigma_embed_dim=args.sigma_embed_dim,94        norm_by_sigma=_get_arg(args, "norm_by_sigma", False),95        ns=args.ns,96        nv=args.nv,97        distance_embed_dim=args.distance_embed_dim,98        cross_distance_embed_dim=args.cross_distance_embed_dim,99        batch_norm=not args.no_batch_norm,100        dropout=args.dropout,101        use_second_order_repr=args.use_second_order_repr,102        cross_max_distance=args.cross_max_distance,103        dynamic_max_cross=args.dynamic_max_cross,104        smooth_edges=_get_arg(args, "smooth_edges", False),105        odd_parity=_get_arg(args, "odd_parity", False),106        lm_embedding_type=lm_embedding_type,107        confidence_mode=confidence_mode,108        confidence_dropout=_get_arg(args, "confidence_dropout", 0.0),109        confidence_no_batchnorm=_get_arg(args, "confidence_no_batchnorm", False),110        affinity_prediction=_get_arg(args, "affinity_prediction", False),111        parallel=_get_arg(args, "parallel", 1),112        num_confidence_outputs=(113            len(args.rmsd_classification_cutoff) + 1114            if isinstance(_get_arg(args, "rmsd_classification_cutoff"), list)115            else 1116        ),117        atom_num_confidence_outputs=(118            len(args.atom_rmsd_classification_cutoff) + 1119            if isinstance(_get_arg(args, "atom_rmsd_classification_cutoff"), list)120            else 1121        ),122        parallel_aggregators=_get_arg(args, "parallel_aggregators", ""),123        fixed_center_conv=not _get_arg(args, "not_fixed_center_conv", False),124        no_aminoacid_identities=_get_arg(args, "no_aminoacid_identities", False),125        include_miscellaneous_atoms=_get_arg(args, "include_miscellaneous_atoms", False),126        sh_lmax=_get_arg(args, "sh_lmax", 2),127        differentiate_convolutions=not _get_arg(args, "no_differentiate_convolutions", False),128        tp_weights_layers=_get_arg(args, "tp_weights_layers", 2),129        num_prot_emb_layers=_get_arg(args, "num_prot_emb_layers", 0),130        reduce_pseudoscalars=_get_arg(args, "reduce_pseudoscalars", False),131        embed_also_ligand=_get_arg(args, "embed_also_ligand", False),132        atom_confidence=_get_arg(args, "atom_confidence_loss_weight", 0.0) > 0.0,133        sidechain_pred=(134            (_has_arg(args, "sidechain_loss_weight") and args.sidechain_loss_weight > 0)135            or (_has_arg(args, "backbone_loss_weight") and args.backbone_loss_weight > 0)136        ),137        depthwise_convolution=_get_arg(args, "depthwise_convolution", False),138    )139    if model_class is AAModel:140        model_kwargs["crop_beyond"] = _get_arg(args, "crop_beyond", None)141    elif model_class is AAOldModel:142        for key in (143            "atom_num_confidence_outputs",144            "differentiate_convolutions",145            "tp_weights_layers",146            "num_prot_emb_layers",147            "reduce_pseudoscalars",148            "embed_also_ligand",149            "atom_confidence",150            "sidechain_pred",151            "depthwise_convolution",152        ):153            model_kwargs.pop(key, None)154        model_kwargs["lm_embedding_type"] = (155            "esm" if _get_arg(args, "esm_embeddings_path") is not None else None156        )157        model_kwargs["use_old_atom_encoder"] = _get_arg(args, "use_old_atom_encoder", True)158 159    model = model_class(**model_kwargs)160 161    if device.type == "cuda" and not no_parallel and _get_arg(args, "dataset") != "torsional":162        model = DataParallel(model)163    model.to(device)164    return model165 166 167def build_score_model(168    model_args,169    device,170    no_parallel=False,171    confidence_mode=False,172    old=False,173):174    t_to_sigma = partial(t_to_sigma_compl, args=model_args)175    model = get_model(176        model_args,177        device,178        t_to_sigma=t_to_sigma,179        no_parallel=no_parallel,180        confidence_mode=confidence_mode,181        old=old,182    )183    return model, t_to_sigma184 185 186def load_score_model(187    model_dir,188    ckpt,189    device=None,190    no_parallel=True,191    confidence_mode=False,192    old=False,193    strict=True,194):195    if device is None:196        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")197 198    model_args = load_model_args(model_dir)199    model, t_to_sigma = build_score_model(200        model_args=model_args,201        device=device,202        no_parallel=no_parallel,203        confidence_mode=confidence_mode,204        old=old,205    )206 207    checkpoint_path = Path(model_dir) / ckpt208    state_dict = torch.load(checkpoint_path, map_location=torch.device("cpu"))209    if isinstance(state_dict, dict) and "model" in state_dict and "optimizer" in state_dict:210        model.load_state_dict(state_dict["model"], strict=strict)211        if "ema_weights" in state_dict and getattr(model_args, "ema_rate", None) is not None:212            ema_weights = ExponentialMovingAverage(model.parameters(), decay=model_args.ema_rate)213            ema_weights.load_state_dict(state_dict["ema_weights"], device=device)214            ema_weights.copy_to(model.parameters())215    else:216        model.load_state_dict(state_dict, strict=strict)217    model = model.to(device)218    model.eval()219    return model, model_args, t_to_sigma220