OneScience-Group/SurfDock

1import os2import subprocess3import warnings4from datetime import datetime5import signal6from contextlib import contextmanager7import numpy as np8import torch9import yaml10from rdkit import Chem11from rdkit.Chem import RemoveHs, MolToPDBFile12from torch_geometric.nn.data_parallel import DataParallel13 14from models.surface_score_model_v3 import TensorProductScoreModel as SurfaceScoreModelV3 15 16from models.mdn_score_model_v6 import TensorProductScoreModelV6 as ConfidenceCGScoreModelV617# from models.score_model_mdn_energy_v1 import TensorProductEnergyModel18from utils.diffusion_utils import get_timestep_embedding19from spyrmsd import rmsd, molecule20 21 22def get_obrmsd(mol1_path, mol2_path, cache_name=None):23    cache_name = datetime.now().strftime('date%d-%m_time%H-%M-%S.%f') if cache_name is None else cache_name24    os.makedirs(".openbabel_cache", exist_ok=True)25    if not isinstance(mol1_path, str):26        MolToPDBFile(mol1_path, '.openbabel_cache/obrmsd_mol1_cache.pdb')27        mol1_path = '.openbabel_cache/obrmsd_mol1_cache.pdb'28    if not isinstance(mol2_path, str):29        MolToPDBFile(mol2_path, '.openbabel_cache/obrmsd_mol2_cache.pdb')30        mol2_path = '.openbabel_cache/obrmsd_mol2_cache.pdb'31    with warnings.catch_warnings():32        warnings.simplefilter("ignore")33        return_code = subprocess.run(f"obrms {mol1_path} {mol2_path} > .openbabel_cache/obrmsd_{cache_name}.rmsd",34                                     shell=True)35        print(return_code)36    obrms_output = read_strings_from_txt(f".openbabel_cache/obrmsd_{cache_name}.rmsd")37    rmsds = [line.split(" ")[-1] for line in obrms_output]38    return np.array(rmsds, dtype=np.float)39 40 41def remove_all_hs(mol,santize=None):42    params = Chem.RemoveHsParameters()43    params.removeAndTrackIsotopes = True44    params.removeDefiningBondStereo = True45    params.removeDegreeZero = True46    params.removeDummyNeighbors = True47    params.removeHigherDegrees = True48    params.removeHydrides = True49    params.removeInSGroups = True50    params.removeIsotopes = True51    params.removeMapped = True52    params.removeNonimplicit = True53    params.removeOnlyHNeighbors = True54    params.removeWithQuery = True55    params.removeWithWedgedBond = True56    if santize is not None:57        params.sanitize = santize58    return RemoveHs(mol, params)59 60 61def read_strings_from_txt(path):62    # every line will be one element of the returned list63    with open(path) as file:64        lines = file.readlines()65        return [line.rstrip() for line in lines]66 67 68def save_yaml_file(path, content):69    assert isinstance(path, str), f'path must be a string, got {path} which is a {type(path)}'70    content = yaml.dump(data=content)71    if '/' in path and os.path.dirname(path) and not os.path.exists(os.path.dirname(path)):72        os.makedirs(os.path.dirname(path),exist_ok=True)73    with open(path, 'w') as f:74        f.write(content)75 76# from accelerate.utils import DummyOptim,DummyScheduler77def get_optimizer_and_scheduler(args, model, accelerator,scheduler_mode='min'):78    optimizer_cls = (79        torch.optim.AdamW80        if accelerator.state.deepspeed_plugin is None81        or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config82        else None #DummyOptim83    )84    optimizer = optimizer_cls(filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr, weight_decay=args.w_decay)85    86    if args.scheduler == 'plateau':87        # scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode=scheduler_mode, factor=0.7,88                                                            #    patience=args.scheduler_patience, min_lr=args.lr / 100)89        90        if (91            accelerator.state.deepspeed_plugin is None92            or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config93        ):94            scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode=scheduler_mode, factor=0.7,95                                                               patience=args.scheduler_patience, min_lr=args.lr / 100)96        # else:97        #     lr_scheduler = DummyScheduler(98        #         optimizer, total_num_steps=args.max_train_steps, warmup_num_steps=args.num_warmup_steps99        #     )100    else:101        print('No scheduler')102        scheduler = None103    return optimizer, scheduler104def get_model(args, device, t_to_sigma, no_parallel=False, model_type='score_model'):105    # ['score_model','mdn_model','energy_score_model']106    timestep_emb_func = get_timestep_embedding(107        embedding_type=args.embedding_type,108        embedding_dim=args.sigma_embed_dim,109        embedding_scale=args.embedding_scale)110    lm_embedding_type = None111    if args.esm_embeddings_path is not None: lm_embedding_type = 'esm'112    if model_type == 'mdn_model':113        model_class = ConfidenceCGScoreModelV6114        model = model_class(args,t_to_sigma=t_to_sigma,115                    device=device,116                    no_torsion=args.no_torsion,117                    timestep_emb_func=timestep_emb_func,118                    num_conv_layers=args.num_conv_layers,119                    lig_max_radius=args.max_radius,120                    scale_by_sigma=args.scale_by_sigma,121                    sigma_embed_dim=args.sigma_embed_dim,122                    ns=args.ns, nv=args.nv,123                    distance_embed_dim=args.distance_embed_dim,124                    cross_distance_embed_dim=args.cross_distance_embed_dim,125                    batch_norm=not args.no_batch_norm,126                    dropout=args.dropout,127                    use_second_order_repr=args.use_second_order_repr,128                    cross_max_distance=args.cross_max_distance,129                    dynamic_max_cross=args.dynamic_max_cross,130                    lm_embedding_type=lm_embedding_type,131                    mdn_dropout=args.mdn_dropout,n_gaussians = args.n_gaussians)132    133    elif model_type == 'surface_score_model':134 135        model_class = SurfaceScoreModelV3136 137        model = model_class(t_to_sigma=t_to_sigma,138                    device=device,139                    no_torsion=args.no_torsion,140                    timestep_emb_func=timestep_emb_func,141                    num_conv_layers=args.num_conv_layers,142                    lig_max_radius=args.max_radius,143                    scale_by_sigma=args.scale_by_sigma,144                    sigma_embed_dim=args.sigma_embed_dim,145                    ns=args.ns, nv=args.nv,146                    distance_embed_dim=args.distance_embed_dim,147                    cross_distance_embed_dim=args.cross_distance_embed_dim,148                    batch_norm=not args.no_batch_norm,149                    dropout=args.dropout,150                    use_second_order_repr=args.use_second_order_repr,151                    cross_max_distance=args.cross_max_distance,152                    dynamic_max_cross=args.dynamic_max_cross,153                    lm_embedding_type=lm_embedding_type,154                   )155    else:156        raise f'not support {model_type} type model setup'157 158    model.to(device)159    return model160 161 162def get_symmetry_rmsd(mol, coords1, coords2, mol2=None):163    with time_limit(10):164        mol = molecule.Molecule.from_rdkit(mol)165        mol2 = molecule.Molecule.from_rdkit(mol2) if mol2 is not None else mol2166        mol2_atomicnums = mol2.atomicnums if mol2 is not None else mol.atomicnums167        mol2_adjacency_matrix = mol2.adjacency_matrix if mol2 is not None else mol.adjacency_matrix168        RMSD = rmsd.symmrmsd(169            coords1,170            coords2,171            mol.atomicnums,172            mol2_atomicnums,173            mol.adjacency_matrix,174            mol2_adjacency_matrix,175        )176        return RMSD177 178 179class TimeoutException(Exception): pass180 181 182@contextmanager183def time_limit(seconds):184    def signal_handler(signum, frame):185        raise TimeoutException("Timed out!")186 187    signal.signal(signal.SIGALRM, signal_handler)188    signal.alarm(seconds)189    try:190        yield191    finally:192        signal.alarm(0)193 194 195class ExponentialMovingAverage:196    """ from https://github.com/yang-song/score_sde_pytorch/blob/main/models/ema.py197    Maintains (exponential) moving average of a set of parameters. """198 199    def __init__(self, parameters, decay, use_num_updates=True):200        """201        Args:202          parameters: Iterable of `torch.nn.Parameter`; usually the result of203            `model.parameters()`.204          decay: The exponential decay.205          use_num_updates: Whether to use number of updates when computing206            averages.207        """208        if decay < 0.0 or decay > 1.0:209            raise ValueError('Decay must be between 0 and 1')210        self.decay = decay211        self.num_updates = 0 if use_num_updates else None212        self.shadow_params = [p.clone().detach()213                              for p in parameters if p.requires_grad]214        self.collected_params = []215 216    def update(self, parameters):217        """218        Update currently maintained parameters.219        Call this every time the parameters are updated, such as the result of220        the `optimizer.step()` call.221        Args:222          parameters: Iterable of `torch.nn.Parameter`; usually the same set of223            parameters used to initialize this object.224        """225        decay = self.decay226        if self.num_updates is not None:227            self.num_updates += 1228            decay = min(decay, (1 + self.num_updates) / (10 + self.num_updates))229        one_minus_decay = 1.0 - decay230        with torch.no_grad():231            parameters = [p for p in parameters if p.requires_grad]232            for s_param, param in zip(self.shadow_params, parameters):233                s_param.sub_(one_minus_decay * (s_param - param))234 235    def copy_to(self, parameters):236        """237        Copy current parameters into given collection of parameters.238        Args:239          parameters: Iterable of `torch.nn.Parameter`; the parameters to be240            updated with the stored moving averages.241        """242        parameters = [p for p in parameters if p.requires_grad]243        for s_param, param in zip(self.shadow_params, parameters):244            if param.requires_grad:245                param.data.copy_(s_param.data)246 247    def store(self, parameters):248        """249        Save the current parameters for restoring later.250        Args:251          parameters: Iterable of `torch.nn.Parameter`; the parameters to be252            temporarily stored.253        """254        self.collected_params = [param.clone() for param in parameters]255 256    def restore(self, parameters):257        """258        Restore the parameters stored with the `store` method.259        Useful to validate the model with EMA parameters without affecting the260        original optimization process. Store the parameters before the261        `copy_to` method. After validation (or model saving), use this to262        restore the former parameters.263        Args:264          parameters: Iterable of `torch.nn.Parameter`; the parameters to be265            updated with the stored parameters.266        """267        for c_param, param in zip(self.collected_params, parameters):268            param.data.copy_(c_param.data)269 270    def state_dict(self):271        return dict(decay=self.decay, num_updates=self.num_updates,272                    shadow_params=self.shadow_params)273 274    def load_state_dict(self, state_dict, device):275        self.decay = state_dict['decay']276        self.num_updates = state_dict['num_updates']277        self.shadow_params = [tensor.to(device) for tensor in state_dict['shadow_params']]278