Fixes to benchmark script (#2982)

examples/benchmark.py

@@ -34,7 +34,7 @@ def downloadAmberSuite():
 
 def runOneTest(testName, options):
     """Perform a single benchmarking simulation."""
-    explicit = (testName in ('rf', 'pme', 'amoebapme'))
+    explicit = (testName not in ('gbsa', 'amoebagk'))
     amoeba = (testName in ('amoebagk', 'amoebapme'))
     apoa1 = testName.startswith('apoa1')
     amber = (testName.startswith('amber'))
@@ -46,10 +46,16 @@ def runOneTest(testName, options):
         print('Test: pme (cutoff=%g)' % options.cutoff)
     else:
         print('Test: %s' % testName)
+    print('Ensemble: %s' % options.ensemble)
     platform = mm.Platform.getPlatformByName(options.platform)
     
     # Create the System.
-    
+
+    temperature = 300*unit.kelvin
+    if explicit:
+        friction = 1*(1/unit.picoseconds)
+    else:
+        friction = 91*(1/unit.picoseconds)
     if amoeba:
         constraints = None
         epsilon = float(options.epsilon)
@@ -69,7 +75,11 @@ def runOneTest(testName, options):
             if isinstance(f, mm.AmoebaMultipoleForce) or isinstance(f, mm.AmoebaVdwForce) or isinstance(f, mm.AmoebaGeneralizedKirkwoodForce) or isinstance(f, mm.AmoebaWcaDispersionForce):
                 f.setForceGroup(1)
         dt = 0.002*unit.picoseconds
-        integ = mm.MTSIntegrator(dt, [(0,2), (1,1)])
+        if options.ensemble == 'NVE':
+            integ = mm.MTSIntegrator(dt, [(0,2), (1,1)])
+        else:
+            integ = mm.MTSLangevinIntegrator(temperature, friction, dt, [(0,2), (1,1)])
+        positions = pdb.positions
     elif amber:
         dirname = downloadAmberSuite()
         names = {'amber20-dhfr':'JAC',  'amber20-factorix':'FactorIX', 'amber20-cellulose':'Cellulose', 'amber20-stmv':'STMV'}
@@ -82,15 +92,11 @@ def runOneTest(testName, options):
         method = app.PME
         cutoff = options.cutoff
         constraints = app.HBonds
-        friction = 1*(1/unit.picoseconds)
         system = prmtop.createSystem(nonbondedMethod=method, nonbondedCutoff=cutoff, constraints=constraints)
-        print('Ensemble: %s' % options.ensemble)
-        if options.ensemble == 'NPT':
-            system.addForce(mm.MonteCarloBarostat(1*unit.bar, 300*unit.kelvin))
         if options.ensemble == 'NVE':
             integ = mm.VerletIntegrator(dt)
         else:
-            integ = mm.LangevinMiddleIntegrator(300*unit.kelvin, friction, dt)
+            integ = mm.LangevinMiddleIntegrator(temperature, friction, dt)
     else:
         if apoa1:
             ff = app.ForceField('amber14/protein.ff14SB.xml', 'amber14/lipid17.xml', 'amber14/tip3p.xml')
@@ -104,7 +110,6 @@ def runOneTest(testName, options):
             else:
                 method = app.CutoffPeriodic
                 cutoff = 1*unit.nanometers
-            friction = 1*(1/unit.picoseconds)
             hydrogenMass = 1.5*unit.amu
         elif explicit:
             ff = app.ForceField('amber99sb.xml', 'tip3p.xml')
@@ -115,24 +120,30 @@ def runOneTest(testName, options):
             else:
                 method = app.CutoffPeriodic
                 cutoff = 1*unit.nanometers
-            friction = 1*(1/unit.picoseconds)
         else:
             ff = app.ForceField('amber99sb.xml', 'amber99_obc.xml')
             pdb = app.PDBFile('5dfr_minimized.pdb')
             method = app.CutoffNonPeriodic
             cutoff = 2*unit.nanometers
-            friction = 91*(1/unit.picoseconds)
         if options.heavy:
             dt = 0.005*unit.picoseconds
             constraints = app.AllBonds
             hydrogenMass = 4*unit.amu
-            integ = mm.LangevinIntegrator(300*unit.kelvin, friction, dt)
+            if options.ensemble == 'NVE':
+                integ = mm.VerletIntegrator(dt)
+            else:
+                integ = mm.LangevinIntegrator(temperature, friction, dt)
         else:
             dt = 0.004*unit.picoseconds
             constraints = app.HBonds
-            integ = mm.LangevinMiddleIntegrator(300*unit.kelvin, friction, dt)
+            if options.ensemble == 'NVE':
+                integ = mm.VerletIntegrator(dt)
+            else:
+                integ = mm.LangevinMiddleIntegrator(temperature, friction, dt)
         positions = pdb.positions
         system = ff.createSystem(pdb.topology, nonbondedMethod=method, nonbondedCutoff=cutoff, constraints=constraints, hydrogenMass=hydrogenMass)
+    if options.ensemble == 'NPT':
+        system.addForce(mm.MonteCarloBarostat(1*unit.bar, temperature, 100))
     print('Step Size: %g fs' % dt.value_in_unit(unit.femtoseconds))
     properties = {}
     initialSteps = 5
@@ -151,11 +162,11 @@ def runOneTest(testName, options):
     else:
         context = mm.Context(system, integ, platform)
     context.setPositions(positions)
-    if not amber:
-        context.setVelocitiesToTemperature(300*unit.kelvin)
     if amber:
         if inpcrd.boxVectors is not None:
             context.setPeriodicBoxVectors(*inpcrd.boxVectors)
+        mm.LocalEnergyMinimizer.minimize(context, 100*unit.kilojoules_per_mole/unit.nanometer)
+    context.setVelocitiesToTemperature(temperature)
 
     steps = 20
     while True:
@@ -176,7 +187,7 @@ platformNames = [mm.Platform.getPlatform(i).getName() for i in range(mm.Platform
 parser.add_argument('--platform', dest='platform', choices=platformNames, help='name of the platform to benchmark')
 parser.add_argument('--test', dest='test', choices=('gbsa', 'rf', 'pme', 'apoa1rf', 'apoa1pme', 'apoa1ljpme', 'amoebagk', 'amoebapme', 'amber20-dhfr',  'amber20-factorix', 'amber20-cellulose', 'amber20-stmv'), 
     help='the test to perform: gbsa, rf, pme, apoa1rf, apoa1pme, apoa1ljpme, amoebagk, amoebapme,  amber20-dhfr,  amber20-factorix, amber20-cellulose, amber20-stmv [default: all except amber-*]')
-parser.add_argument('--ensemble', default='NPT', dest='ensemble', choices=('NPT', 'NVE', 'NVT'), help='the ensemble (for Amber tests only) [default: NPT]')
+parser.add_argument('--ensemble', default='NVT', dest='ensemble', choices=('NPT', 'NVE', 'NVT'), help='the thermodynamic ensemble to simulate [default: NVT]')
 parser.add_argument('--pme-cutoff', default=0.9, dest='cutoff', type=float, help='direct space cutoff for PME in nm [default: 0.9]')
 parser.add_argument('--seconds', default=60, dest='seconds', type=float, help='target simulation length in seconds [default: 60]')
 parser.add_argument('--polarization', default='mutual', dest='polarization', choices=('direct', 'extrapolated', 'mutual'), help='the polarization method for AMOEBA: direct, extrapolated, or mutual [default: mutual]')