Scripts/MMTK: plotMD.py

from MMTK import *
from MMTK.Geometry import SCLattice
from MMTK.ForceFields import LennardJonesForceField
from MMTK.Environment import NoseThermostat, AndersenBarostat
from MMTK.Trajectory import Trajectory, TrajectoryOutput, StandardLogOutput
from MMTK.Dynamics import VelocityVerletIntegrator, VelocityScaler, \
                          TranslationRemover, BarostatReset

edge = 3.2864
n = 9
temperature = 94.4*Units.K
pressure = 1.*Units.atm

ff = LennardJonesForceField(1.5)
universe = CubicPeriodicUniverse(edge, ff)
for point in SCLattice(edge/n, (n, n, n)):
    universe.addObject(Atom('Ar', position=point))

# Create corresponding Chimera molecule
from MMTK2Molecule import convert
m, atomMap = convert(universe)
from chimera import openModels
openModels.add([m])
from chimera import runCommand
runCommand("window ; repr stick")

# Create energy plot that will eventually be updated when
# we integrate
from chimera.mplDialog import MPLDialog
class EnergyPlot(MPLDialog):
        title = "potential energy vs. time step"
        def __init__(self, initialEnergy):
                MPLDialog.__init__(self)
                self.energies = [ initialEnergy ]
                self.subplot = self.add_subplot(1, 1, 1)
                self._displayData()

        def _displayData(self):
                ax = self.subplot
                ax.clear()
                ax.plot(self.energies)
                self.draw()

        def addEnergy(self, energy):
                self.energies.append(energy)
                self._displayData()

energyPlot = EnergyPlot(universe.energy())
def updatePlot(universe, atomMap, state, plot=energyPlot):
        if state is None:
                e = universe.energy()
        else:
                e = state["potential_energy"]
        plot.addEnergy(e)

universe.initializeVelocitiesToTemperature(temperature)

universe.thermostat = NoseThermostat(temperature)
universe.barostat = AndersenBarostat(pressure)

integrator = VelocityVerletIntegrator(universe, delta_t=10*Units.fs,
                                        background=True)

# Note that we do NOT use StandardLogOutput because printing
# in Chimera is not thread-safe (yet)
thread = integrator(steps = 2000,
           actions = [TranslationRemover(0, None, 100),
                      VelocityScaler(temperature, 0.1*temperature,
                                     0, None, 100),
                      BarostatReset(100)])
from MMTK2Molecule import updateChimera
updateChimera(thread, universe, atomMap, 15, callback=updatePlot)