RustOpenMM/openmm
1
0
Fork 0
mirror of https://github.com/openmm/openmm synced 2026-06-03 06:39:48 +09:00
Code Issues Packages Projects Releases Wiki Activity

Add PME optimizer and rigid body scripts

Browse Source
This commit is contained in:
Evan Pretti
2025-03-03 16:10:46 -08:00
parent a6545a1b7a
commit 8a05b45602
3 changed files with 380 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
examples/extras/README.md Normal file
View File

@@ -0,0 +1,10 @@
OpenMM extra utility scripts
============================
This directory contains standalone utility scripts for use with OpenMM.
* `optimizepme.py`: Optimizes parameters for PME simulations by running a series
of simulations with different PME parameters and choosing the combination of
parameters giving the best performance.
* `rigid.py`: Implements rigid bodies by adding constraints between some
particles and converting others to virtual sites.

224
examples/extras/optimizepme.py Normal file
View File

@@ -0,0 +1,224 @@
"""
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.
Portions copyright (c) 2013-2025 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 openmm as mm
import openmm.app as app
import openmm.unit as unit
import itertools
import math
from datetime import datetime
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'
context = _createContext(system, integrator, positions, platform, properties)
steps = 20
time = 0.0
while time < 8.0 or time > 12.0:
time = _timeIntegrator(context, steps)
steps = int(steps*10.0/time)
print(steps, 'steps')
# 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)
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()
# 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):
integrator = mm.XmlSerializer.deserialize(mm.XmlSerializer.serialize(integrator))
context = mm.Context(system, integrator, platform, properties)
context.setPositions(positions)
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)
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

146
examples/extras/rigid.py Normal file
View File

@@ -0,0 +1,146 @@
"""
rigid.py: Implements rigid bodies
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.
Portions copyright (c) 2016-2025 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.
"""
__author__ = "Peter Eastman"
__version__ = "1.0"
import openmm as mm
import openmm.unit as unit
import numpy as np
import numpy.linalg as lin
from itertools import combinations
def createRigidBodies(system, positions, bodies):
"""Modify a System to turn specified sets of particles into rigid bodies.
For every rigid body, four particles are selected as "real" particles whose positions are integrated.
Constraints are added between them to make them move as a rigid body. All other particles in the body
are then turned into virtual sites whose positions are computed based on the "real" particles.
Because virtual sites are massless, the mass properties of the rigid bodies will be slightly different
from the corresponding sets of particles in the original system. The masses of the non-virtual particles
are chosen to guarantee that the total mass and center of mass of each rigid body exactly match those of
the original particles. The moment of inertia will be similar to that of the original particles, but
not identical.
Care is needed when using constraints, since virtual particles cannot participate in constraints. If the
input system includes any constraints, this function will automatically remove ones that connect two
particles in the same rigid body. But if there is a constraint beween a particle in a rigid body and
another particle not in that body, it will likely lead to an exception when you try to create a context.
Parameters:
- system (System) the System to modify
- positions (list) the positions of all particles in the system
- bodies (list) each element of this list defines one rigid body. Each element should itself be a list
of the indices of all particles that make up that rigid body.
"""
# Remove any constraints involving particles in rigid bodies.
for i in range(system.getNumConstraints()-1, -1, -1):
p1, p2, distance = system.getConstraintParameters(i)
if (any(p1 in body and p2 in body for body in bodies)):
system.removeConstraint(i)
# Loop over rigid bodies and process them.
for particles in bodies:
if len(particles) < 5:
# All the particles will be "real" particles.
realParticles = particles
realParticleMasses = [system.getParticleMass(i) for i in particles]
else:
# Select four particles to use as the "real" particles. All others will be virtual sites.
pos = [positions[i] for i in particles]
mass = [system.getParticleMass(i) for i in particles]
cm = unit.sum([p*m for p, m in zip(pos, mass)])/unit.sum(mass)
r = [p-cm for p in pos]
avgR = unit.sqrt(unit.sum([unit.dot(x, x) for x in r])/len(particles))
rank = sorted(range(len(particles)), key=lambda i: abs(unit.norm(r[i])-avgR))
for p in combinations(rank, 4):
# Select masses for the "real" particles. If any is negative, reject this set of particles
# and keep going.
matrix = np.zeros((4, 4))
for i in range(4):
particleR = r[p[i]].value_in_unit(unit.nanometers)
matrix[0][i] = particleR[0]
matrix[1][i] = particleR[1]
matrix[2][i] = particleR[2]
matrix[3][i] = 1.0
rhs = np.array([0.0, 0.0, 0.0, unit.sum(mass).value_in_unit(unit.amu)])
weights = lin.solve(matrix, rhs)
if all(w > 0.0 for w in weights):
# We have a good set of particles.
realParticles = [particles[i] for i in p]
realParticleMasses = [float(w) for w in weights]*unit.amu
break
# Set particle masses.
for i, m in zip(realParticles, realParticleMasses):
system.setParticleMass(i, m)
# Add constraints between the real particles.
for p1, p2 in combinations(realParticles, 2):
distance = unit.norm(positions[p1]-positions[p2])
key = (min(p1, p2), max(p1, p2))
system.addConstraint(p1, p2, distance)
# Select which three particles to use for defining virtual sites.
bestNorm = 0
for p1, p2, p3 in combinations(realParticles, 3):
d12 = (positions[p2]-positions[p1]).value_in_unit(unit.nanometer)
d13 = (positions[p3]-positions[p1]).value_in_unit(unit.nanometer)
crossNorm = unit.norm((d12[1]*d13[2]-d12[2]*d13[1], d12[2]*d13[0]-d12[0]*d13[2], d12[0]*d13[1]-d12[1]*d13[0]))
if crossNorm > bestNorm:
bestNorm = crossNorm
vsiteParticles = (p1, p2, p3)
# Create virtual sites.
d12 = (positions[vsiteParticles[1]]-positions[vsiteParticles[0]]).value_in_unit(unit.nanometer)
d13 = (positions[vsiteParticles[2]]-positions[vsiteParticles[0]]).value_in_unit(unit.nanometer)
cross = mm.Vec3(d12[1]*d13[2]-d12[2]*d13[1], d12[2]*d13[0]-d12[0]*d13[2], d12[0]*d13[1]-d12[1]*d13[0])
matrix = np.zeros((3, 3))
for i in range(3):
matrix[i][0] = d12[i]
matrix[i][1] = d13[i]
matrix[i][2] = cross[i]
for i in particles:
if i not in realParticles:
system.setParticleMass(i, 0)
rhs = np.array((positions[i]-positions[vsiteParticles[0]]).value_in_unit(unit.nanometer))
weights = lin.solve(matrix, rhs)
system.setVirtualSite(i, mm.OutOfPlaneSite(vsiteParticles[0], vsiteParticles[1], vsiteParticles[2], weights[0], weights[1], weights[2]))