Finished reference implementation of CustomManyParticleForce

openmmapi/include/openmm/internal/CustomManyParticleForceImpl.h

@@ -79,11 +79,27 @@ public:
      */
     static Lepton::ParsedExpression prepareExpression(const CustomManyParticleForce& force, const std::map<std::string, Lepton::CustomFunction*>& functions, std::map<std::string, std::vector<int> >& distances,
             std::map<std::string, std::vector<int> >& angles, std::map<std::string, std::vector<int> >& dihedrals);
+    /**
+     * Analyze the type filters for a force and build a set of arrays that can be used for reordering the
+     * particles in an interaction.
+     * 
+     * @param force          the CustomManyParticleForce to process
+     * @param numTypes       on exit, the number of unique particle types
+     * @param particleTypes  on exit, this contains a type code for each particle.  These codes are <i>not</i> necessarily the
+     *                       same as the types assigned by the force.  They are guaranteed to be successive integers starting from 0,
+     *                       whereas the force may have used arbitrary integers.
+     * @param orderIndex     on exit, this contains a lookup table for selecting the particle order for an interaction.
+     *                       orderIndex[t1+numTypes*t2+numTypes*numTypes*t3+...] is the index of the order to use, where t1, t2, etc. are the type codes
+     *                       of the particles involved in the interaction.  If this equals -1, the interaction should be omitted.
+     * @param particleOrder  on exit, particleOrder[i][j] tells which particle to use as the j'th particle, where i is the value found in orderIndex.
+     */
+    static void buildFilterArrays(const CustomManyParticleForce& force, int& numTypes, std::vector<int>& particleTypes, std::vector<int>& orderIndex, std::vector<std::vector<int> >& particleOrder);
 private:
     class FunctionPlaceholder;
     static Lepton::ExpressionTreeNode replaceFunctions(const Lepton::ExpressionTreeNode& node, std::map<std::string, int> atoms,
             std::map<std::string, std::vector<int> >& distances, std::map<std::string, std::vector<int> >& angles,
             std::map<std::string, std::vector<int> >& dihedrals);
+    static void generatePermutations(std::vector<int>& values, int numFixed, std::vector<std::vector<int> >& result);
     const CustomManyParticleForce& owner;
     Kernel kernel;
 };

openmmapi/src/CustomManyParticleForceImpl.cpp

