Continuing to implement CPU version of CustomManyParticleForce

platforms/cpu/include/CpuCustomManyParticleForce.h

@@ -28,7 +28,9 @@
 #include "ReferenceForce.h"
 #include "ReferenceBondIxn.h"
 #include "CompiledExpressionSet.h"
+#include "CpuNeighborList.h"
 #include "openmm/CustomManyParticleForce.h"
+#include "openmm/internal/ThreadPool.h"
 #include "openmm/internal/vectorize.h"
 #include "lepton/CompiledExpression.h"
 #include "lepton/ParsedExpression.h"
@@ -49,6 +51,8 @@ private:
     bool useCutoff, usePeriodic;
     RealOpenMM cutoffDistance;
     RealOpenMM periodicBoxSize[3];
+    CpuNeighborList* neighborList;
+    ThreadPool& threads;
     CompiledExpressionSet expressionSet;
     Lepton::CompiledExpression energyExpression;
     std::vector<std::vector<int> > particleParamIndices;
@@ -61,23 +65,24 @@ private:
     std::vector<AngleTermInfo> angleTerms;
     std::vector<DihedralTermInfo> dihedralTerms;
 
-    void loopOverInteractions(std::vector<int>& particles, int loopIndex, float* posq, std::vector<OpenMM::RealVec>& atomCoordinates,
+    void loopOverInteractions(std::vector<int>& availableParticles, std::vector<int>& particleSet, int loopIndex, int startIndex, float* posq, std::vector<OpenMM::RealVec>& atomCoordinates,
                               RealOpenMM** particleParameters, std::vector<OpenMM::RealVec>& forces, RealOpenMM* totalEnergy, const fvec4& boxSize, const fvec4& invBoxSize);
 
     /**---------------------------------------------------------------------------------------
 
        Calculate custom interaction for one set of particles
 
-       @param particles        the indices of the particles
-       @param posq             atom coordinates in float format
-       @param atomCoordinates  atom coordinates
-       @param forces           force array (forces added)
-       @param totalEnergy      total energy
+       @param particleSet        the indices of the particles
+       @param posq               atom coordinates in float format
+       @param atomCoordinates    atom coordinates
+       @param particleParameters particle parameter values (particleParameters[particleIndex][parameterIndex])
+       @param forces             force array (forces added)
+       @param totalEnergy        total energy
 
        --------------------------------------------------------------------------------------- */
 
-    void calculateOneIxn(const std::vector<int>& particles, float* posq, std::vector<OpenMM::RealVec>& atomCoordinates,
-                         std::vector<OpenMM::RealVec>& forces, RealOpenMM* totalEnergy, const fvec4& boxSize, const fvec4& invBoxSize);
+    void calculateOneIxn(std::vector<int>& particleSet, float* posq, std::vector<OpenMM::RealVec>& atomCoordinates,
+                         RealOpenMM** particleParameters, std::vector<OpenMM::RealVec>& forces, RealOpenMM* totalEnergy, const fvec4& boxSize, const fvec4& invBoxSize);
 
     /**
      * Compute the displacement and squared distance between two points, optionally using
@@ -98,7 +103,7 @@ public:
 
        --------------------------------------------------------------------------------------- */
 
-    CpuCustomManyParticleForce(const OpenMM::CustomManyParticleForce& force);
+    CpuCustomManyParticleForce(const OpenMM::CustomManyParticleForce& force, ThreadPool& threads);
 
     /**---------------------------------------------------------------------------------------
 
@@ -143,7 +148,7 @@ public:
 
        --------------------------------------------------------------------------------------- */
 
-    void calculateIxn(float* posq, std::vector<OpenMM::RealVec>& atomCoordinates, RealOpenMM** particleParameters,
+    void calculateIxn(AlignedArray<float>& posq, std::vector<OpenMM::RealVec>& atomCoordinates, RealOpenMM** particleParameters,
                       const std::map<std::string, double>& globalParameters,
                       std::vector<OpenMM::RealVec>& forces, RealOpenMM* totalEnergy);
 

platforms/cpu/src/CpuCustomManyParticleForce.cpp

@@ -38,7 +38,8 @@
 using namespace OpenMM;
 using namespace std;
 
-CpuCustomManyParticleForce::CpuCustomManyParticleForce(const CustomManyParticleForce& force) : useCutoff(false), usePeriodic(false) {
+CpuCustomManyParticleForce::CpuCustomManyParticleForce(const CustomManyParticleForce& force, ThreadPool& threads) :
+            threads(threads), useCutoff(false), usePeriodic(false), neighborList(NULL) {
     numParticlesPerSet = force.getNumParticlesPerSet();
     numPerParticleParameters = force.getNumPerParticleParameters();
     
@@ -112,23 +113,62 @@ CpuCustomManyParticleForce::CpuCustomManyParticleForce(const CustomManyParticleF
     CustomManyParticleForceImpl::buildFilterArrays(force, numTypes, particleTypes, orderIndex, particleOrder);
 }
 
-CpuCustomManyParticleForce::~CpuCustomManyParticleForce( ){
+CpuCustomManyParticleForce::~CpuCustomManyParticleForce() {
+    if (neighborList != NULL)
+        delete neighborList;
 }
 
-void CpuCustomManyParticleForce::calculateIxn(float* posq, vector<RealVec>& atomCoordinates, RealOpenMM** particleParameters,
+void CpuCustomManyParticleForce::calculateIxn(AlignedArray<float>& posq, vector<RealVec>& atomCoordinates, RealOpenMM** particleParameters,
                                                   const map<string, double>& globalParameters, vector<RealVec>& forces,
                                                   RealOpenMM* totalEnergy) {
     for (map<string, double>::const_iterator iter = globalParameters.begin(); iter != globalParameters.end(); ++iter)
         expressionSet.setVariable(expressionSet.getVariableIndex(iter->first), iter->second);
-    vector<int> particles(numParticlesPerSet);
+    int numParticles = atomCoordinates.size();
+    vector<int> particleIndices(numParticlesPerSet);
     fvec4 boxSize(periodicBoxSize[0], periodicBoxSize[1], periodicBoxSize[2], 0);
     fvec4 invBoxSize((1/periodicBoxSize[0]), (1/periodicBoxSize[1]), (1/periodicBoxSize[2]), 0);
-    loopOverInteractions(particles, 0, posq, atomCoordinates, particleParameters, forces, totalEnergy, boxSize, invBoxSize);
+    if (useCutoff) {
+        // Construct a neighbor list.
+        
+        float boxSizeFloat[] = {(float) periodicBoxSize[0], (float) periodicBoxSize[1], (float) periodicBoxSize[2]};
+        neighborList->computeNeighborList(numParticles, posq, exclusions, boxSizeFloat, usePeriodic, cutoffDistance, threads);
+        for (int blockIndex = 0; blockIndex < neighborList->getNumBlocks(); blockIndex++) {
+            const vector<int>& neighbors = neighborList->getBlockNeighbors(blockIndex);
+            const vector<char>& exclusions = neighborList->getBlockExclusions(blockIndex);
+            int numNeighbors = neighbors.size();
+            for (int i = 0; i < 4; i++) {
+                particleIndices[0] = neighborList->getSortedAtoms()[4*blockIndex+i];
+                
+                // Build a filtered list of neighbors after removing exclusions.  We'll check for actual exclusions
+                // again later, but the neighbor list also includes padding atoms that it marks as exclusions, so
+                // we need to remove those now.
+                
+                vector<int> particles;
+                for (int j = 0; j < numNeighbors; j++)
+                    if ((exclusions[j] & (1<<i)) == 0)
+                        particles.push_back(neighbors[j]);
+                loopOverInteractions(particles, particleIndices, 1, 0, &posq[0], atomCoordinates, particleParameters, forces, totalEnergy, boxSize, invBoxSize);
+            }
+        }
+    }
+    else {
+        // Loop over all possible sets of particles.
+        
+        vector<int> particles(numParticles);
+        for (int i = 0; i < numParticles; i++)
+            particles[i] = i;
+        for (int i = 0; i < numParticles; i++) {
+            particleIndices[0] = i;
+            loopOverInteractions(particles, particleIndices, 1, i+1, &posq[0], atomCoordinates, particleParameters, forces, totalEnergy, boxSize, invBoxSize);
+        }
+    }
 }
 
 void CpuCustomManyParticleForce::setUseCutoff(RealOpenMM distance) {
     useCutoff = true;
     cutoffDistance = distance;
+    if (neighborList == NULL)
+        neighborList = new CpuNeighborList(4);
 }
 
 void CpuCustomManyParticleForce::setPeriodic(RealVec& boxSize) {
@@ -142,65 +182,68 @@ void CpuCustomManyParticleForce::setPeriodic(RealVec& boxSize) {
     periodicBoxSize[2] = boxSize[2];
 }
 
-void CpuCustomManyParticleForce::loopOverInteractions(vector<int>& particles, int loopIndex, float* posq, vector<OpenMM::RealVec>& atomCoordinates,
+void CpuCustomManyParticleForce::loopOverInteractions(vector<int>& availableParticles, vector<int>& particleSet, int loopIndex, int startIndex, float* posq, vector<OpenMM::RealVec>& atomCoordinates,
                                                           RealOpenMM** particleParameters, vector<OpenMM::RealVec>& forces, RealOpenMM* totalEnergy, const fvec4& boxSize, const fvec4& invBoxSize) {
-    int numParticles = atomCoordinates.size();
-    int start = (loopIndex == 0 ? 0 : particles[loopIndex-1]+1);
-    for (int i = start; i < numParticles; i++) {
-        particles[loopIndex] = i;
-        for (int j = 0; j < numPerParticleParameters; j++)
-            expressionSet.setVariable(particleParamIndices[loopIndex][j], particleParameters[i][j]);
-        if (loopIndex == numParticlesPerSet-1)
-            calculateOneIxn(particles, posq, atomCoordinates, forces, totalEnergy, boxSize, invBoxSize);
-        else
-            loopOverInteractions(particles, loopIndex+1, posq, atomCoordinates, particleParameters, forces, totalEnergy, boxSize, invBoxSize);
+    int numParticles = availableParticles.size();
+//    double cutoff2 = cutoffDistance*cutoffDistance;
+    for (int i = startIndex; i < numParticles; i++) {
+        int particle = availableParticles[i];
+        
+        // Check whether this particle can actually participate in interactions with the others found so far.
+        
+        bool include = true;
+        if (useCutoff) {
+//            fvec4 deltaR;
+//            fvec4 pos1(posq+4*particle);
+//            float r2;
+//            for (int j = 0; j < loopIndex && include; j++) {
+//                fvec4 pos2(posq+4*particleSet[j]);
+//                getDeltaR(pos1, pos2, deltaR, r2, boxSize, invBoxSize);
+//                include &= (r2 < cutoff2);
+//            }
+            RealOpenMM delta[ReferenceForce::LastDeltaRIndex];
+            for (int j = 0; j < loopIndex && include; j++) {
+                computeDelta(particle, particleSet[j], delta, atomCoordinates);
+                include &= (delta[ReferenceForce::RIndex] < cutoffDistance);
+            }
+        }
+        for (int j = 0; j < loopIndex && include; j++)
+            include &= (exclusions[particle].find(particleSet[j]) == exclusions[particle].end());
+        if (include) {
+            particleSet[loopIndex] = availableParticles[i];
+            if (loopIndex == numParticlesPerSet-1)
+                calculateOneIxn(particleSet, posq, atomCoordinates, particleParameters, forces, totalEnergy, boxSize, invBoxSize);
+            else
+                loopOverInteractions(availableParticles, particleSet, loopIndex+1, i+1, posq, atomCoordinates, particleParameters, forces, totalEnergy, boxSize, invBoxSize);
+        }
     }
 }
 
-void CpuCustomManyParticleForce::calculateOneIxn(const vector<int>& particles, float* posq, vector<RealVec>& atomCoordinates, vector<RealVec>& forces, RealOpenMM* totalEnergy, const fvec4& boxSize, const fvec4& invBoxSize) {
+void CpuCustomManyParticleForce::calculateOneIxn(vector<int>& particleSet, float* posq, vector<RealVec>& atomCoordinates, RealOpenMM** particleParameters, vector<RealVec>& forces, RealOpenMM* totalEnergy, const fvec4& boxSize, const fvec4& invBoxSize) {
     // Select the ordering to use for the particles.
     
     vector<int> permutedParticles(numParticlesPerSet);
     if (particleOrder.size() == 1) {
         // There are no filters, so we don't need to worry about ordering.
         
-        permutedParticles = particles;
+        permutedParticles = particleSet;
     }
     else {
         int index = 0;
         for (int i = numParticlesPerSet-1; i >= 0; i--)
-            index = particleTypes[particles[i]]+numTypes*index;
+            index = particleTypes[particleSet[i]]+numTypes*index;
         int order = orderIndex[index];
         if (order == -1)
             return;
         for (int i = 0; i < numParticlesPerSet; i++)
-            permutedParticles[i] = particles[particleOrder[order][i]];
+            permutedParticles[i] = particleSet[particleOrder[order][i]];
     }
+
+    // Record per-particle parameters.
     
-    // Decide whether to include this interaction.
-    
-    for (int i = 0; i < (int) permutedParticles.size(); i++) {
-        int p1 = permutedParticles[i];
-        for (int j = i+1; j < (int) permutedParticles.size(); j++) {
-            int p2 = permutedParticles[j];
-            if (exclusions[p1].find(p2) != exclusions[p1].end())
-                return;
-            if (useCutoff) {
-//                fvec4 deltaR;
-//                fvec4 posI(posq+4*p1);
-//                fvec4 posJ(posq+4*p2);
-//                float r2;
-//                getDeltaR(posI, posJ, deltaR, r2, boxSize, invBoxSize);
-//                if (r2 >= cutoffDistance*cutoffDistance)
-//                    return;
-//
-                RealOpenMM delta[ReferenceForce::LastDeltaRIndex];
-                computeDelta(p1, p2, delta, atomCoordinates);
-                if (delta[ReferenceForce::RIndex] >= cutoffDistance)
-                    return;
-            }
-        }
-    }
+    for (int i = 0; i < numParticlesPerSet; i++)
+        for (int j = 0; j < numPerParticleParameters; j++)
+            expressionSet.setVariable(particleParamIndices[i][j], particleParameters[permutedParticles[i]][j]);
     
     // Compute all of the variables the energy can depend on.
 

platforms/cpu/src/CpuKernels.cpp

@@ -863,7 +863,7 @@ void CpuCalcCustomManyParticleForceKernel::initialize(const System& system, cons
     }
     for (int i = 0; i < force.getNumGlobalParameters(); i++)
         globalParameterNames.push_back(force.getGlobalParameterName(i));
-    ixn = new CpuCustomManyParticleForce(force);
+    ixn = new CpuCustomManyParticleForce(force, data.threads);
     nonbondedMethod = CalcCustomManyParticleForceKernel::NonbondedMethod(force.getNonbondedMethod());
     cutoffDistance = force.getCutoffDistance();
 }
@@ -882,7 +882,7 @@ double CpuCalcCustomManyParticleForceKernel::execute(ContextImpl& context, bool
             throw OpenMMException("The periodic box size has decreased to less than twice the nonbonded cutoff.");
         ixn->setPeriodic(box);
     }
-    ixn->calculateIxn(&data.posq[0], posData, particleParamArray, globalParameters, forceData, includeEnergy ? &energy : NULL);
+    ixn->calculateIxn(data.posq, posData, particleParamArray, globalParameters, forceData, includeEnergy ? &energy : NULL);
     return energy;
 }
 

platforms/reference/include/ReferenceCustomManyParticleIxn.h

@@ -66,16 +66,17 @@ class ReferenceCustomManyParticleIxn {
 
          Calculate custom interaction for one set of particles
 
-         @param particles        the indices of the particles
-         @param atomCoordinates  atom coordinates
-         @param variables        the values of variables that may appear in expressions
-         @param forces           force array (forces added)
-         @param totalEnergy      total energy
+         @param particles          the indices of the particles
+         @param atomCoordinates    atom coordinates
+         @param particleParameters particle parameter values (particleParameters[particleIndex][parameterIndex])
+         @param variables          the values of variables that may appear in expressions
+         @param forces             force array (forces added)
+         @param totalEnergy        total energy
 
          --------------------------------------------------------------------------------------- */
 
       void calculateOneIxn(const std::vector<int>& particles, std::vector<OpenMM::RealVec>& atomCoordinates,
-                           std::map<std::string, double>& variables, std::vector<OpenMM::RealVec>& forces,
+                           RealOpenMM** particleParameters, std::map<std::string, double>& variables, std::vector<OpenMM::RealVec>& forces,
                            RealOpenMM* totalEnergy) const;
 
       void computeDelta(int atom1, int atom2, RealOpenMM* delta, std::vector<OpenMM::RealVec>& atomCoordinates) const;

platforms/reference/src/SimTKReference/ReferenceCustomManyParticleIxn.cpp

@@ -137,17 +137,15 @@ void ReferenceCustomManyParticleIxn::loopOverInteractions(vector<int>& particles
     int start = (loopIndex == 0 ? 0 : particles[loopIndex-1]+1);
     for (int i = start; i < numParticles; i++) {
         particles[loopIndex] = i;
-        for (int j = 0; j < numPerParticleParameters; j++)
-            variables[particleParamNames[loopIndex][j]] = particleParameters[i][j];
         if (loopIndex == numParticlesPerSet-1)
-            calculateOneIxn(particles, atomCoordinates, variables, forces, totalEnergy);
+            calculateOneIxn(particles, atomCoordinates, particleParameters, variables, forces, totalEnergy);
         else
             loopOverInteractions(particles, loopIndex+1, atomCoordinates, particleParameters, variables, forces, totalEnergy);
     }
 }
 
 void ReferenceCustomManyParticleIxn::calculateOneIxn(const vector<int>& particles, vector<RealVec>& atomCoordinates,
-                        map<string, double>& variables, vector<RealVec>& forces, RealOpenMM* totalEnergy) const {
+                        RealOpenMM** particleParameters, map<string, double>& variables, vector<RealVec>& forces, RealOpenMM* totalEnergy) const {
     // Select the ordering to use for the particles.
     
     vector<int> permutedParticles(numParticlesPerSet);
@@ -169,9 +167,9 @@ void ReferenceCustomManyParticleIxn::calculateOneIxn(const vector<int>& particle
     
     // Decide whether to include this interaction.
     
-    for (int i = 0; i < (int) permutedParticles.size(); i++) {
+    for (int i = 0; i < numParticlesPerSet; i++) {
         int p1 = permutedParticles[i];
-        for (int j = i+1; j < (int) permutedParticles.size(); j++) {
+        for (int j = i+1; j < numParticlesPerSet; j++) {
             int p2 = permutedParticles[j];
             if (exclusions[p1].find(p2) != exclusions[p1].end())
                 return;
@@ -183,6 +181,12 @@ void ReferenceCustomManyParticleIxn::calculateOneIxn(const vector<int>& particle
             }
         }
     }
+
+    // Record per-particle parameters.
+    
+    for (int i = 0; i < numParticlesPerSet; i++)
+        for (int j = 0; j < numPerParticleParameters; j++)
+            variables[particleParamNames[i][j]] = particleParameters[permutedParticles[i]][j];
     
     // Compute all of the variables the energy can depend on.
 

platforms/reference/tests/TestReferenceCustomManyParticleForce.cpp

@@ -351,13 +351,15 @@ void testTypeFilters() {
     System system;
     for (int i = 0; i < 5; i++)
         system.addParticle(1.0);
-    CustomManyParticleForce* force = new CustomManyParticleForce(3, "distance(p1,p2)+distance(p1,p3)");
-    vector<double> params;
-    force->addParticle(params, 0);
-    force->addParticle(params, 1);
-    force->addParticle(params, 0);
-    force->addParticle(params, 1);
-    force->addParticle(params, 5);
+    CustomManyParticleForce* force = new CustomManyParticleForce(3, "c1*(distance(p1,p2)+distance(p1,p3))");
+    force->addPerParticleParameter("c");
+    double c[] = {1.0, 2.0, 1.3, 1.5, -2.1};
+    int type[] = {0, 1, 0, 1, 5};
+    vector<double> params(1);
+    for (int i = 0; i < 5; i++) {
+        params[0] = c[i];
+        force->addParticle(params, type[i]);
+    }
     vector<Vec3> positions;
     positions.push_back(Vec3(0, 0, 0));
     positions.push_back(Vec3(1, 0, 0));
@@ -390,10 +392,9 @@ void testTypeFilters() {
             Vec3 d13 = positions[p3]-positions[p1];
             double r12 = sqrt(d12.dot(d12));
             double r13 = sqrt(d13.dot(d13));
-            expectedEnergy += r12+r13;
+            expectedEnergy += c[p1]*(r12+r13);
     }
     ASSERT_EQUAL_TOL(expectedEnergy, state.getPotentialEnergy(), 1e-5);
-    
 }
 
 int main() {