Source code for moldesign.interfaces.openmm

# Copyright 2016 Autodesk Inc.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.
import itertools
import imp

import numpy as np

import pyccc
import moldesign as mdt
from moldesign.utils import from_filepath
from moldesign import units as u
from moldesign import compute

except (ImportError, OSError) as exc:
    print 'OpenMM could not be imported; using remote docker container'
    force_remote = True
    force_remote = False

[docs]class OpenMMPickleMixin(object): def __getstate__(self): mystate = self.__dict__.copy() if 'sim' in mystate: assert 'sim_args' not in mystate sim = mystate.pop('sim') mystate['sim_args'] = (sim.topology, sim.system, sim.integrator) return mystate def __setstate__(self, state): from simtk.openmm import app if 'sim_args' in state: assert 'sim' not in state args = state.pop('sim_args') state['sim'] = app.Simulation(*args) self.__dict__.update(state)
# This is a factory for the MdtReporter class. It's here so that we don't have to import # at the module level
[docs]def MdtReporter(mol, report_interval): from import StateDataReporter class MdtReporter(StateDataReporter): """ We'll use this class to capture all the information we need about a trajectory It's pretty basic - the assumption is that there will be more processing on the client side """ LEN = 30 def __init__(self, mol, report_interval): self.mol = mol self.report_interval = report_interval self.trajectory = mdt.Trajectory(mol) self.annotation = None self._row_format = ("{:<%d}" % 10) + 3*("{:>%d}" % self.LEN) self._printed_header = False self.last_report_time = None def report_from_mol(self, **kwargs): self.mol.calculate() if self.annotation is not None: kwargs.setdefault('annotation', self.annotation), self.mol.energy_model.sim.context.getState(getEnergy=True, getForces=True, getPositions=True, getVelocities=True)) def report(self, simulation, state): """ Callback for dynamics after the specified interval :type simulation: :type state: simtk.openmm.State """ # TODO: make sure openmm masses are the same as MDT masses report = dict( positions=simtk2pint(state.getPositions()), momenta=simtk2pint(state.getVelocities())*self.mol.dim_masses, forces=simtk2pint(state.getForces()), time=simtk2pint(state.getTime()), vectors=simtk2pint(state.getPeriodicBoxVectors()), potential_energy=simtk2pint(state.getPotentialEnergy())) if self.annotation is not None: report['annotation'] = self.annotation if not self._printed_header: timeheader = 'time / {units}'.format(units=u.default.time) peheader = 'potential / {units}'.format( keheader = 'kinetic / {units}'.format( temperatureheader = 'T / {units}'.format(units=u.default.temperature) print self._row_format.format(timeheader, peheader, keheader, temperatureheader) self._printed_header = True ke = mdt.helpers.kinetic_energy(report['momenta'], self.mol.dim_masses) t = (2.0 * ke) / (u.k_b * self.mol.dynamic_dof) print self._row_format.format(report['time'].defunits_value(), report['potential_energy'].defunits_value(), ke.defunits_value(), t.defunits_value()) self.last_report_time = self.mol.time self.trajectory.new_frame(properties=report) def describeNextReport(self, simulation): """ :type simulation: :return: report_description A five element tuple. The first element is the number of steps until the next report. The remaining elements specify whether that report will require positions, velocities, forces, and energies respectively. :return type: tuple """ steps = self.report_interval - simulation.currentStep % self.report_interval return (steps, True, True, True, True) return MdtReporter(mol, report_interval)
PINT_NAMES = {'mole': u.avogadro, 'degree': u.degrees, 'radian': u.radians, 'elementary charge': u.q_e}
[docs]def simtk2pint(quantity, flat=False): """ Converts a quantity from the simtk unit system to a quantity from the pint unit system :param quantity: :param flat: if True, flatten 3xN arrays to 3N """ from simtk import unit as stku mag = quantity._value if quantity.unit == stku.radian: return mag * u.radians if quantity.unit == return mag * u.degrees if hasattr(mag, '__getslice__'): mag = np.array(mag[:]) for dim, exp in itertools.chain(quantity.unit.iter_scaled_units(), quantity.unit.iter_top_base_units()): if in PINT_NAMES: pintunit = PINT_NAMES[] else: pintunit = u.ureg.parse_expression( mag = mag * (pintunit**exp) if flat: mag = np.reshape(mag, (np.product(mag.shape),)) return u.default.convert(mag)
[docs]def pint2simtk(quantity): """ Converts a quantity from the pint to simtk unit system. Note SimTK appears limited, esp for energy units. May need to have pint convert to SI first """ from simtk import unit as stku SIMTK_NAMES = {'ang': stku.angstrom, 'fs': stku.femtosecond, 'nm': stku.nanometer, 'ps': stku.picosecond} newvar = quantity._magnitude for dim, exp in quantity._units.iteritems(): if dim in SIMTK_NAMES: stkunit = SIMTK_NAMES[dim] else: stkunit = getattr(stku, dim) newvar = newvar * stkunit ** exp return newvar
@compute.runsremotely(enable=force_remote) def _amber_to_mol(prmtop_file, inpcrd_file): from simtk.openmm import app prmtop = from_filepath(app.AmberPrmtopFile, prmtop_file) inpcrd = from_filepath(app.AmberInpcrdFile, inpcrd_file) mol = topology_to_mol(prmtop.topology, positions=inpcrd.positions, velocities=inpcrd.velocities) return mol if force_remote:
[docs] def amber_to_mol(prmtop_file, inpcrd_file): if not isinstance(prmtop_file, pyccc.FileContainer): prmtop_file = pyccc.LocalFile(prmtop_file) if not isinstance(inpcrd_file, pyccc.FileContainer): inpcrd_file = pyccc.LocalFile(inpcrd_file) return _amber_to_mol(prmtop_file, inpcrd_file)
else: amber_to_mol = _amber_to_mol
[docs]def topology_to_mol(topo, name=None, positions=None, velocities=None, assign_bond_orders=True): """ Convert an OpenMM topology object into an MDT molecule. Args: topo ( topology to convert name (str): name to assign to molecule positions (list): simtk list of atomic positions velocities (list): simtk list of atomic velocities assign_bond_orders (bool): assign bond orders from templates (simtk topologies do not store bond orders) """ from simtk import unit as stku # Atoms atommap = {} newatoms = [] masses = u.amu*[atom.element.mass.value_in_unit(stku.amu) for atom in topo.atoms()] for atom,mass in zip(topo.atoms(), masses): newatom = mdt.Atom(atnum=atom.element.atomic_number,, mass=mass) atommap[atom] = newatom newatoms.append(newatom) # Coordinates if positions is not None: poslist = np.array([p.value_in_unit(stku.nanometer) for p in positions]) * u.nm poslist.ito(u.default.length) for newatom, position in zip(newatoms, poslist): newatom.position = position if velocities is not None: velolist = np.array([v.value_in_unit(stku.nanometer/stku.femtosecond) for v in velocities]) * u.nm/u.fs velolist = u.default.convert(velolist) for newatom, velocity in zip(newatoms, velolist): newatom.momentum = newatom.mass * simtk2pint(velocity) # Biounits chains = {} for chain in topo.chains(): if not in chains: chains[] = mdt.Chain(, index=chain.index) newchain = chains[] for residue in chain.residues(): newresidue = mdt.Residue(name='%s%d' % (, residue.index), chain=newchain, pdbindex=int(, newchain.add(newresidue) for atom in residue.atoms(): newatom = atommap[atom] newatom.residue = newresidue newatom.chain = newchain newresidue.add(newatom) # Bonds bonds = {} for bond in topo.bonds(): a1, a2 = bond na1, na2 = atommap[a1], atommap[a2] if na1 not in bonds: bonds[na1] = {} if na2 not in bonds: bonds[na2] = {} bonds[na1][na2] = 1 bonds[na2][na1] = 1 if name is None: name = 'Unnamed molecule from OpenMM' newmol = mdt.Molecule(newatoms, bond_graph=bonds, name=name) if assign_bond_orders: for residue in newmol.residues: try: residue.assign_template_bonds() except (KeyError, ValueError): pass return newmol
[docs]def mol_to_topology(mol): """ Create an openmm topology object from an MDT molecule Args: mol (moldesign.Molecule): molecule to copy topology from Returns: topology of the molecule """ from simtk.openmm import app top = app.Topology() chainmap = {chain: top.addChain( for chain in mol.chains} resmap = {res: top.addResidue(res.resname, chainmap[res.chain], str(res.pdbindex)) for res in mol.residues} atommap = {atom: top.addAtom(, app.Element.getBySymbol(atom.element), resmap[atom.residue], id=str(atom.pdbindex)) for atom in mol.atoms} for bond in mol.bonds: top.addBond(atommap[bond.a1], atommap[bond.a2]) return top
[docs]def mol_to_modeller(mol): from simtk.openmm import app return app.Modeller(mol_to_topology(mol), pint2simtk(mol.positions))