OneScience-Group/SurfDock

1import os2from openff.toolkit import Molecule3from openmmforcefields.generators import SystemGenerator4from openmm import unit, LangevinIntegrator5from openmm.app import PDBFile, Simulation6from pdbfixer import PDBFixer7import traceback8from rdkit import Chem9from rdkit.Chem import AllChem10import torch11import numpy as np12from rdkit import Chem13import warnings14from openmm import unit, Platform, State15from joblib import wrap_non_picklable_objects16from joblib import delayed17import re18from openmm.app import Modeller19import sys20import loguru21sys.path.append(os.path.dirname(os.path.abspath(__file__)))22from cleaup import clean_structure,fix_pdb23from openmm.app.internal.pdbstructure import PdbStructure24import io25import subprocess26from loguru import logger27def run_command(command: str, cwd_path: str) -> None:28    r"""29    Create a child process and run the command in the cwd_path.30    It is more safe than os.system.31    """32    proc = subprocess.Popen(33        command,34        shell=True,35        cwd=cwd_path,36        executable="/bin/bash",37        stdout=subprocess.PIPE,38        stderr=subprocess.PIPE,39    )40    errorcode = proc.wait()41    if errorcode:42        path = cwd_path43        msg = (44            'Failed with command "{}" failed in '45            ""46            "{} with error code {}"47            "stdout: {}"48            "stderr: {}".format(command, path, errorcode, proc.stdout.read().decode(), proc.stderr.read().decode())49        )50        raise ValueError(msg)51 52def read_molecule(molecule_file, sanitize=False, calc_charges=False, remove_hs=False):53    if molecule_file.endswith('.mol2'):54        mol = Chem.MolFromMol2File(molecule_file, sanitize=False, removeHs=False)55    elif molecule_file.endswith('.sdf'):56        supplier = Chem.SDMolSupplier(molecule_file, sanitize=False, removeHs=False)57        mol = supplier[0]58    elif molecule_file.endswith('.pdbqt'):59        with open(molecule_file) as file:60            pdbqt_data = file.readlines()61        pdb_block = ''62        for line in pdbqt_data:63            pdb_block += '{}\n'.format(line[:66])64        mol = Chem.MolFromPDBBlock(pdb_block, sanitize=False, removeHs=False)65    elif molecule_file.endswith('.pdb'):66        mol = Chem.MolFromPDBFile(molecule_file, sanitize=False, removeHs=False)67    else:68        raise ValueError('Expect the format of the molecule_file to be '69                         'one of .mol2, .sdf, .pdbqt and .pdb, got {}'.format(molecule_file))70 71    try:72        if sanitize or calc_charges:73            Chem.SanitizeMol(mol)74 75        if calc_charges:76            # Compute Gasteiger charges on the molecule.77            try:78                AllChem.ComputeGasteigerCharges(mol)79            except:80                warnings.warn('Unable to compute charges for the molecule.')81        if remove_hs:82            mol = Chem.RemoveHs(mol, sanitize=sanitize)83    except Exception as e:84        logger.info(e)85        logger.info("RDKit was unable to read the molecule.")86        return None87 88    return mol89def read_abs_file_mol(file, remove_hs=False, sanitize=True):90    mol = read_molecule(file, remove_hs=remove_hs, sanitize=True)91    92    if file.endswith(".sdf") and mol is None:93        # mol = read_molecule(file, remove_hs=remove_hs, sanitize=True)94        if os.path.exists(file[:-4] + ".mol2"):95            logger.info('Using the .sdf file failed. We found a .mol2 file instead and are trying to use that.')96            mol = read_molecule(file[:-4] + ".mol2", remove_hs=remove_hs, sanitize=True)97    elif file.endswith(".mol2") and mol is None:98        if os.path.exists(file[:-4] + ".sdf"):99            logger.info('Using the .mol2 file failed. We found a .sdf file instead and are trying to use that.')100            mol = read_molecule(file[:-4] + ".sdf", remove_hs=remove_hs, sanitize=True)101 102    return mol103# from joblib.externals.loky import set_loky_pickler104def trySystem(system_generator,modeller,ligand_mol,lig_path):105   106    max_attempts = 100107    attempts = 0108    success = False109    while attempts < max_attempts and not success:110        try:111            system = system_generator.create_system(modeller.topology, molecules=ligand_mol)112            success = True  # Mark113        except Exception as e:114            # extract the error residue index from the error message115            logger.info(f'Try DELETE THIS ERROE {str(e)}!')116            match = re.search(r"residue (\d+)", str(e))117            if match:118                extracted_index = int(match.group(1)) - 1119                # located and record the residue to delete120                current_index = 0121                residue_to_delete = None122                for residue in modeller.topology.residues():123                    if current_index == extracted_index:124                        residue_to_delete = residue125                        break126                    current_index += 1127 128                modeller.delete([residue_to_delete])129        finally:130            attempts += 1131    if not success:132        logger.info("Try maximum times but cannot create system")133        return None134    else:135        logger.info(f"Try {attempts} times and system is created successfully")136        with open(os.path.join('/home/house/caoduanhua_tmp/DeepLearningForDock/DiffDockForScreen/diffScreen/Screen_dataset/create_system_pdbs',os.path.basename(lig_path).split('_')[0]+'_create_system.pdb'), "w") as f:137            PDBFile.writeFile(modeller.topology, modeller.positions, f)138 139    return modeller140 141def UpdatePose(lig_path,system_generator,modeller,protein_atoms,out_dir,device_num=0):142    try:143        # init save path144        out_base_dir = os.path.join(out_dir,lig_path.split('/')[-2])145        os.makedirs(out_base_dir,exist_ok=True)146        out_file = os.path.join(out_base_dir,os.path.splitext(os.path.basename(lig_path))[0] + '_minimized.sdf')147        if os.path.exists(out_file):148            return 0149        dockingpose = read_molecule(lig_path, remove_hs=True, sanitize=True)150        lig_mol = Molecule.from_rdkit(dockingpose,allow_undefined_stereo=True)151        lig_mol.assign_partial_charges(partial_charge_method='gasteiger')152 153        lig_top = lig_mol.to_topology()154        modeller.add(lig_top.to_openmm(), lig_top.get_positions().to_openmm())155        # create simulation system156        system=system_generator.create_system(modeller.topology,molecules=lig_mol)157        # keep protein atom static in smiulation 158        for atom in protein_atoms:159            system.setParticleMass(atom.index, 0.000*unit.dalton)160        # start simulation161        platform = GetPlatform()162        simulation = EnergyMinimized(modeller,system, platform,verbose=False,device_num=device_num)163        # get energy minimized conformer and modify the graph['ligand'].pos to scoring164        # use conformer mapping165        ligand_atoms = list(filter(lambda atom:  atom.residue.name == 'UNK',list(modeller.topology.atoms())))166        ligand_index = [atom.index for atom in ligand_atoms]167        new_coords = simulation.context.getState(getPositions=True).getPositions(asNumpy=True).value_in_unit(unit.angstrom)[ligand_index]168        lig_mol = lig_mol.to_rdkit()169        conf = lig_mol.GetConformer()170        for i in range(lig_mol.GetNumAtoms()):171            x,y,z = new_coords.astype(np.double)[i]172            conf.SetAtomPosition(i,Point3D(x,y,z))173        try:174            writer = Chem.SDWriter(out_file)175            writer.write(lig_mol)176            writer.close()177        except:178            out_base_dir = os.path.join(out_dir,lig_path.split('/')[-2] + '_tmp')179            os.makedirs(out_base_dir,exist_ok=True)180            out_file = os.path.join(out_base_dir,os.path.splitext(os.path.basename(lig_path))[0] + '_minimized.sdf')181            if os.path.exists(out_file):182                return 0183            writer = Chem.SDWriter(out_file)184            writer.write(lig_mol)185            writer.close()186        return 0187        # return lig_mol188    except Exception as e:189        error_info = traceback.format_exc()190        logger.info(error_info)191        logger.warning(f' : {e}')192        with open('error_sdf.txt','a') as f:193            f.write(lig_path +': error by :' + error_info + '\n')194        return 1195 196def UpdateGrpah(graph,system_generator,modeller,protein_atoms,device_num=0):197    try:198        # raw_position = graph['ligand'].pos199        dockingpose = GetDockingPose(graph)200        lig_mol = Molecule.from_rdkit(dockingpose,allow_undefined_stereo=True)201        lig_mol.assign_partial_charges(partial_charge_method='gasteiger')202        # add ligand to modeller203        lig_top = lig_mol.to_topology()204        modeller.add(lig_top.to_openmm(), lig_top.get_positions().to_openmm())205        # create simulation system206        platform = GetPlatform()207 208        system = system_generator.create_system(modeller.topology,molecules=lig_mol)209        # keep protein atom static in smiulation 210        for atom in protein_atoms:211            system.setParticleMass(atom.index, 0.000*unit.dalton)212        # start simulation213        simulation = EnergyMinimized(modeller,system, platform,verbose=False,device_num=device_num)214        # get energy minimized conformer and modify the graph['ligand'].pos to scoring215        # conformer mapping216        ligand_atoms = list(filter(lambda atom:  atom.residue.name == 'UNK',list(modeller.topology.atoms())))217        ligand_index = [atom.index for atom in ligand_atoms]218        new_coords = simulation.context.getState(getPositions=True).getPositions(asNumpy=True).value_in_unit(unit.angstrom)[ligand_index]219 220        new_coords -= graph.original_center.detach().cpu().numpy()221        lig_mol = lig_mol.to_rdkit()222        conf = lig_mol.GetConformer()223        for i in range(lig_mol.GetNumAtoms()):224            x,y,z = new_coords.astype(np.double)[i]225            conf.SetAtomPosition(i,Point3D(x,y,z))226        lig_mol = Chem.RemoveHs(lig_mol)227      228        graph['ligand'].pos = torch.from_numpy(lig_mol.GetConformer().GetPositions()).to(graph.original_center.device).float()229    230        return graph231    except Exception as e:232        error_info = traceback.format_exc()233        logger.info(error_info)234        warnings.warn(graph['name'][0]+f' : {e}')235        return 1236 237def DescribeState(state: State, name: str):238    """logger.info energy and force information about a simulation state."""239    max_force = max(np.linalg.norm([v.x, v.y, v.z]) for v in state.getForces())240    logger.info(f"{name} has energy {state.getPotentialEnergy().value_in_unit(unit.kilojoule_per_mole):.2f} kJ/mol "241          f"with maximum force {max_force:.2f} kJ/(mol nm)")242def GetFFGenerator(protein_forcefield = 'amber/ff14SB.xml',water_forcefield = 'amber/tip3p_standard.xml',small_molecule_forcefield = 'openff-2.0.0',ignoreExternalBonds=False):243    """244    Get forcefield generator by different forcefield files245    """246    forcefield_kwargs = {'constraints': None, 'rigidWater': True, 'removeCMMotion': False, 'ignoreExternalBonds': ignoreExternalBonds, 'hydrogenMass': 4*unit.amu }247    # forcefield_kwargs = {'constraints': None, 'rigidWater': True, 'removeCMMotion': False, 'hydrogenMass': 4*unit.amu }248    system_generator = SystemGenerator(249                forcefields=[protein_forcefield, water_forcefield ],250                small_molecule_forcefield=small_molecule_forcefield,251                forcefield_kwargs=forcefield_kwargs)252    return system_generator253def GetfixedPDB(receptor_path):254    255    temp_fixd_pdbs = f'{os.path.dirname(os.path.dirname(os.path.abspath(__file__)))}/fixed_pdbs'256    os.makedirs(temp_fixd_pdbs,exist_ok=True)257    if not os.path.exists(os.path.join(temp_fixd_pdbs,os.path.basename(receptor_path).replace('.pdb','_fixer_processed_cleanup.pdb'))):258        alterations_info  = {}259        fixed_pdb = fix_pdb(receptor_path, alterations_info)260        fixed_pdb_file = io.StringIO(fixed_pdb)261        pdb_structure = PdbStructure(fixed_pdb_file)262        clean_structure(pdb_structure, alterations_info)263        fixer = PDBFile(pdb_structure)264        logger.info("Protein loaded with success!")265        PDBFile.writeFile(fixer.topology, fixer.positions, open(os.path.join(temp_fixd_pdbs,os.path.basename(receptor_path).replace('.pdb','_fixer_processed_cleanup.pdb')), 'w'))266        logger.info('Dont have processed by fixer try fix and save in disk')267    else:268        fixer = PDBFixer(os.path.join(temp_fixd_pdbs,os.path.basename(receptor_path).replace('.pdb','_fixer_processed_cleanup.pdb')))269        logger.info('There have a precessed pdb file use it!')270    return fixer271 272import copy273from rdkit.Geometry import Point3D274def GetDockingPose(graph):275    mol = copy.deepcopy(graph.mol[0] if type(graph.mol) == list else graph.mol)276    mol = Chem.RemoveHs(mol)277    docking_position = graph['ligand'].pos.detach().cpu().numpy() # without Hs and dont match with raw pocket278    docking_position = docking_position + graph.original_center.detach().cpu().numpy()279    conf = mol.GetConformer()280    for i in range(mol.GetNumAtoms()):281        x,y,z = docking_position.astype(np.double)[i]282        conf.SetAtomPosition(i,Point3D(x,y,z))283    return mol284 285@delayed286@wrap_non_picklable_objects287def GetPlatformPara():288    """Determine the best simulation platform available."""289    platform_name = os.getenv('PLATFORM')290    # properties = {'CudaDeviceIndex': '0'}291    if platform_name:292        platform = Platform.getPlatformByName(platform_name)293    else:294        platform = max((Platform.getPlatform(i) for i in range(Platform.getNumPlatforms())), key=lambda x: x.getSpeed())295    logger.info(f'Using platform {platform.getName()}')296    if platform.getName() in ['CUDA', 'OpenCL']:297        platform.setPropertyDefaultValue('Precision', 'mixed')298        logger.info(f'Set precision for platform {platform.getName()} to mixed')299    return platform300# @delayed301# @wrap_non_picklable_objects302def GetPlatform():303    """Determine the best simulation platform available."""304    platform_name = os.getenv('PLATFORM')305    # properties = {'CudaDeviceIndex': '0'}306    if platform_name:307        platform = Platform.getPlatformByName(platform_name)308    else:309        platform = max((Platform.getPlatform(i) for i in range(Platform.getNumPlatforms())), key=lambda x: x.getSpeed())310    logger.info(f'Using platform {platform.getName()}')311    if platform.getName() in ['CUDA', 'OpenCL']:312        platform.setPropertyDefaultValue('Precision', 'mixed')313        logger.info(f'Set precision for platform {platform.getName()} to mixed')314    return platform315 316def EnergyMinimized(modeller,system, platform,verbose=False,device_num = 0):317    integrator = LangevinIntegrator(318    300 * unit.kelvin,319    1 / unit.picosecond,320    0.002 * unit.picoseconds,321    )322    properties = {'CudaDeviceIndex': f'{device_num}'}323    simulation = Simulation(modeller.topology, system = system, integrator = integrator, platform=platform,platformProperties=properties)324    simulation.context.setPositions(modeller.positions)325    if verbose:326        DescribeState(327            simulation.context.getState(328                getEnergy=True,329                getForces=True,330            ),331            "Original state",332        )333 334 335    simulation.minimizeEnergy()336    if verbose:337        DescribeState(338            simulation.context.getState(339                getEnergy=True, 340                getForces=True),341            "Minimized state",342        )343    return simulation344