Add correction for self energy of neutralizing plasma (#4907)

* Add correction for self energy of neutralizing plasma

* Fixed compilation errors

* Update total charge in copyParametersToContext()

* Bug fixes

* Fixed compilation errors in HIP

* Bug fix

platforms/common/src/kernels/nonbondedParameters.cc

@@ -3,13 +3,15 @@
  */
 KERNEL void computeParameters(GLOBAL mixed* RESTRICT energyBuffer, int includeSelfEnergy, GLOBAL real* RESTRICT globalParams,
         int numAtoms, GLOBAL const float4* RESTRICT baseParticleParams, GLOBAL real4* RESTRICT posq, GLOBAL real* RESTRICT charge,
-        GLOBAL float2* RESTRICT sigmaEpsilon, GLOBAL float4* RESTRICT particleParamOffsets, GLOBAL int* RESTRICT particleOffsetIndices
+        GLOBAL float2* RESTRICT sigmaEpsilon, GLOBAL float4* RESTRICT particleParamOffsets, GLOBAL int* RESTRICT particleOffsetIndices,
+        GLOBAL real* RESTRICT chargeBuffer
 #ifdef HAS_EXCEPTIONS
         , int numExceptions, GLOBAL const float4* RESTRICT baseExceptionParams, GLOBAL float4* RESTRICT exceptionParams,
         GLOBAL float4* RESTRICT exceptionParamOffsets, GLOBAL int* RESTRICT exceptionOffsetIndices
 #endif
         ) {
     mixed energy = 0;
+    real totalCharge = 0;
 
     // Compute particle parameters.
     
@@ -31,6 +33,7 @@ KERNEL void computeParameters(GLOBAL mixed* RESTRICT energyBuffer, int includeSe
         charge[i] = params.x;
 #endif
         sigmaEpsilon[i] = make_float2(0.5f*params.y, 2*SQRT(params.z));
+        totalCharge += params.x;
 #ifdef HAS_OFFSETS
     #ifdef INCLUDE_EWALD
         energy -= EWALD_SELF_ENERGY_SCALE*params.x*params.x;
@@ -62,6 +65,20 @@ KERNEL void computeParameters(GLOBAL mixed* RESTRICT energyBuffer, int includeSe
 #endif
     if (includeSelfEnergy)
         energyBuffer[GLOBAL_ID] += energy;
+
+    // Record the total charge from particles processed by this block.
+
+#if defined(HAS_OFFSETS) && defined(INCLUDE_EWALD)
+    LOCAL real temp[WORK_GROUP_SIZE];
+    temp[LOCAL_ID] = totalCharge;
+    for (int i = 1; i < WORK_GROUP_SIZE; i *= 2) {
+        SYNC_THREADS;
+        if (LOCAL_ID%(i*2) == 0 && LOCAL_ID+i < WORK_GROUP_SIZE)
+            temp[LOCAL_ID] += temp[LOCAL_ID+i];
+    }
+    if (LOCAL_ID == 0)
+        chargeBuffer[GROUP_ID] = temp[0];
+#endif
 }
 
 /**
@@ -88,3 +105,26 @@ KERNEL void computeExclusionParameters(GLOBAL real4* RESTRICT posq, GLOBAL real*
         exclusionParams[i] = make_float4((float) (ONE_4PI_EPS0*chargeProd), sigma, epsilon, 0);
     }
 }
+
+/**
+ * When using Ewald or PME with parameter offsets, the total charge can change each step.
+ * We therefore need to compute the correction for the neutralizing plasma on the GPU.
+ * This kernel is executed by a single thread block.
+ */
+KERNEL void computePlasmaCorrection(GLOBAL real* RESTRICT chargeBuffer, GLOBAL mixed* RESTRICT energyBuffer,
+        real alpha, real volume) {
+    LOCAL real temp[WORK_GROUP_SIZE];
+    real sum = 0;
+    for (unsigned int index = LOCAL_ID; index < NUM_GROUPS; index += LOCAL_SIZE)
+        sum += chargeBuffer[index];
+    temp[LOCAL_ID] = sum;
+    for (int i = 1; i < WORK_GROUP_SIZE; i *= 2) {
+        SYNC_THREADS;
+        if (LOCAL_ID%(i*2) == 0 && LOCAL_ID+i < WORK_GROUP_SIZE)
+            temp[LOCAL_ID] += temp[LOCAL_ID+i];
+    }
+    if (LOCAL_ID == 0) {
+        real totalCharge = temp[0];
+        energyBuffer[0] -= totalCharge*totalCharge/(8*EPSILON0*volume*alpha*alpha);
+    }
+}
\ No newline at end of file

platforms/cpu/include/CpuKernels.h

@@ -305,7 +305,7 @@ private:
     std::vector<std::vector<int> > bonded14IndexArray;
     std::vector<std::vector<double> > bonded14ParamArray;
     std::map<int, int> nb14Index;
-    double nonbondedCutoff, switchingDistance, rfDielectric, ewaldAlpha, ewaldDispersionAlpha, ewaldSelfEnergy, dispersionCoefficient;
+    double nonbondedCutoff, switchingDistance, rfDielectric, ewaldAlpha, ewaldDispersionAlpha, ewaldSelfEnergy, dispersionCoefficient, totalCharge;
     int kmax[3], gridSize[3], dispersionGridSize[3];
     bool useSwitchingFunction, exceptionsArePeriodic, useOptimizedPme, hasInitializedPme, hasInitializedDispersionPme, hasParticleOffsets, hasExceptionOffsets;
     std::vector<std::set<int> > exclusions;

platforms/cpu/src/CpuKernels.cpp

@@ -712,6 +712,12 @@ double CpuCalcNonbondedForceKernel::execute(ContextImpl& context, bool includeFo
         }
         else
             nonbonded->calculateReciprocalIxn(numParticles, &posq[0], posData, particleParams, C6params, exclusions, forceData, includeEnergy ? &nonbondedEnergy : NULL);
+        if (ewald || pme || ljpme) {
+            // Add the correction for the neutralizing plasma.
+
+            double volume = boxVectors[0][0]*boxVectors[1][1]*boxVectors[2][2];
+            energy -= totalCharge*totalCharge/(8*EPSILON0*volume*ewaldAlpha*ewaldAlpha);
+        }
     }
     energy += nonbondedEnergy;
     if (includeDirect) {
@@ -843,6 +849,7 @@ void CpuCalcNonbondedForceKernel::computeParameters(ContextImpl& context, bool o
 
     if (hasParticleOffsets || !offsetsOnly) {
         double sumSquaredCharges = 0.0;
+        totalCharge = 0.0;
         for (int i = 0; i < numParticles; i++) {
             double charge = baseParticleParams[i][0];
             double sigma = baseParticleParams[i][1];
@@ -857,6 +864,7 @@ void CpuCalcNonbondedForceKernel::computeParameters(ContextImpl& context, bool o
             particleParams[i] = make_pair((float) (0.5*sigma), (float) (2.0*sqrt(epsilon)));
             C6params[i] = 8.0*pow(particleParams[i].first, 3.0) * particleParams[i].second;
             sumSquaredCharges += charge*charge;
+            totalCharge += charge;
         }
         if (nonbondedMethod == Ewald || nonbondedMethod == PME || nonbondedMethod == LJPME) {
             ewaldSelfEnergy = -ONE_4PI_EPS0*ewaldAlpha*sumSquaredCharges/sqrt(M_PI);

platforms/cuda/include/CudaKernels.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-2024 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2025 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -182,6 +182,7 @@ private:
     CudaArray pmeDispersionBsplineModuliZ;
     CudaArray pmeAtomGridIndex;
     CudaArray pmeEnergyBuffer;
+    CudaArray chargeBuffer;
     CudaSort* sort;
     Kernel cpuPme;
     PmeIO* pmeio;
@@ -193,7 +194,7 @@ private:
     CudaFFT3D* dispersionFft;
     cufftHandle dispersionFftForward;
     cufftHandle dispersionFftBackward;
-    CUfunction computeParamsKernel, computeExclusionParamsKernel;
+    CUfunction computeParamsKernel, computeExclusionParamsKernel, computePlasmaCorrectionKernel;
     CUfunction ewaldSumsKernel;
     CUfunction ewaldForcesKernel;
     CUfunction pmeGridIndexKernel;
@@ -212,7 +213,7 @@ private:
     std::vector<std::string> paramNames;
     std::vector<double> paramValues;
     std::map<int, int> exceptionIndex;
-    double ewaldSelfEnergy, dispersionCoefficient, alpha, dispersionAlpha;
+    double ewaldSelfEnergy, dispersionCoefficient, alpha, dispersionAlpha, totalCharge;
     int interpolateForceThreads;
     int gridSizeX, gridSizeY, gridSizeZ;
     int dispersionGridSizeX, dispersionGridSizeY, dispersionGridSizeZ;

platforms/cuda/src/CudaKernels.cpp

@@ -367,8 +367,11 @@ void CudaCalcNonbondedForceKernel::initialize(const System& system, const Nonbon
         dispersionCoefficient = 0.0;
     alpha = 0;
     ewaldSelfEnergy = 0.0;
+    totalCharge = 0.0;
     map<string, string> paramsDefines;
     paramsDefines["ONE_4PI_EPS0"] = cu.doubleToString(ONE_4PI_EPS0);
+    paramsDefines["EPSILON0"] = cu.doubleToString(EPSILON0);
+    paramsDefines["WORK_GROUP_SIZE"] = cu.intToString(CudaContext::ThreadBlockSize);
     hasOffsets = (force.getNumParticleParameterOffsets() > 0 || force.getNumExceptionParameterOffsets() > 0);
     if (hasOffsets)
         paramsDefines["HAS_OFFSETS"] = "1";
@@ -391,8 +394,10 @@ void CudaCalcNonbondedForceKernel::initialize(const System& system, const Nonbon
         if (cu.getContextIndex() == 0) {
             paramsDefines["INCLUDE_EWALD"] = "1";
             paramsDefines["EWALD_SELF_ENERGY_SCALE"] = cu.doubleToString(ONE_4PI_EPS0*alpha/sqrt(M_PI));
-            for (int i = 0; i < numParticles; i++)
+            for (int i = 0; i < numParticles; i++) {
                 ewaldSelfEnergy -= baseParticleParamVec[i].x*baseParticleParamVec[i].x*ONE_4PI_EPS0*alpha/sqrt(M_PI);
+                totalCharge += baseParticleParamVec[i].x;
+            }
 
             // Create the reciprocal space kernels.
 
@@ -444,8 +449,10 @@ void CudaCalcNonbondedForceKernel::initialize(const System& system, const Nonbon
         if (cu.getContextIndex() == 0) {
             paramsDefines["INCLUDE_EWALD"] = "1";
             paramsDefines["EWALD_SELF_ENERGY_SCALE"] = cu.doubleToString(ONE_4PI_EPS0*alpha/sqrt(M_PI));
-            for (int i = 0; i < numParticles; i++)
+            for (int i = 0; i < numParticles; i++) {
                 ewaldSelfEnergy -= baseParticleParamVec[i].x*baseParticleParamVec[i].x*ONE_4PI_EPS0*alpha/sqrt(M_PI);
+                totalCharge += baseParticleParamVec[i].x;
+            }
             if (doLJPME) {
                 paramsDefines["INCLUDE_LJPME"] = "1";
                 paramsDefines["LJPME_SELF_ENERGY_SCALE"] = cu.doubleToString(pow(dispersionAlpha, 6)/3.0);
@@ -831,6 +838,8 @@ void CudaCalcNonbondedForceKernel::initialize(const System& system, const Nonbon
     globalParams.initialize(cu, max((int) paramValues.size(), 1), cu.getUseDoublePrecision() ? sizeof(double) : sizeof(float), "globalParams");
     if (paramValues.size() > 0)
         globalParams.upload(paramValues, true);
+    chargeBuffer.initialize(cu, cu.getNumThreadBlocks(), cu.getUseDoublePrecision() ? sizeof(double) : sizeof(float), "chargeBuffer");
+    cu.clearBuffer(chargeBuffer);
     recomputeParams = true;
     
     // Initialize the kernel for updating parameters.
@@ -838,6 +847,7 @@ void CudaCalcNonbondedForceKernel::initialize(const System& system, const Nonbon
     CUmodule module = cu.createModule(CommonKernelSources::nonbondedParameters, paramsDefines);
     computeParamsKernel = cu.getKernel(module, "computeParameters");
     computeExclusionParamsKernel = cu.getKernel(module, "computeExclusionParameters");
+    computePlasmaCorrectionKernel = cu.getKernel(module, "computePlasmaCorrection");
     info = new ForceInfo(force);
     cu.addForce(info);
 }
@@ -858,13 +868,18 @@ double CudaCalcNonbondedForceKernel::execute(ContextImpl& context, bool includeF
         recomputeParams = true;
         globalParams.upload(paramValues, true);
     }
-    double energy = (includeReciprocal ? ewaldSelfEnergy : 0.0);
+    double energy = 0.0;
+    if (includeReciprocal && (pmeGrid1.isInitialized() || cosSinSums.isInitialized())) {
+        double4 boxSize = cu.getPeriodicBoxSize();
+        double volume = boxSize.x*boxSize.y*boxSize.z;
+        energy = ewaldSelfEnergy - totalCharge*totalCharge/(8*EPSILON0*volume*alpha*alpha);
+    }
     if (recomputeParams || hasOffsets) {
         int computeSelfEnergy = (includeEnergy && includeReciprocal);
-        int numAtoms = cu.getPaddedNumAtoms();
+        int numAtoms = cu.getNumAtoms();
         vector<void*> paramsArgs = {&cu.getEnergyBuffer().getDevicePointer(), &computeSelfEnergy, &globalParams.getDevicePointer(), &numAtoms,
                 &baseParticleParams.getDevicePointer(), &cu.getPosq().getDevicePointer(), &charges.getDevicePointer(), &sigmaEpsilon.getDevicePointer(),
-                &particleParamOffsets.getDevicePointer(), &particleOffsetIndices.getDevicePointer()};
+                &particleParamOffsets.getDevicePointer(), &particleOffsetIndices.getDevicePointer(), &chargeBuffer.getDevicePointer()};
         int numExceptions;
         if (exceptionParams.isInitialized()) {
             numExceptions = exceptionParams.getSize();
@@ -885,8 +900,27 @@ double CudaCalcNonbondedForceKernel::execute(ContextImpl& context, bool includeF
             cuEventRecord(paramsSyncEvent, cu.getCurrentStream());
             cuStreamWaitEvent(pmeStream, paramsSyncEvent, 0);
         }
-        if (hasOffsets)
-            energy = 0.0; // The Ewald self energy was computed in the kernel.
+        if (hasOffsets) {
+            // The Ewald self energy was computed in the kernel.
+
+            energy = 0.0;
+            if (pmeGrid1.isInitialized() || cosSinSums.isInitialized()) {
+                // Invoke a kernel to compute the correction for the neutralizing plasma.
+
+                double4 boxSize = cu.getPeriodicBoxSize();
+                if (cu.getUseDoublePrecision()) {
+                    double volume = boxSize.x*boxSize.y*boxSize.z;
+                    vector<void*> correctionArgs = {&chargeBuffer.getDevicePointer(), &cu.getEnergyBuffer().getDevicePointer(), &alpha, &volume};
+                    cu.executeKernel(computePlasmaCorrectionKernel, &correctionArgs[0], CudaContext::ThreadBlockSize, CudaContext::ThreadBlockSize);
+                }
+                else {
+                    float alphaFloat = (float) alpha;
+                    float volume = boxSize.x*boxSize.y*boxSize.z;
+                    vector<void*> correctionArgs = {&chargeBuffer.getDevicePointer(), &cu.getEnergyBuffer().getDevicePointer(), &alphaFloat, &volume};
+                    cu.executeKernel(computePlasmaCorrectionKernel, &correctionArgs[0], CudaContext::ThreadBlockSize, CudaContext::ThreadBlockSize);
+                }
+            }
+        }
         recomputeParams = false;
     }
     
@@ -1143,10 +1177,12 @@ void CudaCalcNonbondedForceKernel::copyParametersToContext(ContextImpl& context,
         // Compute the self energy.
 
         ewaldSelfEnergy = 0.0;
+        totalCharge = 0.0;
         if (nonbondedMethod == Ewald || nonbondedMethod == PME || nonbondedMethod == LJPME) {
             if (cu.getContextIndex() == 0) {
                 for (int i = 0; i < force.getNumParticles(); i++) {
                     ewaldSelfEnergy -= baseParticleParamVec[i].x*baseParticleParamVec[i].x*ONE_4PI_EPS0*alpha/sqrt(M_PI);
+                    totalCharge += baseParticleParamVec[i].x;
                     if (doLJPME)
                         ewaldSelfEnergy += baseParticleParamVec[i].z*pow(baseParticleParamVec[i].y*dispersionAlpha, 6)/3.0;
                 }

platforms/hip/include/HipKernels.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-2024 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2025 Stanford University and the Authors.      *
  * Portions copyright (c) 2020-2022 Advanced Micro Devices, Inc.              *
  * Authors: Peter Eastman, Nicholas Curtis                                    *
  * Contributors:                                                              *
@@ -182,6 +182,7 @@ private:
     HipArray pmeDispersionBsplineModuliZ;
     HipArray pmeAtomGridIndex;
     HipArray pmeEnergyBuffer;
+    HipArray chargeBuffer;
     HipSort* sort;
     Kernel cpuPme;
     PmeIO* pmeio;
@@ -189,7 +190,7 @@ private:
     hipEvent_t pmeSyncEvent, paramsSyncEvent;
     HipFFT3D* fft;
     HipFFT3D* dispersionFft;
-    hipFunction_t computeParamsKernel, computeExclusionParamsKernel;
+    hipFunction_t computeParamsKernel, computeExclusionParamsKernel, computePlasmaCorrectionKernel;
     hipFunction_t ewaldSumsKernel;
     hipFunction_t ewaldForcesKernel;
     hipFunction_t pmeGridIndexKernel;
@@ -208,7 +209,7 @@ private:
     std::vector<std::string> paramNames;
     std::vector<double> paramValues;
     std::map<int, int> exceptionIndex;
-    double ewaldSelfEnergy, dispersionCoefficient, alpha, dispersionAlpha;
+    double ewaldSelfEnergy, dispersionCoefficient, alpha, dispersionAlpha, totalCharge;
     int interpolateForceThreads;
     int gridSizeX, gridSizeY, gridSizeZ;
     int dispersionGridSizeX, dispersionGridSizeY, dispersionGridSizeZ;

platforms/hip/src/HipKernels.cpp

@@ -361,8 +361,11 @@ void HipCalcNonbondedForceKernel::initialize(const System& system, const Nonbond
         dispersionCoefficient = 0.0;
     alpha = 0;
     ewaldSelfEnergy = 0.0;
+    totalCharge = 0.0;
     map<string, string> paramsDefines;
     paramsDefines["ONE_4PI_EPS0"] = cu.doubleToString(ONE_4PI_EPS0);
+    paramsDefines["EPSILON0"] = cu.doubleToString(EPSILON0);
+    paramsDefines["WORK_GROUP_SIZE"] = cu.intToString(HipContext::ThreadBlockSize);
     hasOffsets = (force.getNumParticleParameterOffsets() > 0 || force.getNumExceptionParameterOffsets() > 0);
     if (hasOffsets)
         paramsDefines["HAS_OFFSETS"] = "1";
@@ -385,8 +388,10 @@ void HipCalcNonbondedForceKernel::initialize(const System& system, const Nonbond
         if (cu.getContextIndex() == 0) {
             paramsDefines["INCLUDE_EWALD"] = "1";
             paramsDefines["EWALD_SELF_ENERGY_SCALE"] = cu.doubleToString(ONE_4PI_EPS0*alpha/sqrt(M_PI));
-            for (int i = 0; i < numParticles; i++)
+            for (int i = 0; i < numParticles; i++) {
                 ewaldSelfEnergy -= baseParticleParamVec[i].x*baseParticleParamVec[i].x*ONE_4PI_EPS0*alpha/sqrt(M_PI);
+                totalCharge += baseParticleParamVec[i].x;
+            }
 
             // Create the reciprocal space kernels.
 
@@ -438,8 +443,10 @@ void HipCalcNonbondedForceKernel::initialize(const System& system, const Nonbond
         if (cu.getContextIndex() == 0) {
             paramsDefines["INCLUDE_EWALD"] = "1";
             paramsDefines["EWALD_SELF_ENERGY_SCALE"] = cu.doubleToString(ONE_4PI_EPS0*alpha/sqrt(M_PI));
-            for (int i = 0; i < numParticles; i++)
+            for (int i = 0; i < numParticles; i++) {
                 ewaldSelfEnergy -= baseParticleParamVec[i].x*baseParticleParamVec[i].x*ONE_4PI_EPS0*alpha/sqrt(M_PI);
+                totalCharge += baseParticleParamVec[i].x;
+            }
             if (doLJPME) {
                 paramsDefines["INCLUDE_LJPME"] = "1";
                 paramsDefines["LJPME_SELF_ENERGY_SCALE"] = cu.doubleToString(pow(dispersionAlpha, 6)/3.0);
@@ -791,6 +798,8 @@ void HipCalcNonbondedForceKernel::initialize(const System& system, const Nonbond
     globalParams.initialize(cu, max((int) paramValues.size(), 1), cu.getUseDoublePrecision() ? sizeof(double) : sizeof(float), "globalParams");
     if (paramValues.size() > 0)
         globalParams.upload(paramValues, true);
+    chargeBuffer.initialize(cu, cu.getNumThreadBlocks(), cu.getUseDoublePrecision() ? sizeof(double) : sizeof(float), "chargeBuffer");
+    cu.clearBuffer(chargeBuffer);
     recomputeParams = true;
 
     // Initialize the kernel for updating parameters.
@@ -798,6 +807,7 @@ void HipCalcNonbondedForceKernel::initialize(const System& system, const Nonbond
     hipModule_t module = cu.createModule(CommonKernelSources::nonbondedParameters, paramsDefines);
     computeParamsKernel = cu.getKernel(module, "computeParameters");
     computeExclusionParamsKernel = cu.getKernel(module, "computeExclusionParameters");
+    computePlasmaCorrectionKernel = cu.getKernel(module, "computePlasmaCorrection");
     info = new ForceInfo(force);
     cu.addForce(info);
 }
@@ -818,13 +828,18 @@ double HipCalcNonbondedForceKernel::execute(ContextImpl& context, bool includeFo
         recomputeParams = true;
         globalParams.upload(paramValues, true);
     }
-    double energy = (includeReciprocal ? ewaldSelfEnergy : 0.0);
+    double energy = 0.0;
+    if (includeReciprocal && (pmeGrid1.isInitialized() || cosSinSums.isInitialized())) {
+        double4 boxSize = cu.getPeriodicBoxSize();
+        double volume = boxSize.x*boxSize.y*boxSize.z;
+        energy = ewaldSelfEnergy - totalCharge*totalCharge/(8*EPSILON0*volume*alpha*alpha);
+    }
     if (recomputeParams || hasOffsets) {
         int computeSelfEnergy = (includeEnergy && includeReciprocal);
-        int numAtoms = cu.getPaddedNumAtoms();
+        int numAtoms = cu.getNumAtoms();
         vector<void*> paramsArgs = {&cu.getEnergyBuffer().getDevicePointer(), &computeSelfEnergy, &globalParams.getDevicePointer(), &numAtoms,
                 &baseParticleParams.getDevicePointer(), &cu.getPosq().getDevicePointer(), &charges.getDevicePointer(), &sigmaEpsilon.getDevicePointer(),
-                &particleParamOffsets.getDevicePointer(), &particleOffsetIndices.getDevicePointer()};
+                &particleParamOffsets.getDevicePointer(), &particleOffsetIndices.getDevicePointer(), &chargeBuffer.getDevicePointer()};
         int numExceptions;
         if (exceptionParams.isInitialized()) {
             numExceptions = exceptionParams.getSize();
@@ -845,8 +860,27 @@ double HipCalcNonbondedForceKernel::execute(ContextImpl& context, bool includeFo
             hipEventRecord(paramsSyncEvent, cu.getCurrentStream());
             hipStreamWaitEvent(pmeStream, paramsSyncEvent, 0);
         }
-        if (hasOffsets)
-            energy = 0.0; // The Ewald self energy was computed in the kernel.
+        if (hasOffsets) {
+            // The Ewald self energy was computed in the kernel.
+
+            energy = 0.0;
+            if (pmeGrid1.isInitialized() || cosSinSums.isInitialized()) {
+                // Invoke a kernel to compute the correction for the neutralizing plasma.
+
+                double4 boxSize = cu.getPeriodicBoxSize();
+                if (cu.getUseDoublePrecision()) {
+                    double volume = boxSize.x*boxSize.y*boxSize.z;
+                    vector<void*> correctionArgs = {&chargeBuffer.getDevicePointer(), &cu.getEnergyBuffer().getDevicePointer(), &alpha, &volume};
+                    cu.executeKernel(computePlasmaCorrectionKernel, &correctionArgs[0], HipContext::ThreadBlockSize, HipContext::ThreadBlockSize);
+                }
+                else {
+                    float alphaFloat = (float) alpha;
+                    float volume = boxSize.x*boxSize.y*boxSize.z;
+                    vector<void*> correctionArgs = {&chargeBuffer.getDevicePointer(), &cu.getEnergyBuffer().getDevicePointer(), &alphaFloat, &volume};
+                    cu.executeKernel(computePlasmaCorrectionKernel, &correctionArgs[0], HipContext::ThreadBlockSize, HipContext::ThreadBlockSize);
+                }
+            }
+        }
         recomputeParams = false;
     }
 
@@ -1051,10 +1085,12 @@ void HipCalcNonbondedForceKernel::copyParametersToContext(ContextImpl& context,
         // Compute the self energy.
 
         ewaldSelfEnergy = 0.0;
+        totalCharge = 0.0;
         if (nonbondedMethod == Ewald || nonbondedMethod == PME || nonbondedMethod == LJPME) {
             if (cu.getContextIndex() == 0) {
                 for (int i = 0; i < force.getNumParticles(); i++) {
                     ewaldSelfEnergy -= baseParticleParamVec[i].x*baseParticleParamVec[i].x*ONE_4PI_EPS0*alpha/sqrt(M_PI);
+                    totalCharge += baseParticleParamVec[i].x;
                     if (doLJPME)
                         ewaldSelfEnergy += baseParticleParamVec[i].z*pow(baseParticleParamVec[i].y*dispersionAlpha, 6)/3.0;
                 }

platforms/opencl/include/OpenCLKernels.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-2024 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2025 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -183,6 +183,7 @@ private:
     OpenCLArray pmeAtomRange;
     OpenCLArray pmeAtomGridIndex;
     OpenCLArray pmeEnergyBuffer;
+    OpenCLArray chargeBuffer;
     OpenCLSort* sort;
     cl::CommandQueue pmeQueue;
     cl::Event pmeSyncEvent;
@@ -191,7 +192,7 @@ private:
     Kernel cpuPme;
     PmeIO* pmeio;
     SyncQueuePostComputation* syncQueue;
-    cl::Kernel computeParamsKernel, computeExclusionParamsKernel;
+    cl::Kernel computeParamsKernel, computeExclusionParamsKernel, computePlasmaCorrectionKernel;
     cl::Kernel ewaldSumsKernel;
     cl::Kernel ewaldForcesKernel;
     cl::Kernel pmeAtomRangeKernel;
@@ -215,7 +216,7 @@ private:
     std::vector<std::string> paramNames;
     std::vector<double> paramValues;
     std::map<int, int> exceptionIndex;
-    double ewaldSelfEnergy, dispersionCoefficient, alpha, dispersionAlpha;
+    double ewaldSelfEnergy, dispersionCoefficient, alpha, dispersionAlpha, totalCharge;
     int gridSizeX, gridSizeY, gridSizeZ;
     int dispersionGridSizeX, dispersionGridSizeY, dispersionGridSizeZ;
     bool hasCoulomb, hasLJ, usePmeQueue, doLJPME, usePosqCharges, recomputeParams, hasOffsets;

platforms/opencl/src/OpenCLKernels.cpp

@@ -370,8 +370,10 @@ void OpenCLCalcNonbondedForceKernel::initialize(const System& system, const Nonb
         dispersionCoefficient = 0.0;
     alpha = 0;
     ewaldSelfEnergy = 0.0;
+    totalCharge = 0.0;
     map<string, string> paramsDefines;
     paramsDefines["ONE_4PI_EPS0"] = cl.doubleToString(ONE_4PI_EPS0);
+    paramsDefines["EPSILON0"] = cl.doubleToString(EPSILON0);
     hasOffsets = (force.getNumParticleParameterOffsets() > 0 || force.getNumExceptionParameterOffsets() > 0);
     if (hasOffsets)
         paramsDefines["HAS_OFFSETS"] = "1";
@@ -394,8 +396,10 @@ void OpenCLCalcNonbondedForceKernel::initialize(const System& system, const Nonb
         if (cl.getContextIndex() == 0) {
             paramsDefines["INCLUDE_EWALD"] = "1";
             paramsDefines["EWALD_SELF_ENERGY_SCALE"] = cl.doubleToString(ONE_4PI_EPS0*alpha/sqrt(M_PI));
-            for (int i = 0; i < numParticles; i++)
+            for (int i = 0; i < numParticles; i++) {
                 ewaldSelfEnergy -= baseParticleParamVec[i].x*baseParticleParamVec[i].x*ONE_4PI_EPS0*alpha/sqrt(M_PI);
+                totalCharge += baseParticleParamVec[i].x;
+            }
 
             // Create the reciprocal space kernels.
 
@@ -446,8 +450,10 @@ void OpenCLCalcNonbondedForceKernel::initialize(const System& system, const Nonb
         if (cl.getContextIndex() == 0) {
             paramsDefines["INCLUDE_EWALD"] = "1";
             paramsDefines["EWALD_SELF_ENERGY_SCALE"] = cl.doubleToString(ONE_4PI_EPS0*alpha/sqrt(M_PI));
-            for (int i = 0; i < numParticles; i++)
+            for (int i = 0; i < numParticles; i++) {
                 ewaldSelfEnergy -= baseParticleParamVec[i].x*baseParticleParamVec[i].x*ONE_4PI_EPS0*alpha/sqrt(M_PI);
+                totalCharge += baseParticleParamVec[i].x;
+            }
             if (doLJPME) {
                 paramsDefines["INCLUDE_LJPME"] = "1";
                 paramsDefines["LJPME_SELF_ENERGY_SCALE"] = cl.doubleToString(pow(dispersionAlpha, 6)/3.0);
@@ -766,6 +772,8 @@ void OpenCLCalcNonbondedForceKernel::initialize(const System& system, const Nonb
     globalParams.initialize(cl, max((int) paramValues.size(), 1), cl.getUseDoublePrecision() ? sizeof(double) : sizeof(float), "globalParams");
     if (paramValues.size() > 0)
         globalParams.upload(paramValues, true);
+    chargeBuffer.initialize(cl, cl.getNumThreadBlocks(), cl.getUseDoublePrecision() ? sizeof(double) : sizeof(float), "chargeBuffer");
+    cl.clearBuffer(chargeBuffer);
     recomputeParams = true;
     
     // Initialize the kernel for updating parameters.
@@ -773,6 +781,7 @@ void OpenCLCalcNonbondedForceKernel::initialize(const System& system, const Nonb
     cl::Program program = cl.createProgram(CommonKernelSources::nonbondedParameters, paramsDefines);
     computeParamsKernel = cl::Kernel(program, "computeParameters");
     computeExclusionParamsKernel = cl::Kernel(program, "computeExclusionParameters");
+    computePlasmaCorrectionKernel = cl::Kernel(program, "computePlasmaCorrection");
     info = new ForceInfo(0, force);
     cl.addForce(info);
 }
@@ -792,6 +801,7 @@ double OpenCLCalcNonbondedForceKernel::execute(ContextImpl& context, bool includ
         computeParamsKernel.setArg<cl::Buffer>(index++, sigmaEpsilon.getDeviceBuffer());
         computeParamsKernel.setArg<cl::Buffer>(index++, particleParamOffsets.getDeviceBuffer());
         computeParamsKernel.setArg<cl::Buffer>(index++, particleOffsetIndices.getDeviceBuffer());
+        computeParamsKernel.setArg<cl::Buffer>(index++, chargeBuffer.getDeviceBuffer());
         if (exceptionParams.isInitialized()) {
             computeParamsKernel.setArg<cl_int>(index++, exceptionParams.getSize());
             computeParamsKernel.setArg<cl::Buffer>(index++, baseExceptionParams.getDeviceBuffer());
@@ -807,6 +817,12 @@ double OpenCLCalcNonbondedForceKernel::execute(ContextImpl& context, bool includ
             computeExclusionParamsKernel.setArg<cl::Buffer>(4, exclusionAtoms.getDeviceBuffer());
             computeExclusionParamsKernel.setArg<cl::Buffer>(5, exclusionParams.getDeviceBuffer());
         }
+        computePlasmaCorrectionKernel.setArg<cl::Buffer>(0, chargeBuffer.getDeviceBuffer());
+        computePlasmaCorrectionKernel.setArg<cl::Buffer>(1, cl.getEnergyBuffer().getDeviceBuffer());
+        if (cl.getUseDoublePrecision())
+            computePlasmaCorrectionKernel.setArg<double>(2, alpha);
+        else
+            computePlasmaCorrectionKernel.setArg<float>(2, alpha);
         if (cosSinSums.isInitialized()) {
             ewaldSumsKernel.setArg<cl::Buffer>(0, cl.getEnergyBuffer().getDeviceBuffer());
             ewaldSumsKernel.setArg<cl::Buffer>(1, cl.getPosq().getDeviceBuffer());
@@ -911,10 +927,15 @@ double OpenCLCalcNonbondedForceKernel::execute(ContextImpl& context, bool includ
         recomputeParams = true;
         globalParams.upload(paramValues, true);
     }
-    double energy = (includeReciprocal ? ewaldSelfEnergy : 0.0);
+    double energy = 0.0;
+    if (includeReciprocal && (pmeGrid1.isInitialized() || cosSinSums.isInitialized())) {
+        mm_double4 boxSize = cl.getPeriodicBoxSizeDouble();
+        double volume = boxSize.x*boxSize.y*boxSize.z;
+        energy = ewaldSelfEnergy - totalCharge*totalCharge/(8*EPSILON0*volume*alpha*alpha);
+    }
     if (recomputeParams || hasOffsets) {
         computeParamsKernel.setArg<cl_int>(1, includeEnergy && includeReciprocal);
-        cl.executeKernel(computeParamsKernel, cl.getPaddedNumAtoms());
+        cl.executeKernel(computeParamsKernel, cl.getNumAtoms());
         if (exclusionParams.isInitialized())
             cl.executeKernel(computeExclusionParamsKernel, exclusionParams.getSize());
         if (usePmeQueue) {
@@ -922,8 +943,22 @@ double OpenCLCalcNonbondedForceKernel::execute(ContextImpl& context, bool includ
             cl.getQueue().enqueueMarkerWithWaitList(NULL, &events[0]);
             pmeQueue.enqueueBarrierWithWaitList(&events);
         }
-        if (hasOffsets)
-            energy = 0.0; // The Ewald self energy was computed in the kernel.
+        if (hasOffsets) {
+            // The Ewald self energy was computed in the kernel.
+
+            energy = 0.0;
+            if (pmeGrid1.isInitialized() || cosSinSums.isInitialized()) {
+                // Invoke a kernel to compute the correction for the neutralizing plasma.
+
+                mm_double4 boxSize = cl.getPeriodicBoxSizeDouble();
+                double volume = boxSize.x*boxSize.y*boxSize.z;
+                if (cl.getUseDoublePrecision())
+                    computePlasmaCorrectionKernel.setArg<double>(3, volume);
+                else
+                    computePlasmaCorrectionKernel.setArg<float>(3, volume);
+                cl.executeKernel(computePlasmaCorrectionKernel, OpenCLContext::ThreadBlockSize, OpenCLContext::ThreadBlockSize);
+            }
+        }
         recomputeParams = false;
     }
     
@@ -1163,10 +1198,12 @@ void OpenCLCalcNonbondedForceKernel::copyParametersToContext(ContextImpl& contex
         // Compute the self energy.
 
         ewaldSelfEnergy = 0.0;
+        totalCharge = 0.0;
         if (nonbondedMethod == Ewald || nonbondedMethod == PME || nonbondedMethod == LJPME) {
             if (cl.getContextIndex() == 0) {
                 for (int i = 0; i < force.getNumParticles(); i++) {
                     ewaldSelfEnergy -= baseParticleParamVec[i].x*baseParticleParamVec[i].x*ONE_4PI_EPS0*alpha/sqrt(M_PI);
+                    totalCharge += baseParticleParamVec[i].x;
                     if (doLJPME)
                         ewaldSelfEnergy += baseParticleParamVec[i].z*pow(baseParticleParamVec[i].y*dispersionAlpha, 6)/3.0;
                 }

platforms/reference/src/SimTKReference/ReferenceLJCoulombIxn.cpp

@@ -218,14 +218,19 @@ void ReferenceLJCoulombIxn::calculateEwaldIxn(int numberOfAtoms, vector<Vec3>& a
     // **************************************************************************************
 
     if (includeReciprocal) {
+        double totalCharge = 0.0;
         for (int atomID = 0; atomID < numberOfAtoms; atomID++) {
             double selfEwaldEnergy = ONE_4PI_EPS0*atomParameters[atomID][QIndex]*atomParameters[atomID][QIndex] * alphaEwald/SQRT_PI;
-            if(ljpme) {
+            totalCharge += atomParameters[atomID][QIndex];
+            if (ljpme) {
                 // Dispersion self term
                 selfEwaldEnergy -= pow(alphaDispersionEwald, 6.0) * 64.0*pow(atomParameters[atomID][SigIndex], 6.0) * pow(atomParameters[atomID][EpsIndex], 2.0) / 12.0;
             }
             totalSelfEwaldEnergy -= selfEwaldEnergy;
         }
+        // Correction for the neutralizing plasma.
+        double volume = periodicBoxVectors[0][0] * periodicBoxVectors[1][1] * periodicBoxVectors[2][2];
+        totalSelfEwaldEnergy -= totalCharge*totalCharge/(8*EPSILON0*volume*alphaEwald*alphaEwald);
     }
 
     if (totalEnergy) {

tests/TestEwald.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-2015 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2025 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -425,6 +425,51 @@ void testPMEParameters() {
     ASSERT(fabs((energy1-energy2)/energy1) > 1e-5);
 }
 
+void testNeutralizingPlasmaCorrection(NonbondedForce::NonbondedMethod method, bool includeOffset) {
+    // Verify that the energy of a system with nonzero charge doesn't depend on alpha.
+
+    System system;
+    NonbondedForce* force = new NonbondedForce();
+    force->setNonbondedMethod(method);
+    system.addForce(force);
+    for (int i = 0; i < 2; i++) {
+        system.addParticle(1.0);
+        force->addParticle(1.0, 0.3, 1.0);
+    }
+    if (includeOffset) {
+        force->addGlobalParameter("a", 0.1);
+        force->addParticleParameterOffset("a", 0, 1.0, 0.1, 0.1);
+    }
+    vector<Vec3> positions(2);
+    positions[0] = Vec3();
+    positions[1] = Vec3(0.3, 0.4, 0.0);
+
+    // Compute the energy.
+
+    VerletIntegrator integrator(0.01);
+    Context context(system, integrator, platform);
+    context.setPositions(positions);
+    double energy1 = context.getState(State::Energy).getPotentialEnergy();
+
+    // Change the cutoff distance, which will change alpha, and see if the energy is the same.
+
+    force->setCutoffDistance(0.7);
+    context.reinitialize(true);
+    double energy2 = context.getState(State::Energy).getPotentialEnergy();
+    ASSERT_EQUAL_TOL(energy1, energy2, 1e-4);
+
+    // Try changing a particle charge with updateParametersInContext() and make sure the
+    // energy changes by the correct amount.
+
+    force->setParticleParameters(0, 2.0, 0.3, 1.0);
+    force->updateParametersInContext(context);
+    double energy3 = context.getState(State::Energy).getPotentialEnergy();
+    force->setCutoffDistance(1.0);
+    context.reinitialize(true);
+    double energy4 = context.getState(State::Energy).getPotentialEnergy();
+    ASSERT_EQUAL_TOL(energy3, energy4, 1e-4);
+}
+
 void runPlatformTests();
 
 int main(int argc, char* argv[]) {
@@ -438,6 +483,11 @@ int main(int argc, char* argv[]) {
         testErrorTolerance(NonbondedForce::Ewald);
         testErrorTolerance(NonbondedForce::PME);
         testPMEParameters();
+        for (bool offset : {false, true}) {
+            testNeutralizingPlasmaCorrection(NonbondedForce::Ewald, offset);
+            testNeutralizingPlasmaCorrection(NonbondedForce::PME, offset);
+            testNeutralizingPlasmaCorrection(NonbondedForce::LJPME, offset);
+        }
         runPlatformTests();
     }
     catch(const exception& e) {