Optimizations to reporters (openmm#4330)

* Optimizations to reporters

* Removed unneeded imports

wrappers/python/openmm/app/dcdfile.py

@@ -121,9 +121,10 @@ def writeModel(self, positions, unitCellDimensions=None, periodicBoxVectors=None
             raise ValueError('The number of positions must match the number of atoms')
         if is_quantity(positions):
             positions = positions.value_in_unit(nanometers)
-        if any(math.isnan(norm(pos)) for pos in positions):
+        import numpy as np
+        if np.isnan(positions).any():
             raise ValueError('Particle position is NaN.  For more information, see https://github.com/openmm/openmm/wiki/Frequently-Asked-Questions#nan')
-        if any(math.isinf(norm(pos)) for pos in positions):
+        if np.isinf(positions).any():
             raise ValueError('Particle position is infinite.  For more information, see https://github.com/openmm/openmm/wiki/Frequently-Asked-Questions#nan')
         file = self._file
 

wrappers/python/openmm/app/dcdreporter.py

@@ -104,7 +104,7 @@ def report(self, simulation, state):
                 self._out, simulation.topology, simulation.integrator.getStepSize(),
                 simulation.currentStep, self._reportInterval, self._append
             )
-        self._dcd.writeModel(state.getPositions(), periodicBoxVectors=state.getPeriodicBoxVectors())
+        self._dcd.writeModel(state.getPositions(asNumpy=True), periodicBoxVectors=state.getPeriodicBoxVectors())
 
     def __del__(self):
         self._out.close()

wrappers/python/openmm/app/pdbfile.py

@@ -343,9 +343,10 @@ def writeModel(topology, positions, file=sys.stdout, modelIndex=None, keepIds=Fa
             raise ValueError('The number of positions must match the number of atoms')
         if is_quantity(positions):
             positions = positions.value_in_unit(angstroms)
-        if any(math.isnan(norm(pos)) for pos in positions):
+        import numpy as np
+        if np.isnan(positions).any():
             raise ValueError('Particle position is NaN.  For more information, see https://github.com/openmm/openmm/wiki/Frequently-Asked-Questions#nan')
-        if any(math.isinf(norm(pos)) for pos in positions):
+        if np.isinf(positions).any():
             raise ValueError('Particle position is infinite.  For more information, see https://github.com/openmm/openmm/wiki/Frequently-Asked-Questions#nan')
         nonHeterogens = PDBFile._standardResidues[:]
         nonHeterogens.remove('HOH')

wrappers/python/openmm/app/pdbreporter.py

@@ -104,12 +104,12 @@ def report(self, simulation, state):
             topology = self._subsetTopology
 
             #PDBFile will convert to angstroms so do it here first instead
-            positions = state.getPositions().value_in_unit(angstroms) 
+            positions = state.getPositions(asNumpy=True).value_in_unit(angstroms)
             positions = [positions[i] for i in self._atomSubset]
 
         else:
             topology = simulation.topology
-            positions = state.getPositions()
+            positions = state.getPositions(asNumpy=True)
 
         if self._nextModel == 0:
             PDBFile.writeHeader(topology, self._out)
@@ -202,12 +202,12 @@ def report(self, simulation, state):
             topology = self._subsetTopology
 
             #PDBFile will convert to angstroms so do it here first instead
-            positions = state.getPositions().value_in_unit(angstroms) 
+            positions = state.getPositions(asNumpy=True).value_in_unit(angstroms)
             positions = [positions[i] for i in self._atomSubset]
 
         else:
             topology = simulation.topology
-            positions = state.getPositions()
+            positions = state.getPositions(asNumpy=True)
 
         if self._nextModel == 0:
             PDBxFile.writeHeader(topology, self._out)

wrappers/python/openmm/app/pdbxfile.py

@@ -418,9 +418,10 @@ def writeModel(topology, positions, file=sys.stdout, modelIndex=1, keepIds=False
             raise ValueError('The number of positions must match the number of atoms')
         if is_quantity(positions):
             positions = positions.value_in_unit(angstroms)
-        if any(math.isnan(norm(pos)) for pos in positions):
+        import numpy as np
+        if np.isnan(positions).any():
             raise ValueError('Particle position is NaN.  For more information, see https://github.com/openmm/openmm/wiki/Frequently-Asked-Questions#nan')
-        if any(math.isinf(norm(pos)) for pos in positions):
+        if np.isinf(positions).any():
             raise ValueError('Particle position is infinite.  For more information, see https://github.com/openmm/openmm/wiki/Frequently-Asked-Questions#nan')
         nonHeterogens = PDBFile._standardResidues[:]
         nonHeterogens.remove('HOH')

wrappers/python/openmm/app/xtcfile.py

@@ -9,11 +9,9 @@
     get_xtc_nframes,
     get_xtc_natoms,
 )
-import numpy as np
 import os
 from openmm import Vec3
 from openmm.unit import nanometers, picoseconds, is_quantity, norm
-import math
 import tempfile
 import shutil
 
@@ -92,11 +90,12 @@ def writeModel(self, positions, unitCellDimensions=None, periodicBoxVectors=None
             raise ValueError("The number of positions must match the number of atoms")
         if is_quantity(positions):
             positions = positions.value_in_unit(nanometers)
-        if any(math.isnan(norm(pos)) for pos in positions):
+        import numpy as np
+        if np.isnan(positions).any():
             raise ValueError(
                 "Particle position is NaN.  For more information, see https://github.com/openmm/openmm/wiki/Frequently-Asked-Questions#nan"
             )
-        if any(math.isinf(norm(pos)) for pos in positions):
+        if np.isinf(positions).any():
             raise ValueError(
                 "Particle position is infinite.  For more information, see https://github.com/openmm/openmm/wiki/Frequently-Asked-Questions#nan"
             )

wrappers/python/openmm/app/xtcreporter.py

@@ -71,5 +71,5 @@ def report(self, simulation, state):
                 self._append,
             )
         self._xtc.writeModel(
-            state.getPositions(), periodicBoxVectors=state.getPeriodicBoxVectors()
+            state.getPositions(asNumpy=True), periodicBoxVectors=state.getPeriodicBoxVectors()
         )