Merge pull request #1120 from peastman/pmeparams

Created getPMEParametersInContext()

olla/include/openmm/kernels.h

@@ -581,6 +581,15 @@ public:
      * @param force      the NonbondedForce to copy the parameters from
      */
     virtual void copyParametersToContext(ContextImpl& context, const NonbondedForce& force) = 0;
+    /**
+     * Get the parameters being used for PME.
+     * 
+     * @param alpha   the separation parameter
+     * @param nx      the number of grid points along the X axis
+     * @param ny      the number of grid points along the Y axis
+     * @param nz      the number of grid points along the Z axis
+     */
+    virtual void getPMEParameters(double& alpha, int& nx, int& ny, int& nz) const = 0;
 };
 
 /**
@@ -1281,6 +1290,15 @@ public:
      * @return the potential energy due to the PME reciprocal space interactions
      */
     virtual double finishComputation(IO& io) = 0;
+    /**
+     * Get the parameters being used for PME.
+     * 
+     * @param alpha   the separation parameter
+     * @param nx      the number of grid points along the X axis
+     * @param ny      the number of grid points along the Y axis
+     * @param nz      the number of grid points along the Z axis
+     */
+    virtual void getPMEParameters(double& alpha, int& nx, int& ny, int& nz) const = 0;
 };
 
 /**

openmmapi/include/openmm/Context.h

@@ -261,6 +261,7 @@ private:
     friend class Force;
     friend class Platform;
     ContextImpl& getImpl();
+    const ContextImpl& getImpl() const;
     ContextImpl* impl;
     std::map<std::string, std::string> properties;
 };

openmmapi/include/openmm/Force.h

@@ -98,6 +98,10 @@ protected:
      * Get the ForceImpl corresponding to this Force in a Context.
      */
     ForceImpl& getImplInContext(Context& context);
+    /**
+     * Get a const reference to the ForceImpl corresponding to this Force in a Context.
+     */
+    const ForceImpl& getImplInContext(const Context& context) const;
     /**
      * Get the ContextImpl corresponding to a Context.
      */

openmmapi/include/openmm/NonbondedForce.h

@@ -217,6 +217,19 @@ public:
      * @param nz      the number of grid points along the Z axis
      */
     void setPMEParameters(double alpha, int nx, int ny, int nz);
