CUDA implementation of GK with extrapolated polarization

plugins/amoeba/platforms/cuda/src/AmoebaCudaKernels.cpp

@@ -1099,6 +1099,8 @@ void CudaCalcAmoebaMultipoleForceKernel::initialize(const System& system, const
     bool useShuffle = (cu.getComputeCapability() >= 3.0 && !cu.getUseDoublePrecision());
     double fixedThreadMemory = 19*elementSize+2*sizeof(float)+3*sizeof(int)/(double) cu.TileSize;
     double inducedThreadMemory = 15*elementSize+2*sizeof(float);
+    if (polarizationType == AmoebaMultipoleForce::Extrapolated)
+        inducedThreadMemory += 12*elementSize;
     double electrostaticsThreadMemory = 0;
     if (!useShuffle)
         fixedThreadMemory += 3*elementSize;
@@ -1175,6 +1177,7 @@ void CudaCalcAmoebaMultipoleForceKernel::initialize(const System& system, const
             prevDipolesGkPolar = new CudaArray(cu, 3*numMultipoles*MaxPrevDIISDipoles, elementSize, "prevDipolesGkPolar");
         }
         else if (polarizationType == AmoebaMultipoleForce::Extrapolated) {
+            inducedThreadMemory += 12*elementSize;
             int numOrders = force.getExtrapolationCoefficients().size();
             extrapolatedDipoleGk = new CudaArray(cu, 3*numMultipoles*numOrders, elementSize, "extrapolatedDipoleGk");
             extrapolatedDipoleGkPolar = new CudaArray(cu, 3*numMultipoles*numOrders, elementSize, "extrapolatedDipoleGkPolar");
@@ -1671,9 +1674,18 @@ double CudaCalcAmoebaMultipoleForceKernel::execute(ContextImpl& context, bool in
     // If using extrapolated polarization, add in force contributions from µ(m) T µ(n).
     
     if (polarizationType == AmoebaMultipoleForce::Extrapolated) {
-        void* extrapolatedArgs[] = {&cu.getForce().getDevicePointer(), &extrapolatedDipole->getDevicePointer(),
-            &extrapolatedDipolePolar->getDevicePointer(), &extrapolatedDipoleFieldGradient->getDevicePointer(), &extrapolatedDipoleFieldGradientPolar->getDevicePointer()};
-        cu.executeKernel(addExtrapolatedGradientKernel, extrapolatedArgs, numMultipoles);
+        if (gkKernel == NULL) {
+            void* extrapolatedArgs[] = {&cu.getForce().getDevicePointer(), &extrapolatedDipole->getDevicePointer(),
+                &extrapolatedDipolePolar->getDevicePointer(), &extrapolatedDipoleFieldGradient->getDevicePointer(), &extrapolatedDipoleFieldGradientPolar->getDevicePointer()};
+            cu.executeKernel(addExtrapolatedGradientKernel, extrapolatedArgs, numMultipoles);
+        }
+        else {
+            void* extrapolatedArgs[] = {&cu.getForce().getDevicePointer(), &extrapolatedDipole->getDevicePointer(),
+                &extrapolatedDipolePolar->getDevicePointer(), &extrapolatedDipoleFieldGradient->getDevicePointer(), &extrapolatedDipoleFieldGradientPolar->getDevicePointer(),
+                &extrapolatedDipoleGk->getDevicePointer(), &extrapolatedDipoleGkPolar->getDevicePointer(),
+                &extrapolatedDipoleFieldGradientGk->getDevicePointer(), &extrapolatedDipoleFieldGradientGkPolar->getDevicePointer()};
+            cu.executeKernel(addExtrapolatedGradientKernel, extrapolatedArgs, numMultipoles);
+        }
     }
 
     // Map torques to force.
@@ -1893,29 +1905,56 @@ bool CudaCalcAmoebaMultipoleForceKernel::iterateDipolesByDIIS(int iteration) {
 void CudaCalcAmoebaMultipoleForceKernel::computeExtrapolatedDipoles(void** recipBoxVectorPointer) {
     // Start by storing the direct dipoles as PT0
 
-    void* initArgs[] = {&inducedDipole->getDevicePointer(), &inducedDipolePolar->getDevicePointer(), &extrapolatedDipole->getDevicePointer(),
-        &extrapolatedDipolePolar->getDevicePointer(), &inducedDipoleFieldGradient->getDevicePointer(), &inducedDipoleFieldGradientPolar->getDevicePointer()};
-    cu.executeKernel(initExtrapolatedKernel, initArgs, extrapolatedDipole->getSize());
+    if (gkKernel == NULL) {
+        void* initArgs[] = {&inducedDipole->getDevicePointer(), &inducedDipolePolar->getDevicePointer(), &extrapolatedDipole->getDevicePointer(),
+            &extrapolatedDipolePolar->getDevicePointer(), &inducedDipoleFieldGradient->getDevicePointer(), &inducedDipoleFieldGradientPolar->getDevicePointer()};
+        cu.executeKernel(initExtrapolatedKernel, initArgs, extrapolatedDipole->getSize());
+    }
+    else {
+        void* initArgs[] = {&inducedDipole->getDevicePointer(), &inducedDipolePolar->getDevicePointer(), &extrapolatedDipole->getDevicePointer(),
+            &extrapolatedDipolePolar->getDevicePointer(), &inducedDipoleFieldGradient->getDevicePointer(), &inducedDipoleFieldGradientPolar->getDevicePointer(),
+            &gkKernel->getInducedDipoles()->getDevicePointer(), &gkKernel->getInducedDipolesPolar()->getDevicePointer(), &extrapolatedDipoleGk->getDevicePointer(),
+            &extrapolatedDipoleGkPolar->getDevicePointer(), &inducedDipoleFieldGradientGk->getDevicePointer(), &inducedDipoleFieldGradientGkPolar->getDevicePointer()};
+        cu.executeKernel(initExtrapolatedKernel, initArgs, extrapolatedDipole->getSize());
+    }
 
     // Recursively apply alpha.Tau to the µ_(n) components to generate µ_(n+1), and store the result
 
-    CudaNonbondedUtilities& nb = cu.getNonbondedUtilities();
     for (int order = 1; order < maxExtrapolationOrder; ++order) {
-        cu.clearBuffer(*inducedField);
-        cu.clearBuffer(*inducedFieldPolar);
         computeInducedField(recipBoxVectorPointer);
-        void* iterateArgs[] = {&order, &inducedDipole->getDevicePointer(), &inducedDipolePolar->getDevicePointer(), &extrapolatedDipole->getDevicePointer(),
-            &extrapolatedDipolePolar->getDevicePointer(), &inducedDipoleFieldGradient->getDevicePointer(), &inducedDipoleFieldGradientPolar->getDevicePointer(),
-            &inducedField->getDevicePointer(), &inducedFieldPolar->getDevicePointer(), &extrapolatedDipoleFieldGradient->getDevicePointer(), &extrapolatedDipoleFieldGradientPolar->getDevicePointer(),
-            &polarizability->getDevicePointer()};
-        cu.executeKernel(iterateExtrapolatedKernel, iterateArgs, extrapolatedDipole->getSize());
+        if (gkKernel == NULL) {
+            void* iterateArgs[] = {&order, &inducedDipole->getDevicePointer(), &inducedDipolePolar->getDevicePointer(), &extrapolatedDipole->getDevicePointer(),
+                &extrapolatedDipolePolar->getDevicePointer(), &inducedDipoleFieldGradient->getDevicePointer(), &inducedDipoleFieldGradientPolar->getDevicePointer(),
+                &inducedField->getDevicePointer(), &inducedFieldPolar->getDevicePointer(), &extrapolatedDipoleFieldGradient->getDevicePointer(), &extrapolatedDipoleFieldGradientPolar->getDevicePointer(),
+                &polarizability->getDevicePointer()};
+            cu.executeKernel(iterateExtrapolatedKernel, iterateArgs, extrapolatedDipole->getSize());
+        }
+        else {
+            void* iterateArgs[] = {&order, &inducedDipole->getDevicePointer(), &inducedDipolePolar->getDevicePointer(), &extrapolatedDipole->getDevicePointer(),
+                &extrapolatedDipolePolar->getDevicePointer(), &inducedDipoleFieldGradient->getDevicePointer(), &inducedDipoleFieldGradientPolar->getDevicePointer(),
+                &inducedField->getDevicePointer(), &inducedFieldPolar->getDevicePointer(), &extrapolatedDipoleFieldGradient->getDevicePointer(), &extrapolatedDipoleFieldGradientPolar->getDevicePointer(),
+                &gkKernel->getInducedDipoles()->getDevicePointer(), &gkKernel->getInducedDipolesPolar()->getDevicePointer(), &extrapolatedDipoleGk->getDevicePointer(),
+                &extrapolatedDipoleGkPolar->getDevicePointer(), &inducedDipoleFieldGradientGk->getDevicePointer(), &inducedDipoleFieldGradientGkPolar->getDevicePointer(),
+                &gkKernel->getInducedField()->getDevicePointer(), &gkKernel->getInducedFieldPolar()->getDevicePointer(),
+                &extrapolatedDipoleFieldGradientGk->getDevicePointer(), &extrapolatedDipoleFieldGradientGkPolar->getDevicePointer(),
+                &polarizability->getDevicePointer()};
+            cu.executeKernel(iterateExtrapolatedKernel, iterateArgs, extrapolatedDipole->getSize());
+        }
     }
     
     // Take a linear combination of the µ_(n) components to form the total dipole
 
-    void* computeArgs[] = {&inducedDipole->getDevicePointer(), &inducedDipolePolar->getDevicePointer(), &extrapolatedDipole->getDevicePointer(),
-            &extrapolatedDipolePolar->getDevicePointer()};
-    cu.executeKernel(computeExtrapolatedKernel, computeArgs, extrapolatedDipole->getSize());
+    if (gkKernel == NULL) {
+        void* computeArgs[] = {&inducedDipole->getDevicePointer(), &inducedDipolePolar->getDevicePointer(), &extrapolatedDipole->getDevicePointer(),
+                &extrapolatedDipolePolar->getDevicePointer()};
+        cu.executeKernel(computeExtrapolatedKernel, computeArgs, extrapolatedDipole->getSize());
+    }
+    else {
+        void* computeArgs[] = {&inducedDipole->getDevicePointer(), &inducedDipolePolar->getDevicePointer(), &extrapolatedDipole->getDevicePointer(),
+                &extrapolatedDipolePolar->getDevicePointer(), &gkKernel->getInducedDipoles()->getDevicePointer(), &gkKernel->getInducedDipolesPolar()->getDevicePointer(),
+                &extrapolatedDipoleGk->getDevicePointer(), &extrapolatedDipoleGkPolar->getDevicePointer()};
+        cu.executeKernel(computeExtrapolatedKernel, computeArgs, extrapolatedDipole->getSize());
+    }
     computeInducedField(recipBoxVectorPointer);
 }
 
@@ -2263,7 +2302,8 @@ void CudaCalcAmoebaGeneralizedKirkwoodForceKernel::initialize(const System& syst
     bornForce = CudaArray::create<long long>(cu, paddedNumAtoms, "bornForce");
     inducedDipoleS = new CudaArray(cu, 3*paddedNumAtoms, elementSize, "inducedDipoleS");
     inducedDipolePolarS = new CudaArray(cu, 3*paddedNumAtoms, elementSize, "inducedDipolePolarS");
-    if (multipoles->getPolarizationType() == AmoebaMultipoleForce::Mutual) {
+    polarizationType = multipoles->getPolarizationType();
+    if (polarizationType != AmoebaMultipoleForce::Direct) {
         inducedField = new CudaArray(cu, 3*paddedNumAtoms, sizeof(long long), "gkInducedField");
         inducedFieldPolar = new CudaArray(cu, 3*paddedNumAtoms, sizeof(long long), "gkInducedFieldPolar");
     }
@@ -2316,8 +2356,12 @@ void CudaCalcAmoebaGeneralizedKirkwoodForceKernel::initialize(const System& syst
     defines["EPSILON_FACTOR"] = cu.doubleToString(138.9354558456);
     defines["M_PI"] = cu.doubleToString(M_PI);
     defines["ENERGY_SCALE_FACTOR"] = cu.doubleToString(138.9354558456/force.getSoluteDielectric());
-    if (multipoles->getPolarizationType() == AmoebaMultipoleForce::Direct)
+    if (polarizationType == AmoebaMultipoleForce::Direct)
         defines["DIRECT_POLARIZATION"] = "";
+    else if (polarizationType == AmoebaMultipoleForce::Mutual)
+        defines["MUTUAL_POLARIZATION"] = "";
+    else if (polarizationType == AmoebaMultipoleForce::Extrapolated)
+        defines["EXTRAPOLATED_POLARIZATION"] = "";
     includeSurfaceArea = force.getIncludeCavityTerm();
     if (includeSurfaceArea) {
         defines["SURFACE_AREA_FACTOR"] = cu.doubleToString(force.getSurfaceAreaFactor());

plugins/amoeba/platforms/cuda/src/AmoebaCudaKernels.h

@@ -528,6 +528,7 @@ private:
     const System& system;
     bool includeSurfaceArea, hasInitializedKernels;
     int computeBornSumThreads, gkForceThreads, chainRuleThreads, ediffThreads;
+    AmoebaMultipoleForce::PolarizationType polarizationType;
     std::map<std::string, std::string> defines;
     CudaArray* params;
     CudaArray* bornSum;

plugins/amoeba/platforms/cuda/src/kernels/gkEDiffPairForce.cu

@@ -365,7 +365,22 @@ __device__ void computeOneEDiffInteractionT3(AtomData4& atom1, volatile AtomData
 
     // correction to convert mutual to direct polarization force
 
-#ifdef DIRECT_POLARIZATION
+#ifdef MUTUAL_POLARIZATION
+    real findmp1 = uscale*(scip2*ddsc3_1 - ddsc5_1*(sci3*scip4+scip3*sci4));
+    real findmp2 = uscale*(scip2*ddsc3_2 - ddsc5_2*(sci3*scip4+scip3*sci4));
+    real findmp3 = uscale*(scip2*ddsc3_3 - ddsc5_3*(sci3*scip4+scip3*sci4));
+    ftm2i1 -= 0.5f*findmp1;
+    ftm2i2 -= 0.5f*findmp2;
+    ftm2i3 -= 0.5f*findmp3;
+
+    real sci3X = sci3 - sci3Y;
+    real sci4X = sci4 - sci4Y;
+    real scip3X = scip3 - scip3Y;
+    real scip4X = scip4 - scip4Y;
+    ftm2i1 += 0.5f*uscale*(-ddsc5_1*(sci3X*scip4X+scip3X*sci4X));
+    ftm2i2 += 0.5f*uscale*(-ddsc5_2*(sci3X*scip4X+scip3X*sci4X));
+    ftm2i3 += 0.5f*uscale*(-ddsc5_3*(sci3X*scip4X+scip3X*sci4X)); 
+#else
     real gfd = (scip2*scale3i - 5*rr2*(scip3*sci4+sci3*scip4)*scale5i);
     real fdir1 = gfd*xr + scale5i* (sci4*atom1.inducedDipolePolarS.x+scip4*atom1.inducedDipoleS.x + sci3*atom2.inducedDipolePolarS.x+scip3*atom2.inducedDipoleS.x);
     real fdir2 = gfd*yr + scale5i* (sci4*atom1.inducedDipolePolarS.y+scip4*atom1.inducedDipoleS.y + sci3*atom2.inducedDipolePolarS.y+scip3*atom2.inducedDipoleS.y);
@@ -385,21 +400,6 @@ __device__ void computeOneEDiffInteractionT3(AtomData4& atom1, volatile AtomData
     ftm2i1 += 0.5f*fdir1;
     ftm2i2 += 0.5f*fdir2;
     ftm2i3 += 0.5f*fdir3;
-#else
-    real findmp1 = uscale*(scip2*ddsc3_1 - ddsc5_1*(sci3*scip4+scip3*sci4));
-    real findmp2 = uscale*(scip2*ddsc3_2 - ddsc5_2*(sci3*scip4+scip3*sci4));
-    real findmp3 = uscale*(scip2*ddsc3_3 - ddsc5_3*(sci3*scip4+scip3*sci4));
-    ftm2i1 -= 0.5f*findmp1;
-    ftm2i2 -= 0.5f*findmp2;
-    ftm2i3 -= 0.5f*findmp3;
-
-    real sci3X = sci3 - sci3Y;
-    real sci4X = sci4 - sci4Y;
-    real scip3X = scip3 - scip3Y;
-    real scip4X = scip4 - scip4Y;
-    ftm2i1 += 0.5f*uscale*(-ddsc5_1*(sci3X*scip4X+scip3X*sci4X));
-    ftm2i2 += 0.5f*uscale*(-ddsc5_2*(sci3X*scip4X+scip3X*sci4X));
-    ftm2i3 += 0.5f*uscale*(-ddsc5_3*(sci3X*scip4X+scip3X*sci4X)); 
 #endif
 #endif
 

plugins/amoeba/platforms/cuda/src/kernels/multipoleInducedField.cu

@@ -721,7 +721,7 @@ extern "C" __global__ void initExtrapolatedDipoles(real* __restrict__ inducedDip
         real* __restrict__ inducedDipoleFieldGradientGk, real* __restrict__ inducedDipoleFieldGradientGkPolar
 #endif
         ) {
-    for (int index = blockIdx.x*blockDim.x+threadIdx.x; index < 3*NUM_ATOMS*MAX_EXTRAPOLATION_ORDER; index += blockDim.x*gridDim.x) {
+    for (int index = blockIdx.x*blockDim.x+threadIdx.x; index < 3*NUM_ATOMS; index += blockDim.x*gridDim.x) {
         extrapolatedDipole[index] = inducedDipole[index];
         extrapolatedDipolePolar[index] = inducedDipolePolar[index];
 #ifdef USE_GK
@@ -729,7 +729,7 @@ extern "C" __global__ void initExtrapolatedDipoles(real* __restrict__ inducedDip
         extrapolatedDipoleGkPolar[index] = inducedDipoleGkPolar[index];
 #endif
     }
-    for (int index = blockIdx.x*blockDim.x+threadIdx.x; index < 6*NUM_ATOMS*MAX_EXTRAPOLATION_ORDER; index += blockDim.x*gridDim.x) {
+    for (int index = blockIdx.x*blockDim.x+threadIdx.x; index < 6*NUM_ATOMS; index += blockDim.x*gridDim.x) {
         inducedDipoleFieldGradient[index] = 0;
         inducedDipoleFieldGradientPolar[index] = 0;
 #ifdef USE_GK
@@ -820,19 +820,35 @@ extern "C" __global__ void addExtrapolatedFieldGradientToForce(long long* __rest
             int index1 = 3*(l*NUM_ATOMS+atom);
             real dipole[] = {extrapolatedDipole[index1], extrapolatedDipole[index1+1], extrapolatedDipole[index1+2]};
             real dipolePolar[] = {extrapolatedDipolePolar[index1], extrapolatedDipolePolar[index1+1], extrapolatedDipolePolar[index1+2]};
+#ifdef USE_GK
+            real dipoleGk[] = {extrapolatedDipoleGk[index1], extrapolatedDipoleGk[index1+1], extrapolatedDipoleGk[index1+2]};
+            real dipoleGkPolar[] = {extrapolatedDipoleGkPolar[index1], extrapolatedDipoleGkPolar[index1+1], extrapolatedDipoleGkPolar[index1+2]};
+#endif
             for (int m = 0; m < MAX_EXTRAPOLATION_ORDER-1-l; m++) {
                 int index2 = 6*(m*NUM_ATOMS+atom);
+                real scale = 0.5f*coeff[l+m+1]*ENERGY_SCALE_FACTOR;
                 real gradient[] = {extrapolatedDipoleFieldGradient[index2], extrapolatedDipoleFieldGradient[index2+1], extrapolatedDipoleFieldGradient[index2+2],
                                    extrapolatedDipoleFieldGradient[index2+3], extrapolatedDipoleFieldGradient[index2+4], extrapolatedDipoleFieldGradient[index2+5]};
                 real gradientPolar[] = {extrapolatedDipoleFieldGradientPolar[index2], extrapolatedDipoleFieldGradientPolar[index2+1], extrapolatedDipoleFieldGradientPolar[index2+2],
                                         extrapolatedDipoleFieldGradientPolar[index2+3], extrapolatedDipoleFieldGradientPolar[index2+4], extrapolatedDipoleFieldGradientPolar[index2+5]};
-                real scale = 0.5f*coeff[l+m+1]*ENERGY_SCALE_FACTOR;
                 fx += scale*(dipole[0]*gradientPolar[0] + dipole[1]*gradientPolar[3] + dipole[2]*gradientPolar[4]);
                 fy += scale*(dipole[0]*gradientPolar[3] + dipole[1]*gradientPolar[1] + dipole[2]*gradientPolar[5]);
                 fz += scale*(dipole[0]*gradientPolar[4] + dipole[1]*gradientPolar[5] + dipole[2]*gradientPolar[2]);
                 fx += scale*(dipolePolar[0]*gradient[0] + dipolePolar[1]*gradient[3] + dipolePolar[2]*gradient[4]);
                 fy += scale*(dipolePolar[0]*gradient[3] + dipolePolar[1]*gradient[1] + dipolePolar[2]*gradient[5]);
                 fz += scale*(dipolePolar[0]*gradient[4] + dipolePolar[1]*gradient[5] + dipolePolar[2]*gradient[2]);
+#ifdef USE_GK
+                real gradientGk[] = {extrapolatedDipoleFieldGradient[index2], extrapolatedDipoleFieldGradient[index2+1], extrapolatedDipoleFieldGradient[index2+2],
+                                   extrapolatedDipoleFieldGradient[index2+3], extrapolatedDipoleFieldGradient[index2+4], extrapolatedDipoleFieldGradient[index2+5]};
+                real gradientGkPolar[] = {extrapolatedDipoleFieldGradientPolar[index2], extrapolatedDipoleFieldGradientPolar[index2+1], extrapolatedDipoleFieldGradientPolar[index2+2],
+                                        extrapolatedDipoleFieldGradientPolar[index2+3], extrapolatedDipoleFieldGradientPolar[index2+4], extrapolatedDipoleFieldGradientPolar[index2+5]};
+                fx += scale*(dipoleGk[0]*gradientGkPolar[0] + dipoleGk[1]*gradientGkPolar[3] + dipoleGk[2]*gradientGkPolar[4]);
+                fy += scale*(dipoleGk[0]*gradientGkPolar[3] + dipoleGk[1]*gradientGkPolar[1] + dipoleGk[2]*gradientGkPolar[5]);
+                fz += scale*(dipoleGk[0]*gradientGkPolar[4] + dipoleGk[1]*gradientGkPolar[5] + dipoleGk[2]*gradientGkPolar[2]);
+                fx += scale*(dipoleGkPolar[0]*gradientGk[0] + dipoleGkPolar[1]*gradientGk[3] + dipoleGkPolar[2]*gradientGk[4]);
+                fy += scale*(dipoleGkPolar[0]*gradientGk[3] + dipoleGkPolar[1]*gradientGk[1] + dipoleGkPolar[2]*gradientGk[5]);
+                fz += scale*(dipoleGkPolar[0]*gradientGk[4] + dipoleGkPolar[1]*gradientGk[5] + dipoleGkPolar[2]*gradientGk[2]);
+#endif
             }
         }
         forceBuffers[atom] += (long long) (fx*0x100000000);

plugins/amoeba/platforms/cuda/tests/TestCudaAmoebaGeneralizedKirkwoodForce.cpp

@@ -50,10 +50,36 @@
 using namespace OpenMM;
+using namespace std;
+
 const double TOL = 1e-4;
 extern "C" void registerAmoebaCudaKernelFactories();
+static void checkFiniteDifferences(vector<Vec3> forces, Context &context, vector<Vec3> positions)
+{
+    // Take a small step in the direction of the energy gradient and see whether the potential energy changes by the expected amount.
+
+    double norm = 0.0;
+    for (int i = 0; i < (int) forces.size(); ++i)
+        norm += forces[i].dot(forces[i]);
+    norm = std::sqrt(norm);
+    const double stepSize = 1e-3;
+    double step = 0.5*stepSize/norm;
+    vector<Vec3> positions2(forces.size()), positions3(forces.size());
+    for (int i = 0; i < (int) positions.size(); ++i) {
+        Vec3 p = positions[i];
+        Vec3 f = forces[i];
+        positions2[i] = Vec3(p[0]-f[0]*step, p[1]-f[1]*step, p[2]-f[2]*step);
+        positions3[i] = Vec3(p[0]+f[0]*step, p[1]+f[1]*step, p[2]+f[2]*step);
+    }
+    context.setPositions(positions2);
+    State state2 = context.getState(State::Energy);
+    context.setPositions(positions3);
+    State state3 = context.getState(State::Energy);
+    ASSERT_EQUAL_TOL(state3.getPotentialEnergy()+norm*stepSize, state2.getPotentialEnergy(), 1e-4);
+}
+
 // setup for 2 ammonia molecules
 static void setupMultipoleAmmonia(System& system, AmoebaGeneralizedKirkwoodForce* amoebaGeneralizedKirkwoodForce,
@@ -289,6 +315,10 @@ static void getForcesEnergyMultipoleAmmonia(Context& context, std::vector<Vec3>&
     State state                      = context.getState(State::Forces | State::Energy);
     forces                           = state.getForces();
     energy                           = state.getPotentialEnergy();
+    
+    // Check that the forces and energy are consistent.
+    
+    checkFiniteDifferences(forces, context, positions);
 }
 // setup for villin
@@ -6972,6 +7002,10 @@ static void setupAndGetForcesEnergyMultipoleVillin(AmoebaMultipoleForce::Polariz
     State state                      = context.getState(State::Forces | State::Energy);
     forces                           = state.getForces();
     energy                           = state.getPotentialEnergy();
+    
+    // Check that the forces and energy are consistent.
+    
+    checkFiniteDifferences(forces, context, positions);
 }
 // compare forces and energies 
@@ -7049,6 +7083,25 @@ static void testGeneralizedKirkwoodAmmoniaDirectPolarization() {
     compareForcesEnergy(testName, expectedEnergy, energy, expectedForces, forces, tolerance);
 }
+static void testGeneralizedKirkwoodAmmoniaExtrapolatedPolarization() {
+
+    std::string testName      = "testGeneralizedKirkwoodAmmoniaExtrapolatedPolarization";
+
+    int numberOfParticles     = 8;
+    std::vector<Vec3> forces;
+    double energy;
+
+    System system;
+    AmoebaGeneralizedKirkwoodForce* amoebaGeneralizedKirkwoodForce  = new AmoebaGeneralizedKirkwoodForce();
+    setupMultipoleAmmonia(system, amoebaGeneralizedKirkwoodForce, AmoebaMultipoleForce::Direct, 0);
+    LangevinIntegrator integrator(0.0, 0.1, 0.01);
+    Context context(system, integrator, Platform::getPlatformByName("CUDA"));
+
+    // We don't have reference values for this case, but at least check that force and energy are consistent.
+
+    getForcesEnergyMultipoleAmmonia(context, forces, energy);
+}
+
 // test GK mutual polarization for system comprised of two ammonia molecules
 static void testGeneralizedKirkwoodAmmoniaMutualPolarization() {
@@ -7765,6 +7818,19 @@ static void testGeneralizedKirkwoodVillinDirectPolarization() {
     compareForceNormsEnergy(testName, expectedEnergy, energy, expectedForces, forces, tolerance);
 }
+static void testGeneralizedKirkwoodVillinExtrapolatedPolarization() {
+
+    std::string testName      = "testGeneralizedKirkwoodVillinExtrapolatedPolarization";
+
+    int numberOfParticles     = 596;
+    std::vector<Vec3> forces;
+    double energy;
+
+    // We don't have reference values for this case, but at least check that force and energy are consistent.
+    
+    setupAndGetForcesEnergyMultipoleVillin(AmoebaMultipoleForce::Extrapolated, 0, forces, energy);
+}
+
 // test GK mutual polarization for villin system
 static void testGeneralizedKirkwoodVillinMutualPolarization() {
@@ -8401,9 +8467,10 @@ int main(int argc, char* argv[]) {
         testGeneralizedKirkwoodAmmoniaDirectPolarization();
         testGeneralizedKirkwoodAmmoniaMutualPolarization();
+        testGeneralizedKirkwoodAmmoniaExtrapolatedPolarization();
         testGeneralizedKirkwoodAmmoniaMutualPolarizationWithCavityTerm();
-
         testGeneralizedKirkwoodVillinDirectPolarization();
+        testGeneralizedKirkwoodVillinExtrapolatedPolarization();
         testGeneralizedKirkwoodVillinMutualPolarization();
     } catch(const std::exception& e) {