Allow disabling the direct space calculation in NonbondedForce (#3173)

openmmapi/include/openmm/NonbondedForce.h

@@ -9,7 +9,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008-2020 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2021 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -543,6 +543,18 @@ public:
      *                 that is specified for direct space.
      */
     void setReciprocalSpaceForceGroup(int group);
+    /**
+     * Get whether to include direct space interactions when calculating forces and energies.  This is useful if you want
+     * to completely replace the direct space calculation, typically with a CustomNonbondedForce that computes it in a
+     * nonstandard way, while still using this object for the reciprocal space calculation.
+     */
+    bool getIncludeDirectSpace() const;
+    /**
+     * Set whether to include direct space interactions when calculating forces and energies.  This is useful if you want
+     * to completely replace the direct space calculation, typically with a CustomNonbondedForce that computes it in a
+     * nonstandard way, while still using this object for the reciprocal space calculation.
+     */
+    void setIncludeDirectSpace(bool include);
     /**
      * Update the particle and exception parameters in a Context to match those stored in this Force object.  This method
      * provides an efficient method to update certain parameters in an existing Context without needing to reinitialize it.
@@ -604,7 +616,7 @@ private:
     class ExceptionOffsetInfo;
     NonbondedMethod nonbondedMethod;
     double cutoffDistance, switchingDistance, rfDielectric, ewaldErrorTol, alpha, dalpha;
-    bool useSwitchingFunction, useDispersionCorrection, exceptionsUsePeriodic;
+    bool useSwitchingFunction, useDispersionCorrection, exceptionsUsePeriodic, includeDirectSpace;
     int recipForceGroup, nx, ny, nz, dnx, dny, dnz;
     void addExclusionsToSet(const std::vector<std::set<int> >& bonded12, std::set<int>& exclusions, int baseParticle, int fromParticle, int currentLevel) const;
     int getGlobalParameterIndex(const std::string& parameter) const;

openmmapi/src/NonbondedForce.cpp

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008-2020 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2021 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -49,7 +49,7 @@ using std::vector;
 
 NonbondedForce::NonbondedForce() : nonbondedMethod(NoCutoff), cutoffDistance(1.0), switchingDistance(-1.0), rfDielectric(78.3),
         ewaldErrorTol(5e-4), alpha(0.0), dalpha(0.0), useSwitchingFunction(false), useDispersionCorrection(true), exceptionsUsePeriodic(false), recipForceGroup(-1),
-        nx(0), ny(0), nz(0), dnx(0), dny(0), dnz(0) {
+        includeDirectSpace(true), nx(0), ny(0), nz(0), dnx(0), dny(0), dnz(0) {
 }
 
 NonbondedForce::NonbondedMethod NonbondedForce::getNonbondedMethod() const {
@@ -344,6 +344,14 @@ void NonbondedForce::setReciprocalSpaceForceGroup(int group) {
     recipForceGroup = group;
 }
 
+bool NonbondedForce::getIncludeDirectSpace() const {
+    return includeDirectSpace;
+}
+
+void NonbondedForce::setIncludeDirectSpace(bool include) {
+    includeDirectSpace = include;
+}
+
 void NonbondedForce::updateParametersInContext(Context& context) {
     dynamic_cast<NonbondedForceImpl&>(getImplInContext(context)).updateParametersInContext(getContextImpl(context));
 }

openmmapi/src/NonbondedForceImpl.cpp

@@ -103,10 +103,11 @@ void NonbondedForceImpl::initialize(ContextImpl& context) {
 }
 
 double NonbondedForceImpl::calcForcesAndEnergy(ContextImpl& context, bool includeForces, bool includeEnergy, int groups) {
-    bool includeDirect = ((groups&(1<<owner.getForceGroup())) != 0);
-    bool includeReciprocal = includeDirect;
-    if (owner.getReciprocalSpaceForceGroup() >= 0)
-        includeReciprocal = ((groups&(1<<owner.getReciprocalSpaceForceGroup())) != 0);
+    bool includeDirect = (owner.getIncludeDirectSpace() && (groups&(1<<owner.getForceGroup())) != 0);
+    int reciprocalGroup = owner.getReciprocalSpaceForceGroup();
+    if (reciprocalGroup < 0)
+        reciprocalGroup = owner.getForceGroup();
+    bool includeReciprocal = ((groups&(1<<reciprocalGroup)) != 0);
     return kernel.getAs<CalcNonbondedForceKernel>().execute(context, includeForces, includeEnergy, includeDirect, includeReciprocal);
 }
 

platforms/cuda/src/CudaKernels.cpp

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008-2019 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2021 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -979,7 +979,8 @@ void CudaCalcNonbondedForceKernel::initialize(const System& system, const Nonbon
             replacements["USE_PERIODIC"] = force.getExceptionsUsePeriodicBoundaryConditions() ? "1" : "0";
             if (doLJPME)
                 replacements["EWALD_DISPERSION_ALPHA"] = cu.doubleToString(dispersionAlpha);
-            cu.getBondedUtilities().addInteraction(atoms, cu.replaceStrings(CommonKernelSources::pmeExclusions, replacements), force.getForceGroup());
+            if (force.getIncludeDirectSpace())
+                cu.getBondedUtilities().addInteraction(atoms, cu.replaceStrings(CommonKernelSources::pmeExclusions, replacements), force.getForceGroup());
         }
     }
 
@@ -1008,7 +1009,8 @@ void CudaCalcNonbondedForceKernel::initialize(const System& system, const Nonbon
         cu.getNonbondedUtilities().addParameter(CudaNonbondedUtilities::ParameterInfo(prefix+"sigmaEpsilon", "float", 2, sizeof(float2), sigmaEpsilon.getDevicePointer()));
     }
     source = cu.replaceStrings(source, replacements);
-    cu.getNonbondedUtilities().addInteraction(useCutoff, usePeriodic, true, force.getCutoffDistance(), exclusionList, source, force.getForceGroup(), true);
+    if (force.getIncludeDirectSpace())
+        cu.getNonbondedUtilities().addInteraction(useCutoff, usePeriodic, true, force.getCutoffDistance(), exclusionList, source, force.getForceGroup(), true);
 
     // Initialize the exceptions.
 
@@ -1033,7 +1035,8 @@ void CudaCalcNonbondedForceKernel::initialize(const System& system, const Nonbon
         map<string, string> replacements;
         replacements["APPLY_PERIODIC"] = (usePeriodic && force.getExceptionsUsePeriodicBoundaryConditions() ? "1" : "0");
         replacements["PARAMS"] = cu.getBondedUtilities().addArgument(exceptionParams.getDevicePointer(), "float4");
-        cu.getBondedUtilities().addInteraction(atoms, cu.replaceStrings(CommonKernelSources::nonbondedExceptions, replacements), force.getForceGroup());
+        if (force.getIncludeDirectSpace())
+            cu.getBondedUtilities().addInteraction(atoms, cu.replaceStrings(CommonKernelSources::nonbondedExceptions, replacements), force.getForceGroup());
     }
     
     // Initialize parameter offsets.

platforms/opencl/src/OpenCLKernels.cpp

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008-2020 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2021 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -927,7 +927,8 @@ void OpenCLCalcNonbondedForceKernel::initialize(const System& system, const Nonb
             replacements["USE_PERIODIC"] = force.getExceptionsUsePeriodicBoundaryConditions() ? "1" : "0";
             if (doLJPME)
                 replacements["EWALD_DISPERSION_ALPHA"] = cl.doubleToString(dispersionAlpha);
-            cl.getBondedUtilities().addInteraction(atoms, cl.replaceStrings(CommonKernelSources::pmeExclusions, replacements), force.getForceGroup());
+            if (force.getIncludeDirectSpace())
+                cl.getBondedUtilities().addInteraction(atoms, cl.replaceStrings(CommonKernelSources::pmeExclusions, replacements), force.getForceGroup());
         }
     }
 
@@ -956,7 +957,8 @@ void OpenCLCalcNonbondedForceKernel::initialize(const System& system, const Nonb
         cl.getNonbondedUtilities().addParameter(OpenCLNonbondedUtilities::ParameterInfo(prefix+"sigmaEpsilon", "float", 2, sizeof(cl_float2), sigmaEpsilon.getDeviceBuffer()));
     }
     source = cl.replaceStrings(source, replacements);
-    cl.getNonbondedUtilities().addInteraction(useCutoff, usePeriodic, true, force.getCutoffDistance(), exclusionList, source, force.getForceGroup());
+    if (force.getIncludeDirectSpace())
+        cl.getNonbondedUtilities().addInteraction(useCutoff, usePeriodic, true, force.getCutoffDistance(), exclusionList, source, force.getForceGroup());
 
     // Initialize the exceptions.
 
@@ -981,7 +983,8 @@ void OpenCLCalcNonbondedForceKernel::initialize(const System& system, const Nonb
         map<string, string> replacements;
         replacements["APPLY_PERIODIC"] = (usePeriodic && force.getExceptionsUsePeriodicBoundaryConditions() ? "1" : "0");
         replacements["PARAMS"] = cl.getBondedUtilities().addArgument(exceptionParams.getDeviceBuffer(), "float4");
-        cl.getBondedUtilities().addInteraction(atoms, cl.replaceStrings(CommonKernelSources::nonbondedExceptions, replacements), force.getForceGroup());
+        if (force.getIncludeDirectSpace())
+            cl.getBondedUtilities().addInteraction(atoms, cl.replaceStrings(CommonKernelSources::nonbondedExceptions, replacements), force.getForceGroup());
     }
     
     // Initialize parameter offsets.

serialization/src/NonbondedForceProxy.cpp

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2010-2020 Stanford University and the Authors.      *
+ * Portions copyright (c) 2010-2021 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -54,6 +54,7 @@ void NonbondedForceProxy::serialize(const void* object, SerializationNode& node)
     node.setDoubleProperty("rfDielectric", force.getReactionFieldDielectric());
     node.setIntProperty("dispersionCorrection", force.getUseDispersionCorrection());
     node.setIntProperty("exceptionsUsePeriodic", force.getExceptionsUsePeriodicBoundaryConditions());
+    node.setBoolProperty("includeDirectSpace", force.getIncludeDirectSpace());
     double alpha;
     int nx, ny, nz;
     force.getPMEParameters(alpha, nx, ny, nz);
@@ -116,6 +117,8 @@ void* NonbondedForceProxy::deserialize(const SerializationNode& node) const {
         force->setEwaldErrorTolerance(node.getDoubleProperty("ewaldTolerance"));
         force->setReactionFieldDielectric(node.getDoubleProperty("rfDielectric"));
         force->setUseDispersionCorrection(node.getIntProperty("dispersionCorrection"));
+        if (node.hasProperty("includeDirectSpace"))
+            force->setIncludeDirectSpace(node.getBoolProperty("includeDirectSpace"));
         double alpha = node.getDoubleProperty("alpha", 0.0);
         int nx = node.getIntProperty("nx", 0);
         int ny = node.getIntProperty("ny", 0);

serialization/tests/TestSerializeNonbondedForce.cpp

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2010-2020 Stanford University and the Authors.      *
+ * Portions copyright (c) 2010-2021 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -52,6 +52,7 @@ void testSerialization() {
     force.setReactionFieldDielectric(50.0);
     force.setUseDispersionCorrection(false);
     force.setExceptionsUsePeriodicBoundaryConditions(true);
+    force.setIncludeDirectSpace(false);
     double alpha = 0.5;
     int nx = 3, ny = 5, nz = 7;
     force.setPMEParameters(alpha, nx, ny, nz);
@@ -92,6 +93,7 @@ void testSerialization() {
     ASSERT_EQUAL(force.getNumGlobalParameters(), force2.getNumGlobalParameters());
     ASSERT_EQUAL(force.getNumParticleParameterOffsets(), force2.getNumParticleParameterOffsets());
     ASSERT_EQUAL(force.getNumExceptionParameterOffsets(), force2.getNumExceptionParameterOffsets());
+    ASSERT_EQUAL(force.getIncludeDirectSpace(), force2.getIncludeDirectSpace());
     double alpha2;
     int nx2, ny2, nz2;
     force2.getPMEParameters(alpha2, nx2, ny2, nz2);

tests/TestNonbondedForce.h

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008-2020 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2021 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -933,6 +933,50 @@ void testEwaldExceptions() {
     ASSERT_EQUAL_TOL(expectedChange, e2-e1, 1e-5);
 }
 
+void testDirectAndReciprocal() {
+    // Create a minimal system with direct space and reciprocal space in different force groups.
+
+    System system;
+    for (int i = 0; i < 4; i++)
+        system.addParticle(1.0);
+    system.setDefaultPeriodicBoxVectors(Vec3(2, 0, 0), Vec3(0, 2, 0), Vec3(0, 0, 2));
+    NonbondedForce* force = new NonbondedForce();
+    system.addForce(force);
+    force->setNonbondedMethod(NonbondedForce::PME);
+    force->setCutoffDistance(1.0);
+    force->setReciprocalSpaceForceGroup(1);
+    force->addParticle(1.0, 0.5, 1.0);
+    force->addParticle(1.0, 0.5, 1.0);
+    force->addParticle(-1.0, 0.5, 1.0);
+    force->addParticle(-1.0, 0.5, 1.0);
+    force->addException(0, 2, -2.0, 0.5, 3.0);
+    vector<Vec3> positions = {
+        Vec3(0, 0, 0),
+        Vec3(1.5, 0, 0),
+        Vec3(0, 0.5, 0.5),
+        Vec3(0.2, 1.3, 0)
+    };
+    VerletIntegrator integrator(0.001);
+    Context context(system, integrator, platform);
+    context.setPositions(positions);
+    
+    // Compute the direct and reciprocal space energies together and separately.
+    
+    double e1 = context.getState(State::Energy).getPotentialEnergy();
+    double e2 = context.getState(State::Energy, true, 1<<0).getPotentialEnergy();
+    double e3 = context.getState(State::Energy, true, 1<<1).getPotentialEnergy();
+    ASSERT_EQUAL_TOL(e1, e2+e3, 1e-5);
+    ASSERT(e2 != 0);
+    ASSERT(e3 != 0);
+    
+    // Completely disable the direct space calculation.
+    
+    force->setIncludeDirectSpace(false);
+    context.reinitialize(true);
+    double e4 = context.getState(State::Energy).getPotentialEnergy();
+    ASSERT_EQUAL_TOL(e3, e4, 1e-5);
+}
+
 void runPlatformTests();
 
 int main(int argc, char* argv[]) {
@@ -954,6 +998,7 @@ int main(int argc, char* argv[]) {
         testTwoForces();
         testParameterOffsets();
         testEwaldExceptions();
+        testDirectAndReciprocal();
         runPlatformTests();
     }
     catch(const exception& e) {