+    /**
+     * Get the parameters being used for PME in a particular Context.  Because some platforms have restrictions
+     * on the allowed grid sizes, the values that are actually used may be slightly different from those
+     * specified with setPMEParameters(), or the standard values calculated based on the Ewald error tolerance.
+     * See the manual for details.
+     * 
+     * @param context the Context for which to get the parameters
+     * @param alpha   the separation parameter
+     * @param nx      the number of grid points along the X axis
+     * @param ny      the number of grid points along the Y axis
+     * @param nz      the number of grid points along the Z axis
+     */
+    void getPMEParametersInContext(const Context& context, double& alpha, int& nx, int& ny, int& nz) const;
     /**
      * Add the nonbonded force parameters for a particle.  This should be called once for each particle
      * in the System.  When it is called for the i'th time, it specifies the parameters for the i'th particle.

openmmapi/include/openmm/internal/NonbondedForceImpl.h

@@ -64,6 +64,7 @@ public:
     }
     std::vector<std::string> getKernelNames();
     void updateParametersInContext(ContextImpl& context);
+    void getPMEParameters(double& alpha, int& nx, int& ny, int& nz) const;
     /**
      * This is a utility routine that calculates the values to use for alpha and kmax when using
      * Ewald summation.

openmmapi/src/Context.cpp

@@ -252,6 +252,10 @@ ContextImpl& Context::getImpl() {
     return *impl;
 }
 
+const ContextImpl& Context::getImpl() const {
+    return *impl;
+}
+
 const vector<vector<int> >& Context::getMolecules() const {
     return impl->getMolecules();
 }

openmmapi/src/Force.cpp

@@ -61,6 +61,14 @@ ForceImpl& Force::getImplInContext(Context& context) {
     throw OpenMMException("getImplInContext: This Force is not present in the Context");
 }
 
+const ForceImpl& Force::getImplInContext(const Context& context) const {
+    const vector<ForceImpl*>& impls = context.getImpl().getForceImpls();
+    for (int i = 0; i < (int) impls.size(); i++)
+        if (&impls[i]->getOwner() == this)
+            return *impls[i];
+    throw OpenMMException("getImplInContext: This Force is not present in the Context");
+}
+
 ContextImpl& Force::getContextImpl(Context& context) {
     return context.getImpl();
 }

openmmapi/src/NonbondedForce.cpp

@@ -113,6 +113,10 @@ void NonbondedForce::setPMEParameters(double alpha, int nx, int ny, int nz) {
     this->nz = nz;
 }
 
+void NonbondedForce::getPMEParametersInContext(const Context& context, double& alpha, int& nx, int& ny, int& nz) const {
+    dynamic_cast<const NonbondedForceImpl&>(getImplInContext(context)).getPMEParameters(alpha, nx, ny, nz);
+}
+
 int NonbondedForce::addParticle(double charge, double sigma, double epsilon) {
     particles.push_back(ParticleInfo(charge, sigma, epsilon));
     return particles.size()-1;

openmmapi/src/NonbondedForceImpl.cpp

@@ -278,3 +278,7 @@ double NonbondedForceImpl::calcDispersionCorrection(const System& system, const
 void NonbondedForceImpl::updateParametersInContext(ContextImpl& context) {
     kernel.getAs<CalcNonbondedForceKernel>().copyParametersToContext(context, owner);
 }
+
+void NonbondedForceImpl::getPMEParameters(double& alpha, int& nx, int& ny, int& nz) const {
+    kernel.getAs<CalcNonbondedForceKernel>().getPMEParameters(alpha, nx, ny, nz);
+}

platforms/cpu/include/CpuKernels.h

@@ -204,6 +204,15 @@ public:
      * @param force      the NonbondedForce to copy the parameters from
      */
     void copyParametersToContext(ContextImpl& context, const NonbondedForce& force);
+    /**
+     * Get the parameters being used for PME.
+     * 
+     * @param alpha   the separation parameter
+     * @param nx      the number of grid points along the X axis
+     * @param ny      the number of grid points along the Y axis
+     * @param nz      the number of grid points along the Z axis
+     */
+    void getPMEParameters(double& alpha, int& nx, int& ny, int& nz) const;
 private:
     class PmeIO;
     CpuPlatform::PlatformData& data;

platforms/cpu/src/CpuKernels.cpp

@@ -654,6 +654,19 @@ void CpuCalcNonbondedForceKernel::copyParametersToContext(ContextImpl& context,
         dispersionCoefficient = NonbondedForceImpl::calcDispersionCorrection(context.getSystem(), force);
 }
 
+void CpuCalcNonbondedForceKernel::getPMEParameters(double& alpha, int& nx, int& ny, int& nz) const {
+    if (nonbondedMethod != PME)
+        throw OpenMMException("getPMEParametersInContext: This Context is not using PME");
+    if (useOptimizedPme)
+        optimizedPme.getAs<const CalcPmeReciprocalForceKernel>().getPMEParameters(alpha, nx, ny, nz);
+    else {
+        alpha = ewaldAlpha;
+        nx = gridSize[0];
+        ny = gridSize[1];
+        nz = gridSize[2];
+    }
+}
+
 CpuCalcCustomNonbondedForceKernel::CpuCalcCustomNonbondedForceKernel(string name, const Platform& platform, CpuPlatform::PlatformData& data) :
             CalcCustomNonbondedForceKernel(name, platform), data(data), forceCopy(NULL), neighborList(NULL), nonbonded(NULL) {
 }

platforms/cpu/tests/TestCpuEwald.cpp

@@ -42,6 +42,7 @@
 #include "openmm/LangevinIntegrator.h"
 #include "openmm/VerletIntegrator.h"
 #include "openmm/internal/ContextImpl.h"
