CpuKernels.cpp 82.4 KB
Newer Older
1
2
3
4
5
6
7
8
/* -------------------------------------------------------------------------- *
 *                                   OpenMM                                   *
 * -------------------------------------------------------------------------- *
 * This is part of the OpenMM molecular simulation toolkit originating from   *
 * Simbios, the NIH National Center for Physics-Based Simulation of           *
 * Biological Structures at Stanford, funded under the NIH Roadmap for        *
 * Medical Research, grant U54 GM072970. See https://simtk.org.               *
 *                                                                            *
9
 * Portions copyright (c) 2013-2025 Stanford University and the Authors.      *
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
 * Authors: Peter Eastman                                                     *
 * Contributors:                                                              *
 *                                                                            *
 * Permission is hereby granted, free of charge, to any person obtaining a    *
 * copy of this software and associated documentation files (the "Software"), *
 * to deal in the Software without restriction, including without limitation  *
 * the rights to use, copy, modify, merge, publish, distribute, sublicense,   *
 * and/or sell copies of the Software, and to permit persons to whom the      *
 * Software is furnished to do so, subject to the following conditions:       *
 *                                                                            *
 * The above copyright notice and this permission notice shall be included in *
 * all copies or substantial portions of the Software.                        *
 *                                                                            *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR *
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,   *
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL    *
 * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,    *
 * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR      *
 * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE  *
 * USE OR OTHER DEALINGS IN THE SOFTWARE.                                     *
 * -------------------------------------------------------------------------- */

#include "CpuKernels.h"
peastman's avatar
peastman committed
33
#include "ReferenceAngleBondIxn.h"
34
#include "ReferenceBondForce.h"
35
#include "ReferenceConstantPotential14.h"
36
#include "ReferenceConstraints.h"
37
38
#include "ReferenceKernelFactory.h"
#include "ReferenceKernels.h"
39
#include "ReferenceLJCoulomb14.h"
40
41
#include "ReferenceProperDihedralBond.h"
#include "ReferenceRbDihedralBond.h"
42
#include "ReferenceTabulatedFunction.h"
43
#include "openmm/Context.h"
44
#include "openmm/OpenMMException.h"
peastman's avatar
peastman committed
45
#include "openmm/Vec3.h"
46
47
#include "openmm/internal/ContextImpl.h"
#include "openmm/internal/NonbondedForceImpl.h"
48
#include "openmm/internal/ConstantPotentialForceImpl.h"
49
#include "openmm/internal/vectorize.h"
50
#include "lepton/CompiledExpression.h"
51
#include "lepton/CustomFunction.h"
52
#include "lepton/Operation.h"
53
#include "lepton/Parser.h"
54
#include <iostream>
55
#include "lepton/ParsedExpression.h"
56
57
58
59

using namespace OpenMM;
using namespace std;

peastman's avatar
peastman committed
60
static vector<Vec3>& extractPositions(ContextImpl& context) {
61
    ReferencePlatform::PlatformData* data = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());
62
    return *data->positions;
63
64
}

peastman's avatar
peastman committed
65
static vector<Vec3>& extractVelocities(ContextImpl& context) {
66
    ReferencePlatform::PlatformData* data = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());
67
    return *data->velocities;
68
69
}

peastman's avatar
peastman committed
70
static vector<Vec3>& extractForces(ContextImpl& context) {
71
    ReferencePlatform::PlatformData* data = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());
72
    return *data->forces;
73
74
}

peastman's avatar
peastman committed
75
static Vec3& extractBoxSize(ContextImpl& context) {
76
    ReferencePlatform::PlatformData* data = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());
77
    return *data->periodicBoxSize;
78
79
}

peastman's avatar
peastman committed
80
static Vec3* extractBoxVectors(ContextImpl& context) {
81
    ReferencePlatform::PlatformData* data = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());
82
    return data->periodicBoxVectors;
83
84
}

85
86
static ReferenceConstraints& extractConstraints(ContextImpl& context) {
    ReferencePlatform::PlatformData* data = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());
87
    return *data->constraints;
88
89
}

90
91
92
93
94
static const ReferenceVirtualSites& extractVirtualSites(ContextImpl& context) {
    ReferencePlatform::PlatformData* data = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());
    return *data->virtualSites;
}

95
96
static map<string, double>& extractEnergyParameterDerivatives(ContextImpl& context) {
    ReferencePlatform::PlatformData* data = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());
97
    return *data->energyParameterDerivatives;
98
99
}

100
101
102
103
104
105
106
/**
 * Make sure an expression doesn't use any undefined variables.
 */
static void validateVariables(const Lepton::ExpressionTreeNode& node, const set<string>& variables) {
    const Lepton::Operation& op = node.getOperation();
    if (op.getId() == Lepton::Operation::VARIABLE && variables.find(op.getName()) == variables.end())
        throw OpenMMException("Unknown variable in expression: "+op.getName());
peastman's avatar
peastman committed
107
108
    for (auto& child : node.getChildren())
        validateVariables(child, variables);
109
110
}

111
112
113
114
115
/**
 * Compute the kinetic energy of the system, possibly shifting the velocities in time to account
 * for a leapfrog integrator.
 */
static double computeShiftedKineticEnergy(ContextImpl& context, vector<double>& masses, double timeShift) {
peastman's avatar
peastman committed
116
117
118
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& velData = extractVelocities(context);
    vector<Vec3>& forceData = extractForces(context);
119
120
121
    int numParticles = context.getSystem().getNumParticles();
    
    // Compute the shifted velocities.
122

peastman's avatar
peastman committed
123
    vector<Vec3> shiftedVel(numParticles);
124
125
126
127
128
129
130
131
    for (int i = 0; i < numParticles; ++i) {
        if (masses[i] > 0)
            shiftedVel[i] = velData[i]+forceData[i]*(timeShift/masses[i]);
        else
            shiftedVel[i] = velData[i];
    }
    
    // Apply constraints to them.
132
133
134
135
136
137
138

    if (timeShift != 0) {
        vector<double> inverseMasses(numParticles);
        for (int i = 0; i < numParticles; i++)
            inverseMasses[i] = (masses[i] == 0 ? 0 : 1/masses[i]);
        extractConstraints(context).applyToVelocities(posData, shiftedVel, inverseMasses, 1e-4);
    }
139
140
141
142
143
144
145
146
147
148
    
    // Compute the kinetic energy.
    
    double energy = 0.0;
    for (int i = 0; i < numParticles; ++i)
        if (masses[i] > 0)
            energy += masses[i]*(shiftedVel[i].dot(shiftedVel[i]));
    return 0.5*energy;
}

149
150
151
152
153
154
155
156
157
158
159
160
161
/**
 * Copy particle charges into the fourth element of the posq array.
 */
static void copyChargesToPosq(ContextImpl& context, const vector<float>& charges, int index) {
    CpuPlatform::PlatformData& data = CpuPlatform::getPlatformData(context);
    if (index == data.currentPosqIndex)
        return;
    data.currentPosqIndex = index;
    AlignedArray<float>& posq = data.posq;
    for (int i = 0; i < charges.size(); i++)
        posq[4*i+3] = charges[i];
}

peastman's avatar
peastman committed
162
163
164
165
166
167
168
CpuCalcForcesAndEnergyKernel::CpuCalcForcesAndEnergyKernel(std::string name, const Platform& platform, CpuPlatform::PlatformData& data, ContextImpl& context) :
        CalcForcesAndEnergyKernel(name, platform), data(data) {
    // Create a Reference platform version of this kernel.
    
    ReferenceKernelFactory referenceFactory;
    referenceKernel = Kernel(referenceFactory.createKernelImpl(name, platform, context));
}
169

peastman's avatar
peastman committed
170
171
172
173
174
175
176
177
178
179
180
181
182
void CpuCalcForcesAndEnergyKernel::initialize(const System& system) {
    referenceKernel.getAs<ReferenceCalcForcesAndEnergyKernel>().initialize(system);
    lastPositions.resize(system.getNumParticles(), Vec3(1e10, 1e10, 1e10));
}

void CpuCalcForcesAndEnergyKernel::beginComputation(ContextImpl& context, bool includeForce, bool includeEnergy, int groups) {
    referenceKernel.getAs<ReferenceCalcForcesAndEnergyKernel>().beginComputation(context, includeForce, includeEnergy, groups);
    
    // Convert positions to single precision and clear the forces.

    int numParticles = context.getSystem().getNumParticles();
    bool positionsValid = true;
    data.threads.execute([&] (ThreadPool& threads, int threadIndex) {
183
184
185
        // Convert the positions to single precision and apply periodic boundary conditions

        AlignedArray<float>& posq = data.posq;
peastman's avatar
peastman committed
186
187
        vector<Vec3>& posData = extractPositions(context);
        Vec3* boxVectors = extractBoxVectors(context);
188
189
190
        double boxSize[3] = {boxVectors[0][0], boxVectors[1][1], boxVectors[2][2]};
        double invBoxSize[3] = {1/boxVectors[0][0], 1/boxVectors[1][1], 1/boxVectors[2][2]};
        bool triclinic = (boxVectors[0][1] != 0 || boxVectors[0][2] != 0 || boxVectors[1][0] != 0 || boxVectors[1][2] != 0 || boxVectors[2][0] != 0 || boxVectors[2][1] != 0);
191
192
193
194
        int numParticles = context.getSystem().getNumParticles();
        int numThreads = threads.getNumThreads();
        int start = threadIndex*numParticles/numThreads;
        int end = (threadIndex+1)*numParticles/numThreads;
195
196
197
        if (data.isPeriodic) {
            if (triclinic) {
                for (int i = start; i < end; i++) {
peastman's avatar
peastman committed
198
                    Vec3 pos = posData[i];
199
200
201
202
203
204
                    pos -= boxVectors[2]*floor(pos[2]*invBoxSize[2]);
                    pos -= boxVectors[1]*floor(pos[1]*invBoxSize[1]);
                    pos -= boxVectors[0]*floor(pos[0]*invBoxSize[0]);
                    posq[4*i] = (float) pos[0];
                    posq[4*i+1] = (float) pos[1];
                    posq[4*i+2] = (float) pos[2];
205
                }
206
207
208
209
            }
            else {
                for (int i = start; i < end; i++) {
                    for (int j = 0; j < 3; j++) {
peastman's avatar
peastman committed
210
                        double x = posData[i][j];
211
212
213
                        double base = floor(x*invBoxSize[j])*boxSize[j];
                        posq[4*i+j] = (float) (x-base);
                    }
214
                }
215
216
            }
        }
217
218
219
220
221
222
        else
            for (int i = start; i < end; i++) {
                posq[4*i] = (float) posData[i][0];
                posq[4*i+1] = (float) posData[i][1];
                posq[4*i+2] = (float) posData[i][2];
            }
223
224
225
        
        // Check for invalid positions.
        
226
227
        for (int i = 4*start; i < 4*end; i += 4)
            if (posq[i] != posq[i] || posq[i+1] != posq[i+1] || posq[i+2] != posq[i+2])
228
                positionsValid = false;
229
230
231
232
233
234

        // Clear the forces.

        fvec4 zero(0.0f);
        for (int j = 0; j < numParticles; j++)
            zero.store(&data.threadForce[threadIndex][j*4]);
peastman's avatar
peastman committed
235
    });
236
    data.threads.waitForThreads();
peastman's avatar
peastman committed
237
    if (!positionsValid)
238
        throw OpenMMException("Particle coordinate is NaN.  For more information, see https://github.com/openmm/openmm/wiki/Frequently-Asked-Questions#nan");
239
240
241

    // Determine whether we need to recompute the neighbor list.
        
242
    if (data.neighborList != NULL && data.cutoff > 0.0) {
243
244
245
246
247
248
        double padding = data.paddedCutoff-data.cutoff;;
        bool needRecompute = false;
        double closeCutoff2 = 0.25*padding*padding;
        double farCutoff2 = 0.5*padding*padding;
        int maxNumMoved = numParticles/10;
        vector<int> moved;
peastman's avatar
peastman committed
249
        vector<Vec3>& posData = extractPositions(context);
250
        for (int i = 0; i < numParticles; i++) {
peastman's avatar
peastman committed
251
            Vec3 delta = posData[i]-lastPositions[i];
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
            double dist2 = delta.dot(delta);
            if (dist2 > closeCutoff2) {
                moved.push_back(i);
                if (dist2 > farCutoff2 || moved.size() > maxNumMoved) {
                    needRecompute = true;
                    break;
                }
            }
        }
        if (!needRecompute && moved.size() > 0) {
            // Some particles have moved further than half the padding distance.  Look for pairs
            // that are missing from the neighbor list.

            int numMoved = moved.size();
            double cutoff2 = data.cutoff*data.cutoff;
            double paddedCutoff2 = data.paddedCutoff*data.paddedCutoff;
            for (int i = 1; i < numMoved && !needRecompute; i++)
                for (int j = 0; j < i; j++) {
peastman's avatar
peastman committed
270
                    Vec3 delta = posData[moved[i]]-posData[moved[j]];
271
272
273
                    if (delta.dot(delta) < cutoff2) {
                        // These particles should interact.  See if they are in the neighbor list.
                        
peastman's avatar
peastman committed
274
                        Vec3 oldDelta = lastPositions[moved[i]]-lastPositions[moved[j]];
275
276
277
278
279
280
281
282
283
284
285
286
                        if (oldDelta.dot(oldDelta) > paddedCutoff2) {
                            needRecompute = true;
                            break;
                        }
                    }
                }
        }
        if (needRecompute) {
            data.neighborList->computeNeighborList(numParticles, data.posq, data.exclusions, extractBoxVectors(context), data.isPeriodic, data.paddedCutoff, data.threads);
            lastPositions = posData;
        }
    }
287
288
}

