-
Notifications
You must be signed in to change notification settings - Fork 1
/
benchmark.py
161 lines (151 loc) · 7.66 KB
/
benchmark.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
from __future__ import print_function
import simtk.openmm.app as app
import simtk.openmm as mm
import simtk.unit as unit
import sys
from datetime import datetime
from optparse import OptionParser
def timeIntegration(context, steps, initialSteps):
"""Integrate a Context for a specified number of steps, then return how many seconds it took."""
context.getIntegrator().step(initialSteps) # Make sure everything is fully initialized
context.getState(getEnergy=True)
start = datetime.now()
context.getIntegrator().step(steps)
context.getState(getEnergy=True)
end = datetime.now()
elapsed = end-start
return elapsed.seconds + elapsed.microseconds*1e-6
def runOneTest(testName, options):
"""Perform a single benchmarking simulation."""
explicit = (testName in ('rf', 'pme', 'amoebapme'))
amoeba = (testName in ('amoebagk', 'amoebapme'))
apoa1 = testName.startswith('apoa1')
hydrogenMass = None
print()
if amoeba:
print('Test: %s (epsilon=%g)' % (testName, options.epsilon))
elif testName == 'pme':
print('Test: pme (cutoff=%g)' % options.cutoff)
else:
print('Test: %s' % testName)
platform = mm.Platform.getPlatformByName(options.platform)
# Create the System.
if amoeba:
constraints = None
epsilon = float(options.epsilon)
if explicit:
ff = app.ForceField('amoeba2009.xml')
pdb = app.PDBFile('5dfr_solv-cube_equil.pdb')
cutoff = 0.7*unit.nanometers
vdwCutoff = 0.9*unit.nanometers
system = ff.createSystem(pdb.topology, nonbondedMethod=app.PME, nonbondedCutoff=cutoff, vdwCutoff=vdwCutoff, constraints=constraints, ewaldErrorTolerance=0.00075, mutualInducedTargetEpsilon=epsilon, polarization=options.polarization)
else:
ff = app.ForceField('amoeba2009.xml', 'amoeba2009_gk.xml')
pdb = app.PDBFile('5dfr_minimized.pdb')
cutoff = 2.0*unit.nanometers
vdwCutoff = 1.2*unit.nanometers
system = ff.createSystem(pdb.topology, nonbondedMethod=app.NoCutoff, constraints=constraints, mutualInducedTargetEpsilon=epsilon, polarization=options.polarization)
for f in system.getForces():
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)])
else:
if apoa1:
ff = app.ForceField('amber14/protein.ff14SB.xml', 'amber14/lipid17.xml', 'amber14/tip3p.xml')
pdb = app.PDBFile('apoa1.pdb')
if testName == 'apoa1pme':
method = app.PME
cutoff = options.cutoff
elif testName == 'apoa1ljpme':
method = app.LJPME
cutoff = options.cutoff
else:
method = app.CutoffPeriodic
cutoff = 1*unit.nanometers
friction = 1*(1/unit.picoseconds)
elif explicit:
ff = app.ForceField('amber99sb.xml', 'tip3p.xml')
pdb = app.PDBFile('5dfr_solv-cube_equil.pdb')
if testName == 'pme':
method = app.PME
cutoff = options.cutoff
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
else:
dt = 0.002*unit.picoseconds
constraints = app.HBonds
hydrogenMass = None
system = ff.createSystem(pdb.topology, nonbondedMethod=method, nonbondedCutoff=cutoff, constraints=constraints, hydrogenMass=hydrogenMass)
integ = mm.LangevinIntegrator(300*unit.kelvin, friction, dt)
print('Step Size: %g fs' % dt.value_in_unit(unit.femtoseconds))
properties = {}
initialSteps = 5
if options.device is not None and platform.getName() in ('CUDA', 'OpenCL'):
properties['DeviceIndex'] = options.device
if ',' in options.device or ' ' in options.device:
initialSteps = 250
if options.precision is not None and platform.getName() in ('CUDA', 'OpenCL'):
properties['Precision'] = options.precision
# Run the simulation.
integ.setConstraintTolerance(1e-5)
if len(properties) > 0:
context = mm.Context(system, integ, platform, properties)
else:
context = mm.Context(system, integ, platform)
context.setPositions(pdb.positions)
context.setVelocitiesToTemperature(300*unit.kelvin)
steps = 20
while True:
time = timeIntegration(context, steps, initialSteps)
if time >= 0.5*options.seconds:
break
if time < 0.5:
steps = int(steps*1.0/time) # Integrate enough steps to get a reasonable estimate for how many we'll need.
else:
steps = int(steps*options.seconds/time)
print('Integrated %d steps in %g seconds' % (steps, time))
print('%g ns/day' % (dt*steps*86400/time).value_in_unit(unit.nanoseconds))
# Parse the command line options.
parser = OptionParser()
platformNames = [mm.Platform.getPlatform(i).getName() for i in range(mm.Platform.getNumPlatforms())]
parser.add_option('--platform', dest='platform', choices=platformNames, help='name of the platform to benchmark')
parser.add_option('--test', dest='test', choices=('gbsa', 'rf', 'pme', 'apoa1rf', 'apoa1pme', 'apoa1ljpme', 'amoebagk', 'amoebapme'), help='the test to perform: gbsa, rf, pme, apoa1rf, apoa1pme, apoa1ljpme, amoebagk, or amoebapme [default: all]')
parser.add_option('--pme-cutoff', default='0.9', dest='cutoff', type='float', help='direct space cutoff for PME in nm [default: 0.9]')
parser.add_option('--seconds', default='60', dest='seconds', type='float', help='target simulation length in seconds [default: 60]')
parser.add_option('--polarization', default='mutual', dest='polarization', choices=('direct', 'extrapolated', 'mutual'), help='the polarization method for AMOEBA: direct, extrapolated, or mutual [default: mutual]')
parser.add_option('--mutual-epsilon', default='1e-5', dest='epsilon', type='float', help='mutual induced epsilon for AMOEBA [default: 1e-5]')
parser.add_option('--heavy-hydrogens', action='store_true', default=False, dest='heavy', help='repartition mass to allow a larger time step')
parser.add_option('--device', default=None, dest='device', help='device index for CUDA or OpenCL')
parser.add_option('--precision', default='single', dest='precision', choices=('single', 'mixed', 'double'), help='precision mode for CUDA or OpenCL: single, mixed, or double [default: single]')
(options, args) = parser.parse_args()
if len(args) > 0:
parser.error('Unknown argument: '+args[0])
if options.platform is None:
parser.error('No platform specified')
print('Platform:', options.platform)
if options.platform in ('CUDA', 'OpenCL'):
print('Precision:', options.precision)
if options.device is not None:
print('Device:', options.device)
# Run the simulations.
if options.test is None:
for test in ('gbsa', 'rf', 'pme', 'apoa1rf', 'apoa1pme', 'apoa1ljpme', 'amoebagk', 'amoebapme'):
try:
runOneTest(test, options)
except Exception as ex:
print('Test failed: %s' % ex.message)
else:
runOneTest(options.test, options)