Merge pull request #4819 from epretti/add-simtk-extras

Add PME optimizer and rigid body scripts

examples/extras/README.md

+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.

examples/extras/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.
+
+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

examples/extras/rigid.py

+"""
+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]))