Spaces:
Sleeping
Sleeping
File size: 3,541 Bytes
dcacefd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | #!/usr/bin/python
# -*- coding:utf-8 -*-
import os
import argparse
from os.path import splitext, basename
from copy import deepcopy
from rdkit import Chem
from rdkit.Chem.rdForceFieldHelpers import UFFOptimizeMolecule
import numpy as np
from openmm import unit
os.environ['OPENMM_CPU_THREADS'] = '4' # prevent openmm from using all cpus available
from openmm import LangevinIntegrator, Platform, CustomExternalForce
from openmm.app import PDBFile, Simulation, ForceField, HBonds, Modeller
from simtk.unit import kilocalories_per_mole, angstroms
from pdbfixer import PDBFixer
import logging
logging.getLogger('openmm').setLevel(logging.ERROR)
FILE_DIR = os.path.abspath(os.path.split(__file__)[0])
CACHE_DIR = '.data/saved/'
def openmm_relax(pdb, out_pdb=None, excluded_chains=None, inverse_exclude=False):
tolerance_in_kj = 2.39 * unit.kilojoules_per_mole / unit.kilocalories_per_mole
stiffness = 10.0 * kilocalories_per_mole / (angstroms ** 2)
if excluded_chains is None:
excluded_chains = []
if out_pdb is None:
out_pdb = pdb[:-4]+'_relaxed'+'.pdb'
fixer = PDBFixer(pdb)
fixer.findNonstandardResidues()
fixer.replaceNonstandardResidues()
fixer.findMissingResidues()
fixer.findMissingAtoms() # [OXT]
fixer.addMissingAtoms()
# force_field = ForceField("amber14/protein.ff14SB.xml")
force_field = ForceField('amber99sb.xml')
modeller = Modeller(fixer.topology, fixer.positions)
modeller.addHydrogens(force_field)
# system = force_field.createSystem(modeller.topology)
system = force_field.createSystem(modeller.topology, constraints=HBonds)
force = CustomExternalForce("0.5 * k * ((x-x0)^2 + (y-y0)^2 + (z-z0)^2)")
force.addGlobalParameter("k", stiffness)
for p in ["x0", "y0", "z0"]:
force.addPerParticleParameter(p)
# add flexible atoms
for residue in modeller.topology.residues():
if (not inverse_exclude and residue.chain.id in excluded_chains) or \
(inverse_exclude and residue.chain.id not in excluded_chains): # antigen
for atom in residue.atoms():
system.setParticleMass(atom.index, 0)
for atom in residue.atoms():
# if atom.name in ['N', 'CA', 'C', 'CB']:
if atom.element.name != 'hydrogen':
force.addParticle(atom.index, modeller.positions[atom.index])
system.addForce(force)
integrator = LangevinIntegrator(0, 0.01, 0.0)
# platform = Platform.getPlatformByName('CPU')
# platform = Platform.getPlatformByName('CUDA')
simulation = Simulation(modeller.topology, system, integrator) # , platform)
simulation.context.setPositions(modeller.positions)
simulation.minimizeEnergy(tolerance_in_kj)
state = simulation.context.getState(getPositions=True, getEnergy=True)
with open(out_pdb, 'w') as fout:
PDBFile.writeFile(simulation.topology, state.getPositions(), fout, keepIds=True)
return out_pdb
def relax_sdf(sdf):
mol = Chem.MolFromMolFile(sdf, sanitize=True)
mol = Chem.AddHs(mol, addCoords=True)
UFFOptimizeMolecule(mol)
# Save to MOL file
Chem.MolToMolFile(mol, sdf)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--pdb', type=str, default='./saved/3upr_C_rec.pdb')
parser.add_argument('--sdf', type=str, default='./saved/3upr_C_rec_3vri_1kx_lig_tt_min_0.sdf')
args = parser.parse_args()
openmm_relax(args.pdb, out_pdb=args.pdb)
relax_sdf(args.sdf, out_sdf=args.sdf)
|