289
double CpuCalcForcesAndEnergyKernel::finishComputation(ContextImpl& context, bool includeForce, bool includeEnergy, int groups, bool& valid) {
290
291
    // Sum the forces from all the threads.
    
peastman's avatar
peastman committed
292
293
294
295
296
297
298
    data.threads.execute([&] (ThreadPool& threads, int threadIndex) {
        // Sum the contributions to forces that have been calculated by different threads.
        
        int numParticles = context.getSystem().getNumParticles();
        int numThreads = threads.getNumThreads();
        int start = threadIndex*numParticles/numThreads;
        int end = (threadIndex+1)*numParticles/numThreads;
peastman's avatar
peastman committed
299
        vector<Vec3>& forceData = extractForces(context);
peastman's avatar
peastman committed
300
301
302
303
304
305
306
307
308
        for (int i = start; i < end; i++) {
            fvec4 f(0.0f);
            for (int j = 0; j < numThreads; j++)
                f += fvec4(&data.threadForce[j][4*i]);
            forceData[i][0] += f[0];
            forceData[i][1] += f[1];
            forceData[i][2] += f[2];
        }
    });
309
    data.threads.waitForThreads();
310
    return referenceKernel.getAs<ReferenceCalcForcesAndEnergyKernel>().finishComputation(context, includeForce, includeEnergy, groups, valid);
311
312
}

313
314
315
316
317
318
319
320
321
322
void CpuUpdateStateDataKernel::createCheckpoint(ContextImpl& context, ostream& stream) {
    ReferenceUpdateStateDataKernel::createCheckpoint(context, stream);
    data.random.createCheckpoint(stream);
}

void CpuUpdateStateDataKernel::loadCheckpoint(ContextImpl& context, istream& stream) {
    ReferenceUpdateStateDataKernel::loadCheckpoint(context, stream);
    data.random.loadCheckpoint(stream);
}

peastman's avatar
peastman committed
323
324
void CpuCalcHarmonicAngleForceKernel::initialize(const System& system, const HarmonicAngleForce& force) {
    numAngles = force.getNumAngles();
325
326
    angleIndexArray.resize(numAngles, vector<int>(3));
    angleParamArray.resize(numAngles, vector<double>(2));
peastman's avatar
peastman committed
327
328
329
330
331
332
333
    for (int i = 0; i < numAngles; ++i) {
        int particle1, particle2, particle3;
        double angle, k;
        force.getAngleParameters(i, particle1, particle2, particle3, angle, k);
        angleIndexArray[i][0] = particle1;
        angleIndexArray[i][1] = particle2;
        angleIndexArray[i][2] = particle3;
peastman's avatar
peastman committed
334
335
        angleParamArray[i][0] = angle;
        angleParamArray[i][1] = k;
peastman's avatar
peastman committed
336
337
    }
    bondForce.initialize(system.getNumParticles(), numAngles, 3, angleIndexArray, data.threads);
338
    usePeriodic = force.usesPeriodicBoundaryConditions();
peastman's avatar
peastman committed
339
340
341
}

double CpuCalcHarmonicAngleForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
342
343
344
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
peastman's avatar
peastman committed
345
    ReferenceAngleBondIxn angleBond;
346
347
    if (usePeriodic)
        angleBond.setPeriodic(extractBoxVectors(context));
peastman's avatar
peastman committed
348
349
350
351
    bondForce.calculateForce(posData, angleParamArray, forceData, includeEnergy ? &energy : NULL, angleBond);
    return energy;
}

352
void CpuCalcHarmonicAngleForceKernel::copyParametersToContext(ContextImpl& context, const HarmonicAngleForce& force, int firstAngle, int lastAngle) {
peastman's avatar
peastman committed
353
354
355
356
357
    if (numAngles != force.getNumAngles())
        throw OpenMMException("updateParametersInContext: The number of angles has changed");

    // Record the values.

358
    for (int i = firstAngle; i <= lastAngle; ++i) {
peastman's avatar
peastman committed
359
360
361
362
363
        int particle1, particle2, particle3;
        double angle, k;
        force.getAngleParameters(i, particle1, particle2, particle3, angle, k);
        if (particle1 != angleIndexArray[i][0] || particle2 != angleIndexArray[i][1] || particle3 != angleIndexArray[i][2])
            throw OpenMMException("updateParametersInContext: The set of particles in an angle has changed");
peastman's avatar
peastman committed
364
365
        angleParamArray[i][0] = angle;
        angleParamArray[i][1] = k;
peastman's avatar
peastman committed
366
367
368
    }
}

369
370
void CpuCalcPeriodicTorsionForceKernel::initialize(const System& system, const PeriodicTorsionForce& force) {
    numTorsions = force.getNumTorsions();
371
372
    torsionIndexArray.resize(numTorsions, vector<int>(4));
    torsionParamArray.resize(numTorsions, vector<double>(3));
373
374
375
376
377
378
379
380
    for (int i = 0; i < numTorsions; ++i) {
        int particle1, particle2, particle3, particle4, periodicity;
        double phase, k;
        force.getTorsionParameters(i, particle1, particle2, particle3, particle4, periodicity, phase, k);
        torsionIndexArray[i][0] = particle1;
        torsionIndexArray[i][1] = particle2;
        torsionIndexArray[i][2] = particle3;
        torsionIndexArray[i][3] = particle4;
peastman's avatar
peastman committed
381
382
383
        torsionParamArray[i][0] = k;
        torsionParamArray[i][1] = phase;
        torsionParamArray[i][2] = periodicity;
384
385
    }
    bondForce.initialize(system.getNumParticles(), numTorsions, 4, torsionIndexArray, data.threads);
386
    usePeriodic = force.usesPeriodicBoundaryConditions();
387
388
389
}

double CpuCalcPeriodicTorsionForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
390
391
392
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
393
    ReferenceProperDihedralBond periodicTorsionBond;
394
395
    if (usePeriodic)
        periodicTorsionBond.setPeriodic(extractBoxVectors(context));
396
397
398
399
    bondForce.calculateForce(posData, torsionParamArray, forceData, includeEnergy ? &energy : NULL, periodicTorsionBond);
    return energy;
}

400
void CpuCalcPeriodicTorsionForceKernel::copyParametersToContext(ContextImpl& context, const PeriodicTorsionForce& force, int firstTorsion, int lastTorsion) {
401
402
403
404
405
    if (numTorsions != force.getNumTorsions())
        throw OpenMMException("updateParametersInContext: The number of torsions has changed");

    // Record the values.

406
    for (int i = firstTorsion; i <= lastTorsion; ++i) {
407
408
409
410
411
        int particle1, particle2, particle3, particle4, periodicity;
        double phase, k;
        force.getTorsionParameters(i, particle1, particle2, particle3, particle4, periodicity, phase, k);
        if (particle1 != torsionIndexArray[i][0] || particle2 != torsionIndexArray[i][1] || particle3 != torsionIndexArray[i][2] || particle4 != torsionIndexArray[i][3])
            throw OpenMMException("updateParametersInContext: The set of particles in a torsion has changed");
peastman's avatar
peastman committed
412
413
414
        torsionParamArray[i][0] = k;
        torsionParamArray[i][1] = phase;
        torsionParamArray[i][2] = periodicity;
415
416
417
418
419
    }
}

void CpuCalcRBTorsionForceKernel::initialize(const System& system, const RBTorsionForce& force) {
    numTorsions = force.getNumTorsions();
420
421
    torsionIndexArray.resize(numTorsions, vector<int>(4));
    torsionParamArray.resize(numTorsions, vector<double>(6));
422
423
424
425
426
427
428
429
    for (int i = 0; i < numTorsions; ++i) {
        int particle1, particle2, particle3, particle4;
        double c0, c1, c2, c3, c4, c5;
        force.getTorsionParameters(i, particle1, particle2, particle3, particle4, c0, c1, c2, c3, c4, c5);
        torsionIndexArray[i][0] = particle1;
        torsionIndexArray[i][1] = particle2;
        torsionIndexArray[i][2] = particle3;
        torsionIndexArray[i][3] = particle4;
peastman's avatar
peastman committed
430
431
432
433
434
435
        torsionParamArray[i][0] = c0;
        torsionParamArray[i][1] = c1;
        torsionParamArray[i][2] = c2;
        torsionParamArray[i][3] = c3;
        torsionParamArray[i][4] = c4;
        torsionParamArray[i][5] = c5;
436
437
    }
    bondForce.initialize(system.getNumParticles(), numTorsions, 4, torsionIndexArray, data.threads);
438
    usePeriodic = force.usesPeriodicBoundaryConditions();
439
440
441
}

double CpuCalcRBTorsionForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
442
443
444
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
445
    ReferenceRbDihedralBond rbTorsionBond;
446
447
    if (usePeriodic)
        rbTorsionBond.setPeriodic(extractBoxVectors(context));
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
    bondForce.calculateForce(posData, torsionParamArray, forceData, includeEnergy ? &energy : NULL, rbTorsionBond);
    return energy;
}