+#include "openmm/internal/NonbondedForceImpl.h"
 #include "SimTKOpenMMRealType.h"
 #include "sfmt/SFMT.h"
 #include <iostream>
@@ -299,6 +300,19 @@ void testErrorTolerance(NonbondedForce::NonbondedMethod method) {
                 diff = sqrt(diff)/norm;
                 ASSERT(diff < 2*tol);
             }
+            if (method == NonbondedForce::PME) {
+                // See if the PME parameters were calculated correctly.
+
+                double expectedAlpha, actualAlpha;
+                int expectedSize[3], actualSize[3];
+                NonbondedForceImpl::calcPMEParameters(system, *force, expectedAlpha, expectedSize[0], expectedSize[1], expectedSize[2]);
+                force->getPMEParametersInContext(context, actualAlpha, actualSize[0], actualSize[1], actualSize[2]);
+                ASSERT_EQUAL_TOL(expectedAlpha, actualAlpha, 1e-5);
+                for (int i = 0; i < 3; i++) {
+                    ASSERT(actualSize[i] >= expectedSize[i]);
+                    ASSERT(actualSize[i] < expectedSize[i]+10);
+                }
+            }
         }
     }
 }

platforms/cuda/include/CudaKernels.h

@@ -620,6 +620,15 @@ public:
      * @param force      the NonbondedForce to copy the parameters from
      */
     void copyParametersToContext(ContextImpl& context, const NonbondedForce& force);
+    /**
+     * Get the parameters being used for PME.
+     * 
+     * @param alpha   the separation parameter
+     * @param nx      the number of grid points along the X axis
+     * @param ny      the number of grid points along the Y axis
+     * @param nz      the number of grid points along the Z axis
+     */
+    void getPMEParameters(double& alpha, int& nx, int& ny, int& nz) const;
 private:
     class SortTrait : public CudaSort::SortTrait {
         int getDataSize() const {return 8;}
@@ -668,7 +677,9 @@ private:
     std::vector<std::pair<int, int> > exceptionAtoms;
     double ewaldSelfEnergy, dispersionCoefficient, alpha;
     int interpolateForceThreads;
+    int gridSizeX, gridSizeY, gridSizeZ;
     bool hasCoulomb, hasLJ, usePmeStream, useCudaFFT;
+    NonbondedMethod nonbondedMethod;
     static const int PmeOrder = 5;
 };
 

platforms/cuda/include/CudaParallelKernels.h

@@ -430,6 +430,15 @@ public:
      * @param force      the NonbondedForce to copy the parameters from
      */
     void copyParametersToContext(ContextImpl& context, const NonbondedForce& force);
+    /**
+     * Get the parameters being used for PME.
+     * 
+     * @param alpha   the separation parameter
+     * @param nx      the number of grid points along the X axis
+     * @param ny      the number of grid points along the Y axis
+     * @param nz      the number of grid points along the Z axis
+     */
+    void getPMEParameters(double& alpha, int& nx, int& ny, int& nz) const;
 private:
     class Task;
     CudaPlatform::PlatformData& data;

platforms/cuda/src/CudaKernels.cpp

@@ -1561,8 +1561,9 @@ void CudaCalcNonbondedForceKernel::initialize(const System& system, const Nonbon
     }
     posq.upload(&temp[0]);
     sigmaEpsilon->upload(sigmaEpsilonVector);
-    bool useCutoff = (force.getNonbondedMethod() != NonbondedForce::NoCutoff);
-    bool usePeriodic = (force.getNonbondedMethod() != NonbondedForce::NoCutoff && force.getNonbondedMethod() != NonbondedForce::CutoffNonPeriodic);
+    nonbondedMethod = CalcNonbondedForceKernel::NonbondedMethod(force.getNonbondedMethod());
+    bool useCutoff = (nonbondedMethod != NoCutoff);
+    bool usePeriodic = (nonbondedMethod != NoCutoff && nonbondedMethod != CutoffNonPeriodic);
     map<string, string> defines;
     defines["HAS_COULOMB"] = (hasCoulomb ? "1" : "0");
     defines["HAS_LENNARD_JONES"] = (hasLJ ? "1" : "0");