@@ -225,3 +225,109 @@ ExpressionTreeNode CustomManyParticleForceImpl::replaceFunctions(const Expressio
 void CustomManyParticleForceImpl::updateParametersInContext(ContextImpl& context) {
     kernel.getAs<CalcCustomManyParticleForceKernel>().copyParametersToContext(context, owner);
 }
+
+void CustomManyParticleForceImpl::buildFilterArrays(const CustomManyParticleForce& force, int& numTypes, vector<int>& particleTypes, vector<int>& orderIndex, vector<vector<int> >& particleOrder) {
+    // Build a canonical list of type codes.
+    
+    int numParticles = force.getNumParticles();
+    int numParticlesPerSet = force.getNumParticlesPerSet();
+    particleTypes.resize(numParticles);
+    map<int, int> typeMap;
+    for (int i = 0; i < numParticles; i++) {
+        vector<double> params;
+        int type;
+        force.getParticleParameters(i, params, type);
+        map<int, int>::const_iterator element = typeMap.find(type);
+        if (element == typeMap.end()) {
+            int newType = typeMap.size();
+            typeMap[type] = newType;
+            particleTypes[i] = newType;
+        }
+        else
+            particleTypes[i] = element->second;
+    }
+    numTypes = typeMap.size();
+    int numIndices = 1;
+    for (int i = 0; i < numParticlesPerSet; i++)
+        numIndices *= numTypes;
+    orderIndex.resize(numIndices, 0);
+    
+    // Find the allowed type codes for each particle in an interaction.
+    
+    vector<set<int> > allowedTypes(numParticlesPerSet);
+    bool anyFilters = false;
+    for (int i = 0; i < numParticlesPerSet; i++) {
+        set<int> types;
+        force.getTypeFilter(i, types);
+        if (types.size() == 0)
+            for (int j = 0; j < numTypes; j++)
+                allowedTypes[i].insert(j);
+        else {
+            for (set<int>::const_iterator iter = types.begin(); iter != types.end(); ++iter)
+                if (typeMap.find(*iter) != typeMap.end())
+                    allowedTypes[i].insert(typeMap[*iter]);
+            if (allowedTypes[i].size() < numTypes)
+                anyFilters = true;
+        }
+    }
+    
+    // If there are no filters, reordering is unnecessary.
+    
+    if (!anyFilters) {
+        particleOrder.resize(1);
+        particleOrder[0].resize(numParticlesPerSet);
+        for (int i = 0; i < numParticlesPerSet; i++)
+            particleOrder[0][i] = i;
+        return;
+    }
+    
+    // Build a list of every possible permutation of the particles.
+    
+    particleOrder.clear();
+    vector<int> values;
+    for (int i = 0; i < numParticlesPerSet; i++)
+        values.push_back(i);
+    generatePermutations(values, 0, particleOrder);
+    int numOrders = particleOrder.size();
+    
+    // Now we need to loop over every possible sequence of type codes, and for each one figure out which order to use.
+    
+    for (int i = 0; i < numIndices; i++) {
+        vector<int> types(numParticlesPerSet);
+        int temp = i;
+        for (int j = 0; j < numParticlesPerSet; j++) {
+            types[j] = temp%numTypes;
+            temp /= numTypes;
+        }
+        
+        // Loop over possible orders until we find one that matches the filters.
+        
+        int order = -1;
+        for (int j = 0; j < numOrders && order == -1; j++) {
+            bool matches = true;
+            for (int k = 0; k < numParticlesPerSet && matches; k++)
+                if (allowedTypes[k].find(types[particleOrder[j][k]]) == allowedTypes[k].end())
+                    matches = false;
+            if (matches)
+                order = j;
+        }
+        orderIndex[i] = order;
+    }
+}
+
+void CustomManyParticleForceImpl::generatePermutations(vector<int>& values, int numFixed, vector<vector<int> >& result) {
+    int numValues = values.size();
+    if (numFixed == numValues) {
+        result.push_back(values);
+        return;
+    }
+    for (int i = numFixed; i < numValues; i++) {
+        int v1 = values[numFixed];
+        int v2 = values[i];
+        values[numFixed] = v2;
+        values[i] = v1;
+        generatePermutations(values, numFixed+1, result);
+        values[numFixed] = v1;
+        values[i] = v2;
+    }
+}
\ No newline at end of file

platforms/reference/include/ReferenceCustomManyParticleIxn.h

@@ -43,13 +43,16 @@ class ReferenceCustomManyParticleIxn {
       class DistanceTermInfo;
       class AngleTermInfo;
       class DihedralTermInfo;
-      int numParticlesPerSet, numPerParticleParameters;
+      int numParticlesPerSet, numPerParticleParameters, numTypes;
       bool useCutoff, usePeriodic;
       RealOpenMM cutoffDistance;
       RealOpenMM periodicBoxSize[3];
       Lepton::ExpressionProgram energyExpression;
       std::vector<std::vector<std::string> > particleParamNames;
       std::vector<std::set<int> > exclusions;
+      std::vector<int> particleTypes;
+      std::vector<int> orderIndex;
+      std::vector<std::vector<int> > particleOrder;
       std::vector<ParticleTermInfo> particleTerms;
       std::vector<DistanceTermInfo> distanceTerms;
       std::vector<AngleTermInfo> angleTerms;

platforms/reference/src/SimTKReference/ReferenceCustomManyParticleIxn.cpp

@@ -97,6 +97,10 @@ ReferenceCustomManyParticleIxn::ReferenceCustomManyParticleIxn(const CustomManyP
         exclusions[p1].insert(p2);
         exclusions[p2].insert(p1);
     }
+    
+    // Record information about type filters.
+    
+    CustomManyParticleForceImpl::buildFilterArrays(force, numTypes, particleTypes, orderIndex, particleOrder);
 }
 
 ReferenceCustomManyParticleIxn::~ReferenceCustomManyParticleIxn( ){
@@ -144,12 +148,31 @@ void ReferenceCustomManyParticleIxn::loopOverInteractions(vector<int>& particles
 
 void ReferenceCustomManyParticleIxn::calculateOneIxn(const vector<int>& particles, vector<RealVec>& atomCoordinates,
                         map<string, double>& variables, vector<RealVec>& forces, RealOpenMM* totalEnergy) const {
+    // 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;
+    }
+    else {
+        int index = 0;
+        for (int i = numParticlesPerSet-1; i >= 0; i--)
+            index = particleTypes[particles[i]]+numTypes*index;
+        int order = orderIndex[index];
+        if (order == -1)
+            return;
+        for (int i = 0; i < numParticlesPerSet; i++)
+            permutedParticles[i] = particles[particleOrder[order][i]];
+    }
+    
     // Decide whether to include this interaction.
     
-    for (int i = 0; i < (int) particles.size(); i++) {
-        int p1 = particles[i];
-        for (int j = i+1; j < (int) particles.size(); j++) {
-            int p2 = particles[j];
+    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) {
@@ -169,20 +192,20 @@ void ReferenceCustomManyParticleIxn::calculateOneIxn(const vector<int>& particle
     }
     for (int i = 0; i < (int) distanceTerms.size(); i++) {
         const DistanceTermInfo& term = distanceTerms[i];
-        computeDelta(particles[term.p1], particles[term.p2], term.delta, atomCoordinates);
+        computeDelta(permutedParticles[term.p1], permutedParticles[term.p2], term.delta, atomCoordinates);
         variables[term.name] = term.delta[ReferenceForce::RIndex];
     }
     for (int i = 0; i < (int) angleTerms.size(); i++) {
         const AngleTermInfo& term = angleTerms[i];
-        computeDelta(particles[term.p1], particles[term.p2], term.delta1, atomCoordinates);
-        computeDelta(particles[term.p3], particles[term.p2], term.delta2, atomCoordinates);
+        computeDelta(permutedParticles[term.p1], permutedParticles[term.p2], term.delta1, atomCoordinates);
+        computeDelta(permutedParticles[term.p3], permutedParticles[term.p2], term.delta2, atomCoordinates);
         variables[term.name] = computeAngle(term.delta1, term.delta2);
     }
     for (int i = 0; i < (int) dihedralTerms.size(); i++) {
         const DihedralTermInfo& term = dihedralTerms[i];
-        computeDelta(particles[term.p2], particles[term.p1], term.delta1, atomCoordinates);
-        computeDelta(particles[term.p2], particles[term.p3], term.delta2, atomCoordinates);
-        computeDelta(particles[term.p4], particles[term.p3], term.delta3, atomCoordinates);
+        computeDelta(permutedParticles[term.p2], permutedParticles[term.p1], term.delta1, atomCoordinates);
+        computeDelta(permutedParticles[term.p2], permutedParticles[term.p3], term.delta2, atomCoordinates);
+        computeDelta(permutedParticles[term.p4], permutedParticles[term.p3], term.delta3, atomCoordinates);
         RealOpenMM dotDihedral, signOfDihedral;
         RealOpenMM* crossProduct[] = {term.cross1, term.cross2};
         variables[term.name] = ReferenceBondIxn::getDihedralAngleBetweenThreeVectors(term.delta1, term.delta2, term.delta3, crossProduct, &dotDihedral, term.delta1, &signOfDihedral, 1);
@@ -192,7 +215,7 @@ void ReferenceCustomManyParticleIxn::calculateOneIxn(const vector<int>& particle
     
     for (int i = 0; i < (int) particleTerms.size(); i++) {
         const ParticleTermInfo& term = particleTerms[i];
-        forces[particles[term.atom]][term.component] -= term.forceExpression.evaluate(variables);
+        forces[permutedParticles[term.atom]][term.component] -= term.forceExpression.evaluate(variables);
     }
 
     // Apply forces based on distances.
@@ -202,8 +225,8 @@ void ReferenceCustomManyParticleIxn::calculateOneIxn(const vector<int>& particle
         RealOpenMM dEdR = (RealOpenMM) (term.forceExpression.evaluate(variables)/(term.delta[ReferenceForce::RIndex]));
         for (int i = 0; i < 3; i++) {
            RealOpenMM force  = -dEdR*term.delta[i];
-           forces[particles[term.p1]][i] -= force;
-           forces[particles[term.p2]][i] += force;
+           forces[permutedParticles[term.p1]][i] -= force;
+           forces[permutedParticles[term.p2]][i] += force;
         }
     }
 
@@ -228,9 +251,9 @@ void ReferenceCustomManyParticleIxn::calculateOneIxn(const vector<int>& particle
             deltaCrossP[1][i] = -(deltaCrossP[0][i]+deltaCrossP[2][i]);
         }
         for (int i = 0; i < 3; i++) {
-            forces[particles[term.p1]][i] += deltaCrossP[0][i];
-            forces[particles[term.p2]][i] += deltaCrossP[1][i];
-            forces[particles[term.p3]][i] += deltaCrossP[2][i];
+            forces[permutedParticles[term.p1]][i] += deltaCrossP[0][i];
+            forces[permutedParticles[term.p2]][i] += deltaCrossP[1][i];
+            forces[permutedParticles[term.p3]][i] += deltaCrossP[2][i];
         }
     }
 
@@ -258,10 +281,10 @@ void ReferenceCustomManyParticleIxn::calculateOneIxn(const vector<int>& particle
             internalF[2][i] = internalF[3][i] + s;
         }
         for (int i = 0; i < 3; i++) {
-            forces[particles[term.p1]][i] += internalF[0][i];
-            forces[particles[term.p2]][i] -= internalF[1][i];
-            forces[particles[term.p3]][i] -= internalF[2][i];
-            forces[particles[term.p4]][i] += internalF[3][i];
+            forces[permutedParticles[term.p1]][i] += internalF[0][i];
+            forces[permutedParticles[term.p2]][i] -= internalF[1][i];
+            forces[permutedParticles[term.p3]][i] -= internalF[2][i];
+            forces[permutedParticles[term.p4]][i] += internalF[3][i];
         }
     }
 

platforms/reference/tests/TestReferenceCustomManyParticleForce.cpp

@@ -345,6 +345,57 @@ void testTabulatedFunctions() {
     ASSERT_EQUAL_TOL(expectedEnergy, state.getPotentialEnergy(), 1e-5);
 }
 
+void testTypeFilters() {
+    // Create a system.
+    
+    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);
+    vector<Vec3> positions;
+    positions.push_back(Vec3(0, 0, 0));
+    positions.push_back(Vec3(1, 0, 0));
+    positions.push_back(Vec3(0, 1.1, 0.3));
+    positions.push_back(Vec3(0.4, 0, -0.8));
+    positions.push_back(Vec3(0.2, 0.5, -0.1));
+    set<int> f1, f2;
+    f1.insert(0);
+    f2.insert(1);
+    f2.insert(5);
+    force->setTypeFilter(0, f1);
+    force->setTypeFilter(1, f2);
+    force->setTypeFilter(2, f2);
+    system.addForce(force);
+    VerletIntegrator integrator(0.001);
+    ReferencePlatform platform;
+    Context context(system, integrator, platform);
+    context.setPositions(positions);
+    
+    // See if the energy is correct.
+
+    State state = context.getState(State::Energy);
+    double expectedEnergy = 0;
+    int sets[6][3] = {{0,1,3}, {0,1,4}, {0,3,4}, {2,1,3}, {2,1,4}, {2,3,4}};
+    for (int i = 0; i < 6; i++) {
+        int p1 = sets[i][0];
+        int p2 = sets[i][1];
+        int p3 = sets[i][2];
+            Vec3 d12 = positions[p2]-positions[p1];
+            Vec3 d13 = positions[p3]-positions[p1];
+            double r12 = sqrt(d12.dot(d12));
+            double r13 = sqrt(d13.dot(d13));
+            expectedEnergy += r12+r13;
+    }
+    ASSERT_EQUAL_TOL(expectedEnergy, state.getPotentialEnergy(), 1e-5);
+    
+}
+
 int main() {
     try {
         testNoCutoff();
@@ -353,6 +404,7 @@ int main() {
         testExclusions();
         testParameters();
         testTabulatedFunctions();
+        testTypeFilters();
     }
     catch(const exception& e) {
         cout << "exception: " << e.what() << endl;