void CpuCalcRBTorsionForceKernel::copyParametersToContext(ContextImpl& context, const RBTorsionForce& force) {
    if (numTorsions != force.getNumTorsions())
        throw OpenMMException("updateParametersInContext: The number of torsions has changed");

    // Record the values.

    for (int i = 0; i < numTorsions; ++i) {
        int particle1, particle2, particle3, particle4;
        double c0, c1, c2, c3, c4, c5;
        force.getTorsionParameters(i, particle1, particle2, particle3, particle4, c0, c1, c2, c3, c4, c5);
        if (particle1 != torsionIndexArray[i][0] || particle2 != torsionIndexArray[i][1] || particle3 != torsionIndexArray[i][2] || particle4 != torsionIndexArray[i][3])
            throw OpenMMException("updateParametersInContext: The set of particles in a torsion has changed");
peastman's avatar
peastman committed
464
465
466
467
468
469
        torsionParamArray[i][0] = c0;
        torsionParamArray[i][1] = c1;
        torsionParamArray[i][2] = c2;
        torsionParamArray[i][3] = c3;
        torsionParamArray[i][4] = c4;
        torsionParamArray[i][5] = c5;
470
471
472
    }
}

473
474
475
476
477
478
479
480
481
482
483
484
485
486
class CpuCalcNonbondedForceKernel::PmeIO : public CalcPmeReciprocalForceKernel::IO {
public:
    PmeIO(float* posq, float* force, int numParticles) : posq(posq), force(force), numParticles(numParticles) {
    }
    float* getPosq() {
        return posq;
    }
    void setForce(float* f) {
        for (int i = 0; i < numParticles; i++) {
            force[4*i] += f[4*i];
            force[4*i+1] += f[4*i+1];
            force[4*i+2] += f[4*i+2];
        }
    }
487
488
    void setChargeDerivatives(float* derivatives) {
    }
489
490
491
492
493
494
private:
    float* posq;
    float* force;
    int numParticles;
};

495
CpuNonbondedForce* createCpuNonbondedForceVec(const CpuNeighborList& neighbors);
496
497

CpuCalcNonbondedForceKernel::CpuCalcNonbondedForceKernel(string name, const Platform& platform, CpuPlatform::PlatformData& data) : CalcNonbondedForceKernel(name, platform),
498
        data(data), hasInitializedPme(false), hasInitializedDispersionPme(false), nonbonded(NULL) {
499
500
}

501
CpuCalcNonbondedForceKernel::~CpuCalcNonbondedForceKernel() {
502
503
    if (nonbonded != NULL)
        delete nonbonded;
504
505
506
}

void CpuCalcNonbondedForceKernel::initialize(const System& system, const NonbondedForce& force) {
507
508
    chargePosqIndex = data.requestPosqIndex();
    ljPosqIndex = data.requestPosqIndex();
509
510
511

    // Identify which exceptions are 1-4 interactions.

512
513
514
515
516
517
518
519
    set<int> exceptionsWithOffsets;
    for (int i = 0; i < force.getNumExceptionParameterOffsets(); i++) {
        string param;
        int exception;
        double charge, sigma, epsilon;
        force.getExceptionParameterOffset(i, param, exception, charge, sigma, epsilon);
        exceptionsWithOffsets.insert(exception);
    }
520
521
522
523
524
525
526
527
528
    numParticles = force.getNumParticles();
    exclusions.resize(numParticles);
    vector<int> nb14s;
    for (int i = 0; i < force.getNumExceptions(); i++) {
        int particle1, particle2;
        double chargeProd, sigma, epsilon;
        force.getExceptionParameters(i, particle1, particle2, chargeProd, sigma, epsilon);
        exclusions[particle1].insert(particle2);
        exclusions[particle2].insert(particle1);
529
530
        if (chargeProd != 0.0 || epsilon != 0.0 || exceptionsWithOffsets.find(i) != exceptionsWithOffsets.end()) {
            nb14Index[i] = nb14s.size();
531
            nb14s.push_back(i);
532
        }
533
534
535
536
537
    }

    // Record the particle parameters.

    num14 = nb14s.size();
538
539
    bonded14IndexArray.resize(num14, vector<int>(2));
    bonded14ParamArray.resize(num14, vector<double>(3));
peastman's avatar
peastman committed
540
    particleParams.resize(numParticles);
541
    charges.resize(numParticles);
542
    C6params.resize(numParticles);
543
544
545
    baseParticleParams.resize(numParticles);
    for (int i = 0; i < numParticles; ++i)
       force.getParticleParameters(i, baseParticleParams[i][0], baseParticleParams[i][1], baseParticleParams[i][2]);
546
    
547
    // Record exception parameters.
548
    
549
    baseExceptionParams.resize(num14);
550
551
    for (int i = 0; i < num14; ++i) {
        int particle1, particle2;
552
        force.getExceptionParameters(nb14s[i], particle1, particle2, baseExceptionParams[i][0], baseExceptionParams[i][1], baseExceptionParams[i][2]);
553
554
        bonded14IndexArray[i][0] = particle1;
        bonded14IndexArray[i][1] = particle2;
555
    }
peastman's avatar
peastman committed
556
    bondForce.initialize(system.getNumParticles(), num14, 2, bonded14IndexArray, data.threads);
557
    
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
    // Record information about parameter offsets.
    
    hasParticleOffsets = (force.getNumParticleParameterOffsets() > 0);
    hasExceptionOffsets = (force.getNumExceptionParameterOffsets() > 0);
    particleParamOffsets.resize(force.getNumParticles());
    exceptionParamOffsets.resize(force.getNumExceptions());
    for (int i = 0; i < force.getNumParticleParameterOffsets(); i++) {
        string param;
        int particle;
        double charge, sigma, epsilon;
        force.getParticleParameterOffset(i, param, particle, charge, sigma, epsilon);
        auto paramPos = find(paramNames.begin(), paramNames.end(), param);
        int paramIndex;
        if (paramPos == paramNames.end()) {
            paramIndex = paramNames.size();
            paramNames.push_back(param);
        }
        else
            paramIndex = paramPos-paramNames.begin();
        particleParamOffsets[particle].push_back(make_tuple(charge, sigma, epsilon, paramIndex));
    }
    for (int i = 0; i < force.getNumExceptionParameterOffsets(); i++) {
        string param;
        int exception;
        double charge, sigma, epsilon;
        force.getExceptionParameterOffset(i, param, exception, charge, sigma, epsilon);
        auto paramPos = find(paramNames.begin(), paramNames.end(), param);
        int paramIndex;
        if (paramPos == paramNames.end()) {
            paramIndex = paramNames.size();
            paramNames.push_back(param);
        }
        else
            paramIndex = paramPos-paramNames.begin();
592
        exceptionParamOffsets[nb14Index[exception]].push_back(make_tuple(charge, sigma, epsilon, paramIndex));
593
    }
594
    paramValues.resize(paramNames.size(), 0.0);
595

596
597
    // Record other parameters.
    
598
599
    nonbondedMethod = CalcNonbondedForceKernel::NonbondedMethod(force.getNonbondedMethod());
    nonbondedCutoff = force.getCutoffDistance();
600
601
    if (nonbondedMethod == NoCutoff) {
        data.requestNeighborList(0.0, 0.0, true, exclusions);
602
        useSwitchingFunction = false;
603
    }
604
    else {
605
        data.requestNeighborList(nonbondedCutoff, 0.25*nonbondedCutoff, true, exclusions);
606
607
608
609
610
611
612
613
614
615
        useSwitchingFunction = force.getUseSwitchingFunction();
        switchingDistance = force.getSwitchingDistance();
    }
    if (nonbondedMethod == Ewald) {
        double alpha;
        NonbondedForceImpl::calcEwaldParameters(system, force, alpha, kmax[0], kmax[1], kmax[2]);
        ewaldAlpha = alpha;
    }
    else if (nonbondedMethod == PME) {
        double alpha;
616
        NonbondedForceImpl::calcPMEParameters(system, force, alpha, gridSize[0], gridSize[1], gridSize[2], false);
617
618
        ewaldAlpha = alpha;
    }
619
620
621
    else if (nonbondedMethod == LJPME) {
        double alpha;
        NonbondedForceImpl::calcPMEParameters(system, force, alpha, gridSize[0], gridSize[1], gridSize[2], false);
peastman's avatar
peastman committed
622
        ewaldAlpha = alpha;
623
        NonbondedForceImpl::calcPMEParameters(system, force, alpha, dispersionGridSize[0], dispersionGridSize[1], dispersionGridSize[2], true);
peastman's avatar
peastman committed
624
        ewaldDispersionAlpha = alpha;
625
626
        useSwitchingFunction = false;
    }
627
628
629
630
    if (nonbondedMethod == NoCutoff || nonbondedMethod == CutoffNonPeriodic)
        exceptionsArePeriodic = false;
    else
        exceptionsArePeriodic = force.getExceptionsUsePeriodicBoundaryConditions();
631
632
633
634
635
    rfDielectric = force.getReactionFieldDielectric();
    if (force.getUseDispersionCorrection())
        dispersionCoefficient = NonbondedForceImpl::calcDispersionCorrection(system, force);
    else
        dispersionCoefficient = 0.0;
636
    data.isPeriodic |= (nonbondedMethod == CutoffPeriodic || nonbondedMethod == Ewald || nonbondedMethod == PME || nonbondedMethod == LJPME);
637
    nonbonded = createCpuNonbondedForceVec(*data.neighborList);
638
639
640
}

double CpuCalcNonbondedForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy, bool includeDirect, bool includeReciprocal) {
641
642
643
    if (!hasInitializedPme) {
        hasInitializedPme = true;
        useOptimizedPme = false;
644
        computeParameters(context, false);
645
646
647
        if (nonbondedMethod == PME) {
            // If available, use the optimized PME implementation.

peastman's avatar
peastman committed
648
649
650
651
            vector<string> kernelNames;
            kernelNames.push_back("CalcPmeReciprocalForce");
            useOptimizedPme = getPlatform().supportsKernels(kernelNames);
            if (useOptimizedPme) {
652
                optimizedPme = getPlatform().createKernel(CalcPmeReciprocalForceKernel::Name(), context);
653
                optimizedPme.getAs<CalcPmeReciprocalForceKernel>().initialize(gridSize[0], gridSize[1], gridSize[2], numParticles, {}, ewaldAlpha, data.deterministicForces);
654
655
            }
        }
656
657
658
659
660
        if (nonbondedMethod == LJPME) {
            // If available, use the optimized PME implementation.

            vector<string> kernelNames;
            kernelNames.push_back("CalcPmeReciprocalForce");
661
            kernelNames.push_back("CalcDispersionPmeReciprocalForce");
662
663
664
            useOptimizedPme = getPlatform().supportsKernels(kernelNames);
            if (useOptimizedPme) {
                optimizedPme = getPlatform().createKernel(CalcPmeReciprocalForceKernel::Name(), context);
665
                optimizedPme.getAs<CalcPmeReciprocalForceKernel>().initialize(gridSize[0], gridSize[1], gridSize[2], numParticles, {}, ewaldAlpha, data.deterministicForces);
666
667
                optimizedDispersionPme = getPlatform().createKernel(CalcDispersionPmeReciprocalForceKernel::Name(), context);
                optimizedDispersionPme.getAs<CalcDispersionPmeReciprocalForceKernel>().initialize(dispersionGridSize[0], dispersionGridSize[1],
668
                                                                                                  dispersionGridSize[2], numParticles, ewaldDispersionAlpha, data.deterministicForces);
669
670
            }
        }
671
    }
672
    computeParameters(context, true);
673
    copyChargesToPosq(context, charges, chargePosqIndex);
674
    AlignedArray<float>& posq = data.posq;
peastman's avatar
peastman committed
675
676
677
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    Vec3* boxVectors = extractBoxVectors(context);
678
    double energy = (includeReciprocal ? ewaldSelfEnergy : 0.0);
679
680
    bool ewald  = (nonbondedMethod == Ewald);
    bool pme  = (nonbondedMethod == PME);
681
    bool ljpme = (nonbondedMethod == LJPME);
682
    if (nonbondedMethod != NoCutoff)
683
        nonbonded->setUseCutoff(nonbondedCutoff, rfDielectric);
684
    if (data.isPeriodic) {
peastman's avatar
peastman committed
685
        Vec3* boxVectors = extractBoxVectors(context);
686
        double minAllowedSize = 1.999999*nonbondedCutoff;
687
        if (boxVectors[0][0] < minAllowedSize || boxVectors[1][1] < minAllowedSize || boxVectors[2][2] < minAllowedSize)
688
            throw OpenMMException("The periodic box size has decreased to less than twice the nonbonded cutoff.");
689
        nonbonded->setPeriodic(boxVectors);
690
        nonbonded->setPeriodicExceptions(exceptionsArePeriodic);
691
692
    }
    if (ewald)
693
        nonbonded->setUseEwald(ewaldAlpha, kmax[0], kmax[1], kmax[2]);
694
    if (pme)
695
        nonbonded->setUsePME(ewaldAlpha, gridSize);
696
    if (useSwitchingFunction)
697
        nonbonded->setUseSwitchingFunction(switchingDistance);
698
699
700
701
    if (ljpme){
        nonbonded->setUsePME(ewaldAlpha, gridSize);
        nonbonded->setUseLJPME(ewaldDispersionAlpha, dispersionGridSize);
    }
702
    double nonbondedEnergy = 0;
peastman's avatar
peastman committed
703
    if (includeDirect)
704
        nonbonded->calculateDirectIxn(numParticles, &posq[0], posData, particleParams, C6params, exclusions, data.threadForce, includeEnergy ? &nonbondedEnergy : NULL, data.threads);
705
706
    if (includeReciprocal) {
        if (useOptimizedPme) {
707
            PmeIO io(&posq[0], &data.threadForce[0][0], numParticles);
peastman's avatar
peastman committed
708
            Vec3 periodicBoxVectors[3] = {boxVectors[0], boxVectors[1], boxVectors[2]};
709
            optimizedPme.getAs<CalcPmeReciprocalForceKernel>().beginComputation(io, periodicBoxVectors, includeEnergy, true, false);
710
            nonbondedEnergy += optimizedPme.getAs<CalcPmeReciprocalForceKernel>().finishComputation(io);
711
712
713
714
715
            if (nonbondedMethod == LJPME) {
                copyChargesToPosq(context, C6params, ljPosqIndex);
                optimizedDispersionPme.getAs<CalcDispersionPmeReciprocalForceKernel>().beginComputation(io, periodicBoxVectors, includeEnergy);
                nonbondedEnergy += optimizedDispersionPme.getAs<CalcDispersionPmeReciprocalForceKernel>().finishComputation(io);
            }
716
717
        }
        else
718
            nonbonded->calculateReciprocalIxn(numParticles, &posq[0], posData, particleParams, C6params, exclusions, forceData, includeEnergy ? &nonbondedEnergy : NULL);
719
720
721
722
723
724
        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);
        }
725
    }
peastman's avatar
peastman committed
726
    energy += nonbondedEnergy;
727
728
    if (includeDirect) {
        ReferenceLJCoulomb14 nonbonded14;
729
730
731
732
        if (exceptionsArePeriodic) {
            Vec3* boxVectors = extractBoxVectors(context);
            nonbonded14.setPeriodic(boxVectors);
        }
peastman's avatar
peastman committed
733
        bondForce.calculateForce(posData, bonded14ParamArray, forceData, includeEnergy ? &energy : NULL, nonbonded14);
734
        if (data.isPeriodic && nonbondedMethod != LJPME)
735
            energy += dispersionCoefficient/(boxVectors[0][0]*boxVectors[1][1]*boxVectors[2][2]);
736
737
    }
    return energy;
738
739
}

740
void CpuCalcNonbondedForceKernel::copyParametersToContext(ContextImpl& context, const NonbondedForce& force, int firstParticle, int lastParticle, int firstException, int lastException) {
741
742
    if (force.getNumParticles() != numParticles)
        throw OpenMMException("updateParametersInContext: The number of particles has changed");
peastman's avatar
peastman committed
743
744
745
746
747
748
749
750
751
752
753

    // Identify which exceptions are 1-4 interactions.

    set<int> exceptionsWithOffsets;
    for (int i = 0; i < force.getNumExceptionParameterOffsets(); i++) {
        string param;
        int exception;
        double charge, sigma, epsilon;
        force.getExceptionParameterOffset(i, param, exception, charge, sigma, epsilon);
        exceptionsWithOffsets.insert(exception);
    }
754
755
756
757
758
    vector<int> nb14s;
    for (int i = 0; i < force.getNumExceptions(); i++) {
        int particle1, particle2;
        double chargeProd, sigma, epsilon;
        force.getExceptionParameters(i, particle1, particle2, chargeProd, sigma, epsilon);
759
760
761
762
763
        if (nb14Index.find(i) == nb14Index.end()) {
            if (chargeProd != 0.0 || epsilon != 0.0 || exceptionsWithOffsets.find(i) != exceptionsWithOffsets.end())
                throw OpenMMException("updateParametersInContext: The set of non-excluded exceptions has changed");
        }
        else
764
765
766
767
768
769
770
            nb14s.push_back(i);
    }
    if (nb14s.size() != num14)
        throw OpenMMException("updateParametersInContext: The number of non-excluded exceptions has changed");

    // Record the values.

771
    for (int i = firstParticle; i <= lastParticle; ++i)
772
       force.getParticleParameters(i, baseParticleParams[i][0], baseParticleParams[i][1], baseParticleParams[i][2]);
773
774
    for (int i = 0; i < num14; ++i) {
        int particle1, particle2;
775
        force.getExceptionParameters(nb14s[i], particle1, particle2, baseExceptionParams[i][0], baseExceptionParams[i][1], baseExceptionParams[i][2]);
776
777
778
        bonded14IndexArray[i][0] = particle1;
        bonded14IndexArray[i][1] = particle2;
    }
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
    particleParamOffsets.clear();
    exceptionParamOffsets.clear();
    particleParamOffsets.resize(force.getNumParticles());
    exceptionParamOffsets.resize(force.getNumExceptions());
    for (int i = 0; i < force.getNumParticleParameterOffsets(); i++) {
        string param;
        int particle;
        double charge, sigma, epsilon;
        force.getParticleParameterOffset(i, param, particle, charge, sigma, epsilon);
        auto paramPos = find(paramNames.begin(), paramNames.end(), param);
        if (paramPos == paramNames.end())
            throw OpenMMException("updateParametersInContext: The parameter of a particle parameter offset has changed");
        int paramIndex = paramPos-paramNames.begin();
        particleParamOffsets[particle].push_back(make_tuple(charge, sigma, epsilon, paramIndex));
    }
    for (int i = 0; i < force.getNumExceptionParameterOffsets(); i++) {
        string param;
        int exception;
        double charge, sigma, epsilon;
        force.getExceptionParameterOffset(i, param, exception, charge, sigma, epsilon);
        auto paramPos = find(paramNames.begin(), paramNames.end(), param);
        if (paramPos == paramNames.end())
            throw OpenMMException("updateParametersInContext: The parameter of an exception parameter offset has changed");
        int paramIndex = paramPos-paramNames.begin();
        exceptionParamOffsets[nb14Index[exception]].push_back(make_tuple(charge, sigma, epsilon, paramIndex));
    }
805
    computeParameters(context, false);
806
807
808
809
810
811
    
    // Recompute the coefficient for the dispersion correction.

    NonbondedForce::NonbondedMethod method = force.getNonbondedMethod();
    if (force.getUseDispersionCorrection() && (method == NonbondedForce::CutoffPeriodic || method == NonbondedForce::Ewald || method == NonbondedForce::PME))
        dispersionCoefficient = NonbondedForceImpl::calcDispersionCorrection(context.getSystem(), force);
812
}
813

814
void CpuCalcNonbondedForceKernel::getPMEParameters(double& alpha, int& nx, int& ny, int& nz) const {
815
    if (nonbondedMethod != PME && nonbondedMethod != LJPME)
816
817
818
819
820
821
822
823
824
825
826
        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];
    }
}

827
void CpuCalcNonbondedForceKernel::getLJPMEParameters(double& alpha, int& nx, int& ny, int& nz) const {
828
829
830
    if (nonbondedMethod != LJPME)
        throw OpenMMException("getPMEParametersInContext: This Context is not using PME");
    if (useOptimizedPme)
831
        optimizedDispersionPme.getAs<const CalcPmeReciprocalForceKernel>().getPMEParameters(alpha, nx, ny, nz);
832
    else {
833
834
835
836
        alpha = ewaldDispersionAlpha;
        nx = dispersionGridSize[0];
        ny = dispersionGridSize[1];
        nz = dispersionGridSize[2];
837
838
839
    }
}

840
void CpuCalcNonbondedForceKernel::computeParameters(ContextImpl& context, bool offsetsOnly) {
841
842
843
844
845
846
847
848
849
850
    bool paramChanged = false;
    for (int i = 0; i < paramNames.size(); i++) {
        double value = context.getParameter(paramNames[i]);
        if (value != paramValues[i]) {
            paramValues[i] = value;;
            paramChanged = true;
        }
    }
    if (!paramChanged && offsetsOnly)
        return;
851
852
853
854
855

    // Compute particle parameters.

    if (hasParticleOffsets || !offsetsOnly) {
        double sumSquaredCharges = 0.0;
856
        totalCharge = 0.0;
857
858
859
860
861
862
863
864
865
866
867
868
869
870
        for (int i = 0; i < numParticles; i++) {
            double charge = baseParticleParams[i][0];
            double sigma = baseParticleParams[i][1];
            double epsilon = baseParticleParams[i][2];
            for (auto& offset : particleParamOffsets[i]) {
                double value = paramValues[get<3>(offset)];
                charge += value*get<0>(offset);
                sigma += value*get<1>(offset);
                epsilon += value*get<2>(offset);
            }
            charges[i] = (float) charge;
            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;
871
            totalCharge += charge;
872
873
874
875
876
877
878
879
880
        }
        if (nonbondedMethod == Ewald || nonbondedMethod == PME || nonbondedMethod == LJPME) {
            ewaldSelfEnergy = -ONE_4PI_EPS0*ewaldAlpha*sumSquaredCharges/sqrt(M_PI);
            if (nonbondedMethod == LJPME)
                for (int atom = 0; atom < numParticles; atom++)
                    ewaldSelfEnergy += pow(ewaldDispersionAlpha, 6.0) * C6params[atom]*C6params[atom] / 12.0;
        }
        else
            ewaldSelfEnergy = 0.0;
881
882
        chargePosqIndex = data.requestPosqIndex();
        ljPosqIndex = data.requestPosqIndex();
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
    }

    // Compute exception parameters.

    if (hasExceptionOffsets || !offsetsOnly) {
        for (int i = 0; i < num14; i++) {
            double chargeProd = baseExceptionParams[i][0];
            double sigma = baseExceptionParams[i][1];
            double epsilon = baseExceptionParams[i][2];
            for (auto& offset : exceptionParamOffsets[i]) {
                double value = paramValues[get<3>(offset)];
                chargeProd += value*get<0>(offset);
                sigma += value*get<1>(offset);
                epsilon += value*get<2>(offset);
            }
            bonded14ParamArray[i][0] = sigma;
            bonded14ParamArray[i][1] = 4.0*epsilon;
            bonded14ParamArray[i][2] = chargeProd;
        }
    }
}

905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
CpuConstantPotentialForce* createCpuConstantPotentialForceVec();

CpuCalcConstantPotentialForceKernel::CpuCalcConstantPotentialForceKernel(string name, const Platform& platform, CpuPlatform::PlatformData& data) : CalcConstantPotentialForceKernel(name, platform),
    data(data), hasInitializedPme(false), constantPotential(NULL), solver(NULL) {
}

CpuCalcConstantPotentialForceKernel::~CpuCalcConstantPotentialForceKernel() {
    if (constantPotential != NULL) {
        delete constantPotential;
    }
    if (solver != NULL) {
        delete solver;
    }
}

void CpuCalcConstantPotentialForceKernel::initialize(const System& system, const ConstantPotentialForce& force) {
    chargePosqIndex = data.requestPosqIndex();

    // Get particle parameters.
    numParticles = force.getNumParticles();
    setCharges.resize(numParticles);
    for (int i = 0; i < numParticles; i++) {
        force.getParticleParameters(i, setCharges[i]);
    }

    // Get "1-4" exceptions (those that don't zero the charge product).
    exclusions.resize(numParticles);
    vector<int> nb14s;
    for (int i = 0; i < force.getNumExceptions(); i++) {
        int particle1, particle2;
        double chargeProd;
        force.getExceptionParameters(i, particle1, particle2, chargeProd);
        exclusions[particle1].insert(particle2);
        exclusions[particle2].insert(particle1);
        if (chargeProd != 0.0) {
            nb14Index[i] = nb14s.size();
            nb14s.push_back(i);
        }
    }

    // Get exception parameters.
    num14 = nb14s.size();
    bonded14ParamArray.resize(num14, vector<double>(1));
    bonded14IndexArray.resize(num14, vector<int>(2));
    for (int i = 0; i < num14; ++i) {
        int particle1, particle2;
        force.getExceptionParameters(nb14s[i], particle1, particle2, bonded14ParamArray[i][0]);
        bonded14IndexArray[i][0] = particle1;
        bonded14IndexArray[i][1] = particle2;
    }
    bondForce.initialize(numParticles, num14, 2, bonded14IndexArray, data.threads);

    // Get electrode parameters.  sysToElec will be a map from system particle
    // indices to electrode particle indices (or -1 if the particle is not an
    // electrode particle), while elecToSys will be a map from electrode
    // particle indices to system particle indices.  sysElec will be a map from
    // system particle indices to electrode indices (or -1 if the particle is
    // not an electrode particle), while elecElec will be a map from electrode
    // particle indices to electrode indices.
    sysToElec.resize(numParticles, -1);
    sysElec.resize(numParticles, -1);
    electrodeParams.resize(force.getNumElectrodes());
    for (int ie = 0; ie < force.getNumElectrodes(); ie++) {
        std::set<int> electrodeParticles;
        double potential;
        double gaussianWidth;
        double thomasFermiScale;
        force.getElectrodeParameters(ie, electrodeParticles, potential, gaussianWidth, thomasFermiScale);
        for (int i : electrodeParticles) {
            sysToElec[i] = elecToSys.size();
            sysElec[i] = ie;
            elecToSys.push_back(i);
            elecElec.push_back(ie);
        }
        electrodeParams[ie] = {potential, gaussianWidth, thomasFermiScale};
    }

    // Clear charges on electrode particles.
    numElectrodeParticles = elecToSys.size();
    for (int ii = 0; ii < numElectrodeParticles; ii++) {
        setCharges[elecToSys[ii]] = 0.0;
    }

    // Initialize single-precision charge vector with initial guesses.
    charges.resize(numParticles);
    for (int i = 0; i < numParticles; i++) {
        charges[i] = (float) setCharges[i];
    }

    // Set options from force.
    nonbondedCutoff = force.getCutoffDistance();
    data.requestNeighborList(nonbondedCutoff, 0.25 * nonbondedCutoff, true, exclusions);
    ConstantPotentialForceImpl::calcPMEParameters(system, force, ewaldAlpha, gridSize[0], gridSize[1], gridSize[2]);
    exceptionsArePeriodic = force.getExceptionsUsePeriodicBoundaryConditions();
    cgErrorTol = force.getCGErrorTolerance();
    useChargeConstraint = force.getUseChargeConstraint();
    chargeTarget = force.getChargeConstraintTarget();
    force.getExternalField(externalField);
    data.isPeriodic = true;

    // Set the charge target to be that on the electrode particles (so that the
    // overall charge is constrained correctly if the non-electrolyte particles
    // are non-neutral).
    for (int i = 0; i < numParticles; i++) {
        if (sysElec[i] == -1) {
            chargeTarget -= setCharges[i];
        }
    }

    // Set up constant potential.
    ConstantPotentialForce::ConstantPotentialMethod method = force.getConstantPotentialMethod();
    if (method == ConstantPotentialForce::Matrix) {
        solver = new CpuConstantPotentialMatrixSolver(numParticles, numElectrodeParticles);
    }
    else if (method == ConstantPotentialForce::CG) {
        solver = new CpuConstantPotentialCGSolver(numParticles, numElectrodeParticles, force.getUsePreconditioner());
    }
    else {
        throw OpenMMException("internal error: invalid constant potential method");
    }

    constantPotential = createCpuConstantPotentialForceVec();
    float externalFieldArray[3] = { (float) externalField[0], (float) externalField[1], (float) externalField[2] };
    constantPotential->initialize(numParticles, numElectrodeParticles, chargePosqIndex, (float) nonbondedCutoff, (float) ewaldAlpha, (float) cgErrorTol,
        gridSize, exceptionsArePeriodic, useChargeConstraint, *data.neighborList, solver, exclusions, sysToElec, elecToSys, sysElec, elecElec, electrodeParams, (float) chargeTarget, externalFieldArray);
}

double CpuCalcConstantPotentialForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
    Vec3* boxVectors = extractBoxVectors(context);
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0.0;

    checkBoxSize(boxVectors);
    ensurePmeInitialized(context);
    copyChargesToPosq(context, charges, chargePosqIndex);

    constantPotential->execute(boxVectors, posData, charges, &data.posq[0], data.threadForce, includeEnergy ? &energy : NULL, data.threads, pmeKernel);

    // Process non-zeroing exceptions.  Since exceptions and electrodes should
    // involve disjoint sets of atoms, this isn't required for the energy
    // minimization with respect to the electrode atom charges.
    ReferenceConstantPotential14 conp14;
    if (exceptionsArePeriodic) {
        conp14.setPeriodic(boxVectors);
    }
    bondForce.calculateForce(posData, bonded14ParamArray, forceData, includeEnergy ? &energy : NULL, conp14);

    return energy;
}

void CpuCalcConstantPotentialForceKernel::copyParametersToContext(ContextImpl& context, const ConstantPotentialForce& force, int firstParticle, int lastParticle, int firstException, int lastException, int firstElectrode, int lastElectrode) {
    // Get particle parameters.
    if (force.getNumParticles() != numParticles) {
        throw OpenMMException("updateParametersInContext: The number of particles has changed");
    }
    for (int i = firstParticle; i <= lastParticle; i++) {
        // Only update charges on non-electrode particles; keep current guesses
        // for electrode particles.
        if (sysElec[i] == -1) {
            force.getParticleParameters(i, setCharges[i]);
            charges[i] = (float) setCharges[i];
        }
    }
    if (firstParticle <= lastParticle) {
        // Signal that charges in posq need to be updated.
        chargePosqIndex = data.requestPosqIndex();
    }

    // Get "1-4" (non-zeroing) exceptions.
    vector<int> nb14s;
    for (int i = 0; i < force.getNumExceptions(); i++) {
        int particle1, particle2;
        double chargeProd;
        force.getExceptionParameters(i, particle1, particle2, chargeProd);
        if (nb14Index.find(i) == nb14Index.end()) {
            if (chargeProd != 0.0) {
                throw OpenMMException("updateParametersInContext: The set of non-excluded exceptions has changed");
            }
        }
        else {
            nb14s.push_back(i);
        }
    }
    if (nb14s.size() != num14) {
        throw OpenMMException("updateParametersInContext: The number of non-excluded exceptions has changed");
    }

    // Get exception parameters.
    for (int i = 0; i < num14; i++) {
        int particle1, particle2;
        force.getExceptionParameters(nb14s[i], particle1, particle2, bonded14ParamArray[i][0]);
        bonded14IndexArray[i][0] = particle1;
        bonded14IndexArray[i][1] = particle2;
    }

    // Get electrode parameters.
    std::set<int> allElectrodeParticles;
    for (int ie = 0; ie < force.getNumElectrodes(); ie++) {
        std::set<int> electrodeParticles;
        double potential;
        double gaussianWidth;
        double thomasFermiScale;
        force.getElectrodeParameters(ie, electrodeParticles, potential, gaussianWidth, thomasFermiScale);
        for (int i : electrodeParticles) {
            if (sysElec[i] != ie) {
                throw OpenMMException("updateParametersInContext: The electrode assignment of a particle has changed");
            }
            allElectrodeParticles.insert(i);
        }
        electrodeParams[ie][0] = potential;
        electrodeParams[ie][1] = gaussianWidth;
        electrodeParams[ie][2] = thomasFermiScale;
    }
    if (allElectrodeParticles.size() != numElectrodeParticles) {
        throw OpenMMException("updateParametersInContext: The electrode state of a particle has changed");
    }

    // Update external field.
    force.getExternalField(externalField);

    // Update charge target.
    chargeTarget = force.getChargeConstraintTarget();
    for (int i = 0; i < numParticles; i++) {
        if (sysElec[i] == -1) {
            chargeTarget -= setCharges[i];
        }
    }

    float externalFieldArray[3] = { (float) externalField[0], (float) externalField[1], (float) externalField[2] };
    constantPotential->update((float) chargeTarget, externalFieldArray, firstElectrode, lastElectrode);
}

void CpuCalcConstantPotentialForceKernel::getPMEParameters(double& alpha, int& nx, int& ny, int& nz) const {
    if (hasInitializedPme) {
        pmeKernel.getAs<const CalcPmeReciprocalForceKernel>().getPMEParameters(alpha, nx, ny, nz);
    } else {
        alpha = ewaldAlpha;
        nx = gridSize[0];
        ny = gridSize[1];
        nz = gridSize[2];
    }
}

void CpuCalcConstantPotentialForceKernel::getCharges(ContextImpl& context, std::vector<double>& chargesOut) {
    // Make sure that positions in posq, and the current neighbor list, are up
    // to date, but don't compute any energies or forces, and exclude all force
    // groups so that no kernels will get executed.
    context.calcForcesAndEnergy(false, false, 0);

    Vec3* boxVectors = extractBoxVectors(context);
    vector<Vec3>& posData = extractPositions(context);

    checkBoxSize(boxVectors);
    ensurePmeInitialized(context);
    copyChargesToPosq(context, charges, chargePosqIndex);

    constantPotential->getCharges(boxVectors, posData, charges, &data.posq[0], data.threadForce, data.threads, pmeKernel);

    // Preserve fixed charges exactly (without single-precision rounding) and
    // load solved values of fluctuating charges.
    chargesOut = setCharges;
    for (int ii = 0; ii < numElectrodeParticles; ii++) {
        int i = elecToSys[ii];
        chargesOut[i] = (double) charges[i];
    }
}