@@ -1590,7 +1591,7 @@ void CudaCalcNonbondedForceKernel::initialize(const System& system, const Nonbon
         dispersionCoefficient = 0.0;
     alpha = 0;
     ewaldSelfEnergy = 0.0;
-    if (force.getNonbondedMethod() == NonbondedForce::Ewald) {
+    if (nonbondedMethod == Ewald) {
         // Compute the Ewald parameters.
 
         int kmaxx, kmaxy, kmaxz;
@@ -1619,11 +1620,9 @@ void CudaCalcNonbondedForceKernel::initialize(const System& system, const Nonbon
             cosSinSums = new CudaArray(cu, (2*kmaxx-1)*(2*kmaxy-1)*(2*kmaxz-1), elementSize, "cosSinSums");
         }
     }
-    else if (force.getNonbondedMethod() == NonbondedForce::PME) {
+    else if (nonbondedMethod == PME) {
         // Compute the PME parameters.
 
-        int gridSizeX, gridSizeY, gridSizeZ;
-
         NonbondedForceImpl::calcPMEParameters(system, force, alpha, gridSizeX, gridSizeY, gridSizeZ);
         gridSizeX = CudaFFT3D::findLegalDimension(gridSizeX);
         gridSizeY = CudaFFT3D::findLegalDimension(gridSizeY);
@@ -1994,14 +1993,26 @@ void CudaCalcNonbondedForceKernel::copyParametersToContext(ContextImpl& context,
     
     // Compute other values.
     
-    NonbondedForce::NonbondedMethod method = force.getNonbondedMethod();
-    if (method == NonbondedForce::Ewald || method == NonbondedForce::PME)
+    if (nonbondedMethod == Ewald || nonbondedMethod == PME)
         ewaldSelfEnergy = (cu.getContextIndex() == 0 ? -ONE_4PI_EPS0*alpha*sumSquaredCharges/sqrt(M_PI) : 0.0);
-    if (force.getUseDispersionCorrection() && cu.getContextIndex() == 0 && (method == NonbondedForce::CutoffPeriodic || method == NonbondedForce::Ewald || method == NonbondedForce::PME))
+    if (force.getUseDispersionCorrection() && cu.getContextIndex() == 0 && (nonbondedMethod == CutoffPeriodic || nonbondedMethod == Ewald || nonbondedMethod == PME))
         dispersionCoefficient = NonbondedForceImpl::calcDispersionCorrection(context.getSystem(), force);
     cu.invalidateMolecules();
 }
 
+void CudaCalcNonbondedForceKernel::getPMEParameters(double& alpha, int& nx, int& ny, int& nz) const {
+    if (nonbondedMethod != PME)
+        throw OpenMMException("getPMEParametersInContext: This Context is not using PME");
+    if (cu.getPlatformData().useCpuPme)
+        cpuPme.getAs<CalcPmeReciprocalForceKernel>().getPMEParameters(alpha, nx, ny, nz);
+    else {
+        alpha = this->alpha;
+        nx = gridSizeX;
+        ny = gridSizeY;
+        nz = gridSizeZ;
+    }
+}
+
 class CudaCustomNonbondedForceInfo : public CudaForceInfo {
 public:
     CudaCustomNonbondedForceInfo(const CustomNonbondedForce& force) : force(force) {

platforms/cuda/src/CudaParallelKernels.cpp

@@ -623,6 +623,10 @@ void CudaParallelCalcNonbondedForceKernel::copyParametersToContext(ContextImpl&
         getKernel(i).copyParametersToContext(context, force);
 }
 
+void CudaParallelCalcNonbondedForceKernel::getPMEParameters(double& alpha, int& nx, int& ny, int& nz) const {
+    dynamic_cast<const CudaCalcNonbondedForceKernel&>(kernels[0].getImpl()).getPMEParameters(alpha, nx, ny, nz);
+}
+
 class CudaParallelCalcCustomNonbondedForceKernel::Task : public CudaContext::WorkTask {
 public:
     Task(ContextImpl& context, CudaCalcCustomNonbondedForceKernel& kernel, bool includeForce,

platforms/cuda/tests/TestCudaEwald.cpp

@@ -42,6 +42,7 @@
 #include "openmm/LangevinIntegrator.h"
 #include "openmm/VerletIntegrator.h"
 #include "openmm/internal/ContextImpl.h"
+#include "openmm/internal/NonbondedForceImpl.h"
 #include "SimTKOpenMMRealType.h"
 #include "sfmt/SFMT.h"
 #include <iostream>
@@ -299,6 +300,19 @@ void testErrorTolerance(NonbondedForce::NonbondedMethod method) {
                 diff = sqrt(diff)/norm;
                 ASSERT(diff < 2*tol);
             }
+            if (method == NonbondedForce::PME) {
+                // See if the PME parameters were calculated correctly.
+
+                double expectedAlpha, actualAlpha;
+                int expectedSize[3], actualSize[3];
+                NonbondedForceImpl::calcPMEParameters(system, *force, expectedAlpha, expectedSize[0], expectedSize[1], expectedSize[2]);
+                force->getPMEParametersInContext(context, actualAlpha, actualSize[0], actualSize[1], actualSize[2]);
+                ASSERT_EQUAL_TOL(expectedAlpha, actualAlpha, 1e-5);
+                for (int i = 0; i < 3; i++) {
+                    ASSERT(actualSize[i] >= expectedSize[i]);
+                    ASSERT(actualSize[i] < expectedSize[i]+10);
+                }
+            }
         }
     }
 }

platforms/opencl/include/OpenCLKernels.h

@@ -598,6 +598,15 @@ public:
      * @param force      the NonbondedForce to copy the parameters from
      */
     void copyParametersToContext(ContextImpl& context, const NonbondedForce& force);
+    /**
+     * Get the parameters being used for PME.
+     * 
+     * @param alpha   the separation parameter
+     * @param nx      the number of grid points along the X axis
+     * @param ny      the number of grid points along the Y axis
+     * @param nz      the number of grid points along the Z axis
+     */
+    void getPMEParameters(double& alpha, int& nx, int& ny, int& nz) const;
 private:
     class SortTrait : public OpenCLSort::SortTrait {
         int getDataSize() const {return 8;}
@@ -647,7 +656,9 @@ private:
     std::map<std::string, std::string> pmeDefines;
     std::vector<std::pair<int, int> > exceptionAtoms;
     double ewaldSelfEnergy, dispersionCoefficient, alpha;
+    int gridSizeX, gridSizeY, gridSizeZ;
     bool hasCoulomb, hasLJ, usePmeQueue;
+    NonbondedMethod nonbondedMethod;
     static const int PmeOrder = 5;
 };
 

platforms/opencl/include/OpenCLParallelKernels.h

@@ -429,6 +429,15 @@ public:
      * @param force      the NonbondedForce to copy the parameters from
      */
     void copyParametersToContext(ContextImpl& context, const NonbondedForce& force);
+    /**
+     * Get the parameters being used for PME.
+     * 
+     * @param alpha   the separation parameter
+     * @param nx      the number of grid points along the X axis
+     * @param ny      the number of grid points along the Y axis
+     * @param nz      the number of grid points along the Z axis
+     */
+    void getPMEParameters(double& alpha, int& nx, int& ny, int& nz) const;
 private:
     class Task;
     OpenCLPlatform::PlatformData& data;

platforms/opencl/src/OpenCLKernels.cpp

@@ -1552,8 +1552,9 @@ void OpenCLCalcNonbondedForceKernel::initialize(const System& system, const Nonb
     else
         cl.getPosq().upload(posqf);
     sigmaEpsilon->upload(sigmaEpsilonVector);
-    bool useCutoff = (force.getNonbondedMethod() != NonbondedForce::NoCutoff);
-    bool usePeriodic = (force.getNonbondedMethod() != NonbondedForce::NoCutoff && force.getNonbondedMethod() != NonbondedForce::CutoffNonPeriodic);
+    nonbondedMethod = CalcNonbondedForceKernel::NonbondedMethod(force.getNonbondedMethod());
+    bool useCutoff = (nonbondedMethod != NoCutoff);
+    bool usePeriodic = (nonbondedMethod != NoCutoff && nonbondedMethod != CutoffNonPeriodic);
     map<string, string> defines;
     defines["HAS_COULOMB"] = (hasCoulomb ? "1" : "0");
     defines["HAS_LENNARD_JONES"] = (hasLJ ? "1" : "0");
@@ -1581,7 +1582,7 @@ void OpenCLCalcNonbondedForceKernel::initialize(const System& system, const Nonb
         dispersionCoefficient = 0.0;
     alpha = 0;
     ewaldSelfEnergy = 0.0;
-    if (force.getNonbondedMethod() == NonbondedForce::Ewald) {
+    if (nonbondedMethod == Ewald) {
         // Compute the Ewald parameters.
 
         int kmaxx, kmaxy, kmaxz;
@@ -1607,10 +1608,9 @@ void OpenCLCalcNonbondedForceKernel::initialize(const System& system, const Nonb
             cosSinSums = new OpenCLArray(cl, (2*kmaxx-1)*(2*kmaxy-1)*(2*kmaxz-1), elementSize, "cosSinSums");
         }
     }
-    else if (force.getNonbondedMethod() == NonbondedForce::PME) {
+    else if (nonbondedMethod == PME) {
         // Compute the PME parameters.
 
-        int gridSizeX, gridSizeY, gridSizeZ;
         NonbondedForceImpl::calcPMEParameters(system, force, alpha, gridSizeX, gridSizeY, gridSizeZ);
         gridSizeX = OpenCLFFT3D::findLegalDimension(gridSizeX);
         gridSizeY = OpenCLFFT3D::findLegalDimension(gridSizeY);
@@ -2057,14 +2057,26 @@ void OpenCLCalcNonbondedForceKernel::copyParametersToContext(ContextImpl& contex
     
     // Compute other values.
     
-    NonbondedForce::NonbondedMethod method = force.getNonbondedMethod();
-    if (method == NonbondedForce::Ewald || method == NonbondedForce::PME)
+    if (nonbondedMethod == Ewald || nonbondedMethod == PME)
         ewaldSelfEnergy = (cl.getContextIndex() == 0 ? -ONE_4PI_EPS0*alpha*sumSquaredCharges/sqrt(M_PI) : 0.0);
-    if (force.getUseDispersionCorrection() && cl.getContextIndex() == 0 && (method == NonbondedForce::CutoffPeriodic || method == NonbondedForce::Ewald || method == NonbondedForce::PME))
+    if (force.getUseDispersionCorrection() && cl.getContextIndex() == 0 && (nonbondedMethod == CutoffPeriodic || nonbondedMethod == Ewald || nonbondedMethod == PME))
         dispersionCoefficient = NonbondedForceImpl::calcDispersionCorrection(context.getSystem(), force);
     cl.invalidateMolecules();
 }
 
+void OpenCLCalcNonbondedForceKernel::getPMEParameters(double& alpha, int& nx, int& ny, int& nz) const {
+    if (nonbondedMethod != PME)
+        throw OpenMMException("getPMEParametersInContext: This Context is not using PME");
+    if (cl.getPlatformData().useCpuPme)
+        cpuPme.getAs<CalcPmeReciprocalForceKernel>().getPMEParameters(alpha, nx, ny, nz);
+    else {
+        alpha = this->alpha;
+        nx = gridSizeX;
+        ny = gridSizeY;
+        nz = gridSizeZ;
+    }
+}
+
 class OpenCLCustomNonbondedForceInfo : public OpenCLForceInfo {
 public:
     OpenCLCustomNonbondedForceInfo(int requiredBuffers, const CustomNonbondedForce& force) : OpenCLForceInfo(requiredBuffers), force(force) {