isolate contribution of random number generation (result: not a major contribution)

quick_speed_test_customintegrator.py

#!/bin/env python

"""
Benchmark DHFR in explicit solvent with minimal CustomIntegrators that include Gaussian random numbers or not. 
"""

import time
from simtk import openmm, unit
from openmmtools import testsystems

from simtk import unit
from simtk import openmm as mm

class GradientDescent(mm.CustomIntegrator):
    def __init__(self, timestep=0.01 * unit.femtoseconds):
        super(GradientDescent, self).__init__(timestep)
        self.addUpdateContextState()
        self.addComputePerDof("x", "x + dt*f/m")
        self.addConstrainPositions()

class NoisyGradientDescent(mm.CustomIntegrator):
    def __init__(self, timestep=0.01 * unit.femtoseconds):
        super(NoisyGradientDescent, self).__init__(timestep)
        self.addUpdateContextState()
        self.addComputePerDof("x", "x + dt*((f/m) + gaussian)")
        self.addConstrainPositions()

print('OpenMM version: ', mm.version.full_version)
timestep = 2.0 * unit.femtoseconds
nsteps = 5000

testsystem = testsystems.DHFRExplicit()
system, positions = testsystem.system, testsystem.positions

platform = mm.Platform.getPlatformByName('CUDA')
properties = {'CudaPrecision' : 'mixed'}

integrator = GradientDescent()
context = mm.Context(system, integrator, platform, properties)
context.setPositions(positions)
integrator.step(2) # compile the kernels
print('Running %d steps GradientDescent (no random numbers) on %s...' % (nsteps, context.getPlatform().getName()))
context.getState(getEnergy=True)
initial_time = time.time()
integrator.step(nsteps)
context.getState(getEnergy=True)
final_time = time.time()
del context, integrator
elapsed_time = final_time - initial_time
print('%.3f s elapsed.' % elapsed_time)

integrator = NoisyGradientDescent()
context = mm.Context(system, integrator, platform, properties)
context.setPositions(positions)
integrator.step(2) # compile the kernels
print('Running %d steps NoisyGradientDescent (with random numbers) on %s...' % (nsteps, context.getPlatform().getName()))
context.getState(getEnergy=True)
initial_time = time.time()
integrator.step(nsteps)
context.getState(getEnergy=True)
final_time = time.time()
del context, integrator
elapsed_time = final_time - initial_time
print('%.3f s elapsed.' % elapsed_time)