void CpuCalcConstantPotentialForceKernel::checkBoxSize(const Vec3* boxVectors) {
    double minAllowedSize = 1.999999*nonbondedCutoff;
    if (boxVectors[0][0] < minAllowedSize || boxVectors[1][1] < minAllowedSize || boxVectors[2][2] < minAllowedSize) {
        throw OpenMMException("The periodic box size has decreased to less than twice the nonbonded cutoff.");
    }
}

void CpuCalcConstantPotentialForceKernel::ensurePmeInitialized(ContextImpl& context) {
    if (!hasInitializedPme) {
        vector<string> kernelNames;
        kernelNames.push_back("CalcPmeReciprocalForce");
        if (!getPlatform().supportsKernels(kernelNames)) {
            throw OpenMMException("ConstantPotentialForce unsupported on CPU platform without PME kernel");
        }
        pmeKernel = getPlatform().createKernel(CalcPmeReciprocalForceKernel::Name(), context);
        pmeKernel.getAs<CalcPmeReciprocalForceKernel>().initialize(gridSize[0], gridSize[1], gridSize[2], numParticles, elecToSys, ewaldAlpha, data.deterministicForces);
        hasInitializedPme = true;
    }
}

1193
CpuCalcCustomNonbondedForceKernel::CpuCalcCustomNonbondedForceKernel(string name, const Platform& platform, CpuPlatform::PlatformData& data) :
1194
            CalcCustomNonbondedForceKernel(name, platform), data(data), forceCopy(NULL), nonbonded(NULL) {
1195
1196
}

1197
CpuCalcCustomNonbondedForceKernel::~CpuCalcCustomNonbondedForceKernel() {
1198
1199
    if (nonbonded != NULL)
        delete nonbonded;
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
    if (forceCopy != NULL)
        delete forceCopy;
}

void CpuCalcCustomNonbondedForceKernel::initialize(const System& system, const CustomNonbondedForce& force) {

    // Record the exclusions.

    numParticles = force.getNumParticles();
    exclusions.resize(numParticles);
    for (int i = 0; i < force.getNumExclusions(); i++) {
        int particle1, particle2;
        force.getExclusionParticles(i, particle1, particle2);
        exclusions[particle1].insert(particle2);
        exclusions[particle2].insert(particle1);
    }

    // Build the arrays.

1219
1220
1221
    particleParamArray.resize(numParticles);
    for (int i = 0; i < numParticles; ++i)
        force.getParticleParameters(i, particleParamArray[i]);
1222
1223
    nonbondedMethod = CalcCustomNonbondedForceKernel::NonbondedMethod(force.getNonbondedMethod());
    nonbondedCutoff = force.getCutoffDistance();
1224
1225
    if (nonbondedMethod == NoCutoff) {
        data.requestNeighborList(0.0, 0.0, true, exclusions);
1226
        useSwitchingFunction = false;
1227
    }
1228
    else {
1229
        data.requestNeighborList(nonbondedCutoff, 0.25*nonbondedCutoff, true, exclusions);
1230
1231
1232
1233
        useSwitchingFunction = force.getUseSwitchingFunction();
        switchingDistance = force.getSwitchingDistance();
    }

1234
    // Record the tabulated function update counts for future reference.
1235
1236

    for (int i = 0; i < force.getNumTabulatedFunctions(); i++)
1237
        tabulatedFunctionUpdateCount[force.getTabulatedFunctionName(i)] = force.getTabulatedFunction(i).getUpdateCount();
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264

    // Record information for the long range correction.

    if (force.getNonbondedMethod() == CustomNonbondedForce::CutoffPeriodic && force.getUseLongRangeCorrection()) {
        forceCopy = new CustomNonbondedForce(force);
        hasInitializedLongRangeCorrection = false;
    }
    else {
        longRangeCoefficient = 0.0;
        hasInitializedLongRangeCorrection = true;
    }

    // Record the interaction groups.

    for (int i = 0; i < force.getNumInteractionGroups(); i++) {
        set<int> set1, set2;
        force.getInteractionGroupParameters(i, set1, set2);
        interactionGroups.push_back(make_pair(set1, set2));
    }
    data.isPeriodic |= (nonbondedMethod == CutoffPeriodic);

    // Create the interaction.

    createInteraction(force);
}

void CpuCalcCustomNonbondedForceKernel::createInteraction(const CustomNonbondedForce& force) {
1265
1266
1267
    // Create custom functions for the tabulated functions.

    map<string, Lepton::CustomFunction*> functions;
1268
    for (int i = 0; i < force.getNumTabulatedFunctions(); i++)
1269
1270
1271
1272
        functions[force.getTabulatedFunctionName(i)] = createReferenceTabulatedFunction(force.getTabulatedFunction(i));

    // Parse the various expressions used to calculate the force.

1273
1274
    Lepton::ParsedExpression energyExpression = Lepton::Parser::parse(force.getEnergyFunction(), functions).optimize();
    Lepton::ParsedExpression forceExpression = energyExpression.differentiate("r");
1275
    for (int i = 0; i < force.getNumPerParticleParameters(); i++)
1276
1277
1278
1279
1280
        parameterNames.push_back(force.getPerParticleParameterName(i));
    for (int i = 0; i < force.getNumGlobalParameters(); i++) {
        globalParameterNames.push_back(force.getGlobalParameterName(i));
        globalParamValues[force.getGlobalParameterName(i)] = force.getGlobalParameterDefaultValue(i);
    }
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
    set<string> particleVariables, pairVariables;
    particleVariables.insert("r");
    pairVariables.insert("r");
    for (auto& name : parameterNames) {
        particleVariables.insert(name);
        pairVariables.insert(name+"1");
        pairVariables.insert(name+"2");
    }
    particleVariables.insert(globalParameterNames.begin(), globalParameterNames.end());
    pairVariables.insert(globalParameterNames.begin(), globalParameterNames.end());
1291
    vector<Lepton::ParsedExpression> computedValueExpressions, energyParamDerivExpressions;
1292
1293
1294
1295
1296
1297
    for (int i = 0; i < force.getNumComputedValues(); i++) {
        string name, exp;
        force.getComputedValueParameters(i, name, exp);
        Lepton::ParsedExpression parsed = Lepton::Parser::parse(exp, functions);
        validateVariables(parsed.getRootNode(), particleVariables);
        computedValueNames.push_back(name);
1298
        computedValueExpressions.push_back(parsed);
1299
    }
1300
1301
1302
    for (int i = 0; i < force.getNumEnergyParameterDerivatives(); i++) {
        string param = force.getEnergyParameterDerivativeName(i);
        energyParamDerivNames.push_back(param);
1303
        energyParamDerivExpressions.push_back(energyExpression.differentiate(param));
1304
    }
1305
1306
1307
    for (auto& name : computedValueNames) {
        pairVariables.insert(name+"1");
        pairVariables.insert(name+"2");
1308
    }
1309
    validateVariables(energyExpression.getRootNode(), pairVariables);
1310
1311
1312

    // Delete the custom functions.

peastman's avatar
peastman committed
1313
1314
    for (auto& function : functions)
        delete function.second;
1315
1316
1317

    // Create the object that computes the interaction.

1318
    nonbonded = createCpuCustomNonbondedForce(data.threads, *data.neighborList);
1319
1320
    nonbonded->initialize(energyExpression, forceExpression, parameterNames, exclusions, energyParamDerivExpressions,
            computedValueNames, computedValueExpressions);
peastman's avatar
Bug fix  
peastman committed
1321
1322
    if (interactionGroups.size() > 0)
        nonbonded->setInteractionGroups(interactionGroups);
1323
1324
1325
}

double CpuCalcCustomNonbondedForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
1326
1327
1328
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    Vec3* boxVectors = extractBoxVectors(context);
1329
1330
    double energy = 0;
    bool periodic = (nonbondedMethod == CutoffPeriodic);
1331
    if (nonbondedMethod != NoCutoff)
1332
        nonbonded->setUseCutoff(nonbondedCutoff);
1333
1334
    if (periodic) {
        double minAllowedSize = 2*nonbondedCutoff;
1335
        if (boxVectors[0][0] < minAllowedSize || boxVectors[1][1] < minAllowedSize || boxVectors[2][2] < minAllowedSize)
1336
            throw OpenMMException("The periodic box size has decreased to less than twice the nonbonded cutoff.");
1337
        nonbonded->setPeriodic(boxVectors);
1338
1339
    }
    bool globalParamsChanged = false;
peastman's avatar
peastman committed
1340
1341
1342
    for (auto& name : globalParameterNames) {
        double value = context.getParameter(name);
        if (globalParamValues[name] != value)
1343
            globalParamsChanged = true;
peastman's avatar
peastman committed
1344
        globalParamValues[name] = value;
1345
1346
    }
    if (useSwitchingFunction)
1347
        nonbonded->setUseSwitchingFunction(switchingDistance);
1348
    vector<double> energyParamDerivValues(energyParamDerivNames.size()+1, 0.0);
1349
    nonbonded->calculatePairIxn(numParticles, &data.posq[0], posData, particleParamArray, globalParamValues, data.threadForce, includeForces, includeEnergy, energy, &energyParamDerivValues[0]);
1350
1351
1352
    map<string, double>& energyParamDerivs = extractEnergyParameterDerivatives(context);
    for (int i = 0; i < energyParamDerivNames.size(); i++)
        energyParamDerivs[energyParamDerivNames[i]] += energyParamDerivValues[i];
1353
1354
1355
    
    // Add in the long range correction.
    
1356
    if (!hasInitializedLongRangeCorrection) {
1357
        longRangeCorrectionData = CustomNonbondedForceImpl::prepareLongRangeCorrection(*forceCopy, data.threads.getNumThreads());
1358
        CustomNonbondedForceImpl::calcLongRangeCorrection(*forceCopy, longRangeCorrectionData, context.getOwner(), longRangeCoefficient, longRangeCoefficientDerivs, data.threads);
1359
1360
        hasInitializedLongRangeCorrection = true;
    }
1361
1362
    else if (globalParamsChanged && forceCopy != NULL)
        CustomNonbondedForceImpl::calcLongRangeCorrection(*forceCopy, longRangeCorrectionData, context.getOwner(), longRangeCoefficient, longRangeCoefficientDerivs, data.threads);
1363
1364
1365
1366
    double volume = boxVectors[0][0]*boxVectors[1][1]*boxVectors[2][2];
    energy += longRangeCoefficient/volume;
    for (int i = 0; i < longRangeCoefficientDerivs.size(); i++)
        energyParamDerivs[energyParamDerivNames[i]] += longRangeCoefficientDerivs[i]/volume;
1367
1368
1369
    return energy;
}

