Skip to content

GitLab

Unverified Commit 4c2f3323 authored by Peter Eastman's avatar Peter Eastman Committed by GitHub
Browse files

Overhaul benchmarking script to add lots of new features (#3887) 

* Overhaul benchmarking script to add lots of new features

* Added another example
parent bf1cae4f
Show whitespace changes
Inline Side-by-side
Showing with 441 additions and 90 deletions
examples/benchmark.py
View file @ 4c2f3323
...@@ -3,8 +3,98 @@ import openmm.app as app ...@@ -3,8 +3,98 @@ import openmm.app as app
import openmm as mm import openmm as mm
import openmm.unit as unit import openmm.unit as unit
from datetime import datetime from datetime import datetime
import os import argparse
from argparse import ArgumentParser
def cpuinfo():
"""Return CPU info"""
import os, platform, subprocess, re
if platform.system() == "Windows":
return platform.processor()
elif platform.system() == "Darwin":
os.environ['PATH'] = os.environ['PATH'] + os.pathsep + '/usr/sbin'
command ="sysctl -n machdep.cpu.brand_string"
return subprocess.check_output(command, shell=True, text=True).strip()
elif platform.system() == "Linux":
command = "cat /proc/cpuinfo"
all_info = subprocess.check_output(command, shell=True, text=True).strip()
for line in all_info.split("\n"):
if "model name" in line:
return re.sub( ".*model name.*: ", "", line,1)
return ""
def appendTestResult(filename=None, test_result=None, system_info=None):
"""Append a test result to a JSON or YAML file.
TODO: The efficiency of this could be improved.
Parameters
----------
filename : str, optional, default=None
The filename to append a result to, ending either in .yaml or .json
If None, no result is written
test_result : dict, optional, default=None
The test result to append to the 'benchmarks' blcok
system_info : dict, optional, default=None
System info to append to the 'system' block
"""
# Do nothing if filename is None
if filename is None:
return
import os
all_results = { 'benchmarks' : list() }
if system_info is not None:
all_results['system'] = system_info
if filename.endswith('.yaml'):
# Append test result to a YAML file, creating if file does not exist.
import yaml
if os.path.exists(filename):
with open(filename, 'rt') as infile:
all_results = yaml.safe_load(infile)
if test_result is not None:
all_results['benchmarks'].append(test_result)
with open(filename, 'wt') as outfile:
yaml.dump(all_results, outfile, sort_keys=False)
elif filename.endswith('.json'):
# Append test result to a JSON file, creating if file does not exist.
import json
if os.path.exists(filename):
with open(filename, 'rt') as infile:
all_results = json.load(infile)
if test_result is not None:
all_results['benchmarks'].append(test_result)
with open(filename, 'wt') as outfile:
json.dump(all_results, outfile, sort_keys=False, indent=4)
else:
raise ValueError('--output filename must end with .json or .yaml')
def printTestResult(test_result, options):
"""Render a test result
Parameters
----------
test_result : dict
The test result
options :
Options structure
"""
if options.style == 'simple':
for (key, value) in test_result.items():
print(f'{key}: {value}')
print('')
elif options.style == 'table':
print('%-18s%-12s%-14s%-15s%-10g%-11s%-11s%-g' %
(test_result['test'],
test_result['precision'],
test_result['constraints'],
test_result['hydrogen_mass'],
test_result['timestep_in_fs'],
test_result['ensemble'],
test_result['platform'],
test_result['ns_per_day']))
else:
raise ValueError(f"style '{style}' must be one of ['simple', 'table']")
def timeIntegration(context, steps, initialSteps): def timeIntegration(context, steps, initialSteps):
"""Integrate a Context for a specified number of steps, then return how many seconds it took.""" """Integrate a Context for a specified number of steps, then return how many seconds it took."""
...@@ -31,52 +121,68 @@ def downloadAmberSuite(): ...@@ -31,52 +121,68 @@ def downloadAmberSuite():
tarfh.extractall(path=dirname) tarfh.extractall(path=dirname)
return dirname return dirname
def runOneTest(testName, options): import functools
"""Perform a single benchmarking simulation.""" @functools.lru_cache(maxsize=None)
def retrieveTestSystem(testName, pme_cutoff=0.9, bond_constraints='hbonds', polarization='mutual', epsilon=1e-5):
"""Retrieve a benchmark system
Parameters
----------
testName : str
The name of the test
pme_cutoff : float
PME cutoff, in nm
bond_constraints : str, optional, default='hbonds'
'hbonds' : use app.HBonds and set H mass to 1.5*amu
'allbonds' : use app.AllBonds and set H mass to 4*amu
polarization : str, optional, default='mutual'
Polarization scheme for Amoeba
epsilon : str or float, optional, default=1e-5
mutualInducedTargetEpsilon for Amoeba
Returns
-------
system : openmm.System
The test system object
positions : openmm.unit.Quantity with shape (natoms,3)
The initial positions
test_parameters : dict of str : str
Special test parameters to report in test results
"""
explicit = (testName not in ('gbsa', 'amoebagk')) explicit = (testName not in ('gbsa', 'amoebagk'))
amoeba = (testName in ('amoebagk', 'amoebapme')) amoeba = (testName in ('amoebagk', 'amoebapme'))
apoa1 = testName.startswith('apoa1') apoa1 = testName.startswith('apoa1')
amber = (testName.startswith('amber')) amber = (testName.startswith('amber'))
hydrogenMass = None 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)
print('Ensemble: %s' % options.ensemble)
platform = mm.Platform.getPlatformByName(options.platform)
# Create the System. # Create dictionary of test parameters
test_parameters = dict()
temperature = 300*unit.kelvin # Store the test name
if explicit: test_parameters['test'] = testName
friction = 1*(1/unit.picoseconds)
else: # Create the System.
friction = 91*(1/unit.picoseconds)
if amoeba: if amoeba:
constraints = None constraints = None
epsilon = float(options.epsilon) test_parameters['epsilon'] = epsilon
epsilon = float(epsilon)
if explicit: if explicit:
ff = app.ForceField('amoeba2009.xml') ff = app.ForceField('amoeba2009.xml')
pdb = app.PDBFile('5dfr_solv-cube_equil.pdb') pdb = app.PDBFile('5dfr_solv-cube_equil.pdb')
cutoff = 0.7*unit.nanometers cutoff = 0.7*unit.nanometers
vdwCutoff = 0.9*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) system = ff.createSystem(pdb.topology, nonbondedMethod=app.PME, nonbondedCutoff=cutoff, vdwCutoff=vdwCutoff, constraints=constraints, ewaldErrorTolerance=0.00075, mutualInducedTargetEpsilon=epsilon, polarization=polarization)
else: else:
ff = app.ForceField('amoeba2009.xml', 'amoeba2009_gk.xml') ff = app.ForceField('amoeba2009.xml', 'amoeba2009_gk.xml')
pdb = app.PDBFile('5dfr_minimized.pdb') pdb = app.PDBFile('5dfr_minimized.pdb')
system = ff.createSystem(pdb.topology, nonbondedMethod=app.NoCutoff, constraints=constraints, mutualInducedTargetEpsilon=epsilon, polarization=options.polarization) system = ff.createSystem(pdb.topology, nonbondedMethod=app.NoCutoff, constraints=constraints, mutualInducedTargetEpsilon=epsilon, polarization=polarization)
for f in system.getForces(): 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): if isinstance(f, mm.AmoebaMultipoleForce) or isinstance(f, mm.AmoebaVdwForce) or isinstance(f, mm.AmoebaGeneralizedKirkwoodForce) or isinstance(f, mm.AmoebaWcaDispersionForce):
f.setForceGroup(1) f.setForceGroup(1)
dt = 0.002*unit.picoseconds
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 positions = pdb.positions
test_parameters['constraints'] = 'None'
test_parameters['hydrogen_mass'] = '1'
elif amber: elif amber:
dirname = downloadAmberSuite() dirname = downloadAmberSuite()
names = {'amber20-dhfr':'JAC', 'amber20-factorix':'FactorIX', 'amber20-cellulose':'Cellulose', 'amber20-stmv':'STMV'} names = {'amber20-dhfr':'JAC', 'amber20-factorix':'FactorIX', 'amber20-cellulose':'Cellulose', 'amber20-stmv':'STMV'}
...@@ -84,130 +190,375 @@ def runOneTest(testName, options): ...@@ -84,130 +190,375 @@ def runOneTest(testName, options):
prmtop = app.AmberPrmtopFile(os.path.join(dirname, f'PME/Topologies/{fileName}.prmtop')) prmtop = app.AmberPrmtopFile(os.path.join(dirname, f'PME/Topologies/{fileName}.prmtop'))
inpcrd = app.AmberInpcrdFile(os.path.join(dirname, f'PME/Coordinates/{fileName}.inpcrd')) inpcrd = app.AmberInpcrdFile(os.path.join(dirname, f'PME/Coordinates/{fileName}.inpcrd'))
positions = inpcrd.positions positions = inpcrd.positions
dt = 0.004*unit.picoseconds
method = app.PME method = app.PME
cutoff = options.cutoff
constraints = app.HBonds constraints = app.HBonds
hydrogenMass = 1.5*unit.amu cutoff = pme_cutoff
system = prmtop.createSystem(nonbondedMethod=method, nonbondedCutoff=cutoff, constraints=constraints) system = prmtop.createSystem(nonbondedMethod=method, nonbondedCutoff=cutoff, constraints=constraints)
if options.ensemble == 'NVE': if inpcrd.boxVectors is not None:
integ = mm.VerletIntegrator(dt) system.setDefaultPeriodicBoxVectors(*inpcrd.boxVectors)
else: test_parameters['cutoff'] = cutoff
integ = mm.LangevinMiddleIntegrator(temperature, friction, dt) test_parameters['constraints'] = 'HBonds'
test_parameters['hydrogen_mass'] = '1.5'
else: else:
if apoa1: if apoa1:
ff = app.ForceField('amber14/protein.ff14SB.xml', 'amber14/lipid17.xml', 'amber14/tip3p.xml') ff = app.ForceField('amber14/protein.ff14SB.xml', 'amber14/lipid17.xml', 'amber14/tip3p.xml')
pdb = app.PDBFile('apoa1.pdb') pdb = app.PDBFile('apoa1.pdb')
if testName == 'apoa1pme': if testName == 'apoa1pme':
method = app.PME method = app.PME
cutoff = options.cutoff cutoff = pme_cutoff
elif testName == 'apoa1ljpme': elif testName == 'apoa1ljpme':
method = app.LJPME method = app.LJPME
cutoff = options.cutoff cutoff = pme_cutoff
else: else:
# Reaction field uses hard-coded cutoff
method = app.CutoffPeriodic method = app.CutoffPeriodic
cutoff = 1*unit.nanometers cutoff = 1.0 # nanometers ; JDC: Shouldn't this be larger for reaction field?
elif explicit: elif explicit:
ff = app.ForceField('amber99sb.xml', 'tip3p.xml') ff = app.ForceField('amber99sb.xml', 'tip3p.xml')
pdb = app.PDBFile('5dfr_solv-cube_equil.pdb') pdb = app.PDBFile('5dfr_solv-cube_equil.pdb')
if testName == 'pme': if testName == 'pme':
method = app.PME method = app.PME
cutoff = options.cutoff cutoff = pme_cutoff
else: else:
# Reaction field uses hard-coded cutoff
method = app.CutoffPeriodic method = app.CutoffPeriodic
cutoff = 1*unit.nanometers cutoff = 1.0 # nanometers; JDC: Shouldn't this be larger for reaction field?
else: else:
ff = app.ForceField('amber99sb.xml', 'amber99_obc.xml') ff = app.ForceField('amber99sb.xml', 'amber99_obc.xml')
pdb = app.PDBFile('5dfr_minimized.pdb') pdb = app.PDBFile('5dfr_minimized.pdb')
method = app.CutoffNonPeriodic method = app.CutoffNonPeriodic
cutoff = 2*unit.nanometers cutoff = 2.0 # nanometers
if options.heavy: if bond_constraints == 'hbonds':
dt = 0.005*unit.picoseconds constraints = app.HBonds
constraints = app.AllBonds hydrogenMass = 1.5*unit.amu
test_parameters['constraints'] = 'HBonds'
test_parameters['hydrogen_mass'] = '1.5'
elif bond_constraints == 'allbonds':
constraints = app.AllBonds # CAUTION: Constraining all bonds will perturb the effective dihedral marginal distributions.
hydrogenMass = 4*unit.amu hydrogenMass = 4*unit.amu
test_parameters['constraints'] = 'AllBonds'
test_parameters['hydrogen_mass'] = '4'
else:
raise ValueError(f"bond_constraints must be one of 'hbonds', 'allbonds': found {bond_constraints}")
test_parameters['cutoff'] = cutoff
positions = pdb.positions
system = ff.createSystem(pdb.topology, nonbondedMethod=method, nonbondedCutoff=cutoff, constraints=constraints, hydrogenMass=hydrogenMass)
return system, positions, test_parameters
def serializeTest(directory=None, system=None, integrator=None, state=None, coredata=None, metadata=None):
"""Serialize test system for benchmarking on Folding@home.
Parameters
----------
directory : str
Directory to write serialized test XML files
system : openmm.System
The System to serialize
integrator : openmm.Integrator
The Integrator to serialize
state : openmm.State
The State to serialize
coredata : dict
Core parameters to serialize to XML
metadata : dict
Metadata to dump to YAML
"""
import os
os.makedirs(directory, exist_ok=True)
import openmm
def serialize(obj, filename):
if filename.endswith('.bz2'):
import bz2
with bz2.open(filename, 'wt') as outfile:
outfile.write(openmm.XmlSerializer.serialize(obj))
elif filename.endswith('.gz'):
import gzip
with gzip.open(filename, 'wt') as outfile:
outfile.write(openmm.XmlSerializer.serialize(obj))
else:
with open(filename, 'wt') as outfile:
outfile.write(openmm.XmlSerializer.serialize(obj))
serialize(system, os.path.join(directory, 'system.xml.bz2'))
serialize(integrator, os.path.join(directory, 'integrator.xml.bz2'))
serialize(state, os.path.join(directory, 'state.xml.bz2'))
with open(os.path.join(directory, 'core.xml'), 'wt') as outfile:
outfile.write('<config>\n')
for key, value in coredata.items():
outfile.write(f' <{key} v="{value}"/>\n')
outfile.write('</config>\n')
import yaml
with open(os.path.join(directory, 'metadata.yaml'), 'wt') as outfile:
outfile.write(yaml.dump(metadata))
def runOneTest(testName, options):
"""Perform a single benchmarking simulation."""
system, positions, test_parameters = retrieveTestSystem(testName, pme_cutoff=options.pme_cutoff, bond_constraints=options.bond_constraints, polarization=options.polarization, epsilon=options.epsilon)
# Create a copy of the basic test_parameters dict (which may be cached) to report the test results
test_result = test_parameters.copy()
test_result['ensemble'] = options.ensemble
test_result['precision'] = options.precision
explicit = (testName not in ('gbsa', 'amoebagk'))
amoeba = (testName in ('amoebagk', 'amoebapme'))
apoa1 = testName.startswith('apoa1')
amber = (testName.startswith('amber'))
# Create the integrator
temperature = 300*unit.kelvin
if explicit:
friction = 1*(1/unit.picoseconds)
else:
friction = 91*(1/unit.picoseconds)
if amoeba:
dt = 0.002*unit.picoseconds
if options.ensemble == 'NVE':
integ = mm.MTSIntegrator(dt, [(0,2), (1,1)])
else:
integ = mm.MTSLangevinIntegrator(temperature, friction, dt, [(0,2), (1,1)])
elif amber:
dt = 0.004*unit.picoseconds
if options.ensemble == 'NVE': if options.ensemble == 'NVE':
integ = mm.VerletIntegrator(dt) integ = mm.VerletIntegrator(dt)
else: else:
integ = mm.LangevinIntegrator(temperature, friction, dt) integ = mm.LangevinMiddleIntegrator(temperature, friction, dt)
else: else:
if options.bond_constraints == 'hbonds':
dt = 0.004*unit.picoseconds dt = 0.004*unit.picoseconds
constraints = app.HBonds
hydrogenMass = 1.5*unit.amu
if options.ensemble == 'NVE': if options.ensemble == 'NVE':
integ = mm.VerletIntegrator(dt) integ = mm.VerletIntegrator(dt)
else: else:
integ = mm.LangevinMiddleIntegrator(temperature, friction, dt) integ = mm.LangevinMiddleIntegrator(temperature, friction, dt)
positions = pdb.positions elif options.bond_constraints == 'allbonds':
system = ff.createSystem(pdb.topology, nonbondedMethod=method, nonbondedCutoff=cutoff, constraints=constraints, hydrogenMass=hydrogenMass) dt = 0.005*unit.picoseconds
if options.ensemble == 'NPT': if options.ensemble == 'NVE':
system.addForce(mm.MonteCarloBarostat(1*unit.bar, temperature, 100)) integ = mm.VerletIntegrator(dt)
print('Step Size: %g fs' % dt.value_in_unit(unit.femtoseconds)) else:
integ = mm.LangevinIntegrator(temperature, friction, dt)
else:
raise ValueError(f'Unknown bond_constraints {options.bond_constraints}')
test_result['timestep_in_fs'] = dt.value_in_unit(unit.femtoseconds)
properties = {} properties = {}
initialSteps = 5 initialSteps = 5
if options.device is not None and platform.getName() in ('CUDA', 'OpenCL'): platform = mm.Platform.getPlatformByName(options.platform)
if options.device is not None and 'DeviceIndex' in platform.getPropertyNames():
properties['DeviceIndex'] = options.device properties['DeviceIndex'] = options.device
if ',' in options.device or ' ' in options.device: if ',' in options.device or ' ' in options.device:
initialSteps = 250 initialSteps = 250
if options.precision is not None and platform.getName() in ('CUDA', 'OpenCL'): if (options.precision is not None) and ('Precision' in platform.getPropertyNames()):
properties['Precision'] = options.precision properties['Precision'] = options.precision
# Run the simulation. # Add barostat if requested
if options.ensemble == 'NPT':
system.addForce(mm.MonteCarloBarostat(1*unit.bar, temperature, 100))
# Create the Context
integ.setConstraintTolerance(1e-5) integ.setConstraintTolerance(1e-5)
if len(properties) > 0: if len(properties) > 0:
context = mm.Context(system, integ, platform, properties) context = mm.Context(system, integ, platform, properties)
else: else:
context = mm.Context(system, integ, platform) context = mm.Context(system, integ, platform)
# Store information about the Platform used by the Context
platform = context.getPlatform()
test_result['platform'] = platform.getName()
test_result['platform_properties'] = { property_name : platform.getPropertyValue(context, property_name) for property_name in platform.getPropertyNames() }
# Prepare the simulation
context.setPositions(positions) context.setPositions(positions)
if amber: if amber:
if inpcrd.boxVectors is not None:
context.setPeriodicBoxVectors(*inpcrd.boxVectors)
mm.LocalEnergyMinimizer.minimize(context, 100*unit.kilojoules_per_mole/unit.nanometer) mm.LocalEnergyMinimizer.minimize(context, 100*unit.kilojoules_per_mole/unit.nanometer)
context.setVelocitiesToTemperature(temperature) context.setVelocitiesToTemperature(temperature)
if options.serialize:
# Apply constraints on positions and velocities if needing to serialize for Folding@home
tol = 1.0e-8
context.applyConstraints(tol)
context.applyVelocityConstraints(tol)
state = context.getState(getPositions=True, getVelocities=True, getEnergy=True, getForces=True, getParameters=True)
# Time integration, ensuring we trigger kernel compilation before we start timing
steps = 20 steps = 20
while True: while True:
time = timeIntegration(context, steps, initialSteps) elapsed_time = timeIntegration(context, steps, initialSteps)
if time >= 0.5*options.seconds: if elapsed_time >= 0.5*options.seconds:
break break
if time < 0.5: if elapsed_time < 0.5:
steps = int(steps*1.0/time) # Integrate enough steps to get a reasonable estimate for how many we'll need. steps = int(steps*1.0/elapsed_time) # Integrate enough steps to get a reasonable estimate for how many we'll need.
else:
steps = int(steps*options.seconds/elapsed_time)
test_result['steps'] = steps
test_result['elapsed_time'] = elapsed_time
time_per_step = elapsed_time * unit.seconds / steps
ns_per_day = (integ.getStepSize() / time_per_step) / (unit.nanoseconds/unit.day)
test_result['ns_per_day'] = ns_per_day
# Serialize XML files for Folding@home benchmark if requested
if options.serialize:
wu_duration = 5*unit.minutes
ncheckpoints_per_wu = 2
nsteps_per_wu = int(wu_duration / time_per_step)
coredata = dict()
coredata['checkpointFreq'] = int(nsteps_per_wu / ncheckpoints_per_wu)
coredata['numSteps'] = ncheckpoints_per_wu * coredata['checkpointFreq']
coredata['xtcFreq'] = coredata['numSteps']
coredata['precision'] = options.precision
coredata['xtcAtoms'] = 'solute'
coredata['disableCheckpointStateTests'] = 1 # disable checkpoint state tests
coredata['DisablePmeStream'] = 0 # don't disable separate PME streams
coredata['forceTolerance'] = 200000 # for some reason, we have to make this large (TODO: Investigate why)
coredata['energyTolerance'] = 200000 # for some reason, we have to make this large (TODO: Investigate why)
forces = { force.getName() : force for force in system.getForces() }
NONBONDED_METHODS = { 0 : 'NoCutoff', 1 : 'CutoffNonPeriodic', 2 : 'CutoffPeriodic', 3 : 'Ewald', 4 : 'PME', 5 : 'LJPME' }
if 'NonbondedForce' in forces:
nonbonded_method = NONBONDED_METHODS[forces['NonbondedForce'].getNonbondedMethod()]
else: else:
steps = int(steps*options.seconds/time) nonbonded_method = 'unknown'
print('Integrated %d steps in %g seconds' % (steps, time))
print('%g ns/day' % (dt*steps*86400/time).value_in_unit(unit.nanoseconds)) metadata = {
'name' : testName,
'description' : f'OpenMM Benchmark Suite : {testName}',
'precision' : options.precision,
'num_atoms' : context.getSystem().getNumParticles(),
'nonbonded_method' : f'{nonbonded_method}',
'timestep' : integ.getStepSize() / unit.picoseconds,
'integrator' : f'{integ.__class__.__name__}',
}
directory = os.path.join(options.serialize, f'{testName}-{options.precision}')
serializeTest(directory=directory, system=context.getSystem(), integrator=context.getIntegrator(), state=state, coredata=coredata, metadata=metadata)
# Clean up
del context, integ
printTestResult(test_result, options)
appendTestResult(test_result=test_result, filename=options.outfile)
# Parse the command line options. # Parse the command line options.
parser = ArgumentParser() platform_speeds = { mm.Platform.getPlatform(i).getName() : mm.Platform.getPlatform(i).getSpeed() for i in range(mm.Platform.getNumPlatforms()) }
platformNames = [mm.Platform.getPlatform(i).getName() for i in range(mm.Platform.getNumPlatforms())] PLATFORMS = [platform for platform, speed in sorted(platform_speeds.items(), key=lambda item: item[1], reverse=True)]
parser.add_argument('--platform', dest='platform', choices=platformNames, help='name of the platform to benchmark') TESTS = ('gbsa', 'rf', 'pme', 'apoa1rf', 'apoa1pme', 'apoa1ljpme', 'amoebagk', 'amoebapme', 'amber20-dhfr', 'amber20-factorix', 'amber20-cellulose', 'amber20-stmv')
parser.add_argument('--test', dest='test', choices=('gbsa', 'rf', 'pme', 'apoa1rf', 'apoa1pme', 'apoa1ljpme', 'amoebagk', 'amoebapme', 'amber20-dhfr', 'amber20-factorix', 'amber20-cellulose', 'amber20-stmv'), ENSEMBLES = ('NVE', 'NVT', 'NPT')
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-*]') BOND_CONSTRAINTS = ('hbonds', 'allbonds')
parser.add_argument('--ensemble', default='NVT', dest='ensemble', choices=('NPT', 'NVE', 'NVT'), help='the thermodynamic ensemble to simulate [default: NVT]') PRECISIONS = ('single', 'mixed', 'double')
parser.add_argument('--pme-cutoff', default=0.9, dest='cutoff', type=float, help='direct space cutoff for PME in nm [default: 0.9]') POLARIZATION_MODES = ('direct', 'extrapolated', 'mutual')
STYLES = ('simple', 'table')
parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter,
description="Run one or more benchmarks of OpenMM",
epilog="""
Example: run the full suite of benchmarks for the CUDA platform, printing the results as a table
python benchmark.py --platform=CUDA --style=table
Example: run the apoa1pme benchmark for the CPU platform with a reduced cutoff distance
python benchmark.py --platform=CPU --test=apoa1pme --pme-cutoff=0.8
Example: run the full suite in mixed precision mode, saving the results to a YAML file
python benchmark.py --platform=CUDA --precision=mixed --outfile=benchmark.yaml""")
parser.add_argument('--platform', dest='platform', choices=PLATFORMS, help='name of the platform to benchmark')
parser.add_argument('--test', default=','.join(TESTS), dest='test', help=f'the test to perform, or comma-separated list: {TESTS} [default: all]')
parser.add_argument('--ensemble', default='NVT', dest='ensemble', help=f'the thermodynamic ensemble to simulate: {ENSEMBLES} [default: NVT]')
parser.add_argument('--pme-cutoff', default=0.9, dest='pme_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('--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]') parser.add_argument('--polarization', default='mutual', dest='polarization', choices=POLARIZATION_MODES, help=f'the polarization method for AMOEBA: {POLARIZATION_MODES} [default: mutual]')
parser.add_argument('--mutual-epsilon', default=1e-5, dest='epsilon', type=float, help='mutual induced epsilon for AMOEBA [default: 1e-5]') parser.add_argument('--mutual-epsilon', default=1e-5, dest='epsilon', type=float, help='mutual induced epsilon for AMOEBA [default: 1e-5]')
parser.add_argument('--heavy-hydrogens', action='store_true', default=False, dest='heavy', help='repartition mass to allow a larger time step') parser.add_argument('--bond-constraints', default='hbonds', dest='bond_constraints', help=f'hbonds: constrain bonds to hydrogen, use 1.5*amu H mass; allbonds: constrain all bonds, use 4*amu H mass, and use larger timestep. This option is ignored for AMOEBA: {BOND_CONSTRAINTS} [default: hbonds]')
parser.add_argument('--device', default=None, dest='device', help='device index for CUDA or OpenCL') parser.add_argument('--device', default=None, dest='device', help='device index for CUDA or OpenCL')
parser.add_argument('--precision', default='single', dest='precision', choices=('single', 'mixed', 'double'), help='precision mode for CUDA or OpenCL: single, mixed, or double [default: single]') parser.add_argument('--precision', default='single', dest='precision', help=f'precision modes for CUDA or OpenCL: {PRECISIONS} [default: single]')
parser.add_argument('--style', default='simple', dest='style', choices=STYLES, help=f'output style: {STYLES} [default: simple]')
parser.add_argument('--outfile', default=None, dest='outfile', help='output filename for benchmark logging (must end with .yaml or .json)')
parser.add_argument('--serialize', default=None, dest='serialize', help='if specified, output serialized test systems for Folding@home or other uses')
parser.add_argument('--verbose', default=False, action='store_true', dest='verbose', help='if specified, print verbose output')
args = parser.parse_args() args = parser.parse_args()
if args.platform is None: if args.platform is None:
parser.error('No platform specified') parser.error('No platform specified')
print('Platform:', args.platform)
if args.platform in ('CUDA', 'OpenCL'):
print('Precision:', args.precision)
if args.device is not None:
print('Device:', args.device)
# Run the simulations. # Collect system information
system_info = dict()
import socket, platform
system_info['hostname'] = socket.gethostname()
system_info['timestamp'] = datetime.utcnow().isoformat()
system_info['openmm_version'] = mm.version.version
system_info['cpuinfo'] = cpuinfo()
system_info['cpuarch'] = platform.processor()
system_info['system'] = platform.system()
# TODO: Capture information on how many CPU threads will be used
# Attempt to get GPU info
try:
import subprocess
cmd = 'nvidia-smi --query-gpu=driver_version,gpu_name --format=csv,noheader'
output = subprocess.check_output(cmd, shell=True, text=True)
system_info['nvidia_driver'], system_info['gpu'] = output.strip().split(', ')
except Exception as e:
pass
for key, value in system_info.items():
print(f'{key}: {value}')
if args.outfile is not None:
# Remove output file
import os
if os.path.exists(args.outfile):
os.unlink(args.outfile)
# Write system info
appendTestResult(args.outfile, system_info=system_info)
tests = args.test.split(',')
if not set(tests).issubset(TESTS):
parser.error(f'Available tests: {TESTS}')
if args.test is None: precisions = args.precision.split(',')
for test in ('gbsa', 'rf', 'pme', 'apoa1rf', 'apoa1pme', 'apoa1ljpme', 'amoebagk', 'amoebapme'): if args.platform == 'Reference':
precisions = ['double']
if args.platform == 'CPU':
precisions = ['mixed']
if not set(precisions).issubset(PRECISIONS):
parser.error(f'Available precisions: {PRECISIONS}')
ensembles = args.ensemble.split(',')
if not set(ensembles).issubset(ENSEMBLES):
parser.error(f'Available ensembles: {ENSEMBLES}')
bond_constraints = args.bond_constraints.split(',')
if not set(bond_constraints).issubset(BOND_CONSTRAINTS):
parser.error(f'Available bond constraints: {BOND_CONSTRAINTS}')
# Combinatorially run all requested benchmarks, ignoring combinations that cannot be run
from openmm import OpenMMException
from itertools import product
if args.style == 'simple':
for (test, args.bond_constraints, args.ensemble, args.precision) in product(tests, bond_constraints, ensembles, precisions):
try:
runOneTest(test, args)
except OpenMMException as e:
if args.verbose:
print(e)
pass
elif args.style == 'table':
print()
print('Test Precision Constraints H mass (amu) dt (fs) Ensemble Platform ns/day')
for (test, args.bond_constraints, args.ensemble, args.precision) in product(tests, bond_constraints, ensembles, precisions):
try: try:
runOneTest(test, args) runOneTest(test, args)
except Exception as ex: except OpenMMException as e:
print('Test failed: %s' % ex.message) if args.verbose:
print(e)
pass
else: else:
runOneTest(args.test, args) raise ValueError(f"style {args.style} unkown; must be one of ['simple', 'table']")
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment