optimizepme.py

"""
optimizepme.py: Optimizes parameters for PME simulations

This is part of the OpenMM molecular simulation toolkit originating from
Simbios, the NIH National Center for Physics-Based Simulation of
Biological Structures at Stanford, funded under the NIH Roadmap for
Medical Research, grant U54 GM072970. See https://simtk.org.

Portions copyright (c) 2013 Stanford University and the Authors.
Authors: Peter Eastman
Contributors:

Permission is hereby granted, free of charge, to any person obtaining a
copy of this software and associated documentation files (the "Software"),
to deal in the Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, sublicense,
and/or sell copies of the Software, and to permit persons to whom the
Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
USE OR OTHER DEALINGS IN THE SOFTWARE.
"""
import simtk.openmm as mm
import simtk.openmm.app as app
import simtk.unit as unit
import itertools
import math
from datetime import datetime

def calc_pme_parameters(system):
    """Calculate PME parameters using scheme similar to OpenMM OpenCL platform.
    
    Parameters
    ----------
    system : simtk.openmm.System
        The system for which parameters are to be computed.

    Returns
    -------
    alpha : float
        The PME alpha parameter
    nx, ny, nz : int
        The grid numbers in each dimension

    """

    # Find nonbonded force.
    forces = { system.getForce(index).__class__.__name__ : system.getForce(index) for index in range(system.getNumForces()) }
    force = forces['NonbondedForce']
    tol = force.getEwaldErrorTolerance()
    boxVectors = system.getDefaultPeriodicBoxVectors()

    from numpy import sqrt, log, ceil
    from math import pow
    alpha = (1.0/force.getCutoffDistance())*sqrt(-log(2.0*tol))
    xsize = int(ceil(2*alpha*boxVectors[0][0]/(3*pow(tol, 0.2))))
    ysize = int(ceil(2*alpha*boxVectors[1][1]/(3*pow(tol, 0.2))))
    zsize = int(ceil(2*alpha*boxVectors[2][2]/(3*pow(tol, 0.2))))

    print (xsize,ysize,zsize)
    def findLegalDimension(minimum):
        while (True):
            # Attempt to factor the current value.
            unfactored = minimum
            for factor in range(2, 8):
                while (unfactored > 1) and (unfactored%factor == 0):
                    unfactored /= factor
        
            if (unfactored == 1):
                return minimum

            minimum += 1
    
    nx = findLegalDimension(xsize)
    ny = findLegalDimension(ysize)
    nz = findLegalDimension(zsize)

    return (alpha, nx, ny, nz)

def optimizePME(system, integrator, positions, platform, properties, minCutoff, maxCutoff):
    """Run a series of simulations using different parameters to see which give the best performance.
    
    When running a simulation with PME, different combinations of parameters may give equivalent accuracy
    but differ in performance.  In particular:
    
    1. The nonbonded cutoff does not affect the accuracy of the Coulomb interaction with PME.  You can
    freely vary the cutoff distance, and OpenMM will automatically select internal parameters to give
    whatever accuracy has been selected with the ewaldErrorTolerance parameter.  (The cutoff does affect
    other nonbonded interactions, such as Lennard-Jones, so this generally places a lower limit on the
    cutoffs you consider acceptable.)
    2. In some cases, OpenMM can perform reciprocal space calculations on the CPU at the same time it is
    doing direct space calculations on the GPU.  Depending on your hardware, this might or might not
    be faster.
    
    This function runs a series of simulations to measure the performance of simulating a particular system
    on the current hardware.  This allows you to choose the combination of parameters that give the
    best performance while still providing the required accuracy.  The function prints out the results of
    each simulation, along with a final recommendation of the best parameters to use.  On exit, the
    system and properties arguments will have been modified to use the recommended parameters.
    
    Parameters:
     - system (System) the System to simulate
     - integrator (Integrator) the Integrator to use for simulating it
     - positions (list) the initial particle positions
     - platform (Platform) the Platform to use for running the simulation
     - properties (dict) any platform-specific properties you want to specify
     - minCutoff (distance) the minimum cutoff distance to try
     - maxCutoff (distance) the maximum cutoff distance to try
    """
    if unit.is_quantity(minCutoff):
        minCutoff = minCutoff.value_in_unit(unit.nanometers)
    if unit.is_quantity(maxCutoff):
        maxCutoff = maxCutoff.value_in_unit(unit.nanometers)
    
    # Find the NonbondedForce or AmoebaMultipoleForce to optimize.
    
    nonbonded = None
    for force in system.getForces():
        if isinstance(force, mm.NonbondedForce):
            nonbonded = force
            if nonbonded.getNonbondedMethod() != mm.NonbondedForce.PME:
                raise ValueError('The System does not use PME')
            break
        if isinstance(force, mm.AmoebaMultipoleForce):
            nonbonded = force
            if nonbonded.getNonbondedMethod() != mm.AmoebaMultipoleForce.PME:
                raise ValueError('The System does not use PME')
            nonbonded.setAEwald(0)
            break
    if nonbonded is None:
        raise ValueError('The System does not include a NonbondedForce or AmoebaMultipoleForce')
    errorTolerance = nonbonded.getEwaldErrorTolerance()
    canUseCpuPme = (isinstance(nonbonded, mm.NonbondedForce) and platform.supportsKernels(['CalcPmeReciprocalForce']))
    if platform.getName() == 'CUDA':
        cpuPmeProperty = 'CudaUseCpuPme'
    else:
        cpuPmeProperty = 'OpenCLUseCpuPme'

    # Build a list of cutoff distances to try.
    
    gpuCutoffs = set()
    cpuCutoffs = set()
    gpuCutoffs.add(minCutoff)
    cpuCutoffs.add(minCutoff)
    vec1, vec2, vec3 = system.getDefaultPeriodicBoxVectors()
    errorTolerance5 = math.pow(errorTolerance, 0.2)
    boxDimensions = [x.value_in_unit(unit.nanometers) for x in (vec1[0], vec2[1], vec3[2])]
    for boxSize in boxDimensions: # Loop over the three dimensions of the periodic box.
        for gridSize in itertools.count(start=5, step=1): # Loop over possible sizes of the PME grid.
            # Determine whether this is a legal size for the FFT.
            
            unfactored = gridSize
            for factor in (2, 3, 5, 7):
                while unfactored > 1 and unfactored%factor == 0:
                    unfactored /= factor
            if unfactored not in (1, 11, 13):
                continue
            
            # Compute the smallest cutoff that will give this grid size.
            
            alpha = 1.5*gridSize*errorTolerance5/boxSize
            cutoff = math.sqrt(-math.log(2*errorTolerance))/alpha
            cutoff = 0.001*int(cutoff*1000) # Round up to the next picometer to avoid roundoff errors.
            if cutoff < minCutoff:
                break
            if cutoff < maxCutoff:
                cpuCutoffs.add(cutoff)
                if unfactored == 1:
                    gpuCutoffs.add(cutoff)
    gpuCutoffs = sorted(gpuCutoffs)
    cpuCutoffs = sorted(cpuCutoffs)
    
    # Select a length for the simulations so they will each take about 10 seconds.
    
    print()
    print('Selecting a length for the test simulations... ')
    nonbonded.setCutoffDistance(math.sqrt(minCutoff*maxCutoff))
    properties[cpuPmeProperty] = 'false'
    print("Creating Context...")
    context = _createContext(system, integrator, positions, platform, properties)
    steps = 20
    time = 0.0
    while time < 8.0 or time > 12.0:
        print("Trying %d steps..." % steps)
        time = _timeIntegrator(context, steps)
        steps = int(steps*10.0/time)
    print(steps, 'steps')
    del context
    
    # Run the simulations.
    
    print()
    print('Running simulations with standard PME')
    print()
    results = []
    properties[cpuPmeProperty] = 'false'
    gpuTimes = _timeWithCutoffs(system, integrator, positions, platform, properties, nonbonded, gpuCutoffs, steps)
    for time, cutoff in zip(gpuTimes, gpuCutoffs):
        results.append((time, cutoff, 'false'))
    if canUseCpuPme:
        print()
        print('Running simulations with CPU based PME')
        print()
        properties[cpuPmeProperty] = 'true'
        cpuTimes = _timeWithCutoffs(system, integrator, positions, platform, properties, nonbonded, cpuCutoffs, steps)
        for time, cutoff in zip(cpuTimes, cpuCutoffs):
            results.append((time, cutoff, 'true'))

    # Rerun the fastest configurations to make sure the results are consistent.
    
    print()
    print('Confirming results for best configurations')
    print()
    results.sort(key=lambda x: x[0])
    finalResults = []
    for time, cutoff, useCpu in results[:5]:
        nonbonded.setCutoffDistance(cutoff)
        properties[cpuPmeProperty] = useCpu
        context = _createContext(system, integrator, positions, platform, properties)
        print calc_pme_parameters(system)
        time2 = _timeIntegrator(context, steps)
        time3 = _timeIntegrator(context, steps)
        medianTime = sorted((time, time2, time3))[1]
        finalResults.append((medianTime, cutoff, useCpu))
        print('Cutoff=%g, %s=%s' % (cutoff, cpuPmeProperty, useCpu))
        print('Times: %g, %g, %g' % (time, time2, time3))
        print('Median time: %g' % medianTime)
        print()
        del context

    # Select the best configuration.
    
    finalResults.sort(key=lambda x: x[0])
    best = finalResults[0]
    nonbonded.setCutoffDistance(best[1])
    properties[cpuPmeProperty] = best[2]
    print('Best configuration:')
    print()
    print('Cutoff=%g nm, %s=%s' % (best[1], cpuPmeProperty, best[2]))
    print()