1370
void CpuCalcCustomNonbondedForceKernel::copyParametersToContext(ContextImpl& context, const CustomNonbondedForce& force, int firstParticle, int lastParticle) {
1371
1372
1373
1374
1375
1376
    if (numParticles != force.getNumParticles())
        throw OpenMMException("updateParametersInContext: The number of particles has changed");

    // Record the values.

    int numParameters = force.getNumPerParticleParameters();
1377
1378
    vector<double> parameters;
    for (int i = firstParticle; i <= lastParticle; ++i) {
1379
1380
1381
1382
1383
1384
1385
1386
        force.getParticleParameters(i, parameters);
        for (int j = 0; j < numParameters; j++)
            particleParamArray[i][j] = parameters[j];
    }
    
    // If necessary, recompute the long range correction.
    
    if (forceCopy != NULL) {
1387
        longRangeCorrectionData = CustomNonbondedForceImpl::prepareLongRangeCorrection(force, data.threads.getNumThreads());
1388
        CustomNonbondedForceImpl::calcLongRangeCorrection(force, longRangeCorrectionData, context.getOwner(), longRangeCoefficient, longRangeCoefficientDerivs, data.threads);
1389
1390
1391
        hasInitializedLongRangeCorrection = true;
        *forceCopy = force;
    }
1392
1393
1394
1395
1396
1397

    // See if any tabulated functions have changed.

    bool changed = false;
    for (int i = 0; i < force.getNumTabulatedFunctions(); i++) {
        string name = force.getTabulatedFunctionName(i);
1398
1399
        if (force.getTabulatedFunction(i).getUpdateCount() != tabulatedFunctionUpdateCount[name]) {
            tabulatedFunctionUpdateCount[name] = force.getTabulatedFunction(i).getUpdateCount();
1400
1401
1402
1403
1404
1405
1406
1407
            changed = true;
        }
    }
    if (changed) {
        delete nonbonded;
        nonbonded = NULL;
        createInteraction(force);
    }
1408
1409
}

1410
1411
1412
1413
CpuCalcGBSAOBCForceKernel::~CpuCalcGBSAOBCForceKernel() {
}

void CpuCalcGBSAOBCForceKernel::initialize(const System& system, const GBSAOBCForce& force) {
1414
    posqIndex = data.requestPosqIndex();
1415
1416
    int numParticles = system.getNumParticles();
    particleParams.resize(numParticles);
1417
    charges.resize(numParticles);
1418
1419
1420
    for (int i = 0; i < numParticles; ++i) {
        double charge, radius, scalingFactor;
        force.getParticleParameters(i, charge, radius, scalingFactor);
1421
        charges[i] = (float) charge;
1422
1423
        radius -= 0.009;
        particleParams[i] = make_pair((float) radius, (float) (scalingFactor*radius));
1424
1425
1426
1427
    }
    obc.setParticleParameters(particleParams);
    obc.setSolventDielectric((float) force.getSolventDielectric());
    obc.setSoluteDielectric((float) force.getSoluteDielectric());
1428
    obc.setSurfaceAreaEnergy((float) force.getSurfaceAreaEnergy());
1429
1430
    if (force.getNonbondedMethod() != GBSAOBCForce::NoCutoff)
        obc.setUseCutoff((float) force.getCutoffDistance());
1431
    data.isPeriodic |= (force.getNonbondedMethod() == GBSAOBCForce::CutoffPeriodic);
1432
1433
1434
}

double CpuCalcGBSAOBCForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
1435
    copyChargesToPosq(context, charges, posqIndex);
1436
    if (data.isPeriodic) {
peastman's avatar
peastman committed
1437
        Vec3& boxSize = extractBoxSize(context);
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
        float floatBoxSize[3] = {(float) boxSize[0], (float) boxSize[1], (float) boxSize[2]};
        obc.setPeriodic(floatBoxSize);
    }
    double energy = 0.0;
    obc.computeForce(data.posq, data.threadForce, includeEnergy ? &energy : NULL, data.threads);
    return energy;
}

void CpuCalcGBSAOBCForceKernel::copyParametersToContext(ContextImpl& context, const GBSAOBCForce& force) {
    int numParticles = force.getNumParticles();
    if (numParticles != obc.getParticleParameters().size())
        throw OpenMMException("updateParametersInContext: The number of particles has changed");

    // Record the values.

1453
    posqIndex = data.requestPosqIndex();
1454
1455
1456
    for (int i = 0; i < numParticles; ++i) {
        double charge, radius, scalingFactor;
        force.getParticleParameters(i, charge, radius, scalingFactor);
1457
        charges[i] = (float) charge;
1458
1459
        radius -= 0.009;
        particleParams[i] = make_pair((float) radius, (float) (scalingFactor*radius));
1460
1461
1462
    }
    obc.setParticleParameters(particleParams);
}
1463

1464
CpuCalcCustomGBForceKernel::~CpuCalcCustomGBForceKernel() {
1465
1466
    if (ixn != NULL)
        delete ixn;
Peter Eastman's avatar
Peter Eastman committed
1467
1468
    if (neighborList != NULL)
        delete neighborList;
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
}

void CpuCalcCustomGBForceKernel::initialize(const System& system, const CustomGBForce& force) {
    if (force.getNumComputedValues() > 0) {
        string name, expression;
        CustomGBForce::ComputationType type;
        force.getComputedValueParameters(0, name, expression, type);
        if (type == CustomGBForce::SingleParticle)
            throw OpenMMException("CpuPlatform requires that the first computed value for a CustomGBForce be of type ParticlePair or ParticlePairNoExclusions.");
        for (int i = 1; i < force.getNumComputedValues(); i++) {
            force.getComputedValueParameters(i, name, expression, type);
            if (type != CustomGBForce::SingleParticle)
                throw OpenMMException("CpuPlatform requires that a CustomGBForce only have one computed value of type ParticlePair or ParticlePairNoExclusions.");
        }
    }

    // Record the exclusions.

    numParticles = force.getNumParticles();
    exclusions.resize(numParticles);
    for (int i = 0; i < force.getNumExclusions(); i++) {
        int particle1, particle2;
        force.getExclusionParticles(i, particle1, particle2);
        exclusions[particle1].insert(particle2);
        exclusions[particle2].insert(particle1);
    }

    // Build the arrays.

1498
1499
1500
    particleParamArray.resize(numParticles);
    for (int i = 0; i < numParticles; ++i)
        force.getParticleParameters(i, particleParamArray[i]);
1501
    for (int i = 0; i < force.getNumPerParticleParameters(); i++)
1502
1503
1504
1505
        particleParameterNames.push_back(force.getPerParticleParameterName(i));
    for (int i = 0; i < force.getNumGlobalParameters(); i++)
        globalParameterNames.push_back(force.getGlobalParameterName(i));
    nonbondedMethod = CalcCustomGBForceKernel::NonbondedMethod(force.getNonbondedMethod());
peastman's avatar
peastman committed
1506
    nonbondedCutoff = force.getCutoffDistance();
1507
    if (nonbondedMethod != NoCutoff)
Peter Eastman's avatar
Peter Eastman committed
1508
        neighborList = new CpuNeighborList(4);
1509
1510
    data.isPeriodic |= (force.getNonbondedMethod() == CustomGBForce::CutoffPeriodic);

1511
    // Record the tabulated function update counts for future reference.
1512
1513

    for (int i = 0; i < force.getNumTabulatedFunctions(); i++)
1514
        tabulatedFunctionUpdateCount[force.getTabulatedFunctionName(i)] = force.getTabulatedFunction(i).getUpdateCount();
1515
1516

    // Create the interaction.
1517

1518
1519
1520
1521
    createInteraction(force);
}

void CpuCalcCustomGBForceKernel::createInteraction(const CustomGBForce& force) {
1522
1523
1524
    // Create custom functions for the tabulated functions.

    map<string, Lepton::CustomFunction*> functions;
1525
    for (int i = 0; i < force.getNumTabulatedFunctions(); i++)
1526
1527
1528
1529
        functions[force.getTabulatedFunctionName(i)] = createReferenceTabulatedFunction(force.getTabulatedFunction(i));

    // Parse the expressions for computed values.

1530
1531
1532
    valueTypes.clear();
    valueNames.clear();
    energyParamDerivNames.clear();
1533
1534
    vector<vector<Lepton::CompiledExpression> > valueDerivExpressions(force.getNumComputedValues());
    vector<vector<Lepton::CompiledExpression> > valueGradientExpressions(force.getNumComputedValues());
1535
    vector<vector<Lepton::CompiledExpression> > valueParamDerivExpressions(force.getNumComputedValues());
1536
1537
    vector<Lepton::CompiledExpression> valueExpressions;
    vector<Lepton::CompiledExpression> energyExpressions;
1538
1539
1540
1541
1542
    set<string> particleVariables, pairVariables;
    pairVariables.insert("r");
    particleVariables.insert("x");
    particleVariables.insert("y");
    particleVariables.insert("z");
1543
    for (int i = 0; i < force.getNumPerParticleParameters(); i++) {
1544
1545
1546
1547
1548
1549
        particleVariables.insert(particleParameterNames[i]);
        pairVariables.insert(particleParameterNames[i]+"1");
        pairVariables.insert(particleParameterNames[i]+"2");
    }
    particleVariables.insert(globalParameterNames.begin(), globalParameterNames.end());
    pairVariables.insert(globalParameterNames.begin(), globalParameterNames.end());
1550
1551
1552
1553
1554
1555
1556
1557
    for (int i = 0; i < force.getNumComputedValues(); i++) {
        string name, expression;
        CustomGBForce::ComputationType type;
        force.getComputedValueParameters(i, name, expression, type);
        Lepton::ParsedExpression ex = Lepton::Parser::parse(expression, functions).optimize();
        valueExpressions.push_back(ex.createCompiledExpression());
        valueTypes.push_back(type);
        valueNames.push_back(name);
1558
        if (i == 0) {
1559
            valueDerivExpressions[i].push_back(ex.differentiate("r").createCompiledExpression());
1560
1561
            validateVariables(ex.getRootNode(), pairVariables);
        }
1562
1563
1564
1565
1566
1567
        else {
            valueGradientExpressions[i].push_back(ex.differentiate("x").createCompiledExpression());
            valueGradientExpressions[i].push_back(ex.differentiate("y").createCompiledExpression());
            valueGradientExpressions[i].push_back(ex.differentiate("z").createCompiledExpression());
            for (int j = 0; j < i; j++)
                valueDerivExpressions[i].push_back(ex.differentiate(valueNames[j]).createCompiledExpression());
1568
            validateVariables(ex.getRootNode(), particleVariables);
1569
        }
1570
1571
1572
1573
1574
        for (int j = 0; j < force.getNumEnergyParameterDerivatives(); j++) {
            string param = force.getEnergyParameterDerivativeName(j);
            energyParamDerivNames.push_back(param);
            valueParamDerivExpressions[i].push_back(ex.differentiate(param).createCompiledExpression());
        }
1575
1576
1577
        particleVariables.insert(name);
        pairVariables.insert(name+"1");
        pairVariables.insert(name+"2");
1578
1579
1580
1581
    }

    // Parse the expressions for energy terms.

1582
    energyTypes.clear();
1583
1584
    vector<vector<Lepton::CompiledExpression> > energyDerivExpressions(force.getNumEnergyTerms());
    vector<vector<Lepton::CompiledExpression> > energyGradientExpressions(force.getNumEnergyTerms());
1585
    vector<vector<Lepton::CompiledExpression> > energyParamDerivExpressions(force.getNumEnergyTerms());
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
    for (int i = 0; i < force.getNumEnergyTerms(); i++) {
        string expression;
        CustomGBForce::ComputationType type;
        force.getEnergyTermParameters(i, expression, type);
        Lepton::ParsedExpression ex = Lepton::Parser::parse(expression, functions).optimize();
        energyExpressions.push_back(ex.createCompiledExpression());
        energyTypes.push_back(type);
        if (type != CustomGBForce::SingleParticle)
            energyDerivExpressions[i].push_back(ex.differentiate("r").createCompiledExpression());
        for (int j = 0; j < force.getNumComputedValues(); j++) {
            if (type == CustomGBForce::SingleParticle) {
                energyDerivExpressions[i].push_back(ex.differentiate(valueNames[j]).createCompiledExpression());
                energyGradientExpressions[i].push_back(ex.differentiate("x").createCompiledExpression());
                energyGradientExpressions[i].push_back(ex.differentiate("y").createCompiledExpression());
                energyGradientExpressions[i].push_back(ex.differentiate("z").createCompiledExpression());
1601
                validateVariables(ex.getRootNode(), particleVariables);
1602
1603
1604
1605
            }
            else {
                energyDerivExpressions[i].push_back(ex.differentiate(valueNames[j]+"1").createCompiledExpression());
                energyDerivExpressions[i].push_back(ex.differentiate(valueNames[j]+"2").createCompiledExpression());
1606
                validateVariables(ex.getRootNode(), pairVariables);
1607
1608
            }
        }
1609
1610
        for (int j = 0; j < force.getNumEnergyParameterDerivatives(); j++)
            energyParamDerivExpressions[i].push_back(ex.differentiate(force.getEnergyParameterDerivativeName(j)).createCompiledExpression());
1611
1612
1613
1614
    }

    // Delete the custom functions.

peastman's avatar
peastman committed
1615
1616
    for (auto& function : functions)
        delete function.second;
1617
1618
1619
    ixn = new CpuCustomGBForce(numParticles, exclusions, valueExpressions, valueDerivExpressions, valueGradientExpressions, valueParamDerivExpressions,
        valueNames, valueTypes, energyExpressions, energyDerivExpressions, energyGradientExpressions, energyParamDerivExpressions, energyTypes,
        particleParameterNames, data.threads);
1620
1621
1622
}

double CpuCalcCustomGBForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
1623
1624
1625
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
    Vec3* boxVectors = extractBoxVectors(context);
1626
1627
    if (data.isPeriodic)
        ixn->setPeriodic(extractBoxSize(context));
1628
    if (nonbondedMethod != NoCutoff) {
1629
        vector<set<int> > noExclusions(numParticles);
Peter Eastman's avatar
Peter Eastman committed
1630
1631
        neighborList->computeNeighborList(numParticles, data.posq, noExclusions, boxVectors, data.isPeriodic, nonbondedCutoff, data.threads);
        ixn->setUseCutoff(nonbondedCutoff, *neighborList);
1632
1633
    }
    map<string, double> globalParameters;
peastman's avatar
peastman committed
1634
1635
    for (auto& name : globalParameterNames)
        globalParameters[name] = context.getParameter(name);
1636
1637
1638
1639
1640
    vector<double> energyParamDerivValues(energyParamDerivNames.size()+1, 0.0);
    ixn->calculateIxn(numParticles, &data.posq[0], particleParamArray, globalParameters, data.threadForce, includeForces, includeEnergy, energy, &energyParamDerivValues[0]);
    map<string, double>& energyParamDerivs = extractEnergyParameterDerivatives(context);
    for (int i = 0; i < energyParamDerivNames.size(); i++)
        energyParamDerivs[energyParamDerivNames[i]] += energyParamDerivValues[i];
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
    return energy;
}

void CpuCalcCustomGBForceKernel::copyParametersToContext(ContextImpl& context, const CustomGBForce& force) {
    if (numParticles != force.getNumParticles())
        throw OpenMMException("updateParametersInContext: The number of particles has changed");

    // Record the values.

    int numParameters = force.getNumPerParticleParameters();
    vector<double> params;
    for (int i = 0; i < numParticles; ++i) {
        vector<double> parameters;
        force.getParticleParameters(i, parameters);
        for (int j = 0; j < numParameters; j++)
peastman's avatar
peastman committed
1656
            particleParamArray[i][j] = static_cast<double>(parameters[j]);
1657
    }
1658
1659
1660
1661
1662
1663

    // See if any tabulated functions have changed.

    bool changed = false;
    for (int i = 0; i < force.getNumTabulatedFunctions(); i++) {
        string name = force.getTabulatedFunctionName(i);
1664
1665
        if (force.getTabulatedFunction(i).getUpdateCount() != tabulatedFunctionUpdateCount[name]) {
            tabulatedFunctionUpdateCount[name] = force.getTabulatedFunction(i).getUpdateCount();
1666
1667
1668
1669
1670
1671
1672
1673
            changed = true;
        }
    }
    if (changed) {
        delete ixn;
        ixn = NULL;
        createInteraction(force);
    }
1674
1675
}

1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
CpuCalcCustomManyParticleForceKernel::~CpuCalcCustomManyParticleForceKernel() {
    if (ixn != NULL)
        delete ixn;
}

void CpuCalcCustomManyParticleForceKernel::initialize(const System& system, const CustomManyParticleForce& force) {

    // Build the arrays.

    numParticles = system.getNumParticles();
1686
    particleParamArray.resize(numParticles);
1687
1688
    for (int i = 0; i < numParticles; ++i) {
        int type;
1689
        force.getParticleParameters(i, particleParamArray[i], type);
1690
1691
1692
    }
    for (int i = 0; i < force.getNumGlobalParameters(); i++)
        globalParameterNames.push_back(force.getGlobalParameterName(i));
1693

1694
    // Record the tabulated function update counts for future reference.
1695
1696

    for (int i = 0; i < force.getNumTabulatedFunctions(); i++)
1697
        tabulatedFunctionUpdateCount[force.getTabulatedFunctionName(i)] = force.getTabulatedFunction(i).getUpdateCount();
1698
1699
1700

    // Create the interaction.

1701
    ixn = new CpuCustomManyParticleForce(force, data.threads);
1702
1703
    nonbondedMethod = CalcCustomManyParticleForceKernel::NonbondedMethod(force.getNonbondedMethod());
    cutoffDistance = force.getCutoffDistance();
1704
    data.isPeriodic |= (nonbondedMethod == CutoffPeriodic);
1705
1706
1707
1708
}

double CpuCalcCustomManyParticleForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
    map<string, double> globalParameters;
peastman's avatar
peastman committed
1709
1710
    for (auto& name : globalParameterNames)
        globalParameters[name] = context.getParameter(name);
1711
    if (nonbondedMethod == CutoffPeriodic) {
peastman's avatar
peastman committed
1712
        Vec3* boxVectors = extractBoxVectors(context);
1713
        double minAllowedSize = 2*cutoffDistance;
1714
        if (boxVectors[0][0] < minAllowedSize || boxVectors[1][1] < minAllowedSize || boxVectors[2][2] < minAllowedSize)
1715
            throw OpenMMException("The periodic box size has decreased to less than twice the nonbonded cutoff.");
1716
        ixn->setPeriodic(boxVectors);
1717
    }
1718
    double energy = 0;
1719
    ixn->calculateIxn(data.posq, particleParamArray, globalParameters, data.threadForce, includeForces, includeEnergy, energy);
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
    return energy;
}

void CpuCalcCustomManyParticleForceKernel::copyParametersToContext(ContextImpl& context, const CustomManyParticleForce& force) {
    if (numParticles != force.getNumParticles())
        throw OpenMMException("updateParametersInContext: The number of particles has changed");

    // Record the values.

    int numParameters = force.getNumPerParticleParameters();
    vector<double> params;
    for (int i = 0; i < numParticles; ++i) {
        vector<double> parameters;
        int type;
        force.getParticleParameters(i, parameters, type);
        for (int j = 0; j < numParameters; j++)
peastman's avatar
peastman committed
1736
            particleParamArray[i][j] = static_cast<double>(parameters[j]);
1737
    }
1738
1739
1740
1741
1742
1743

    // See if any tabulated functions have changed.

    bool changed = false;
    for (int i = 0; i < force.getNumTabulatedFunctions(); i++) {
        string name = force.getTabulatedFunctionName(i);
1744
1745
        if (force.getTabulatedFunction(i).getUpdateCount() != tabulatedFunctionUpdateCount[name]) {
            tabulatedFunctionUpdateCount[name] = force.getTabulatedFunction(i).getUpdateCount();
1746
1747
1748
1749
1750
1751
1752
1753
            changed = true;
        }
    }
    if (changed) {
        delete ixn;
        ixn = NULL;
        ixn = new CpuCustomManyParticleForce(force, data.threads);
    }
1754
1755
}

1756
1757
1758
1759
1760
1761
1762
CpuCalcGayBerneForceKernel::~CpuCalcGayBerneForceKernel() {
    if (ixn != NULL)
        delete ixn;
}

void CpuCalcGayBerneForceKernel::initialize(const System& system, const GayBerneForce& force) {
    ixn = new CpuGayBerneForce(force);
1763
    data.isPeriodic |= (force.getNonbondedMethod() == GayBerneForce::CutoffPeriodic);
1764
1765
1766
1767
    if (force.getNonbondedMethod() != GayBerneForce::NoCutoff) {
        double cutoff = force.getCutoffDistance();
        data.requestNeighborList(cutoff, 0.1*cutoff, true, ixn->getExclusions());
    }
1768
1769
1770
}

double CpuCalcGayBerneForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
1771
    return ixn->calculateForce(extractPositions(context), extractForces(context), data.threadForce, extractBoxVectors(context), data);
1772
1773
1774
1775
1776
1777
1778
1779
}

void CpuCalcGayBerneForceKernel::copyParametersToContext(ContextImpl& context, const GayBerneForce& force) {
    delete ixn;
    ixn = NULL;
    ixn = new CpuGayBerneForce(force);
}

1780
CpuIntegrateLangevinMiddleStepKernel::~CpuIntegrateLangevinMiddleStepKernel() {
1781
1782
1783
1784
    if (dynamics)
        delete dynamics;
}

1785
void CpuIntegrateLangevinMiddleStepKernel::initialize(const System& system, const LangevinMiddleIntegrator& integrator) {
1786
1787
1788
1789
1790
1791
1792
    int numParticles = system.getNumParticles();
    masses.resize(numParticles);
    for (int i = 0; i < numParticles; ++i)
        masses[i] = system.getParticleMass(i);
    data.random.initialize(integrator.getRandomNumberSeed(), data.threads.getNumThreads());
}

1793
void CpuIntegrateLangevinMiddleStepKernel::execute(ContextImpl& context, const LangevinMiddleIntegrator& integrator) {
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
    double temperature = integrator.getTemperature();
    double friction = integrator.getFriction();
    double stepSize = integrator.getStepSize();
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& velData = extractVelocities(context);
    if (dynamics == 0 || temperature != prevTemp || friction != prevFriction || stepSize != prevStepSize) {
        // Recreate the computation objects with the new parameters.
        
        if (dynamics)
            delete dynamics;
1804
        dynamics = new CpuLangevinMiddleDynamics(context.getSystem().getNumParticles(), stepSize, friction, temperature, data.threads, data.random);
1805
        dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
1806
        dynamics->setVirtualSites(extractVirtualSites(context));
1807
1808
1809
1810
        prevTemp = temperature;
        prevFriction = friction;
        prevStepSize = stepSize;
    }
1811
    dynamics->update(context, posData, velData, masses, integrator.getConstraintTolerance(), extractBoxVectors(context));
1812
1813
1814
1815
1816
    ReferencePlatform::PlatformData* refData = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());
    refData->time += stepSize;
    refData->stepCount++;
}

1817
double CpuIntegrateLangevinMiddleStepKernel::computeKineticEnergy(ContextImpl& context, const LangevinMiddleIntegrator& integrator) {
1818
1819
    return computeShiftedKineticEnergy(context, masses, 0.0);
}