simtk.org

Code: Select all

class XyzReporter(object):
    def __init__(self, file, reportInterval):
        self._reportInterval = reportInterval
        self._out = None
        self.fname = file

    def describeNextReport(self, simulation):
        stepsleft = simulation.currentStep % self._reportInterval
        steps = self._reportInterval - stepsleft
        return (steps, True, True, False, False)

    def report(self, simulation, state):
        crds = state.getPositions().value_in_unit(unit.angstrom)
        vels = state.getVelocities().value_in_unit(unit.angstrom / unit.femtosecond)
        if self._out is None:
            self.atom = len(crds)
            self._out = open(self.fname, 'w')

        self._out.write(f"{self.atom}\nOpenMM Trajectory\n")
        # Add the coordinates, velocities, and/or forces as needed
        atoms=simulation.topology.atoms()
        for atom, crd, vel in zip(atoms, crds, vels):
            self._out.write(f"{atom.name} {crd[0]} {crd[1]} {crd[2]} {vel[0]} {vel[1]} {vel[2]}\n")