Merge pull request #1895 from peastman/periodic

Reporters can request wrapped or unwrapped coordinates

docs-source/usersguide/application.rst

@@ -1794,7 +1794,7 @@ Here is the definition of the :class:`ForceReporter` class:
 
             def describeNextReport(self, simulation):
                 steps = self._reportInterval - simulation.currentStep%self._reportInterval
-                return (steps, False, False, True, False)
+                return (steps, False, False, True, False, None)
 
             def report(self, simulation, state):
                 forces = state.getForces().value_in_unit(kilojoules/mole/nanometer)
@@ -1814,7 +1814,7 @@ We then have two methods that every reporter must implement:
 :meth:`describeNextReport()` and :meth:`report()`.  A Simulation object
 periodically calls :meth:`describeNextReport()` on each of its reporters to
 find out when that reporter will next generate a report, and what information
-will be needed to generate it.  The return value should be a five element tuple,
+will be needed to generate it.  The return value should be a six element tuple,
 whose elements are as follows:
 
 * The number of time steps until the next report.  We calculate this as
@@ -1825,6 +1825,9 @@ whose elements are as follows:
 * Whether the next report will need particle velocities.
 * Whether the next report will need forces.
 * Whether the next report will need energies.
+* Whether the positions should be wrapped to the periodic box.  If None, it will
+  automatically decide whether to wrap positions based on whether the System uses
+  periodic boundary conditions.
 
 
 When the time comes for the next scheduled report, the :class:`Simulation` calls

wrappers/python/simtk/openmm/app/dcdreporter.py

@@ -42,7 +42,7 @@ class DCDReporter(object):
     To use it, create a DCDReporter, then add it to the Simulation's list of reporters.
     """
 
-    def __init__(self, file, reportInterval, append=False):
+    def __init__(self, file, reportInterval, append=False, enforcePeriodicBox=None):
         """Create a DCDReporter.
 
         Parameters
@@ -53,9 +53,15 @@ class DCDReporter(object):
             The interval (in time steps) at which to write frames
         append : bool=False
             If True, open an existing DCD file to append to.  If False, create a new file.
+        enforcePeriodicBox: bool
+            Specifies whether particle positions should be translated so the center of every molecule
+            lies in the same periodic box.  If None (the default), it will automatically decide whether
+            to translate molecules based on whether the system being simulated uses periodic boundary
+            conditions.
         """
         self._reportInterval = reportInterval
         self._append = append
+        self._enforcePeriodicBox = enforcePeriodicBox
         if append:
             mode = 'r+b'
         else:
@@ -74,13 +80,14 @@ class DCDReporter(object):
         Returns
         -------
         tuple
-            A five element tuple. The first element is the number of steps
-            until the next report. The remaining elements specify whether
+            A six element tuple. The first element is the number of steps
+            until the next report. The next four elements specify whether
             that report will require positions, velocities, forces, and
-            energies respectively.
+            energies respectively.  The final element specifies whether
+            positions should be wrapped to lie in a single periodic box.
         """
         steps = self._reportInterval - simulation.currentStep%self._reportInterval
-        return (steps, True, False, False, False)
+        return (steps, True, False, False, False, self._enforcePeriodicBox)
 
     def report(self, simulation, state):
         """Generate a report.

wrappers/python/simtk/openmm/app/pdbreporter.py

@@ -41,7 +41,7 @@ class PDBReporter(object):
     To use it, create a PDBReporter, then add it to the Simulation's list of reporters.
     """
 
-    def __init__(self, file, reportInterval):
+    def __init__(self, file, reportInterval, enforcePeriodicBox=None):
         """Create a PDBReporter.
 
         Parameters
@@ -50,8 +50,14 @@ class PDBReporter(object):
             The file to write to
         reportInterval : int
             The interval (in time steps) at which to write frames
+        enforcePeriodicBox: bool
+            Specifies whether particle positions should be translated so the center of every molecule
+            lies in the same periodic box.  If None (the default), it will automatically decide whether
+            to translate molecules based on whether the system being simulated uses periodic boundary
+            conditions.
         """
         self._reportInterval = reportInterval
+        self._enforcePeriodicBox = enforcePeriodicBox
         self._out = open(file, 'w')
         self._topology = None
         self._nextModel = 0
@@ -67,13 +73,14 @@ class PDBReporter(object):
         Returns
         -------
         tuple
-            A five element tuple. The first element is the number of steps
-            until the next report. The remaining elements specify whether
+            A six element tuple. The first element is the number of steps
+            until the next report. The next four elements specify whether
             that report will require positions, velocities, forces, and
-            energies respectively.
+            energies respectively.  The final element specifies whether
+            positions should be wrapped to lie in a single periodic box.
         """
         steps = self._reportInterval - simulation.currentStep%self._reportInterval
-        return (steps, True, False, False, False)
+        return (steps, True, False, False, False, self._enforcePeriodicBox)
 
     def report(self, simulation, state):
         """Generate a report.

