Loop over trajectory to decompose OpenMM energies

[Marcelo D Poleto]

Hi,

I want to loop over a MD trajectory written by OpenMM to get the decomposed energy groups and write them to a file. I am using ParmEd's (pmd.openmm.energy_decomposition_system()) to read a single coordinate file as shown below:

Code: Select all

top   = 'ala.c36.psf'
coord = 'ala.c36.crd'
traj  = '100frames.dcd'

psf  = CharmmPsfFile(top)
crd  = CharmmCrdFile(coord)

box = 2.9*nanometer #
psf.setBox(box,box,box)
system = psf.createSystem(charmm_params, nonbondedMethod=PME, nonbondedCutoff=1.2*nanometers, switchDistance=1.0*nanometer, ewaldErrorTolerance = 0.0005, constraints=HBonds)

integrator = LangevinIntegrator(298, 1/picosecond, 0.002)
simulation = Simulation(psf.topology, system, integrator)
simulation.context.setPositions(crd.positions)

# print out energy and forces
positions = pmd.load_file(coord)
energies = pmd.openmm.energy_decomposition_system(positions,system)

for term, ener in energies:
    print(str(term) + ": " + str(ener))

Does anyone know a way to use MDtraj or MDAnalysis to update the coordinates of the systems and use them to update parmed positions object? I am trying to connect the dots here but I feel I am missing something.

[Peter Eastman]

Your code is a little misleading:

positions = pmd.load_file(coord)

The value returned by load_file(), which you pass as the first argument to energy_decomposition_system(), is actually a Structure, not a set of positions. It assumes you have already set the positions by calling setPositions() on the Context.

You can just loop over frames, call setPositions() for each one, and then ask ParmEd to produce an energy decomposition. For example, if you have a MDTraj trajectory you can do something like this.

Code: Select all

for frame in range(traj.n_frames):
    simulation.context.setPositions(traj.xyz[frame])
    energies = pmd.openmm.energy_decomposition_system(structure, system)
    # Do something with the result

Make sure the positions are in the correct units.

[Marcelo D Poleto]

Hi Peter,

Thanks for your input and correction. However, I don't see how the Structure object is updated with the new positions. The simulation context is being updated, but I don't think that is being passed to Structure. I tested your suggestion with the code below:

Code: Select all

top   = 'ala.c36.psf'
coord = 'ala.c36.crd'
traj  = '100frames.dcd'

psf  = CharmmPsfFile(top)
crd  = CharmmCrdFile(coord)

box = 2.9*nanometer #
psf.setBox(box,box,box)
system = psf.createSystem(charmm_params, nonbondedMethod=PME, nonbondedCutoff=1.2*nanometers, switchDistance=1.0*nanometer, ewaldErrorTolerance = 0.0005, constraints=HBonds)

integrator = LangevinIntegrator(298, 1/picosecond, 0.002)
simulation = Simulation(psf.topology, system, integrator)
simulation.context.setPositions(crd.positions)

#####
# Initial evaluation
state = simulation.context.getState(getEnergy=True)
potenergy = state.getPotentialEnergy().value_in_unit(kilojoules_per_mole)

energies = pmd.openmm.energy_decomposition_system(stru, system)

totalpotenergy_array  = []
for term, energy in energies:
    energy = energy*4.184
    totalpotenergy_array.append(energy)
    
totalpotenergy  = np.sum(totalpotenergy_array)
print("Initial", potenergy,totalpotenergy)

####
# Loop over trajectory
traj = md.load(dcd_file, top=psf_file)
for frame in range(traj.n_frames):
    newpositions = traj.xyz[frame]
    simulation.context.setPositions(newpositions)

    state = simulation.context.getState(getEnergy=True)
    potenergy = state.getPotentialEnergy().value_in_unit(kilojoules_per_mole)

    energies = pmd.openmm.energy_decomposition_system(stru, system)

    totalpotenergy_array  = []

    for term, energy in energies:
        energy = energy*4.184
        totalpotenergy_array.append(energy)
    
    totalpotenergy  = np.sum(totalpotenergy_array)
    print(frame, potenergy, totalpotenergy)

If I run this, what I get is this comparison:

Code: Select all

