Implemented long range correction for interaction groups

openmmapi/src/CustomNonbondedForceImpl.cpp

@@ -42,14 +42,10 @@
 #include "lepton/Parser.h"
 #include <cmath>
 #include <sstream>
+#include <utility>
 
 using namespace OpenMM;
-using std::map;
-using std::pair;
-using std::vector;
-using std::set;
-using std::string;
-using std::stringstream;
+using namespace std;
 
 CustomNonbondedForceImpl::CustomNonbondedForceImpl(const CustomNonbondedForce& owner) : owner(owner) {
 }
@@ -176,20 +172,6 @@ double CustomNonbondedForceImpl::calcLongRangeCorrection(const CustomNonbondedFo
     if (force.getNonbondedMethod() == CustomNonbondedForce::NoCutoff || force.getNonbondedMethod() == CustomNonbondedForce::CutoffNonPeriodic)
         return 0.0;
     
-    // Identify all particle classes (defined by parameters), and count the number of
-    // particles in each class.
-
-    map<vector<double>, int> classCounts;
-    for (int i = 0; i < force.getNumParticles(); i++) {
-        vector<double> parameters;
-        force.getParticleParameters(i, parameters);
-        map<vector<double>, int>::iterator entry = classCounts.find(parameters);
-        if (entry == classCounts.end())
-            classCounts[parameters] = 1;
-        else
-            entry->second++;
-    }
-    
     // Parse the energy expression.
     
     map<string, Lepton::CustomFunction*> functions;
@@ -201,20 +183,69 @@ double CustomNonbondedForceImpl::calcLongRangeCorrection(const CustomNonbondedFo
         functions[name] = new TabulatedFunction(min, max, values);
     }
     Lepton::ExpressionProgram expression = Lepton::Parser::parse(force.getEnergyFunction(), functions).optimize().createProgram();
+    
+    // Identify all particle classes (defined by parameters), and record the class of each particle.
+    
+    int numParticles = force.getNumParticles();
+    vector<vector<double> > classes;
+    map<vector<double>, int> classIndex;
+    vector<int> atomClass(numParticles);
+    for (int i = 0; i < numParticles; i++) {
+        vector<double> parameters;
+        force.getParticleParameters(i, parameters);
+        if (classIndex.find(parameters) == classIndex.end()) {
+            classIndex[parameters] = classes.size();
+            classes.push_back(parameters);
+        }
+        atomClass[i] = classIndex[parameters];
+    }
+    int numClasses = classes.size();
+    
+    // Count the total number of particle pairs for each pair of classes.
+    
+    map<pair<int, int>, int> interactionCount;
+    if (force.getNumInteractionGroups() == 0) {
+        // Count the particles of each class.
+        
+        vector<int> classCounts(numClasses, 0);
+        for (int i = 0; i < numParticles; i++)
+            classCounts[atomClass[i]]++;
+        for (int i = 0; i < numClasses; i++) {
+            interactionCount[make_pair(i, i)] = (classCounts[i]*(classCounts[i]+1))/2;
+            for (int j = i+1; j < numClasses; j++)
+                interactionCount[make_pair(i, j)] = classCounts[i]*classCounts[j];
+        }
+    }
+    else {
+        // Initialize the counts to 0.
+        
+        for (int i = 0; i < numClasses; i++) {
+            for (int j = i; j < numClasses; j++)
+                interactionCount[make_pair(i, j)] = 0;
+        }
+        
+        // Loop over interaction groups and count the interactions in each one.
+        
+        for (int group = 0; group < force.getNumInteractionGroups(); group++) {
+            set<int> set1, set2;
+            force.getInteractionGroupParameters(group, set1, set2);
+            for (set<int>::const_iterator a1 = set1.begin(); a1 != set1.end(); ++a1)
+                for (set<int>::const_iterator a2 = set2.begin(); a2 != set2.end(); ++a2) {
+                    if (*a1 >= *a2 && set1.find(*a2) != set1.end() && set2.find(*a1) != set2.end())
+                        continue;
+                    int class1 = atomClass[*a1];
+                    int class2 = atomClass[*a2];
+                    interactionCount[make_pair(min(class1, class2), max(class1, class2))]++;
+                }
+        }
+    }
 
     // Loop over all pairs of classes to compute the coefficient.
 
     double sum = 0;
-    for (map<vector<double>, int>::const_iterator entry = classCounts.begin(); entry != classCounts.end(); ++entry) {
-        int count = (entry->second*(entry->second+1))/2;
-        sum += count*integrateInteraction(expression, entry->first, entry->first, force, context);
-    }
-    for (map<vector<double>, int>::const_iterator class1 = classCounts.begin(); class1 != classCounts.end(); ++class1)
-        for (map<vector<double>, int>::const_iterator class2 = classCounts.begin(); class2 != class1; ++class2) {
-            int count = class1->second*class2->second;
-            sum += count*integrateInteraction(expression, class1->first, class2->first, force, context);
-        }
-    int numParticles = force.getNumParticles();
+    for (int i = 0; i < numClasses; i++)
+        for (int j = i; j < numClasses; j++)
+            sum += interactionCount[make_pair(i, j)]*integrateInteraction(expression, classes[i], classes[j], force, context);
     int numInteractions = (numParticles*(numParticles+1))/2;
     sum /= numInteractions;
     return 2*M_PI*numParticles*numParticles*sum;

platforms/cuda/tests/TestCudaCustomNonbondedForce.cpp

@@ -654,6 +654,63 @@ void testLargeInteractionGroup() {
         ASSERT_EQUAL_VEC(state1.getForces()[i], state3.getForces()[i], 1e-4);
 }
 
+void testInteractionGroupLongRangeCorrection() {
+    const int numParticles = 10;
+    const double boxSize = 10.0;
+    const double cutoff = 0.5;
+    System system;
+    system.setDefaultPeriodicBoxVectors(Vec3(boxSize, 0, 0), Vec3(0, boxSize, 0), Vec3(0, 0, boxSize));
+    CustomNonbondedForce* nonbonded = new CustomNonbondedForce("c1*c2*r^-4");
+    nonbonded->addPerParticleParameter("c");
+    vector<Vec3> positions(numParticles);
+    vector<double> params(1);
+    for (int i = 0; i < numParticles; i++) {
+        system.addParticle(1.0);
+        params[0] = (i%2 == 0 ? 1.1 : 2.0);
+        nonbonded->addParticle(params);
+        positions[i] = Vec3(0.5*i, 0, 0);
+    }
+    nonbonded->setNonbondedMethod(CustomNonbondedForce::CutoffPeriodic);
+    nonbonded->setCutoffDistance(cutoff);
+    system.addForce(nonbonded);
+    
+    // Setup nonbonded groups.  They involve 1 interaction of type AA,
+    // 2 of type BB, and 5 of type AB.
+    
+    set<int> set1, set2, set3, set4, set5;
+    set1.insert(0);
+    set1.insert(1);
+    set1.insert(2);
+    nonbonded->addInteractionGroup(set1, set1);
+    set2.insert(3);
+    set3.insert(4);
+    set3.insert(6);
+    set3.insert(8);
+    nonbonded->addInteractionGroup(set2, set3);
+    set4.insert(5);
+    set5.insert(7);
+    set5.insert(9);
+    nonbonded->addInteractionGroup(set4, set5);
+    
+    // Compute energy with and without the correction.
+    
+    VerletIntegrator integrator(0.01);
+    Context context(system, integrator, platform);
+    context.setPositions(positions);
+    double energy1 = context.getState(State::Energy).getPotentialEnergy();
+    nonbonded->setUseLongRangeCorrection(true);
+    context.reinitialize();
+    context.setPositions(positions);
+    double energy2 = context.getState(State::Energy).getPotentialEnergy();
+    
+    // Check the result.
+    
+    double sum = (1.1*1.1 + 2*2.0*2.0 + 5*1.1*2.0)*2.0;
+    int numPairs = (numParticles*(numParticles+1))/2;
+    double expected = 2*M_PI*numParticles*numParticles*sum/(numPairs*boxSize*boxSize*boxSize);
+    ASSERT_EQUAL_TOL(expected, energy2-energy1, 1e-4);
+}
+
 int main(int argc, char* argv[]) {
     try {
         if (argc > 1)
@@ -671,6 +728,7 @@ int main(int argc, char* argv[]) {
         testLongRangeCorrection();
         testInteractionGroups();
         testLargeInteractionGroup();
+        testInteractionGroupLongRangeCorrection();
     }
     catch(const exception& e) {
         cout << "exception: " << e.what() << endl;

platforms/opencl/tests/TestOpenCLCustomNonbondedForce.cpp

@@ -654,6 +654,63 @@ void testLargeInteractionGroup() {
         ASSERT_EQUAL_VEC(state1.getForces()[i], state3.getForces()[i], 1e-4);
 }
 
+void testInteractionGroupLongRangeCorrection() {
+    const int numParticles = 10;
+    const double boxSize = 10.0;
+    const double cutoff = 0.5;
+    System system;
+    system.setDefaultPeriodicBoxVectors(Vec3(boxSize, 0, 0), Vec3(0, boxSize, 0), Vec3(0, 0, boxSize));
+    CustomNonbondedForce* nonbonded = new CustomNonbondedForce("c1*c2*r^-4");
+    nonbonded->addPerParticleParameter("c");
+    vector<Vec3> positions(numParticles);
+    vector<double> params(1);
+    for (int i = 0; i < numParticles; i++) {
+        system.addParticle(1.0);
+        params[0] = (i%2 == 0 ? 1.1 : 2.0);
+        nonbonded->addParticle(params);
+        positions[i] = Vec3(0.5*i, 0, 0);
+    }
+    nonbonded->setNonbondedMethod(CustomNonbondedForce::CutoffPeriodic);
+    nonbonded->setCutoffDistance(cutoff);
+    system.addForce(nonbonded);
+    
+    // Setup nonbonded groups.  They involve 1 interaction of type AA,
+    // 2 of type BB, and 5 of type AB.
+    
+    set<int> set1, set2, set3, set4, set5;
+    set1.insert(0);
+    set1.insert(1);
+    set1.insert(2);
+    nonbonded->addInteractionGroup(set1, set1);
+    set2.insert(3);
+    set3.insert(4);
+    set3.insert(6);
+    set3.insert(8);
+    nonbonded->addInteractionGroup(set2, set3);
+    set4.insert(5);
+    set5.insert(7);
+    set5.insert(9);
+    nonbonded->addInteractionGroup(set4, set5);
+    
+    // Compute energy with and without the correction.
+    
+    VerletIntegrator integrator(0.01);
+    Context context(system, integrator, platform);
+    context.setPositions(positions);
+    double energy1 = context.getState(State::Energy).getPotentialEnergy();
+    nonbonded->setUseLongRangeCorrection(true);
+    context.reinitialize();
+    context.setPositions(positions);
+    double energy2 = context.getState(State::Energy).getPotentialEnergy();
+    
+    // Check the result.
+    
+    double sum = (1.1*1.1 + 2*2.0*2.0 + 5*1.1*2.0)*2.0;
+    int numPairs = (numParticles*(numParticles+1))/2;
+    double expected = 2*M_PI*numParticles*numParticles*sum/(numPairs*boxSize*boxSize*boxSize);
+    ASSERT_EQUAL_TOL(expected, energy2-energy1, 1e-4);
+}
+
 int main(int argc, char* argv[]) {
     try {
         if (argc > 1)
@@ -671,6 +728,7 @@ int main(int argc, char* argv[]) {
         testLongRangeCorrection();
         testInteractionGroups();
         testLargeInteractionGroup();
+        testInteractionGroupLongRangeCorrection();
     }
     catch(const exception& e) {
         cout << "exception: " << e.what() << endl;

platforms/reference/tests/TestReferenceCustomNonbondedForce.cpp

@@ -589,6 +589,64 @@ void testLargeInteractionGroup() {
         ASSERT_EQUAL_VEC(state1.getForces()[i], state3.getForces()[i], 1e-4);
 }
 
+void testInteractionGroupLongRangeCorrection() {
+    const int numParticles = 10;
+    const double boxSize = 10.0;
+    const double cutoff = 0.5;
+    ReferencePlatform platform;
+    System system;
+    system.setDefaultPeriodicBoxVectors(Vec3(boxSize, 0, 0), Vec3(0, boxSize, 0), Vec3(0, 0, boxSize));
+    CustomNonbondedForce* nonbonded = new CustomNonbondedForce("c1*c2*r^-4");
+    nonbonded->addPerParticleParameter("c");
+    vector<Vec3> positions(numParticles);
+    vector<double> params(1);
+    for (int i = 0; i < numParticles; i++) {
+        system.addParticle(1.0);
+        params[0] = (i%2 == 0 ? 1.1 : 2.0);
+        nonbonded->addParticle(params);
+        positions[i] = Vec3(0.5*i, 0, 0);
+    }
+    nonbonded->setNonbondedMethod(CustomNonbondedForce::CutoffPeriodic);
+    nonbonded->setCutoffDistance(cutoff);
+    system.addForce(nonbonded);
+    
+    // Setup nonbonded groups.  They involve 1 interaction of type AA,
+    // 2 of type BB, and 5 of type AB.
+    
+    set<int> set1, set2, set3, set4, set5;
+    set1.insert(0);
+    set1.insert(1);
+    set1.insert(2);
+    nonbonded->addInteractionGroup(set1, set1);
+    set2.insert(3);
+    set3.insert(4);
+    set3.insert(6);
+    set3.insert(8);
+    nonbonded->addInteractionGroup(set2, set3);
+    set4.insert(5);
+    set5.insert(7);
+    set5.insert(9);
+    nonbonded->addInteractionGroup(set4, set5);
+    
+    // Compute energy with and without the correction.
+    
+    VerletIntegrator integrator(0.01);
+    Context context(system, integrator, platform);
+    context.setPositions(positions);
+    double energy1 = context.getState(State::Energy).getPotentialEnergy();
+    nonbonded->setUseLongRangeCorrection(true);
+    context.reinitialize();
+    context.setPositions(positions);
+    double energy2 = context.getState(State::Energy).getPotentialEnergy();
+    
+    // Check the result.
+    
+    double sum = (1.1*1.1 + 2*2.0*2.0 + 5*1.1*2.0)*2.0;
+    int numPairs = (numParticles*(numParticles+1))/2;
+    double expected = 2*M_PI*numParticles*numParticles*sum/(numPairs*boxSize*boxSize*boxSize);
+    ASSERT_EQUAL_TOL(expected, energy2-energy1, 1e-4);
+}
+
 int main() {
     try {
         testSimpleExpression();
@@ -602,6 +660,7 @@ int main() {
         testLongRangeCorrection();
         testInteractionGroups();
         testLargeInteractionGroup();
+        testInteractionGroupLongRangeCorrection();
     }
     catch(const exception& e) {
         cout << "exception: " << e.what() << endl;