wrappers/python/simtk/openmm/app/simulation.py

@@ -183,6 +183,9 @@ class Simulation(object):
         nextReport = [None]*len(self.reporters)
         while self.currentStep < endStep and (endTime is None or datetime.now() < endTime):
             nextSteps = endStep-self.currentStep
+            
+            # Find when the next report will happen.
+            
             anyReport = False
             for i, reporter in enumerate(self.reporters):
                 nextReport[i] = reporter.describeNextReport(self)
@@ -198,25 +201,57 @@ class Simulation(object):
             self.integrator.step(stepsToGo)
             self.currentStep += nextSteps
             if anyReport:
-                getPositions = False
-                getVelocities = False
-                getForces = False
-                getEnergy = False
-                for reporter, next in zip(self.reporters, nextReport):
-                    if next[0] == nextSteps:
-                        if next[1]:
-                            getPositions = True
-                        if next[2]:
-                            getVelocities = True
-                        if next[3]:
-                            getForces = True
-                        if next[4]:
-                            getEnergy = True
-                state = self.context.getState(getPositions=getPositions, getVelocities=getVelocities, getForces=getForces,
-                                              getEnergy=getEnergy, getParameters=True, enforcePeriodicBox=self._usesPBC)
-                for reporter, next in zip(self.reporters, nextReport):
-                    if next[0] == nextSteps:
-                        reporter.report(self, state)
+                # One or more reporters are ready to generate reports.  Organize them into three
+                # groups: ones that want wrapped positions, ones that want unwrapped positions,
+                # and ones that don't care about positions.
+                
+                wrapped = []
+                unwrapped = []
+                either = []
+                for reporter, report in zip(self.reporters, nextReport):
+                    if report[0] == nextSteps:
+                        if len(report) > 5:
+                            wantWrap = report[5]
+                            if wantWrap is None:
+                                wantWrap = self._usesPBC
+                        else:
+                            wantWrap = self._usesPBC
+                        if not report[1]:
+                            either.append((reporter, report))
+                        elif wantWrap:
+                            wrapped.append((reporter, report))
+                        else:
+                            unwrapped.append((reporter, report))
+                if len(wrapped) > len(unwrapped):
+                    wrapped += either
+                else:
+                    unwrapped += either
+                
+                # Generate the reports.
+
+                if len(wrapped) > 0:
+                    self._generate_reports(wrapped, True)
+                if len(unwrapped) > 0:
+                    self._generate_reports(unwrapped, False)
+    
+    def _generate_reports(self, reports, periodic):
+        getPositions = False
+        getVelocities = False
+        getForces = False
+        getEnergy = False
+        for reporter, next in reports:
+            if next[1]:
+                getPositions = True
+            if next[2]:
+                getVelocities = True
+            if next[3]:
+                getForces = True
+            if next[4]:
+                getEnergy = True
+        state = self.context.getState(getPositions=getPositions, getVelocities=getVelocities, getForces=getForces,
+                                      getEnergy=getEnergy, getParameters=True, enforcePeriodicBox=periodic)
+        for reporter, next in reports:
+            reporter.report(self, state)
 
     def saveCheckpoint(self, file):
         """Save a checkpoint of the simulation to a file.

wrappers/python/tests/TestSimulation.py

@@ -157,6 +157,38 @@ class TestSimulation(unittest.TestCase):
         simulation.loadState(stateFile)
         self.assertEqual(velocities, simulation.context.getState(getVelocities=True).getVelocities())
 
+    def testWrappedCoordinates(self):
+        """Test generating reports with and without wrapped coordinates."""
+        pdb = PDBFile('systems/alanine-dipeptide-explicit.pdb')
+        ff = ForceField('amber99sb.xml', 'tip3p.xml')
+        system = ff.createSystem(pdb.topology, nonbondedMethod=CutoffPeriodic, constraints=HBonds)
+        integrator = LangevinIntegrator(300*kelvin, 1/picosecond, 0.002*picoseconds)
+
+        class CompareCoordinatesReporter(object):
+            def __init__(self, periodic):
+                self.periodic = periodic
+                self.interval = 100
+                
+            def describeNextReport(self, simulation):
+                steps = self.interval - simulation.currentStep%self.interval
+                return (steps, True, False, False, False, self.periodic)
+        
+            def report(self, simulation, state):
+                state2 = simulation.context.getState(getPositions=True, enforcePeriodicBox=self.periodic)
+                assert state.getPositions() == state2.getPositions()
+
+        # Create a Simulation.
+
+        simulation = Simulation(pdb.topology, system, integrator)
+        simulation.context.setPositions(pdb.positions)
+        simulation.context.setVelocitiesToTemperature(300*kelvin)
+        simulation.reporters.append(CompareCoordinatesReporter(False))
+        simulation.reporters.append(CompareCoordinatesReporter(True))
+        
+        # Run for a little while and make sure the reporters don't find any problems.
+        
+        simulation.step(500)
+
 
 if __name__ == '__main__':
     unittest.main()