File size: 1,653 Bytes
1fdc49a | 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 | import pdbfixer
import openmm
ENERGY = openmm.unit.kilocalories_per_mole
LENGTH = openmm.unit.angstroms
def refine(pdb_file, stiffness=10., tolerance=2.39, use_gpu=False):
tolerance = tolerance * ENERGY
stiffness = stiffness * ENERGY / (LENGTH**2)
fixer = pdbfixer.PDBFixer(pdb_file)
fixer.findMissingResidues()
fixer.findMissingAtoms()
fixer.addMissingAtoms()
force_field = openmm.app.ForceField("amber14/protein.ff14SB.xml")
modeller = openmm.app.Modeller(fixer.topology, fixer.positions)
modeller.addHydrogens(force_field)
system = force_field.createSystem(modeller.topology)
force = openmm.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)
for residue in modeller.topology.residues():
for atom in residue.atoms():
if atom.name in ["N", "CA", "C", "CB"]:
force.addParticle(atom.index,
modeller.positions[atom.index])
system.addForce(force)
integrator = openmm.LangevinIntegrator(0, 0.01, 1.0)
platform = openmm.Platform.getPlatformByName("CUDA" if use_gpu else "CPU")
simulation = openmm.app.Simulation(modeller.topology, system, integrator, platform)
simulation.context.setPositions(modeller.positions)
simulation.minimizeEnergy(tolerance)
with open(pdb_file, "w") as f:
openmm.app.PDBFile.writeFile(
simulation.topology,
simulation.context.getState(getPositions=True).getPositions(),
f,
keepIds=True,) |