def _createContext(system, integrator, positions, platform, properties):
    print("copying integrator...")
    #integrator = mm.XmlSerializer.deserialize(mm.XmlSerializer.serialize(integrator))
    integrator = mm.VerletIntegrator(1.0 * unit.femtoseconds)
    print("Creating context...")
    context = mm.Context(system, integrator, platform, properties)
    print("Setting positions...")
    context.setPositions(positions)
    print("Done...")
    return context


def _timeIntegrator(context, steps):
    context.getIntegrator().step(5) # Make sure everything is fully initialized
    context.getState(getEnergy=True)
    start = datetime.now()
    context.getIntegrator().step(steps)
    context.getState(getEnergy=True)
    end = datetime.now()
    return (end-start).total_seconds()


def _timeWithCutoffs(system, integrator, positions, platform, properties, nonbonded, cutoffs, steps):
    times = []
    for cutoff in cutoffs:
        nonbonded.setCutoffDistance(cutoff)
        context = _createContext(system, integrator, positions, platform, properties)
        time = _timeIntegrator(context, steps)
        del context
        print calc_pme_parameters(system)
        print('cutoff=%g, time=%g' % (cutoff, time))
        times.append(time)
        if len(times) > 3 and times[-1] > times[-2] > times[-3] > times[-4]:
            # It's steadily getting slower as we increase the cutoff, so stop now.
            break
        
    return times

if __name__ == '__main__':
    import gzip
    try:
        system = mm.XmlSerializer.deserialize(gzip.open('system.xml.gz').read().decode('utf-8'))
        integrator = mm.XmlSerializer.deserialize(gzip.open('integrator.xml.gz').read().decode('utf-8'))
        state = mm.XmlSerializer.deserialize(gzip.open('state0.xml.gz').read().decode('utf-8'))
    except:
        system = mm.XmlSerializer.deserialize(gzip.open('system.xml.gz').read())
        integrator = mm.XmlSerializer.deserialize(gzip.open('integrator.xml.gz').read())
        state = mm.XmlSerializer.deserialize(gzip.open('state0.xml.gz').read())
    platform = mm.Platform.getPlatformByName('OpenCL')
    properties = {'OpenCLPrecision': 'single'}
    optimizePME(system, integrator, state.getPositions(), platform, properties, 0.8*unit.nanometers, 1.2*unit.nanometers)