Initial  -29602.05819440112   -29602.055461647367
0        -32901.96444440112   -29602.047649147367
1        -31255.77694440112   -29602.047649147367
2        -30556.41756940112   -29602.063274147367
3        -29760.02694440112   -29602.055705787992
4        -29603.23006940112   -29602.055461647367
5        -29794.88631940112   -29602.055705787992
6        -29388.67538190112   -29602.055705787992
7        -29687.15975690112   -29602.055461647367
8        -29786.55819440112   -29602.063518287992
9        -29441.23006940112   -29602.063274147367

Am I missing something?

[Peter Eastman]

It looks like there are two different functions: energy_decomposition(), which uses the positions already stored in the Context, and energy_decomposition_system(), which takes them from the Structure. Can you just set them directly?

Code: Select all

structure.positions = traj.xyz[frame]

[Marcelo D Poleto]

Not really.

Code: Select all

stru.positions = traj.xyz[frame]
AttributeError: can't set attribute

I also tried to use:

Code: Select all

energies = pmd.openmm.energy_decomposition(stru, simulation.context)

But I get the following:

Code: Select all

File "/home/mdpoleto/Documents/Softwares/anaconda3/envs/openmm7.7/lib/python3.9/site-packages/parmed/openmm/utils.py", line 43, in energy_decomposition
    force_group_names[gp] = all_names[gp]
KeyError: 0

[Peter Eastman]

I don't know what the intended idiom is then. Since this is a ParmEd question, I suggest asking at https://github.com/ParmEd/ParmEd/issues.

Alternatively you could decompose the energies directly through OpenMM rather than using ParmEd. See this cookbook entry for a description of how to do that.

[Marcelo D Poleto]

Hi,

I just want to update the topic with my final code. I checked that forces from this code matches the forces calculated with the CHARMM software.

Code: Select all

import simtk.openmm as mm
from simtk.openmm import *
from simtk.openmm.app import *
from simtk.unit import *
import mdtraj as md

import os, argparse
import numpy as np

def forcegroupify(system):
    forcegroups = {}
    for i in range(system.getNumForces()):
        force = system.getForce(i)
        force.setForceGroup(i)
        forcename = force.getName()
        forcegroups[forcename] = i
    return forcegroups


psffile = "ala.c36.psf"
crdfile = "ala.c36.crd"
dcdfile = "ala.c36.dcd"

# reading and parsing topology and coordinates
psf = CharmmPsfFile(psffile)
crd = CharmmCrdFile(crdfile)

# applying boxvectors to psf
box = 2.9*nanometer #
psf.setBox(box,box,box)

# reading parameters and creating system
system = psf.createSystem(charmm_params, nonbondedMethod=PME, nonbondedCutoff=1.2*nanometers, switchDistance=1.0*nanometer, ewaldErrorTolerance = 0.0005, constraints=HBonds)

# making sure each ForceGroup has its own index
forcegroups = {}
for i in range(system.getNumForces()):
    force = system.getForce(i)
    force.setForceGroup(i)
    forcename = force.getName()
    forcegroups[forcename] = i

integrator = LangevinIntegrator(temperature, 1/picosecond, dt)
simulation = Simulation(psf.topology, system, integrator)
simulation.context.setPositions(crd.positions)

######################################
# Calculating Energies for entire trajectory
traj = md.load(cmd.dcd, top=cmd.psf)

print("\n>>> Reading and analyzing:")
for frame in range(traj.n_frames):
    print("> Frame: " + str(frame))
    
    # update simulation context with new positions
    simulation.context.setPositions(traj.xyz[frame])

    # update simulation context with new box vectors
    a, b, c = traj.unitcell_vectors[frame]
    simulation.context.setPeriodicBoxVectors(*(a,b,c))

    # get updated energies
    state = simulation.context.getState(getEnergy=True)
    potenergy = getPotentialEnergy(state)
    
    for f, i in forcegroups.items():
        energy = simulation.context.getState(getEnergy=True, groups=2**i).getPotentialEnergy()
        energy = energy.value_in_unit(kilojoules_per_mole) # remove unit
        print(f,i)

[Peter Eastman]

Thanks! By the way, your code can be simplified a bit. Instead of

Code: Select all

    for i in range(system.getNumForces()):
        force = system.getForce(i)
        force.setForceGroup(i)
        forcename = force.getName()
        forcegroups[forcename] = i

you can write

Code: Select all

    for i, force in enumerate(system.getForces()):
        force.setForceGroup(i)
        forcegroups[force.getName()] = i