ReferenceKernels.cpp 155 KB
Newer Older
1
2
3
/* -------------------------------------------------------------------------- *
 *                                   OpenMM                                   *
 * -------------------------------------------------------------------------- *
Evan Pretti's avatar
Evan Pretti committed
4
5
 * This is part of the OpenMM molecular simulation toolkit.                   *
 * See https://openmm.org/development.                                        *
6
 *                                                                            *
7
 * Portions copyright (c) 2008-2025 Stanford University and the Authors.      *
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
 * 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 "ReferenceKernels.h"
31
#include "ReferenceObc.h"
32
33
34
35
36
#include "ReferenceAndersenThermostat.h"
#include "ReferenceAngleBondIxn.h"
#include "ReferenceBondForce.h"
#include "ReferenceBrownianDynamics.h"
#include "ReferenceCMAPTorsionIxn.h"
37
38
#include "ReferenceConstantPotential.h"
#include "ReferenceConstantPotential14.h"
39
#include "ReferenceConstraints.h"
40
41
#include "ReferenceCustomAngleIxn.h"
#include "ReferenceCustomBondIxn.h"
42
#include "ReferenceCustomCentroidBondIxn.h"
43
#include "ReferenceCustomCompoundBondIxn.h"
44
#include "ReferenceCustomCVForce.h"
45
46
47
48
49
#include "ReferenceCustomDynamics.h"
#include "ReferenceCustomExternalIxn.h"
#include "ReferenceCustomGBIxn.h"
#include "ReferenceCustomHbondIxn.h"
#include "ReferenceCustomNonbondedIxn.h"
50
#include "ReferenceCustomManyParticleIxn.h"
51
#include "ReferenceCustomTorsionIxn.h"
Peter Eastman's avatar
Peter Eastman committed
52
#include "ReferenceDPDDynamics.h"
53
#include "ReferenceGayBerneForce.h"
54
#include "ReferenceHarmonicBondIxn.h"
55
#include "ReferenceLangevinMiddleDynamics.h"
56
57
58
#include "ReferenceLJCoulomb14.h"
#include "ReferenceLJCoulombIxn.h"
#include "ReferenceMonteCarloBarostat.h"
59
#include "ReferenceNoseHooverChain.h"
60
#include "ReferenceNoseHooverDynamics.h"
61
#include "ReferenceOrientationRestraintForce.h"
62
#include "ReferencePointFunctions.h"
63
#include "ReferenceProperDihedralBond.h"
64
#include "ReferenceQTBDynamics.h"
65
#include "ReferenceRbDihedralBond.h"
66
#include "ReferenceRGForce.h"
67
#include "ReferenceRMSDForce.h"
68
#include "ReferenceTabulatedFunction.h"
69
70
71
72
#include "ReferenceVariableStochasticDynamics.h"
#include "ReferenceVariableVerletDynamics.h"
#include "ReferenceVerletDynamics.h"
#include "ReferenceVirtualSites.h"
73
#include "openmm/CMMotionRemover.h"
74
#include "openmm/Context.h"
75
#include "openmm/System.h"
76
#include "openmm/internal/AndersenThermostatImpl.h"
77
#include "openmm/internal/ContextImpl.h"
78
#include "openmm/internal/CustomCentroidBondForceImpl.h"
79
#include "openmm/internal/CustomCompoundBondForceImpl.h"
80
#include "openmm/internal/CustomHbondForceImpl.h"
81
#include "openmm/internal/CMAPTorsionForceImpl.h"
82
#include "openmm/internal/NonbondedForceImpl.h"
83
#include "openmm/internal/ConstantPotentialForceImpl.h"
84
#include "openmm/Integrator.h"
85
#include "openmm/OpenMMException.h"
86
#include "SimTKOpenMMUtilities.h"
87
#include "lepton/CustomFunction.h"
88
#include "lepton/Operation.h"
89
90
#include "lepton/Parser.h"
#include "lepton/ParsedExpression.h"
91
#include <cmath>
Peter Eastman's avatar
Peter Eastman committed
92
#include <iostream>
93
#include <limits>
94
95
96
97

using namespace OpenMM;
using namespace std;

peastman's avatar
peastman committed
98
static vector<Vec3>& extractPositions(ContextImpl& context) {
99
    ReferencePlatform::PlatformData* data = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());
100
    return *data->positions;
101
102
}

peastman's avatar
peastman committed
103
static vector<Vec3>& extractVelocities(ContextImpl& context) {
104
    ReferencePlatform::PlatformData* data = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());
105
    return *data->velocities;
106
107
}

peastman's avatar
peastman committed
108
static vector<Vec3>& extractForces(ContextImpl& context) {
109
    ReferencePlatform::PlatformData* data = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());
110
    return *data->forces;
111
112
}

peastman's avatar
peastman committed
113
static Vec3& extractBoxSize(ContextImpl& context) {
114
    ReferencePlatform::PlatformData* data = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());
115
    return *data->periodicBoxSize;
116
117
}

peastman's avatar
peastman committed
118
static Vec3* extractBoxVectors(ContextImpl& context) {
119
    ReferencePlatform::PlatformData* data = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());
120
    return data->periodicBoxVectors;
121
122
}

123
124
static ReferenceConstraints& extractConstraints(ContextImpl& context) {
    ReferencePlatform::PlatformData* data = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());
125
    return *data->constraints;
126
127
}

128
129
130
131
132
static const ReferenceVirtualSites& extractVirtualSites(ContextImpl& context) {
    ReferencePlatform::PlatformData* data = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());
    return *data->virtualSites;
}

133
134
static map<string, double>& extractEnergyParameterDerivatives(ContextImpl& context) {
    ReferencePlatform::PlatformData* data = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());
135
    return *data->energyParameterDerivatives;
136
137
}

138
139
140
141
142
static ThreadPool& extractThreadPool(ContextImpl& context) {
    ReferencePlatform::PlatformData* data = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());
    return data->threads;
}

143
144
145
146
147
148
149
/**
 * 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
150
151
    for (auto& child : node.getChildren())
        validateVariables(child, variables);
152
153
}

154
155
156
157
/**
 * Compute the kinetic energy of the system, possibly shifting the velocities in time to account
 * for a leapfrog integrator.
 */
158
static double computeShiftedKineticEnergy(ContextImpl& context, vector<double>& masses, double timeShift) {
159
    int numParticles = context.getSystem().getNumParticles();
peastman's avatar
peastman committed
160
    vector<Vec3> shiftedVel(numParticles);
161
    context.computeShiftedVelocities(timeShift, shiftedVel);
162
163
164
165
    double energy = 0.0;
    for (int i = 0; i < numParticles; ++i)
        if (masses[i] > 0)
            energy += masses[i]*(shiftedVel[i].dot(shiftedVel[i]));
166
167
168
    return 0.5*energy;
}

169
void ReferenceCalcForcesAndEnergyKernel::initialize(const System& system) {
170
171
}

172
void ReferenceCalcForcesAndEnergyKernel::beginComputation(ContextImpl& context, bool includeForces, bool includeEnergy, int groups) {
peastman's avatar
peastman committed
173
    vector<Vec3>& forceData = extractForces(context);
174
175
176
    if (includeForces) {
        int numParticles = context.getSystem().getNumParticles();
        for (int i = 0; i < numParticles; ++i) {
peastman's avatar
peastman committed
177
178
179
            forceData[i][0] = 0.0;
            forceData[i][1] = 0.0;
            forceData[i][2] = 0.0;
180
        }
181
    }
182
183
    else
        savedForces = forceData;
peastman's avatar
peastman committed
184
185
    for (auto& param : context.getParameters())
        extractEnergyParameterDerivatives(context)[param.first] = 0;
186
187
}

188
double ReferenceCalcForcesAndEnergyKernel::finishComputation(ContextImpl& context, bool includeForces, bool includeEnergy, int groups, bool& valid) {
189
190
    if (!includeForces)
        extractForces(context) = savedForces; // Restore the forces so computing the energy doesn't overwrite the forces with incorrect values.
191
    else
192
        extractVirtualSites(context).distributeForces(context.getSystem(), extractPositions(context), extractForces(context), extractBoxVectors(context));
193
194
195
    return 0.0;
}

196
void ReferenceUpdateStateDataKernel::initialize(const System& system) {
197
198
199
200
    int numParticles = system.getNumParticles();
    masses.resize(numParticles);
    for (int i = 0; i < numParticles; ++i)
        masses[i] = system.getParticleMass(i);
201
202
}

203
double ReferenceUpdateStateDataKernel::getTime(const ContextImpl& context) const {
204
205
206
    return data.time;
}

207
void ReferenceUpdateStateDataKernel::setTime(ContextImpl& context, double time) {
208
209
210
    data.time = time;
}

211
212
213
214
215
216
217
218
long long ReferenceUpdateStateDataKernel::getStepCount(const ContextImpl& context) const {
    return data.stepCount;
}

void ReferenceUpdateStateDataKernel::setStepCount(const ContextImpl& context, long long count) {
    data.stepCount = count;
}

219
220
void ReferenceUpdateStateDataKernel::getPositions(ContextImpl& context, std::vector<Vec3>& positions) {
    int numParticles = context.getSystem().getNumParticles();
peastman's avatar
peastman committed
221
    vector<Vec3>& posData = extractPositions(context);
222
223
224
225
226
227
228
    positions.resize(numParticles);
    for (int i = 0; i < numParticles; ++i)
        positions[i] = Vec3(posData[i][0], posData[i][1], posData[i][2]);
}

void ReferenceUpdateStateDataKernel::setPositions(ContextImpl& context, const std::vector<Vec3>& positions) {
    int numParticles = context.getSystem().getNumParticles();
peastman's avatar
peastman committed
229
    vector<Vec3>& posData = extractPositions(context);
230
    for (int i = 0; i < numParticles; ++i) {
peastman's avatar
peastman committed
231
232
233
        posData[i][0] = positions[i][0];
        posData[i][1] = positions[i][1];
        posData[i][2] = positions[i][2];
234
235
236
237
238
    }
}

void ReferenceUpdateStateDataKernel::getVelocities(ContextImpl& context, std::vector<Vec3>& velocities) {
    int numParticles = context.getSystem().getNumParticles();
peastman's avatar
peastman committed
239
    vector<Vec3>& velData = extractVelocities(context);
240
241
242
243
244
245
246
    velocities.resize(numParticles);
    for (int i = 0; i < numParticles; ++i)
        velocities[i] = Vec3(velData[i][0], velData[i][1], velData[i][2]);
}

void ReferenceUpdateStateDataKernel::setVelocities(ContextImpl& context, const std::vector<Vec3>& velocities) {
    int numParticles = context.getSystem().getNumParticles();
peastman's avatar
peastman committed
247
    vector<Vec3>& velData = extractVelocities(context);
248
    for (int i = 0; i < numParticles; ++i) {
peastman's avatar
peastman committed
249
250
251
        velData[i][0] = velocities[i][0];
        velData[i][1] = velocities[i][1];
        velData[i][2] = velocities[i][2];
252
253
254
    }
}

255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
void ReferenceUpdateStateDataKernel::computeShiftedVelocities(ContextImpl& context, double timeShift, std::vector<Vec3>& velocities) {
    int numParticles = context.getSystem().getNumParticles();
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& velData = extractVelocities(context);
    vector<Vec3>& forceData = extractForces(context);
    
    // Compute the shifted velocities.
    
    velocities.resize(numParticles);
    vector<double> inverseMasses(numParticles);
    for (int i = 0; i < numParticles; ++i) {
        if (masses[i] == 0) {
            velocities[i] = velData[i];
            inverseMasses[i] = 0;
        }
        else {
            velocities[i] = velData[i]+forceData[i]*(timeShift/masses[i]);
            inverseMasses[i] = 1/masses[i];
        }
    }
    
    // Apply constraints to them.
277
278
279

    if (timeShift != 0)
        extractConstraints(context).applyToVelocities(posData, velocities, inverseMasses, 1e-4);
280
281
}

282
283
void ReferenceUpdateStateDataKernel::getForces(ContextImpl& context, std::vector<Vec3>& forces) {
    int numParticles = context.getSystem().getNumParticles();
peastman's avatar
peastman committed
284
    vector<Vec3>& forceData = extractForces(context);
285
286
287
288
289
    forces.resize(numParticles);
    for (int i = 0; i < numParticles; ++i)
        forces[i] = Vec3(forceData[i][0], forceData[i][1], forceData[i][2]);
}

290
291
292
293
void ReferenceUpdateStateDataKernel::getEnergyParameterDerivatives(ContextImpl& context, map<string, double>& derivs) {
    derivs = extractEnergyParameterDerivatives(context);
}

294
void ReferenceUpdateStateDataKernel::getPeriodicBoxVectors(ContextImpl& context, Vec3& a, Vec3& b, Vec3& c) const {
peastman's avatar
peastman committed
295
    Vec3* vectors = extractBoxVectors(context);
296
297
298
    a = vectors[0];
    b = vectors[1];
    c = vectors[2];
299
300
}

301
void ReferenceUpdateStateDataKernel::setPeriodicBoxVectors(ContextImpl& context, const Vec3& a, const Vec3& b, const Vec3& c) {
peastman's avatar
peastman committed
302
303
304
305
306
    Vec3& box = extractBoxSize(context);
    box[0] = a[0];
    box[1] = b[1];
    box[2] = c[2];
    Vec3* vectors = extractBoxVectors(context);
307
308
309
    vectors[0] = a;
    vectors[1] = b;
    vectors[2] = c;
310
311
}

Peter Eastman's avatar
Peter Eastman committed
312
void ReferenceUpdateStateDataKernel::createCheckpoint(ContextImpl& context, ostream& stream) {
313
    int version = 3;
Peter Eastman's avatar
Peter Eastman committed
314
315
    stream.write((char*) &version, sizeof(int));
    stream.write((char*) &data.time, sizeof(data.time));
316
    stream.write((char*) &data.stepCount, sizeof(long long));
peastman's avatar
peastman committed
317
318
319
320
321
322
    vector<Vec3>& posData = extractPositions(context);
    stream.write((char*) &posData[0], sizeof(Vec3)*posData.size());
    vector<Vec3>& velData = extractVelocities(context);
    stream.write((char*) &velData[0], sizeof(Vec3)*velData.size());
    Vec3* vectors = extractBoxVectors(context);
    stream.write((char*) vectors, 3*sizeof(Vec3));
Peter Eastman's avatar
Peter Eastman committed
323
324
325
326
327
328
    SimTKOpenMMUtilities::createCheckpoint(stream);
}

void ReferenceUpdateStateDataKernel::loadCheckpoint(ContextImpl& context, istream& stream) {
    int version;
    stream.read((char*) &version, sizeof(int));
329
    if (version != 3)
Peter Eastman's avatar
Peter Eastman committed
330
331
        throw OpenMMException("Checkpoint was created with a different version of OpenMM");
    stream.read((char*) &data.time, sizeof(data.time));
332
    stream.read((char*) &data.stepCount, sizeof(long long));
peastman's avatar
peastman committed
333
334
335
336
337
338
    vector<Vec3>& posData = extractPositions(context);
    stream.read((char*) &posData[0], sizeof(Vec3)*posData.size());
    vector<Vec3>& velData = extractVelocities(context);
    stream.read((char*) &velData[0], sizeof(Vec3)*velData.size());
    Vec3* vectors = extractBoxVectors(context);
    stream.read((char*) vectors, 3*sizeof(Vec3));
Peter Eastman's avatar
Peter Eastman committed
339
340
341
    SimTKOpenMMUtilities::loadCheckpoint(stream);
}

342
343
void ReferenceApplyConstraintsKernel::initialize(const System& system) {
    int numParticles = system.getNumParticles();
344
345
    masses.resize(numParticles);
    inverseMasses.resize(numParticles);
346
    for (int i = 0; i < numParticles; ++i) {
peastman's avatar
peastman committed
347
        masses[i] = system.getParticleMass(i);
348
349
350
351
352
353
354
355
        inverseMasses[i] = 1.0/masses[i];
    }
}

ReferenceApplyConstraintsKernel::~ReferenceApplyConstraintsKernel() {
}

void ReferenceApplyConstraintsKernel::apply(ContextImpl& context, double tol) {
peastman's avatar
peastman committed
356
    vector<Vec3>& positions = extractPositions(context);
357
    extractConstraints(context).apply(positions, positions, inverseMasses, tol);
358
    extractVirtualSites(context).computePositions(context.getSystem(), positions, extractBoxVectors(context));
359
360
}

361
void ReferenceApplyConstraintsKernel::applyToVelocities(ContextImpl& context, double tol) {
peastman's avatar
peastman committed
362
363
    vector<Vec3>& positions = extractPositions(context);
    vector<Vec3>& velocities = extractVelocities(context);
364
    extractConstraints(context).applyToVelocities(positions, velocities, inverseMasses, tol);
365
366
}

367
368
369
370
void ReferenceVirtualSitesKernel::initialize(const System& system) {
}

void ReferenceVirtualSitesKernel::computePositions(ContextImpl& context) {
peastman's avatar
peastman committed
371
    vector<Vec3>& positions = extractPositions(context);
372
    extractVirtualSites(context).computePositions(context.getSystem(), positions, extractBoxVectors(context));
373
374
}

375
void ReferenceCalcHarmonicBondForceKernel::initialize(const System& system, const HarmonicBondForce& force) {
376
    numBonds = force.getNumBonds();
377
378
    bondIndexArray.resize(numBonds, vector<int>(2));
    bondParamArray.resize(numBonds, vector<double>(2));
379
    for (int i = 0; i < numBonds; ++i) {
Peter Eastman's avatar
Peter Eastman committed
380
        int particle1, particle2;
381
        double length, k;
Peter Eastman's avatar
Peter Eastman committed
382
383
384
        force.getBondParameters(i, particle1, particle2, length, k);
        bondIndexArray[i][0] = particle1;
        bondIndexArray[i][1] = particle2;
peastman's avatar
peastman committed
385
386
        bondParamArray[i][0] = length;
        bondParamArray[i][1] = k;
387
    }
388
    usePeriodic = force.usesPeriodicBoundaryConditions();
389
390
}

391
double ReferenceCalcHarmonicBondForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
392
393
394
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
395
396
    ReferenceBondForce refBondForce;
    ReferenceHarmonicBondIxn harmonicBond;
397
398
    if (usePeriodic)
        harmonicBond.setPeriodic(extractBoxVectors(context));
399
    refBondForce.calculateForce(numBonds, bondIndexArray, posData, bondParamArray, forceData, includeEnergy ? &energy : NULL, harmonicBond);
400
401
402
    return energy;
}

403
void ReferenceCalcHarmonicBondForceKernel::copyParametersToContext(ContextImpl& context, const HarmonicBondForce& force, int firstBond, int lastBond) {
404
405
406
407
408
    if (numBonds != force.getNumBonds())
        throw OpenMMException("updateParametersInContext: The number of bonds has changed");

    // Record the values.

409
    for (int i = firstBond; i <= lastBond; ++i) {
410
411
412
413
414
415
416
        int particle1, particle2;
        double length, k;
        force.getBondParameters(i, particle1, particle2, length, k);
        if (particle1 != bondIndexArray[i][0] || particle2 != bondIndexArray[i][1])
            throw OpenMMException("updateParametersInContext: The set of particles in a bond has changed");
        bondIndexArray[i][0] = particle1;
        bondIndexArray[i][1] = particle2;
peastman's avatar
peastman committed
417
418
        bondParamArray[i][0] = length;
        bondParamArray[i][1] = k;
419
420
421
    }
}

422
423
424
425
426
ReferenceCalcCustomBondForceKernel::~ReferenceCalcCustomBondForceKernel() {
    if (ixn != NULL)
        delete ixn;
}

427
428
429
void ReferenceCalcCustomBondForceKernel::initialize(const System& system, const CustomBondForce& force) {
    numBonds = force.getNumBonds();
    int numParameters = force.getNumPerBondParameters();
430
    usePeriodic = force.usesPeriodicBoundaryConditions();
431
432
433

    // Build the arrays.

434
435
    bondIndexArray.resize(numBonds, vector<int>(2));
    bondParamArray.resize(numBonds, vector<double>(numParameters));
436
    vector<double> params;
437
    for (int i = 0; i < numBonds; ++i) {
438
439
440
441
442
        int particle1, particle2;
        force.getBondParameters(i, particle1, particle2, params);
        bondIndexArray[i][0] = particle1;
        bondIndexArray[i][1] = particle2;
        for (int j = 0; j < numParameters; j++)
peastman's avatar
peastman committed
443
            bondParamArray[i][j] = params[j];
444
445
446
447
448
    }

    // Parse the expression used to calculate the force.

    Lepton::ParsedExpression expression = Lepton::Parser::parse(force.getEnergyFunction()).optimize();
449
450
    energyExpression = expression.createCompiledExpression();
    forceExpression = expression.differentiate("r").createCompiledExpression();
451
452
453
454
    for (int i = 0; i < numParameters; i++)
        parameterNames.push_back(force.getPerBondParameterName(i));
    for (int i = 0; i < force.getNumGlobalParameters(); i++)
        globalParameterNames.push_back(force.getGlobalParameterName(i));
455
456
457
458
459
    for (int i = 0; i < force.getNumEnergyParameterDerivatives(); i++) {
        string param = force.getEnergyParameterDerivativeName(i);
        energyParamDerivNames.push_back(param);
        energyParamDerivExpressions.push_back(expression.differentiate(param).createCompiledExpression());
    }
460
461
462
463
464
    set<string> variables;
    variables.insert("r");
    variables.insert(parameterNames.begin(), parameterNames.end());
    variables.insert(globalParameterNames.begin(), globalParameterNames.end());
    validateVariables(expression.getRootNode(), variables);
465
    ixn = new ReferenceCustomBondIxn(energyExpression, forceExpression, parameterNames, energyParamDerivExpressions);
466
467
}

468
double ReferenceCalcCustomBondForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
469
470
471
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
472
    map<string, double> globalParameters;
peastman's avatar
peastman committed
473
474
    for (auto& name : globalParameterNames)
        globalParameters[name] = context.getParameter(name);
475
    ixn->setGlobalParameters(globalParameters);
476
    if (usePeriodic)
477
        ixn->setPeriodic(extractBoxVectors(context));
478
479
    vector<double> energyParamDerivValues(energyParamDerivNames.size()+1, 0.0);
    for (int i = 0; i < numBonds; i++)
480
        ixn->calculateBondIxn(bondIndexArray[i], posData, bondParamArray[i], forceData, includeEnergy ? &energy : NULL, &energyParamDerivValues[0]);
481
482
483
    map<string, double>& energyParamDerivs = extractEnergyParameterDerivatives(context);
    for (int i = 0; i < energyParamDerivNames.size(); i++)
        energyParamDerivs[energyParamDerivNames[i]] += energyParamDerivValues[i];
484
485
486
    return energy;
}

487
void ReferenceCalcCustomBondForceKernel::copyParametersToContext(ContextImpl& context, const CustomBondForce& force, int firstBond, int lastBond) {
488
489
490
491
492
493
494
    if (numBonds != force.getNumBonds())
        throw OpenMMException("updateParametersInContext: The number of bonds has changed");

    // Record the values.

    int numParameters = force.getNumPerBondParameters();
    vector<double> params;
495
    for (int i = firstBond; i <= lastBond; ++i) {
496
497
498
499
500
        int particle1, particle2;
        force.getBondParameters(i, particle1, particle2, params);
        if (particle1 != bondIndexArray[i][0] || particle2 != bondIndexArray[i][1])
            throw OpenMMException("updateParametersInContext: The set of particles in a bond has changed");
        for (int j = 0; j < numParameters; j++)
peastman's avatar
peastman committed
501
            bondParamArray[i][j] = params[j];
502
503
504
    }
}

505
506
void ReferenceCalcHarmonicAngleForceKernel::initialize(const System& system, const HarmonicAngleForce& force) {
    numAngles = force.getNumAngles();
507
508
    angleIndexArray.resize(numAngles, vector<int>(3));
    angleParamArray.resize(numAngles, vector<double>(2));
509
    for (int i = 0; i < numAngles; ++i) {
Peter Eastman's avatar
Peter Eastman committed
510
        int particle1, particle2, particle3;
511
        double angle, k;
Peter Eastman's avatar
Peter Eastman committed
512
513
514
515
        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
516
517
        angleParamArray[i][0] = angle;
        angleParamArray[i][1] = k;
518
    }
519
    usePeriodic = force.usesPeriodicBoundaryConditions();
520
521
}

522
double ReferenceCalcHarmonicAngleForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
523
524
525
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
526
527
    ReferenceBondForce refBondForce;
    ReferenceAngleBondIxn angleBond;
528
529
    if (usePeriodic)
        angleBond.setPeriodic(extractBoxVectors(context));
530
    refBondForce.calculateForce(numAngles, angleIndexArray, posData, angleParamArray, forceData, includeEnergy ? &energy : NULL, angleBond);
531
532
533
    return energy;
}

534
void ReferenceCalcHarmonicAngleForceKernel::copyParametersToContext(ContextImpl& context, const HarmonicAngleForce& force, int firstAngle, int lastAngle) {
535
536
537
538
539
    if (numAngles != force.getNumAngles())
        throw OpenMMException("updateParametersInContext: The number of angles has changed");

    // Record the values.

540
    for (int i = firstAngle; i <= lastAngle; ++i) {
541
542
543
544
545
        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
546
547
        angleParamArray[i][0] = angle;
        angleParamArray[i][1] = k;
548
549
550
    }
}

551
552
553
554
555
ReferenceCalcCustomAngleForceKernel::~ReferenceCalcCustomAngleForceKernel() {
    if (ixn != NULL)
        delete ixn;
}

556
557
558
void ReferenceCalcCustomAngleForceKernel::initialize(const System& system, const CustomAngleForce& force) {
    numAngles = force.getNumAngles();
    int numParameters = force.getNumPerAngleParameters();
559
    usePeriodic = force.usesPeriodicBoundaryConditions();
560
561
562

    // Build the arrays.

563
564
    angleIndexArray.resize(numAngles, vector<int>(3));
    angleParamArray.resize(numAngles, vector<double>(numParameters));
565
    vector<double> params;
566
    for (int i = 0; i < numAngles; ++i) {
567
568
569
570
571
572
        int particle1, particle2, particle3;
        force.getAngleParameters(i, particle1, particle2, particle3, params);
        angleIndexArray[i][0] = particle1;
        angleIndexArray[i][1] = particle2;
        angleIndexArray[i][2] = particle3;
        for (int j = 0; j < numParameters; j++)
peastman's avatar
peastman committed
573
            angleParamArray[i][j] = params[j];
574
575
576
577
578
    }

    // Parse the expression used to calculate the force.

    Lepton::ParsedExpression expression = Lepton::Parser::parse(force.getEnergyFunction()).optimize();
579
580
    energyExpression = expression.createCompiledExpression();
    forceExpression = expression.differentiate("theta").createCompiledExpression();
581
582
583
584
    for (int i = 0; i < numParameters; i++)
        parameterNames.push_back(force.getPerAngleParameterName(i));
    for (int i = 0; i < force.getNumGlobalParameters(); i++)
        globalParameterNames.push_back(force.getGlobalParameterName(i));
585
586
587
588
589
    for (int i = 0; i < force.getNumEnergyParameterDerivatives(); i++) {
        string param = force.getEnergyParameterDerivativeName(i);
        energyParamDerivNames.push_back(param);
        energyParamDerivExpressions.push_back(expression.differentiate(param).createCompiledExpression());
    }
590
591
592
593
594
    set<string> variables;
    variables.insert("theta");
    variables.insert(parameterNames.begin(), parameterNames.end());
    variables.insert(globalParameterNames.begin(), globalParameterNames.end());
    validateVariables(expression.getRootNode(), variables);
595
    ixn = new ReferenceCustomAngleIxn(energyExpression, forceExpression, parameterNames, energyParamDerivExpressions);
596
597
}

598
double ReferenceCalcCustomAngleForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
599
600
601
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
602
    map<string, double> globalParameters;
peastman's avatar
peastman committed
603
604
    for (auto& name : globalParameterNames)
        globalParameters[name] = context.getParameter(name);
605
    ixn->setGlobalParameters(globalParameters);
606
    if (usePeriodic)
607
        ixn->setPeriodic(extractBoxVectors(context));
608
609
    vector<double> energyParamDerivValues(energyParamDerivNames.size()+1, 0.0);
    for (int i = 0; i < numAngles; i++)
610
        ixn->calculateBondIxn(angleIndexArray[i], posData, angleParamArray[i], forceData, includeEnergy ? &energy : NULL, &energyParamDerivValues[0]);
611
612
613
    map<string, double>& energyParamDerivs = extractEnergyParameterDerivatives(context);
    for (int i = 0; i < energyParamDerivNames.size(); i++)
        energyParamDerivs[energyParamDerivNames[i]] += energyParamDerivValues[i];
614
615
616
    return energy;
}

617
void ReferenceCalcCustomAngleForceKernel::copyParametersToContext(ContextImpl& context, const CustomAngleForce& force, int firstAngle, int lastAngle) {
618
619
620
621
622
623
624
    if (numAngles != force.getNumAngles())
        throw OpenMMException("updateParametersInContext: The number of angles has changed");

    // Record the values.

    int numParameters = force.getNumPerAngleParameters();
    vector<double> params;
625
    for (int i = firstAngle; i <= lastAngle; ++i) {
626
627
628
629
630
        int particle1, particle2, particle3;
        force.getAngleParameters(i, particle1, particle2, particle3, params);
        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");
        for (int j = 0; j < numParameters; j++)
peastman's avatar
peastman committed
631
            angleParamArray[i][j] = params[j];
632
633
634
    }
}

635
636
void ReferenceCalcPeriodicTorsionForceKernel::initialize(const System& system, const PeriodicTorsionForce& force) {
    numTorsions = force.getNumTorsions();
637
638
    torsionIndexArray.resize(numTorsions, vector<int>(4));
    torsionParamArray.resize(numTorsions, vector<double>(3));
639
    for (int i = 0; i < numTorsions; ++i) {
Peter Eastman's avatar
Peter Eastman committed
640
        int particle1, particle2, particle3, particle4, periodicity;
641
        double phase, k;
Peter Eastman's avatar
Peter Eastman committed
642
643
644
645
646
        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
647
648
649
        torsionParamArray[i][0] = k;
        torsionParamArray[i][1] = phase;
        torsionParamArray[i][2] = periodicity;
650
    }
651
    usePeriodic = force.usesPeriodicBoundaryConditions();
652
653
}

654
double ReferenceCalcPeriodicTorsionForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
655
656
657
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
658
659
    ReferenceBondForce refBondForce;
    ReferenceProperDihedralBond periodicTorsionBond;
660
661
    if (usePeriodic)
        periodicTorsionBond.setPeriodic(extractBoxVectors(context));
662
    refBondForce.calculateForce(numTorsions, torsionIndexArray, posData, torsionParamArray, forceData, includeEnergy ? &energy : NULL, periodicTorsionBond);
663
664
665
    return energy;
}

666
void ReferenceCalcPeriodicTorsionForceKernel::copyParametersToContext(ContextImpl& context, const PeriodicTorsionForce& force, int firstTorsion, int lastTorsion) {
667
668
669
670
671
    if (numTorsions != force.getNumTorsions())
        throw OpenMMException("updateParametersInContext: The number of torsions has changed");

    // Record the values.

672
    for (int i = firstTorsion; i <= lastTorsion; ++i) {
673
674
675
676
677
        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
678
679
680
        torsionParamArray[i][0] = k;
        torsionParamArray[i][1] = phase;
        torsionParamArray[i][2] = periodicity;
681
682
683
    }
}

684
685
void ReferenceCalcRBTorsionForceKernel::initialize(const System& system, const RBTorsionForce& force) {
    numTorsions = force.getNumTorsions();
686
687
    torsionIndexArray.resize(numTorsions, vector<int>(4));
    torsionParamArray.resize(numTorsions, vector<double>(6));
688
    for (int i = 0; i < numTorsions; ++i) {
Peter Eastman's avatar
Peter Eastman committed
689
        int particle1, particle2, particle3, particle4;
690
        double c0, c1, c2, c3, c4, c5;
Peter Eastman's avatar
Peter Eastman committed
691
692
693
694
695
        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
696
697
698
699
700
701
        torsionParamArray[i][0] = c0;
        torsionParamArray[i][1] = c1;
        torsionParamArray[i][2] = c2;
        torsionParamArray[i][3] = c3;
        torsionParamArray[i][4] = c4;
        torsionParamArray[i][5] = c5;
702
    }
703
    usePeriodic = force.usesPeriodicBoundaryConditions();
704
705
}

706
double ReferenceCalcRBTorsionForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
707
708
709
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
710
711
    ReferenceBondForce refBondForce;
    ReferenceRbDihedralBond rbTorsionBond;
712
713
    if (usePeriodic)
        rbTorsionBond.setPeriodic(extractBoxVectors(context));
714
    refBondForce.calculateForce(numTorsions, torsionIndexArray, posData, torsionParamArray, forceData, includeEnergy ? &energy : NULL, rbTorsionBond);
715
716
717
    return energy;
}

718
719
720
721
722
723
724
725
726
727
728
729
void ReferenceCalcRBTorsionForceKernel::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
730
731
732
733
734
735
        torsionParamArray[i][0] = c0;
        torsionParamArray[i][1] = c1;
        torsionParamArray[i][2] = c2;
        torsionParamArray[i][3] = c3;
        torsionParamArray[i][4] = c4;
        torsionParamArray[i][5] = c5;
736
737
738
    }
}

739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
void ReferenceCalcCMAPTorsionForceKernel::initialize(const System& system, const CMAPTorsionForce& force) {
    int numMaps = force.getNumMaps();
    int numTorsions = force.getNumTorsions();
    coeff.resize(numMaps);
    vector<double> energy;
    vector<vector<double> > c;
    for (int i = 0; i < numMaps; i++) {
        int size;
        force.getMapParameters(i, size, energy);
        CMAPTorsionForceImpl::calcMapDerivatives(size, energy, c);
        coeff[i].resize(size*size);
        for (int j = 0; j < size*size; j++) {
            coeff[i][j].resize(16);
            for (int k = 0; k < 16; k++)
                coeff[i][j][k] = c[j][k];
        }
    }
    torsionMaps.resize(numTorsions);
    torsionIndices.resize(numTorsions);
    for (int i = 0; i < numTorsions; i++) {
        torsionIndices[i].resize(8);
        force.getTorsionParameters(i, torsionMaps[i], torsionIndices[i][0], torsionIndices[i][1], torsionIndices[i][2],
            torsionIndices[i][3], torsionIndices[i][4], torsionIndices[i][5], torsionIndices[i][6], torsionIndices[i][7]);
    }
763
    usePeriodic = force.usesPeriodicBoundaryConditions();
764
765
}

766
double ReferenceCalcCMAPTorsionForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
767
768
769
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double totalEnergy = 0;
770
    ReferenceCMAPTorsionIxn torsion(coeff, torsionMaps, torsionIndices);
771
772
    if (usePeriodic)
        torsion.setPeriodic(extractBoxVectors(context));
773
774
775
776
    torsion.calculateIxn(posData, forceData, &totalEnergy);
    return totalEnergy;
}

777
778
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
805
806
807
808
809
810
void ReferenceCalcCMAPTorsionForceKernel::copyParametersToContext(ContextImpl& context, const CMAPTorsionForce& force) {
    int numMaps = force.getNumMaps();
    int numTorsions = force.getNumTorsions();
    if (coeff.size() != numMaps)
        throw OpenMMException("updateParametersInContext: The number of maps has changed");
    if (torsionMaps.size() != numTorsions)
        throw OpenMMException("updateParametersInContext: The number of CMAP torsions has changed");

    // Update the maps.

    vector<double> energy;
    vector<vector<double> > c;
    for (int i = 0; i < numMaps; i++) {
        int size;
        force.getMapParameters(i, size, energy);
        if (coeff[i].size() != size*size)
            throw OpenMMException("updateParametersInContext: The size of a map has changed");
        CMAPTorsionForceImpl::calcMapDerivatives(size, energy, c);
        for (int j = 0; j < size*size; j++)
            for (int k = 0; k < 16; k++)
                coeff[i][j][k] = c[j][k];
    }

    // Update the indices.

    for (int i = 0; i < numTorsions; i++) {
        int index[8];
        force.getTorsionParameters(i, torsionMaps[i], index[0], index[1], index[2], index[3], index[4], index[5], index[6], index[7]);
        for (int j = 0; j < 8; j++)
            if (index[j] != torsionIndices[i][j])
                throw OpenMMException("updateParametersInContext: The set of particles in a CMAP torsion has changed");
    }
}

811
812
813
814
815
ReferenceCalcCustomTorsionForceKernel::~ReferenceCalcCustomTorsionForceKernel() {
    if (ixn != NULL)
        delete ixn;
}

816
817
818
void ReferenceCalcCustomTorsionForceKernel::initialize(const System& system, const CustomTorsionForce& force) {
    numTorsions = force.getNumTorsions();
    int numParameters = force.getNumPerTorsionParameters();
819
    usePeriodic = force.usesPeriodicBoundaryConditions();
820
821
822

    // Build the arrays.

823
824
    torsionIndexArray.resize(numTorsions, vector<int>(4));
    torsionParamArray.resize(numTorsions, vector<double>(numParameters));
825
    vector<double> params;
826
    for (int i = 0; i < numTorsions; ++i) {
827
828
829
830
831
832
833
        int particle1, particle2, particle3, particle4;
        force.getTorsionParameters(i, particle1, particle2, particle3, particle4, params);
        torsionIndexArray[i][0] = particle1;
        torsionIndexArray[i][1] = particle2;
        torsionIndexArray[i][2] = particle3;
        torsionIndexArray[i][3] = particle4;
        for (int j = 0; j < numParameters; j++)
peastman's avatar
peastman committed
834
            torsionParamArray[i][j] = params[j];
835
836
837
838
839
    }

    // Parse the expression used to calculate the force.

    Lepton::ParsedExpression expression = Lepton::Parser::parse(force.getEnergyFunction()).optimize();
840
841
    energyExpression = expression.createCompiledExpression();
    forceExpression = expression.differentiate("theta").createCompiledExpression();
842
843
844
845
    for (int i = 0; i < numParameters; i++)
        parameterNames.push_back(force.getPerTorsionParameterName(i));
    for (int i = 0; i < force.getNumGlobalParameters(); i++)
        globalParameterNames.push_back(force.getGlobalParameterName(i));
846
847
848
849
850
    for (int i = 0; i < force.getNumEnergyParameterDerivatives(); i++) {
        string param = force.getEnergyParameterDerivativeName(i);
        energyParamDerivNames.push_back(param);
        energyParamDerivExpressions.push_back(expression.differentiate(param).createCompiledExpression());
    }
851
852
853
854
855
    set<string> variables;
    variables.insert("theta");
    variables.insert(parameterNames.begin(), parameterNames.end());
    variables.insert(globalParameterNames.begin(), globalParameterNames.end());
    validateVariables(expression.getRootNode(), variables);
856
    ixn = new ReferenceCustomTorsionIxn(energyExpression, forceExpression, parameterNames, energyParamDerivExpressions);
857
858
}

859
double ReferenceCalcCustomTorsionForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
860
861
862
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
863
    map<string, double> globalParameters;
peastman's avatar
peastman committed
864
865
    for (auto& name : globalParameterNames)
        globalParameters[name] = context.getParameter(name);
866
    ixn->setGlobalParameters(globalParameters);
867
    if (usePeriodic)
868
        ixn->setPeriodic(extractBoxVectors(context));
869
870
    vector<double> energyParamDerivValues(energyParamDerivNames.size()+1, 0.0);
    for (int i = 0; i < numTorsions; i++)
871
        ixn->calculateBondIxn(torsionIndexArray[i], posData, torsionParamArray[i], forceData, includeEnergy ? &energy : NULL, &energyParamDerivValues[0]);
872
873
874
    map<string, double>& energyParamDerivs = extractEnergyParameterDerivatives(context);
    for (int i = 0; i < energyParamDerivNames.size(); i++)
        energyParamDerivs[energyParamDerivNames[i]] += energyParamDerivValues[i];
875
876
877
    return energy;
}

878
void ReferenceCalcCustomTorsionForceKernel::copyParametersToContext(ContextImpl& context, const CustomTorsionForce& force, int firstTorsion, int lastTorsion) {
879
880
881
882
883
884
885
    if (numTorsions != force.getNumTorsions())
        throw OpenMMException("updateParametersInContext: The number of torsions has changed");

    // Record the values.

    int numParameters = force.getNumPerTorsionParameters();
    vector<double> params;
886
    for (int i = firstTorsion; i <= lastTorsion; ++i) {
887
888
889
890
891
        int particle1, particle2, particle3, particle4;
        force.getTorsionParameters(i, particle1, particle2, particle3, particle4, params);
        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");
        for (int j = 0; j < numParameters; j++)
peastman's avatar
peastman committed
892
            torsionParamArray[i][j] = params[j];
893
894
895
    }
}

896
897
898
899
900
ReferenceCalcNonbondedForceKernel::~ReferenceCalcNonbondedForceKernel() {
    if (neighborList != NULL)
        delete neighborList;
}

901
902
903
904
void ReferenceCalcNonbondedForceKernel::initialize(const System& system, const NonbondedForce& force) {

    // Identify which exceptions are 1-4 interactions.

905
906
907
908
909
910
911
912
    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);
    }
Peter Eastman's avatar
Peter Eastman committed
913
    numParticles = force.getNumParticles();
914
915
916
917
918
919
920
921
    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);
922
923
        if (chargeProd != 0.0 || epsilon != 0.0 || exceptionsWithOffsets.find(i) != exceptionsWithOffsets.end()) {
            nb14Index[i] = nb14s.size();
924
            nb14s.push_back(i);
925
        }
926
927
928
929
930
    }

    // Build the arrays.

    num14 = nb14s.size();
931
932
933
    bonded14IndexArray.resize(num14, vector<int>(2));
    bonded14ParamArray.resize(num14, vector<double>(3));
    particleParamArray.resize(numParticles, vector<double>(3));
934
935
936
937
    baseParticleParams.resize(numParticles);
    baseExceptionParams.resize(num14);
    for (int i = 0; i < numParticles; ++i)
       force.getParticleParameters(i, baseParticleParams[i][0], baseParticleParams[i][1], baseParticleParams[i][2]);
938
    for (int i = 0; i < num14; ++i) {
Peter Eastman's avatar
Peter Eastman committed
939
        int particle1, particle2;
940
        force.getExceptionParameters(nb14s[i], particle1, particle2, baseExceptionParams[i][0], baseExceptionParams[i][1], baseExceptionParams[i][2]);
Peter Eastman's avatar
Peter Eastman committed
941
942
        bonded14IndexArray[i][0] = particle1;
        bonded14IndexArray[i][1] = particle2;
943
944
945
946
947
948
949
950
951
952
953
954
955
    }
    for (int i = 0; i < force.getNumParticleParameterOffsets(); i++) {
        string param;
        int particle;
        double charge, sigma, epsilon;
        force.getParticleParameterOffset(i, param, particle, charge, sigma, epsilon);
        particleParamOffsets[make_pair(param, particle)] = {charge, sigma, epsilon};
    }
    for (int i = 0; i < force.getNumExceptionParameterOffsets(); i++) {
        string param;
        int exception;
        double charge, sigma, epsilon;
        force.getExceptionParameterOffset(i, param, exception, charge, sigma, epsilon);
956
        exceptionParamOffsets[make_pair(param, nb14Index[exception])] = {charge, sigma, epsilon};
957
    }
958
    nonbondedMethod = CalcNonbondedForceKernel::NonbondedMethod(force.getNonbondedMethod());
peastman's avatar
peastman committed
959
    nonbondedCutoff = force.getCutoffDistance();
960
    if (nonbondedMethod == NoCutoff) {
961
        neighborList = NULL;
962
963
964
        useSwitchingFunction = false;
    }
    else {
965
        neighborList = new NeighborList();
966
967
968
        useSwitchingFunction = force.getUseSwitchingFunction();
        switchingDistance = force.getSwitchingDistance();
    }
969
970
971
    if (nonbondedMethod == Ewald) {
        double alpha;
        NonbondedForceImpl::calcEwaldParameters(system, force, alpha, kmax[0], kmax[1], kmax[2]);
peastman's avatar
peastman committed
972
        ewaldAlpha = alpha;
973
974
975
    }
    else if (nonbondedMethod == PME) {
        double alpha;
976
        NonbondedForceImpl::calcPMEParameters(system, force, alpha, gridSize[0], gridSize[1], gridSize[2], false);
peastman's avatar
peastman committed
977
        ewaldAlpha = alpha;
978
    }
979
980
981
    else if (nonbondedMethod == LJPME) {
        double alpha;
        NonbondedForceImpl::calcPMEParameters(system, force, alpha, gridSize[0], gridSize[1], gridSize[2], false);
peastman's avatar
peastman committed
982
        ewaldAlpha = alpha;
983
        NonbondedForceImpl::calcPMEParameters(system, force, alpha, dispersionGridSize[0], dispersionGridSize[1], dispersionGridSize[2], true);
peastman's avatar
peastman committed
984
        ewaldDispersionAlpha = alpha;
985
986
        useSwitchingFunction = false;
    }
987
988
989
990
    if (nonbondedMethod == NoCutoff || nonbondedMethod == CutoffNonPeriodic)
        exceptionsArePeriodic = false;
    else
        exceptionsArePeriodic = force.getExceptionsUsePeriodicBoundaryConditions();
peastman's avatar
peastman committed
991
    rfDielectric = force.getReactionFieldDielectric();
992
993
994
995
    if (force.getUseDispersionCorrection())
        dispersionCoefficient = NonbondedForceImpl::calcDispersionCorrection(system, force);
    else
        dispersionCoefficient = 0.0;
996
997
}

998
double ReferenceCalcNonbondedForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy, bool includeDirect, bool includeReciprocal) {
999
    computeParameters(context);
peastman's avatar
peastman committed
1000
1001
1002
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
1003
    ReferenceLJCoulombIxn clj;
1004
    bool periodic = (nonbondedMethod == CutoffPeriodic);
1005
    bool ewald  = (nonbondedMethod == Ewald);
1006
    bool pme  = (nonbondedMethod == PME);
1007
    bool ljpme = (nonbondedMethod == LJPME);
1008
    if (nonbondedMethod != NoCutoff) {
1009
        computeNeighborListVoxelHash(*neighborList, numParticles, posData, exclusions, extractBoxVectors(context), periodic || ewald || pme || ljpme, nonbondedCutoff, 0.0);
1010
        clj.setUseCutoff(nonbondedCutoff, *neighborList, rfDielectric);
1011
    }
1012
    if (periodic || ewald || pme || ljpme) {
peastman's avatar
peastman committed
1013
        Vec3* boxVectors = extractBoxVectors(context);
1014
        double minAllowedSize = 1.999999*nonbondedCutoff;
1015
        if (boxVectors[0][0] < minAllowedSize || boxVectors[1][1] < minAllowedSize || boxVectors[2][2] < minAllowedSize)
1016
            throw OpenMMException("The periodic box size has decreased to less than twice the nonbonded cutoff.");
1017
        clj.setPeriodic(boxVectors);
1018
        clj.setPeriodicExceptions(exceptionsArePeriodic);
1019
    }
1020
1021
    if (ewald)
        clj.setUseEwald(ewaldAlpha, kmax[0], kmax[1], kmax[2]);
1022
    if (pme)
1023
        clj.setUsePME(ewaldAlpha, gridSize);
1024
1025
1026
1027
    if (ljpme){
        clj.setUsePME(ewaldAlpha, gridSize);
        clj.setUseLJPME(ewaldDispersionAlpha, dispersionGridSize);
    }
1028
1029
    if (useSwitchingFunction)
        clj.setUseSwitchingFunction(switchingDistance);
1030
    clj.calculatePairIxn(numParticles, posData, particleParamArray, exclusions, forceData, includeEnergy ? &energy : NULL, includeDirect, includeReciprocal);
1031
1032
1033
    if (includeDirect) {
        ReferenceBondForce refBondForce;
        ReferenceLJCoulomb14 nonbonded14;
1034
1035
1036
1037
        if (exceptionsArePeriodic) {
            Vec3* boxVectors = extractBoxVectors(context);
            nonbonded14.setPeriodic(boxVectors);
        }
1038
        refBondForce.calculateForce(num14, bonded14IndexArray, posData, bonded14ParamArray, forceData, includeEnergy ? &energy : NULL, nonbonded14);
1039
        if (periodic || ewald || pme) {
peastman's avatar
peastman committed
1040
            Vec3* boxVectors = extractBoxVectors(context);
1041
            energy += dispersionCoefficient/(boxVectors[0][0]*boxVectors[1][1]*boxVectors[2][2]);
1042
        }
1043
    }
1044
1045
1046
    return energy;
}

1047
void ReferenceCalcNonbondedForceKernel::copyParametersToContext(ContextImpl& context, const NonbondedForce& force, int firstParticle, int lastParticle, int firstException, int lastException) {
1048
1049
    if (force.getNumParticles() != numParticles)
        throw OpenMMException("updateParametersInContext: The number of particles has changed");
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
    if (force.getNumParticleParameterOffsets() != particleParamOffsets.size())
        throw OpenMMException("updateParametersInContext: The number of particle parameter offsets has changed");
    if (force.getNumExceptionParameterOffsets() != exceptionParamOffsets.size())
        throw OpenMMException("updateParametersInContext: The number of exception parameter offsets has changed");
    for (int i = 0; i < force.getNumParticleParameterOffsets(); i++) {
        string param;
        int particle;
        double charge, sigma, epsilon;
        force.getParticleParameterOffset(i, param, particle, charge, sigma, epsilon);
        if (particleParamOffsets.find(make_pair(param, particle)) == particleParamOffsets.end())
            throw OpenMMException("updateParametersInContext: The parameter or particle index of a particle parameter offset has changed");
    }
peastman's avatar
peastman committed
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071

    // 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);
1072
1073
        if (exceptionParamOffsets.find(make_pair(param, nb14Index[exception])) == exceptionParamOffsets.end())
            throw OpenMMException("updateParametersInContext: The parameter or exception index of an exception parameter offset has changed");
peastman's avatar
peastman committed
1074
    }
1075
1076
1077
1078
1079
    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);
1080
1081
1082
1083
1084
        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
1085
1086
1087
1088
1089
1090
1091
            nb14s.push_back(i);
    }
    if (nb14s.size() != num14)
        throw OpenMMException("updateParametersInContext: The number of non-excluded exceptions has changed");

    // Record the values.

1092
    for (int i = firstParticle; i <= lastParticle; ++i)
1093
        force.getParticleParameters(i, baseParticleParams[i][0], baseParticleParams[i][1], baseParticleParams[i][2]);
1094
1095
    for (int i = 0; i < num14; ++i) {
        int particle1, particle2;
1096
        force.getExceptionParameters(nb14s[i], particle1, particle2, baseExceptionParams[i][0], baseExceptionParams[i][1], baseExceptionParams[i][2]);
1097
1098
1099
        bonded14IndexArray[i][0] = particle1;
        bonded14IndexArray[i][1] = particle2;
    }
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
    particleParamOffsets.clear();
    exceptionParamOffsets.clear();
    for (int i = 0; i < force.getNumParticleParameterOffsets(); i++) {
        string param;
        int particle;
        double charge, sigma, epsilon;
        force.getParticleParameterOffset(i, param, particle, charge, sigma, epsilon);
        particleParamOffsets[make_pair(param, particle)] = {charge, sigma, epsilon};
    }
    for (int i = 0; i < force.getNumExceptionParameterOffsets(); i++) {
        string param;
        int exception;
        double charge, sigma, epsilon;
        force.getExceptionParameterOffset(i, param, exception, charge, sigma, epsilon);
        exceptionParamOffsets[make_pair(param, nb14Index[exception])] = {charge, sigma, epsilon};
    }
1116
1117
1118
1119
1120
1121
1122
1123
    
    // 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);
}

1124
void ReferenceCalcNonbondedForceKernel::getPMEParameters(double& alpha, int& nx, int& ny, int& nz) const {
1125
1126
    if (nonbondedMethod != PME && nonbondedMethod != LJPME)
        throw OpenMMException("getPMEParametersInContext: This Context is not using PME or LJPME");
1127
1128
1129
1130
1131
1132
    alpha = ewaldAlpha;
    nx = gridSize[0];
    ny = gridSize[1];
    nz = gridSize[2];
}

1133
void ReferenceCalcNonbondedForceKernel::getLJPMEParameters(double& alpha, int& nx, int& ny, int& nz) const {
1134
1135
    if (nonbondedMethod != LJPME)
        throw OpenMMException("getPMEParametersInContext: This Context is not using LJPME");
1136
1137
1138
1139
    alpha = ewaldDispersionAlpha;
    nx = dispersionGridSize[0];
    ny = dispersionGridSize[1];
    nz = dispersionGridSize[2];
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
void ReferenceCalcNonbondedForceKernel::computeParameters(ContextImpl& context) {
    // Compute particle parameters.

    vector<double> charges(numParticles), sigmas(numParticles), epsilons(numParticles);
    for (int i = 0; i < numParticles; i++) {
        charges[i] = baseParticleParams[i][0];
        sigmas[i] = baseParticleParams[i][1];
        epsilons[i] = baseParticleParams[i][2];
    }
    for (auto& offset : particleParamOffsets) {
        double value = context.getParameter(offset.first.first);
        int index = offset.first.second;
        charges[index] += value*offset.second[0];
        sigmas[index] += value*offset.second[1];
        epsilons[index] += value*offset.second[2];
    }
    for (int i = 0; i < numParticles; i++) {
        particleParamArray[i][0] = 0.5*sigmas[i];
        particleParamArray[i][1] = 2.0*sqrt(epsilons[i]);
        particleParamArray[i][2] = charges[i];
    }

    // Compute exception parameters.

    charges.resize(num14);
    sigmas.resize(num14);
    epsilons.resize(num14);
    for (int i = 0; i < num14; i++) {
        charges[i] = baseExceptionParams[i][0];
        sigmas[i] = baseExceptionParams[i][1];
        epsilons[i] = baseExceptionParams[i][2];
    }
    for (auto& offset : exceptionParamOffsets) {
        double value = context.getParameter(offset.first.first);
        int index = offset.first.second;
        charges[index] += value*offset.second[0];
        sigmas[index] += value*offset.second[1];
        epsilons[index] += value*offset.second[2];
    }
    for (int i = 0; i < num14; i++) {
        bonded14ParamArray[i][0] = sigmas[i];
        bonded14ParamArray[i][1] = 4.0*epsilons[i];
        bonded14ParamArray[i][2] = charges[i];
    }
}

1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
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
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
ReferenceCalcConstantPotentialForceKernel::~ReferenceCalcConstantPotentialForceKernel() {
    if (neighborList != NULL) {
        delete neighborList;
    }
    if (solver != NULL) {
        delete solver;
    }
}

void ReferenceCalcConstantPotentialForceKernel::initialize(const System& system, const ConstantPotentialForce& force) {
    // Get particle parameters.
    numParticles = force.getNumParticles();
    charges.resize(numParticles);
    for (int i = 0; i < numParticles; i++) {
        force.getParticleParameters(i, charges[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;
    }

    // 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.
    sysToElec.resize(numParticles, -1);
    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] = electrodeParamArray.size();
            elecToSys.push_back(i);
            electrodeParamArray.push_back({potential, gaussianWidth, thomasFermiScale});
        }
    }

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

    // Set options from force.
    nonbondedCutoff = force.getCutoffDistance();
    neighborList = new NeighborList();
    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);

    // 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 (sysToElec[i] == -1) {
            chargeTarget -= charges[i];
        }
    }

    ConstantPotentialForce::ConstantPotentialMethod method = force.getConstantPotentialMethod();
    if (method == ConstantPotentialForce::Matrix) {
        solver = new ReferenceConstantPotentialMatrixSolver(numElectrodeParticles);
    }
    else if (method == ConstantPotentialForce::CG) {
        solver = new ReferenceConstantPotentialCGSolver(numElectrodeParticles, force.getUsePreconditioner());
    }
    else {
        throw OpenMMException("internal error: invalid constant potential method");
    }
}

double ReferenceCalcConstantPotentialForceKernel::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;

    // Solve for charges, then calculate forces and energy.
    updateNeighborList(boxVectors, posData);
    ReferenceConstantPotential conp(nonbondedCutoff, neighborList, boxVectors, exceptionsArePeriodic, ewaldAlpha, gridSize, cgErrorTol, useChargeConstraint, chargeTarget, externalField);
    solver->update(numParticles, numElectrodeParticles, posData, charges, exclusions, sysToElec, elecToSys, electrodeParamArray, conp);
    conp.execute(numParticles, numElectrodeParticles, posData, forceData, charges, exclusions, sysToElec, elecToSys, electrodeParamArray, includeEnergy ? &energy : NULL, solver);

    // 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.
    ReferenceBondForce refBondForce;
    ReferenceConstantPotential14 conp14;
    if (exceptionsArePeriodic) {
        conp14.setPeriodic(boxVectors);
    }
    refBondForce.calculateForce(num14, bonded14IndexArray, posData, bonded14ParamArray, forceData, includeEnergy ? &energy : NULL, conp14);

    return energy;
}

void ReferenceCalcConstantPotentialForceKernel::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 (sysToElec[i] == -1) {
            force.getParticleParameters(i, charges[i]);
        }
    }

    // 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) {
            int ii = sysToElec[i];
            if (ii == -1) {
                // Particle was not an electrode particle but is now.
                throw OpenMMException("updateParametersInContext: The electrode state of a particle has changed");
            }
            electrodeParamArray[ii][0] = potential;
            electrodeParamArray[ii][1] = gaussianWidth;
            electrodeParamArray[ii][2] = thomasFermiScale;
            allElectrodeParticles.insert(i);
        }
    }
    if (allElectrodeParticles.size() != numElectrodeParticles) {
        // Particle that was an electrode particle might not be now.
        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 (sysToElec[i] == -1) {
            chargeTarget -= charges[i];
        }
    }

    // Invalidate matrix or CG data if electrode parameters changed.
    if (firstElectrode <= lastElectrode) {
        solver->invalidate();
    }
}

void ReferenceCalcConstantPotentialForceKernel::getPMEParameters(double& alpha, int& nx, int& ny, int& nz) const {
    alpha = ewaldAlpha;
    nx = gridSize[0];
    ny = gridSize[1];
    nz = gridSize[2];
}

void ReferenceCalcConstantPotentialForceKernel::getCharges(ContextImpl& context, std::vector<double>& chargesOut) {
    Vec3* boxVectors = extractBoxVectors(context);
    vector<Vec3>& posData = extractPositions(context);
    
    // Solve for charges only.
    updateNeighborList(boxVectors, posData);
    ReferenceConstantPotential conp(nonbondedCutoff, neighborList, boxVectors, exceptionsArePeriodic, ewaldAlpha, gridSize, cgErrorTol, useChargeConstraint, chargeTarget, externalField);
    solver->update(numParticles, numElectrodeParticles, posData, charges, exclusions, sysToElec, elecToSys, electrodeParamArray, conp);
    conp.getCharges(numParticles, numElectrodeParticles, posData, charges, exclusions, sysToElec, elecToSys, electrodeParamArray, solver);

    chargesOut = charges;
}

void ReferenceCalcConstantPotentialForceKernel::updateNeighborList(const Vec3* boxVectors, const std::vector<Vec3>& posData) {
    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.");
    }
    computeNeighborListVoxelHash(*neighborList, numParticles, posData, exclusions, boxVectors, true, nonbondedCutoff, 0.0);
}

1421
1422
1423
ReferenceCalcCustomNonbondedForceKernel::~ReferenceCalcCustomNonbondedForceKernel() {
    if (neighborList != NULL)
        delete neighborList;
1424
1425
    if (forceCopy != NULL)
        delete forceCopy;
1426
1427
1428
1429
}

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

1430
    // Record the exclusions.
1431
1432
1433

    numParticles = force.getNumParticles();
    exclusions.resize(numParticles);
1434
    for (int i = 0; i < force.getNumExclusions(); i++) {
1435
        int particle1, particle2;
1436
        force.getExclusionParticles(i, particle1, particle2);
1437
1438
1439
1440
1441
1442
        exclusions[particle1].insert(particle2);
        exclusions[particle2].insert(particle1);
    }

    // Build the arrays.

1443
1444
1445
    particleParamArray.resize(numParticles);
    for (int i = 0; i < numParticles; ++i)
        force.getParticleParameters(i, particleParamArray[i]);
1446
    nonbondedMethod = CalcCustomNonbondedForceKernel::NonbondedMethod(force.getNonbondedMethod());
peastman's avatar
peastman committed
1447
    nonbondedCutoff = force.getCutoffDistance();
1448
    if (nonbondedMethod == NoCutoff) {
1449
        neighborList = NULL;
1450
1451
1452
        useSwitchingFunction = false;
    }
    else {
1453
        neighborList = new NeighborList();
1454
1455
1456
        useSwitchingFunction = force.getUseSwitchingFunction();
        switchingDistance = force.getSwitchingDistance();
    }
1457

1458
    // Record the tabulated function update counts for future reference.
1459
1460

    for (int i = 0; i < force.getNumTabulatedFunctions(); i++)
1461
        tabulatedFunctionUpdateCount[force.getTabulatedFunctionName(i)] = force.getTabulatedFunction(i).getUpdateCount();
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487

    // Create the expressions.
    
    createExpressions(force);
    
    // 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));
    }
}

void ReferenceCalcCustomNonbondedForceKernel::createExpressions(const CustomNonbondedForce& force) {
1488
1489
1490
    // Create custom functions for the tabulated functions.

    map<string, Lepton::CustomFunction*> functions;
1491
    for (int i = 0; i < force.getNumTabulatedFunctions(); i++)
1492
        functions[force.getTabulatedFunctionName(i)] = createReferenceTabulatedFunction(force.getTabulatedFunction(i));
1493

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

1496
    Lepton::ParsedExpression expression = Lepton::Parser::parse(force.getEnergyFunction(), functions).optimize();
1497
1498
    energyExpression = expression.createCompiledExpression();
    forceExpression = expression.differentiate("r").createCompiledExpression();
1499
1500
1501
    parameterNames.clear();
    globalParameterNames.clear();
    globalParamValues.clear();
1502
1503
    computedValueNames.clear();
    computedValueExpressions.clear();
1504
1505
1506
    energyParamDerivNames.clear();
    energyParamDerivExpressions.clear();
    for (int i = 0; i < force.getNumPerParticleParameters(); i++)
1507
        parameterNames.push_back(force.getPerParticleParameterName(i));
1508
    for (int i = 0; i < force.getNumGlobalParameters(); i++) {
1509
        globalParameterNames.push_back(force.getGlobalParameterName(i));
1510
1511
        globalParamValues[force.getGlobalParameterName(i)] = force.getGlobalParameterDefaultValue(i);
    }
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
    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());
    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);
        computedValueExpressions.push_back(parsed.createCompiledExpression());
    }
1530
1531
1532
1533
1534
    for (int i = 0; i < force.getNumEnergyParameterDerivatives(); i++) {
        string param = force.getEnergyParameterDerivativeName(i);
        energyParamDerivNames.push_back(param);
        energyParamDerivExpressions.push_back(expression.differentiate(param).createCompiledExpression());
    }
1535
1536
1537
    for (auto& name : computedValueNames) {
        pairVariables.insert(name+"1");
        pairVariables.insert(name+"2");
1538
    }
1539
    validateVariables(expression.getRootNode(), pairVariables);
1540
1541
1542

    // Delete the custom functions.

peastman's avatar
peastman committed
1543
1544
    for (auto& function : functions)
        delete function.second;
1545
1546
}

1547
double ReferenceCalcCustomNonbondedForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
1548
1549
1550
1551
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    Vec3* boxVectors = extractBoxVectors(context);
    double energy = 0;
1552
    ReferenceCustomNonbondedIxn ixn(energyExpression, forceExpression, parameterNames, energyParamDerivExpressions, computedValueNames, computedValueExpressions);
1553
1554
    bool periodic = (nonbondedMethod == CutoffPeriodic);
    if (nonbondedMethod != NoCutoff) {
1555
        computeNeighborListVoxelHash(*neighborList, numParticles, posData, exclusions, extractBoxVectors(context), periodic, nonbondedCutoff, 0.0);
1556
1557
        ixn.setUseCutoff(nonbondedCutoff, *neighborList);
    }
1558
1559
    if (periodic) {
        double minAllowedSize = 2*nonbondedCutoff;
1560
        if (boxVectors[0][0] < minAllowedSize || boxVectors[1][1] < minAllowedSize || boxVectors[2][2] < minAllowedSize)
1561
            throw OpenMMException("The periodic box size has decreased to less than twice the nonbonded cutoff.");
1562
        ixn.setPeriodic(boxVectors);
1563
    }
1564
1565
    if (interactionGroups.size() > 0)
        ixn.setInteractionGroups(interactionGroups);
1566
    bool globalParamsChanged = false;
peastman's avatar
peastman committed
1567
1568
1569
    for (auto& name : globalParameterNames) {
        double value = context.getParameter(name);
        if (globalParamValues[name] != value)
1570
            globalParamsChanged = true;
peastman's avatar
peastman committed
1571
        globalParamValues[name] = value;
1572
    }
1573
1574
    if (useSwitchingFunction)
        ixn.setUseSwitchingFunction(switchingDistance);
1575
    vector<double> energyParamDerivValues(energyParamDerivNames.size()+1, 0.0);
1576
    ixn.calculatePairIxn(numParticles, posData, particleParamArray, exclusions, globalParamValues, forceData, includeEnergy ? &energy : NULL, &energyParamDerivValues[0]);
1577
1578
1579
    map<string, double>& energyParamDerivs = extractEnergyParameterDerivatives(context);
    for (int i = 0; i < energyParamDerivNames.size(); i++)
        energyParamDerivs[energyParamDerivNames[i]] += energyParamDerivValues[i];
1580
1581
1582
    
    // Add in the long range correction.
    
1583
    if (!hasInitializedLongRangeCorrection) {
1584
        ThreadPool& threads = extractThreadPool(context);
1585
        longRangeCorrectionData = CustomNonbondedForceImpl::prepareLongRangeCorrection(*forceCopy, threads.getNumThreads());
1586
        CustomNonbondedForceImpl::calcLongRangeCorrection(*forceCopy, longRangeCorrectionData, context.getOwner(), longRangeCoefficient, longRangeCoefficientDerivs, threads);
1587
1588
        hasInitializedLongRangeCorrection = true;
    }
1589
    else if (globalParamsChanged && forceCopy != NULL) {
1590
        ThreadPool& threads = extractThreadPool(context);
1591
1592
        CustomNonbondedForceImpl::calcLongRangeCorrection(*forceCopy, longRangeCorrectionData, context.getOwner(), longRangeCoefficient, longRangeCoefficientDerivs, threads);
    }
1593
1594
1595
1596
    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;
1597
1598
1599
    return energy;
}

1600
void ReferenceCalcCustomNonbondedForceKernel::copyParametersToContext(ContextImpl& context, const CustomNonbondedForce& force, int firstParticle, int lastParticle) {
1601
1602
1603
1604
1605
1606
    if (numParticles != force.getNumParticles())
        throw OpenMMException("updateParametersInContext: The number of particles has changed");

    // Record the values.

    int numParameters = force.getNumPerParticleParameters();
1607
1608
    vector<double> parameters;
    for (int i = firstParticle; i <= lastParticle; ++i) {
1609
1610
        force.getParticleParameters(i, parameters);
        for (int j = 0; j < numParameters; j++)
peastman's avatar
peastman committed
1611
            particleParamArray[i][j] = parameters[j];
1612
    }
1613
1614
1615
1616
    
    // If necessary, recompute the long range correction.
    
    if (forceCopy != NULL) {
1617
        ThreadPool& threads = extractThreadPool(context);
1618
        longRangeCorrectionData = CustomNonbondedForceImpl::prepareLongRangeCorrection(force, threads.getNumThreads());
1619
        CustomNonbondedForceImpl::calcLongRangeCorrection(force, longRangeCorrectionData, context.getOwner(), longRangeCoefficient, longRangeCoefficientDerivs, threads);
1620
1621
1622
        hasInitializedLongRangeCorrection = true;
        *forceCopy = force;
    }
1623
1624
1625
1626
1627
1628

    // See if any tabulated functions have changed.

    bool changed = false;
    for (int i = 0; i < force.getNumTabulatedFunctions(); i++) {
        string name = force.getTabulatedFunctionName(i);
1629
1630
        if (force.getTabulatedFunction(i).getUpdateCount() != tabulatedFunctionUpdateCount[name]) {
            tabulatedFunctionUpdateCount[name] = force.getTabulatedFunction(i).getUpdateCount();
1631
1632
1633
1634
1635
            changed = true;
        }
    }
    if (changed)
        createExpressions(force);
1636
1637
}

1638
ReferenceCalcGBSAOBCForceKernel::~ReferenceCalcGBSAOBCForceKernel() {
1639
    if (obc) {
Peter Eastman's avatar
Peter Eastman committed
1640
        delete obc->getObcParameters();
1641
1642
1643
1644
        delete obc;
    }
}

1645
void ReferenceCalcGBSAOBCForceKernel::initialize(const System& system, const GBSAOBCForce& force) {
Peter Eastman's avatar
Peter Eastman committed
1646
1647
    int numParticles = system.getNumParticles();
    charges.resize(numParticles);
peastman's avatar
peastman committed
1648
1649
    vector<double> atomicRadii(numParticles);
    vector<double> scaleFactors(numParticles);
Peter Eastman's avatar
Peter Eastman committed
1650
    for (int i = 0; i < numParticles; ++i) {
1651
        double charge, radius, scalingFactor;
Peter Eastman's avatar
Peter Eastman committed
1652
        force.getParticleParameters(i, charge, radius, scalingFactor);
peastman's avatar
peastman committed
1653
1654
1655
        charges[i] = charge;
        atomicRadii[i] = radius;
        scaleFactors[i] = scalingFactor;
1656
    }
1657
    ObcParameters* obcParameters = new ObcParameters(numParticles, ObcParameters::ObcTypeII);
1658
    obcParameters->setAtomicRadii(atomicRadii);
1659
    obcParameters->setScaledRadiusFactors(scaleFactors);
peastman's avatar
peastman committed
1660
1661
    obcParameters->setSolventDielectric(force.getSolventDielectric());
    obcParameters->setSoluteDielectric(force.getSoluteDielectric());
1662
    obcParameters->setPi4Asolv(4*M_PI*force.getSurfaceAreaEnergy());
1663
    if (force.getNonbondedMethod() != GBSAOBCForce::NoCutoff)
peastman's avatar
peastman committed
1664
        obcParameters->setUseCutoff(force.getCutoffDistance());
1665
    isPeriodic = (force.getNonbondedMethod() == GBSAOBCForce::CutoffPeriodic);
1666
    obc = new ReferenceObc(obcParameters);
1667
    obc->setIncludeAceApproximation(true);
1668
1669
}

1670
double ReferenceCalcGBSAOBCForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
1671
1672
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
1673
    if (isPeriodic)
1674
        obc->getObcParameters()->setPeriodic(extractBoxVectors(context));
Mark Friedrichs's avatar
Mark Friedrichs committed
1675
    return obc->computeBornEnergyForces(posData, charges, forceData);
1676
1677
}

1678
1679
1680
1681
1682
1683
1684
1685
void ReferenceCalcGBSAOBCForceKernel::copyParametersToContext(ContextImpl& context, const GBSAOBCForce& force) {
    int numParticles = force.getNumParticles();
    ObcParameters* obcParameters = obc->getObcParameters();
    if (numParticles != obcParameters->getAtomicRadii().size())
        throw OpenMMException("updateParametersInContext: The number of particles has changed");

    // Record the values.

peastman's avatar
peastman committed
1686
1687
    vector<double> atomicRadii(numParticles);
    vector<double> scaleFactors(numParticles);
1688
1689
1690
    for (int i = 0; i < numParticles; ++i) {
        double charge, radius, scalingFactor;
        force.getParticleParameters(i, charge, radius, scalingFactor);
peastman's avatar
peastman committed
1691
1692
1693
        charges[i] = charge;
        atomicRadii[i] = radius;
        scaleFactors[i] = scalingFactor;
1694
1695
1696
1697
1698
    }
    obcParameters->setAtomicRadii(atomicRadii);
    obcParameters->setScaledRadiusFactors(scaleFactors);
}

1699
1700
1701
1702
1703
1704
ReferenceCalcCustomGBForceKernel::~ReferenceCalcCustomGBForceKernel() {
    if (neighborList != NULL)
        delete neighborList;
}

void ReferenceCalcCustomGBForceKernel::initialize(const System& system, const CustomGBForce& force) {
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
    if (force.getNumComputedValues() > 0) {
        string name, expression;
        CustomGBForce::ComputationType type;
        force.getComputedValueParameters(0, name, expression, type);
        if (type == CustomGBForce::SingleParticle)
            throw OpenMMException("ReferencePlatform 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("ReferencePlatform requires that a CustomGBForce only have one computed value of type ParticlePair or ParticlePairNoExclusions.");
        }
    }
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730

    // 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.

1731
1732
1733
    particleParamArray.resize(numParticles);
    for (int i = 0; i < numParticles; ++i)
        force.getParticleParameters(i, particleParamArray[i]);
1734
    for (int i = 0; i < force.getNumPerParticleParameters(); i++)
1735
1736
1737
1738
        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
1739
    nonbondedCutoff = force.getCutoffDistance();
1740
1741
1742
1743
1744
    if (nonbondedMethod == NoCutoff)
        neighborList = NULL;
    else
        neighborList = new NeighborList();

1745
    // Record the tabulated function update counts for future reference.
1746
1747

    for (int i = 0; i < force.getNumTabulatedFunctions(); i++)
1748
        tabulatedFunctionUpdateCount[force.getTabulatedFunctionName(i)] = force.getTabulatedFunction(i).getUpdateCount();
1749
1750
1751
1752
1753
1754
1755

    // Create the expressions.
    
    createExpressions(force);
}

void ReferenceCalcCustomGBForceKernel::createExpressions(const CustomGBForce& force) {
1756
1757
1758
    // Create custom functions for the tabulated functions.

    map<string, Lepton::CustomFunction*> functions;
1759
    for (int i = 0; i < force.getNumTabulatedFunctions(); i++)
1760
        functions[force.getTabulatedFunctionName(i)] = createReferenceTabulatedFunction(force.getTabulatedFunction(i));
1761
1762
1763

    // Parse the expressions for computed values.

1764
1765
1766
1767
1768
1769
1770
    valueExpressions.clear();
    valueTypes.clear();
    valueNames.clear();
    energyParamDerivNames.clear();
    valueDerivExpressions.clear();
    valueGradientExpressions.clear();
    valueParamDerivExpressions.clear();
1771
    valueDerivExpressions.resize(force.getNumComputedValues());
1772
    valueGradientExpressions.resize(force.getNumComputedValues());
1773
    valueParamDerivExpressions.resize(force.getNumComputedValues());
1774
1775
1776
1777
1778
    set<string> particleVariables, pairVariables;
    pairVariables.insert("r");
    particleVariables.insert("x");
    particleVariables.insert("y");
    particleVariables.insert("z");
1779
    for (int i = 0; i < force.getNumPerParticleParameters(); i++) {
1780
1781
1782
1783
1784
1785
        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());
1786
1787
1788
1789
1790
    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();
1791
        valueExpressions.push_back(ex.createCompiledExpression());
1792
1793
        valueTypes.push_back(type);
        valueNames.push_back(name);
1794
        if (i == 0) {
1795
            valueDerivExpressions[i].push_back(ex.differentiate("r").createCompiledExpression());
1796
1797
            validateVariables(ex.getRootNode(), pairVariables);
        }
1798
        else {
1799
1800
1801
            valueGradientExpressions[i].push_back(ex.differentiate("x").createCompiledExpression());
            valueGradientExpressions[i].push_back(ex.differentiate("y").createCompiledExpression());
            valueGradientExpressions[i].push_back(ex.differentiate("z").createCompiledExpression());
1802
            for (int j = 0; j < i; j++)
1803
                valueDerivExpressions[i].push_back(ex.differentiate(valueNames[j]).createCompiledExpression());
1804
            validateVariables(ex.getRootNode(), particleVariables);
1805
        }
1806
1807
1808
1809
1810
        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());
        }
1811
1812
1813
        particleVariables.insert(name);
        pairVariables.insert(name+"1");
        pairVariables.insert(name+"2");
1814
1815
    }

1816
    // Parse the expressions for energy terms.
1817

1818
1819
1820
1821
1822
    energyExpressions.clear();
    energyTypes.clear();
    energyDerivExpressions.clear();
    energyGradientExpressions.clear();
    energyParamDerivExpressions.clear();
1823
    energyDerivExpressions.resize(force.getNumEnergyTerms());
1824
    energyGradientExpressions.resize(force.getNumEnergyTerms());
1825
    energyParamDerivExpressions.resize(force.getNumEnergyTerms());
1826
1827
1828
1829
1830
    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();
1831
        energyExpressions.push_back(ex.createCompiledExpression());
1832
1833
        energyTypes.push_back(type);
        if (type != CustomGBForce::SingleParticle)
1834
            energyDerivExpressions[i].push_back(ex.differentiate("r").createCompiledExpression());
1835
        for (int j = 0; j < force.getNumComputedValues(); j++) {
1836
            if (type == CustomGBForce::SingleParticle) {
1837
1838
1839
1840
                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());
1841
                validateVariables(ex.getRootNode(), particleVariables);
1842
            }
1843
            else {
1844
1845
                energyDerivExpressions[i].push_back(ex.differentiate(valueNames[j]+"1").createCompiledExpression());
                energyDerivExpressions[i].push_back(ex.differentiate(valueNames[j]+"2").createCompiledExpression());
1846
                validateVariables(ex.getRootNode(), pairVariables);
1847
1848
            }
        }
1849
1850
        for (int j = 0; j < force.getNumEnergyParameterDerivatives(); j++)
            energyParamDerivExpressions[i].push_back(ex.differentiate(force.getEnergyParameterDerivativeName(j)).createCompiledExpression());
1851
1852
1853
1854
    }

    // Delete the custom functions.

peastman's avatar
peastman committed
1855
1856
    for (auto& function : functions)
        delete function.second;
1857
1858
}

1859
double ReferenceCalcCustomGBForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
1860
1861
1862
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
1863
1864
    ReferenceCustomGBIxn ixn(valueExpressions, valueDerivExpressions, valueGradientExpressions, valueParamDerivExpressions, valueNames, valueTypes,
        energyExpressions, energyDerivExpressions, energyGradientExpressions, energyParamDerivExpressions, energyTypes, particleParameterNames);
1865
    bool periodic = (nonbondedMethod == CutoffPeriodic);
1866
    if (periodic)
1867
        ixn.setPeriodic(extractBoxVectors(context));
1868
    if (nonbondedMethod != NoCutoff) {
Peter Eastman's avatar
Peter Eastman committed
1869
1870
        vector<set<int> > empty(context.getSystem().getNumParticles()); // Don't omit exclusions from the neighbor list
        computeNeighborListVoxelHash(*neighborList, numParticles, posData, empty, extractBoxVectors(context), periodic, nonbondedCutoff, 0.0);
1871
1872
1873
        ixn.setUseCutoff(nonbondedCutoff, *neighborList);
    }
    map<string, double> globalParameters;
peastman's avatar
peastman committed
1874
1875
    for (auto& name : globalParameterNames)
        globalParameters[name] = context.getParameter(name);
1876
1877
1878
1879
1880
    vector<double> energyParamDerivValues(energyParamDerivNames.size()+1, 0.0);
    ixn.calculateIxn(numParticles, posData, particleParamArray, exclusions, globalParameters, forceData, includeEnergy ? &energy : NULL, &energyParamDerivValues[0]);
    map<string, double>& energyParamDerivs = extractEnergyParameterDerivatives(context);
    for (int i = 0; i < energyParamDerivNames.size(); i++)
        energyParamDerivs[energyParamDerivNames[i]] += energyParamDerivValues[i];
1881
1882
1883
    return energy;
}

1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
void ReferenceCalcCustomGBForceKernel::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
1896
            particleParamArray[i][j] = parameters[j];
1897
    }
1898
1899
1900
1901
1902
1903

    // See if any tabulated functions have changed.

    bool changed = false;
    for (int i = 0; i < force.getNumTabulatedFunctions(); i++) {
        string name = force.getTabulatedFunctionName(i);
1904
1905
        if (force.getTabulatedFunction(i).getUpdateCount() != tabulatedFunctionUpdateCount[name]) {
            tabulatedFunctionUpdateCount[name] = force.getTabulatedFunction(i).getUpdateCount();
1906
1907
1908
1909
1910
            changed = true;
        }
    }
    if (changed)
        createExpressions(force);
1911
1912
}

1913
1914
1915
1916
1917
ReferenceCalcCustomExternalForceKernel::~ReferenceCalcCustomExternalForceKernel() {
    if (ixn != NULL)
        delete ixn;
}

1918
1919
1920
1921
1922
1923
1924
void ReferenceCalcCustomExternalForceKernel::initialize(const System& system, const CustomExternalForce& force) {
    numParticles = force.getNumParticles();
    int numParameters = force.getNumPerParticleParameters();

    // Build the arrays.

    particles.resize(numParticles);
1925
1926
1927
    particleParamArray.resize(numParticles);
    for (int i = 0; i < numParticles; ++i)
        force.getParticleParameters(i, particles[i], particleParamArray[i]);
1928
1929
1930

    // Parse the expression used to calculate the force.

1931
    map<string, Lepton::CustomFunction*> functions;
1932
    ReferencePointDistanceFunction periodicDistance(true, &boxVectors);
1933
1934
    functions["periodicdistance"] = &periodicDistance;
    Lepton::ParsedExpression expression = Lepton::Parser::parse(force.getEnergyFunction(), functions).optimize();
1935
1936
1937
1938
    energyExpression = expression.createCompiledExpression();
    forceExpressionX = expression.differentiate("x").createCompiledExpression();
    forceExpressionY = expression.differentiate("y").createCompiledExpression();
    forceExpressionZ = expression.differentiate("z").createCompiledExpression();
1939
1940
1941
1942
    for (int i = 0; i < numParameters; i++)
        parameterNames.push_back(force.getPerParticleParameterName(i));
    for (int i = 0; i < force.getNumGlobalParameters(); i++)
        globalParameterNames.push_back(force.getGlobalParameterName(i));
1943
1944
1945
1946
1947
1948
1949
    set<string> variables;
    variables.insert("x");
    variables.insert("y");
    variables.insert("z");
    variables.insert(parameterNames.begin(), parameterNames.end());
    variables.insert(globalParameterNames.begin(), globalParameterNames.end());
    validateVariables(expression.getRootNode(), variables);
1950
1951
    ixn = new ReferenceCustomExternalIxn(energyExpression, forceExpressionX, forceExpressionY, forceExpressionZ, parameterNames);

1952
1953
}

1954
double ReferenceCalcCustomExternalForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
1955
1956
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
1957
    boxVectors = extractBoxVectors(context);
peastman's avatar
peastman committed
1958
    double energy = 0;
1959
    map<string, double> globalParameters;
peastman's avatar
peastman committed
1960
1961
    for (auto& name : globalParameterNames)
        globalParameters[name] = context.getParameter(name);
1962
    ixn->setGlobalParameters(globalParameters);
1963
    for (int i = 0; i < numParticles; ++i)
1964
        ixn->calculateForce(particles[i], posData, particleParamArray[i], forceData, includeEnergy ? &energy : NULL);
1965
1966
1967
    return energy;
}

1968
void ReferenceCalcCustomExternalForceKernel::copyParametersToContext(ContextImpl& context, const CustomExternalForce& force, int firstParticle, int lastParticle) {
1969
1970
1971
1972
1973
1974
1975
    if (numParticles != force.getNumParticles())
        throw OpenMMException("updateParametersInContext: The number of particles has changed");

    // Record the values.

    int numParameters = force.getNumPerParticleParameters();
    vector<double> params;
1976
    for (int i = firstParticle; i <= lastParticle; ++i) {
1977
        int particle;
1978
        force.getParticleParameters(i, particle, params);
1979
1980
1981
        if (particle != particles[i])
            throw OpenMMException("updateParametersInContext: A particle index has changed");
        for (int j = 0; j < numParameters; j++)
1982
            particleParamArray[i][j] = params[j];
1983
1984
1985
    }
}

1986
ReferenceCalcCustomHbondForceKernel::~ReferenceCalcCustomHbondForceKernel() {
1987
1988
    if (ixn != NULL)
        delete ixn;
1989
1990
1991
1992
1993
1994
1995
1996
1997
1998
1999
2000
2001
2002
2003
2004
2005
2006
}

void ReferenceCalcCustomHbondForceKernel::initialize(const System& system, const CustomHbondForce& force) {

    // Record the exclusions.

    numDonors = force.getNumDonors();
    numAcceptors = force.getNumAcceptors();
    numParticles = system.getNumParticles();
    exclusions.resize(numDonors);
    for (int i = 0; i < force.getNumExclusions(); i++) {
        int donor, acceptor;
        force.getExclusionParticles(i, donor, acceptor);
        exclusions[donor].insert(acceptor);
    }

    // Build the arrays.

2007
    donorParticles.resize(numDonors);
2008
    donorParamArray.resize(numDonors);
2009
    for (int i = 0; i < numDonors; ++i) {
2010
        int d1, d2, d3;
2011
        force.getDonorParameters(i, d1, d2, d3, donorParamArray[i]);
2012
2013
2014
        donorParticles[i].push_back(d1);
        donorParticles[i].push_back(d2);
        donorParticles[i].push_back(d3);
2015
    }
2016
    acceptorParticles.resize(numAcceptors);
2017
    acceptorParamArray.resize(numAcceptors);
2018
    for (int i = 0; i < numAcceptors; ++i) {
2019
        int a1, a2, a3;
2020
        force.getAcceptorParameters(i, a1, a2, a3, acceptorParamArray[i]);
2021
2022
2023
        acceptorParticles[i].push_back(a1);
        acceptorParticles[i].push_back(a2);
        acceptorParticles[i].push_back(a3);
2024
    }
2025
2026
    for (int i = 0; i < force.getNumGlobalParameters(); i++)
        globalParameterNames.push_back(force.getGlobalParameterName(i));
peastman's avatar
peastman committed
2027
    nonbondedCutoff = force.getCutoffDistance();
2028

2029
    // Record the tabulated function update counts for future reference.
2030
2031

    for (int i = 0; i < force.getNumTabulatedFunctions(); i++)
2032
        tabulatedFunctionUpdateCount[force.getTabulatedFunctionName(i)] = force.getTabulatedFunction(i).getUpdateCount();
2033
2034
2035
2036
2037
2038
2039

    // Create the interaction.
    
    createInteraction(force);
}

void ReferenceCalcCustomHbondForceKernel::createInteraction(const CustomHbondForce& force) {
2040
2041
2042
    // Create custom functions for the tabulated functions.

    map<string, Lepton::CustomFunction*> functions;
2043
    for (int i = 0; i < force.getNumTabulatedFunctions(); i++)
2044
        functions[force.getTabulatedFunctionName(i)] = createReferenceTabulatedFunction(force.getTabulatedFunction(i));
2045

2046
    // Parse the expression and create the object used to calculate the interaction.
2047

2048
2049
2050
    map<string, vector<int> > distances;
    map<string, vector<int> > angles;
    map<string, vector<int> > dihedrals;
2051
    Lepton::ParsedExpression energyExpression = CustomHbondForceImpl::prepareExpression(force, functions, distances, angles, dihedrals);
2052
2053
    vector<string> donorParameterNames;
    vector<string> acceptorParameterNames;
2054
    for (int i = 0; i < force.getNumPerDonorParameters(); i++)
2055
        donorParameterNames.push_back(force.getPerDonorParameterName(i));
2056
    for (int i = 0; i < force.getNumPerAcceptorParameters(); i++)
2057
        acceptorParameterNames.push_back(force.getPerAcceptorParameterName(i));
2058
    ixn = new ReferenceCustomHbondIxn(donorParticles, acceptorParticles, energyExpression, donorParameterNames, acceptorParameterNames, distances, angles, dihedrals);
2059
    NonbondedMethod nonbondedMethod = CalcCustomHbondForceKernel::NonbondedMethod(force.getNonbondedMethod());
2060
    isPeriodic = (nonbondedMethod == CutoffPeriodic);
2061
2062
    if (nonbondedMethod != NoCutoff)
        ixn->setUseCutoff(nonbondedCutoff);
2063
2064
2065

    // Delete the custom functions.

peastman's avatar
peastman committed
2066
2067
    for (auto& function : functions)
        delete function.second;
2068
2069
}

2070
double ReferenceCalcCustomHbondForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
2071
2072
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
2073
    if (isPeriodic)
2074
        ixn->setPeriodic(extractBoxVectors(context));
peastman's avatar
peastman committed
2075
    double energy = 0;
2076
    map<string, double> globalParameters;
peastman's avatar
peastman committed
2077
2078
    for (auto& name : globalParameterNames)
        globalParameters[name] = context.getParameter(name);
2079
    ixn->calculatePairIxn(posData, donorParamArray, acceptorParamArray, exclusions, globalParameters, forceData, includeEnergy ? &energy : NULL);
2080
2081
2082
    return energy;
}

2083
2084
2085
2086
2087
2088
2089
2090
2091
2092
2093
2094
2095
2096
2097
2098
2099
void ReferenceCalcCustomHbondForceKernel::copyParametersToContext(ContextImpl& context, const CustomHbondForce& force) {
    if (numDonors != force.getNumDonors())
        throw OpenMMException("updateParametersInContext: The number of donors has changed");
    if (numAcceptors != force.getNumAcceptors())
        throw OpenMMException("updateParametersInContext: The number of acceptors has changed");

    // Record the values.

    vector<double> parameters;
    int numDonorParameters = force.getNumPerDonorParameters();
    const vector<vector<int> >& donorAtoms = ixn->getDonorAtoms();
    for (int i = 0; i < numDonors; ++i) {
        int d1, d2, d3;
        force.getDonorParameters(i, d1, d2, d3, parameters);
        if (d1 != donorAtoms[i][0] || d2 != donorAtoms[i][1] || d3 != donorAtoms[i][2])
            throw OpenMMException("updateParametersInContext: The set of particles in a donor group has changed");
        for (int j = 0; j < numDonorParameters; j++)
peastman's avatar
peastman committed
2100
            donorParamArray[i][j] = parameters[j];
2101
2102
2103
2104
2105
2106
2107
2108
2109
    }
    int numAcceptorParameters = force.getNumPerAcceptorParameters();
    const vector<vector<int> >& acceptorAtoms = ixn->getAcceptorAtoms();
    for (int i = 0; i < numAcceptors; ++i) {
        int a1, a2, a3;
        force.getAcceptorParameters(i, a1, a2, a3, parameters);
        if (a1 != acceptorAtoms[i][0] || a2 != acceptorAtoms[i][1] || a3 != acceptorAtoms[i][2])
            throw OpenMMException("updateParametersInContext: The set of particles in an acceptor group has changed");
        for (int j = 0; j < numAcceptorParameters; j++)
peastman's avatar
peastman committed
2110
            acceptorParamArray[i][j] = parameters[j];
2111
    }
2112
2113
2114
2115
2116
2117

    // See if any tabulated functions have changed.

    bool changed = false;
    for (int i = 0; i < force.getNumTabulatedFunctions(); i++) {
        string name = force.getTabulatedFunctionName(i);
2118
2119
        if (force.getTabulatedFunction(i).getUpdateCount() != tabulatedFunctionUpdateCount[name]) {
            tabulatedFunctionUpdateCount[name] = force.getTabulatedFunction(i).getUpdateCount();
2120
2121
2122
2123
2124
2125
2126
2127
            changed = true;
        }
    }
    if (changed) {
        delete ixn;
        ixn = NULL;
        createInteraction(force);
    }
2128
2129
}

2130
2131
2132
2133
2134
2135
ReferenceCalcCustomCentroidBondForceKernel::~ReferenceCalcCustomCentroidBondForceKernel() {
    if (ixn != NULL)
        delete ixn;
}

void ReferenceCalcCustomCentroidBondForceKernel::initialize(const System& system, const CustomCentroidBondForce& force) {
2136
    usePeriodic = force.usesPeriodicBoundaryConditions();
2137
2138
2139
2140

    // Build the arrays.

    int numGroups = force.getNumGroups();
2141
    groupAtoms.resize(numGroups);
2142
2143
2144
2145
2146
    vector<double> ignored;
    for (int i = 0; i < numGroups; i++)
        force.getGroupParameters(i, groupAtoms[i], ignored);
    CustomCentroidBondForceImpl::computeNormalizedWeights(force, system, normalizedWeights);
    numBonds = force.getNumBonds();
2147
    bondGroups.resize(numBonds);
2148
2149
2150
    bondParamArray.resize(numBonds);
    for (int i = 0; i < numBonds; ++i)
        force.getBondParameters(i, bondGroups[i], bondParamArray[i]);
2151

2152
    // Record the tabulated function update counts for future reference.
2153
2154

    for (int i = 0; i < force.getNumTabulatedFunctions(); i++)
2155
        tabulatedFunctionUpdateCount[force.getTabulatedFunctionName(i)] = force.getTabulatedFunction(i).getUpdateCount();
2156
2157
2158
2159
2160
2161
2162

    // Create the interaction.
    
    createInteraction(force);
}

void ReferenceCalcCustomCentroidBondForceKernel::createInteraction(const CustomCentroidBondForce& force) {
2163
2164
2165
    // Create custom functions for the tabulated functions.

    map<string, Lepton::CustomFunction*> functions;
2166
    for (int i = 0; i < force.getNumTabulatedFunctions(); i++)
2167
2168
        functions[force.getTabulatedFunctionName(i)] = createReferenceTabulatedFunction(force.getTabulatedFunction(i));

2169
2170
2171
2172
2173
2174
    // Create implementations of point functions.

    functions["pointdistance"] = new ReferencePointDistanceFunction(usePeriodic, &boxVectors);
    functions["pointangle"] = new ReferencePointAngleFunction(usePeriodic, &boxVectors);
    functions["pointdihedral"] = new ReferencePointDihedralFunction(usePeriodic, &boxVectors);

2175
2176
    // Parse the expression and create the object used to calculate the interaction.

2177
    int numGroups = force.getNumGroups();
2178
    Lepton::ParsedExpression energyExpression = CustomCentroidBondForceImpl::prepareExpression(force, functions);
2179
    vector<string> bondParameterNames;
2180
    for (int i = 0; i < force.getNumPerBondParameters(); i++)
2181
2182
2183
        bondParameterNames.push_back(force.getPerBondParameterName(i));
    for (int i = 0; i < force.getNumGlobalParameters(); i++)
        globalParameterNames.push_back(force.getGlobalParameterName(i));
2184
2185
2186
2187
2188
2189
    vector<Lepton::CompiledExpression> energyParamDerivExpressions;
    for (int i = 0; i < force.getNumEnergyParameterDerivatives(); i++) {
        string param = force.getEnergyParameterDerivativeName(i);
        energyParamDerivNames.push_back(param);
        energyParamDerivExpressions.push_back(energyExpression.differentiate(param).createCompiledExpression());
    }
2190
    ixn = new ReferenceCustomCentroidBondIxn(force.getNumGroupsPerBond(), groupAtoms, normalizedWeights, bondGroups, energyExpression, bondParameterNames, energyParamDerivExpressions);
2191
2192
2193

    // Delete the custom functions.

peastman's avatar
peastman committed
2194
2195
    for (auto& function : functions)
        delete function.second;
2196
2197
2198
}

double ReferenceCalcCustomCentroidBondForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
2199
2200
2201
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
2202
    map<string, double> globalParameters;
peastman's avatar
peastman committed
2203
2204
    for (auto& name : globalParameterNames)
        globalParameters[name] = context.getParameter(name);
2205
2206
2207
2208
    if (usePeriodic) {
        boxVectors = extractBoxVectors(context);
        ixn->setPeriodic(boxVectors);
    }
2209
2210
2211
2212
2213
    vector<double> energyParamDerivValues(energyParamDerivNames.size()+1, 0.0);
    ixn->calculatePairIxn(posData, bondParamArray, globalParameters, forceData, includeEnergy ? &energy : NULL, &energyParamDerivValues[0]);
    map<string, double>& energyParamDerivs = extractEnergyParameterDerivatives(context);
    for (int i = 0; i < energyParamDerivNames.size(); i++)
        energyParamDerivs[energyParamDerivNames[i]] += energyParamDerivValues[i];
2214
2215
2216
2217
2218
2219
2220
2221
2222
2223
2224
2225
2226
2227
2228
2229
2230
2231
2232
    return energy;
}

void ReferenceCalcCustomCentroidBondForceKernel::copyParametersToContext(ContextImpl& context, const CustomCentroidBondForce& force) {
    if (numBonds != force.getNumBonds())
        throw OpenMMException("updateParametersInContext: The number of bonds has changed");

    // Record the values.

    int numParameters = force.getNumPerBondParameters();
    const vector<vector<int> >& bondGroups = ixn->getBondGroups();
    vector<int> groups;
    vector<double> params;
    for (int i = 0; i < numBonds; ++i) {
        force.getBondParameters(i, groups, params);
        for (int j = 0; j < groups.size(); j++)
            if (groups[j] != bondGroups[i][j])
                throw OpenMMException("updateParametersInContext: The set of groups in a bond has changed");
        for (int j = 0; j < numParameters; j++)
peastman's avatar
peastman committed
2233
            bondParamArray[i][j] = params[j];
2234
    }
2235
2236
2237
2238
2239
2240

    // See if any tabulated functions have changed.

    bool changed = false;
    for (int i = 0; i < force.getNumTabulatedFunctions(); i++) {
        string name = force.getTabulatedFunctionName(i);
2241
2242
        if (force.getTabulatedFunction(i).getUpdateCount() != tabulatedFunctionUpdateCount[name]) {
            tabulatedFunctionUpdateCount[name] = force.getTabulatedFunction(i).getUpdateCount();
2243
2244
2245
2246
2247
2248
2249
2250
            changed = true;
        }
    }
    if (changed) {
        delete ixn;
        ixn = NULL;
        createInteraction(force);
    }
2251
2252
}

2253
2254
2255
2256
2257
2258
ReferenceCalcCustomCompoundBondForceKernel::~ReferenceCalcCustomCompoundBondForceKernel() {
    if (ixn != NULL)
        delete ixn;
}

void ReferenceCalcCustomCompoundBondForceKernel::initialize(const System& system, const CustomCompoundBondForce& force) {
2259
    usePeriodic = force.usesPeriodicBoundaryConditions();
2260
2261
2262
2263

    // Build the arrays.

    numBonds = force.getNumBonds();
2264
    bondParticles.resize(numBonds);
2265
2266
2267
    bondParamArray.resize(numBonds);
    for (int i = 0; i < numBonds; ++i)
        force.getBondParameters(i, bondParticles[i], bondParamArray[i]);
2268

2269
    // Record the tabulated function update counts for future reference.
2270
2271

    for (int i = 0; i < force.getNumTabulatedFunctions(); i++)
2272
        tabulatedFunctionUpdateCount[force.getTabulatedFunctionName(i)] = force.getTabulatedFunction(i).getUpdateCount();
2273
2274
2275
2276
2277
2278
2279

    // Create the interaction.

    createInteraction(force);
}

void ReferenceCalcCustomCompoundBondForceKernel::createInteraction(const CustomCompoundBondForce& force) {
2280
2281
2282
    // Create custom functions for the tabulated functions.

    map<string, Lepton::CustomFunction*> functions;
2283
    for (int i = 0; i < force.getNumTabulatedFunctions(); i++)
2284
        functions[force.getTabulatedFunctionName(i)] = createReferenceTabulatedFunction(force.getTabulatedFunction(i));
2285

2286
2287
2288
2289
2290
2291
    // Create implementations of point functions.

    functions["pointdistance"] = new ReferencePointDistanceFunction(usePeriodic, &boxVectors);
    functions["pointangle"] = new ReferencePointAngleFunction(usePeriodic, &boxVectors);
    functions["pointdihedral"] = new ReferencePointDihedralFunction(usePeriodic, &boxVectors);

2292
2293
    // Parse the expression and create the object used to calculate the interaction.

2294
    Lepton::ParsedExpression energyExpression = CustomCompoundBondForceImpl::prepareExpression(force, functions);
2295
    vector<string> bondParameterNames;
2296
    for (int i = 0; i < force.getNumPerBondParameters(); i++)
2297
2298
2299
        bondParameterNames.push_back(force.getPerBondParameterName(i));
    for (int i = 0; i < force.getNumGlobalParameters(); i++)
        globalParameterNames.push_back(force.getGlobalParameterName(i));
2300
2301
2302
2303
2304
2305
    vector<Lepton::CompiledExpression> energyParamDerivExpressions;
    for (int i = 0; i < force.getNumEnergyParameterDerivatives(); i++) {
        string param = force.getEnergyParameterDerivativeName(i);
        energyParamDerivNames.push_back(param);
        energyParamDerivExpressions.push_back(energyExpression.differentiate(param).createCompiledExpression());
    }
2306
    ixn = new ReferenceCustomCompoundBondIxn(force.getNumParticlesPerBond(), bondParticles, energyExpression, bondParameterNames, energyParamDerivExpressions);
2307
2308
2309

    // Delete the custom functions.

peastman's avatar
peastman committed
2310
2311
    for (auto& function : functions)
        delete function.second;
2312
2313
2314
}

double ReferenceCalcCustomCompoundBondForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
2315
2316
2317
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
2318
    map<string, double> globalParameters;
peastman's avatar
peastman committed
2319
2320
    for (auto& name : globalParameterNames)
        globalParameters[name] = context.getParameter(name);
2321
2322
2323
2324
    if (usePeriodic) {
        boxVectors = extractBoxVectors(context);
        ixn->setPeriodic(boxVectors);
    }
2325
2326
2327
2328
2329
    vector<double> energyParamDerivValues(energyParamDerivNames.size()+1, 0.0);
    ixn->calculatePairIxn(posData, bondParamArray, globalParameters, forceData, includeEnergy ? &energy : NULL, &energyParamDerivValues[0]);
    map<string, double>& energyParamDerivs = extractEnergyParameterDerivatives(context);
    for (int i = 0; i < energyParamDerivNames.size(); i++)
        energyParamDerivs[energyParamDerivNames[i]] += energyParamDerivValues[i];
2330
2331
2332
    return energy;
}

2333
2334
2335
2336
2337
2338
2339
2340
2341
2342
2343
2344
void ReferenceCalcCustomCompoundBondForceKernel::copyParametersToContext(ContextImpl& context, const CustomCompoundBondForce& force) {
    if (numBonds != force.getNumBonds())
        throw OpenMMException("updateParametersInContext: The number of bonds has changed");

    // Record the values.

    int numParameters = force.getNumPerBondParameters();
    const vector<vector<int> >& bondAtoms = ixn->getBondAtoms();
    vector<int> particles;
    vector<double> params;
    for (int i = 0; i < numBonds; ++i) {
        force.getBondParameters(i, particles, params);
2345
        for (int j = 0; j < particles.size(); j++)
2346
2347
2348
            if (particles[j] != bondAtoms[i][j])
                throw OpenMMException("updateParametersInContext: The set of particles in a bond has changed");
        for (int j = 0; j < numParameters; j++)
peastman's avatar
peastman committed
2349
            bondParamArray[i][j] = params[j];
2350
    }
2351
2352
2353
2354
2355
2356

    // See if any tabulated functions have changed.

    bool changed = false;
    for (int i = 0; i < force.getNumTabulatedFunctions(); i++) {
        string name = force.getTabulatedFunctionName(i);
2357
2358
        if (force.getTabulatedFunction(i).getUpdateCount() != tabulatedFunctionUpdateCount[name]) {
            tabulatedFunctionUpdateCount[name] = force.getTabulatedFunction(i).getUpdateCount();
2359
2360
2361
2362
2363
2364
2365
2366
            changed = true;
        }
    }
    if (changed) {
        delete ixn;
        ixn = NULL;
        createInteraction(force);
    }
2367
2368
}

2369
2370
2371
2372
2373
2374
2375
2376
2377
ReferenceCalcCustomManyParticleForceKernel::~ReferenceCalcCustomManyParticleForceKernel() {
    if (ixn != NULL)
        delete ixn;
}

void ReferenceCalcCustomManyParticleForceKernel::initialize(const System& system, const CustomManyParticleForce& force) {
    // Build the arrays.

    numParticles = system.getNumParticles();
2378
    particleParamArray.resize(numParticles);
2379
2380
    for (int i = 0; i < numParticles; ++i) {
        int type;
2381
        force.getParticleParameters(i, particleParamArray[i], type);
2382
2383
2384
    }
    for (int i = 0; i < force.getNumGlobalParameters(); i++)
        globalParameterNames.push_back(force.getGlobalParameterName(i));
2385

2386
    // Record the tabulated function update counts for future reference.
2387
2388

    for (int i = 0; i < force.getNumTabulatedFunctions(); i++)
2389
        tabulatedFunctionUpdateCount[force.getTabulatedFunctionName(i)] = force.getTabulatedFunction(i).getUpdateCount();
2390
2391
2392

    // Create the interaction.

2393
    ixn = new ReferenceCustomManyParticleIxn(force);
2394
2395
2396
2397
2398
    nonbondedMethod = CalcCustomManyParticleForceKernel::NonbondedMethod(force.getNonbondedMethod());
    cutoffDistance = force.getCutoffDistance();
}

double ReferenceCalcCustomManyParticleForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
2399
2400
2401
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
2402
    map<string, double> globalParameters;
peastman's avatar
peastman committed
2403
2404
    for (auto& name : globalParameterNames)
        globalParameters[name] = context.getParameter(name);
2405
    if (nonbondedMethod == CutoffPeriodic) {
peastman's avatar
peastman committed
2406
        Vec3* boxVectors = extractBoxVectors(context);
2407
        double minAllowedSize = 2*cutoffDistance;
2408
        if (boxVectors[0][0] < minAllowedSize || boxVectors[1][1] < minAllowedSize || boxVectors[2][2] < minAllowedSize)
2409
            throw OpenMMException("The periodic box size has decreased to less than twice the nonbonded cutoff.");
2410
        ixn->setPeriodic(boxVectors);
2411
2412
2413
2414
2415
2416
2417
2418
2419
2420
2421
2422
2423
2424
2425
2426
2427
2428
    }
    ixn->calculateIxn(posData, particleParamArray, globalParameters, forceData, includeEnergy ? &energy : NULL);
    return energy;
}

void ReferenceCalcCustomManyParticleForceKernel::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
2429
            particleParamArray[i][j] = parameters[j];
2430
    }
2431
2432
2433
2434
2435
2436

    // See if any tabulated functions have changed.

    bool changed = false;
    for (int i = 0; i < force.getNumTabulatedFunctions(); i++) {
        string name = force.getTabulatedFunctionName(i);
2437
2438
        if (force.getTabulatedFunction(i).getUpdateCount() != tabulatedFunctionUpdateCount[name]) {
            tabulatedFunctionUpdateCount[name] = force.getTabulatedFunction(i).getUpdateCount();
2439
2440
2441
2442
2443
2444
2445
2446
            changed = true;
        }
    }
    if (changed) {
        delete ixn;
        ixn = NULL;
        ixn = new ReferenceCustomManyParticleIxn(force);
    }
2447
2448
}

2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
2461
2462
2463
2464
2465
2466
2467
ReferenceCalcGayBerneForceKernel::~ReferenceCalcGayBerneForceKernel() {
    if (ixn != NULL)
        delete ixn;
}

void ReferenceCalcGayBerneForceKernel::initialize(const System& system, const GayBerneForce& force) {
    ixn = new ReferenceGayBerneForce(force);
}

double ReferenceCalcGayBerneForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
    return ixn->calculateForce(extractPositions(context), extractForces(context), extractBoxVectors(context));
}

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

2468
2469
2470
2471
2472
ReferenceCalcCustomCVForceKernel::~ReferenceCalcCustomCVForceKernel() {
    if (ixn != NULL)
        delete ixn;
}

2473
void ReferenceCalcCustomCVForceKernel::initialize(const System& system, const CustomCVForce& force, ContextImpl& innerContext) {
2474
2475
2476
2477
2478
2479
2480
2481
2482
2483
2484
2485
2486
2487
2488
2489
2490
2491
2492
2493
2494
2495
2496
    for (int i = 0; i < force.getNumGlobalParameters(); i++)
        globalParameterNames.push_back(force.getGlobalParameterName(i));
    for (int i = 0; i < force.getNumEnergyParameterDerivatives(); i++)
        energyParamDerivNames.push_back(force.getEnergyParameterDerivativeName(i));
    ixn = new ReferenceCustomCVForce(force);
}

double ReferenceCalcCustomCVForceKernel::execute(ContextImpl& context, ContextImpl& innerContext, bool includeForces, bool includeEnergy) {
    copyState(context, innerContext);
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
    map<string, double> globalParameters;
    for (auto& name : globalParameterNames)
        globalParameters[name] = context.getParameter(name);
    map<string, double>& energyParamDerivs = extractEnergyParameterDerivatives(context);
    ixn->calculateIxn(innerContext, posData, globalParameters, forceData, includeEnergy ? &energy : NULL, energyParamDerivs);
    return energy;
}

void ReferenceCalcCustomCVForceKernel::copyState(ContextImpl& context, ContextImpl& innerContext) {
    extractPositions(innerContext) = extractPositions(context);
    extractVelocities(innerContext) = extractVelocities(context);
2497
2498
2499
2500
    Vec3 a, b, c;
    context.getPeriodicBoxVectors(a, b, c);
    innerContext.setPeriodicBoxVectors(a, b, c);
    innerContext.setTime(context.getTime());
2501
2502
2503
2504
2505
    map<string, double> innerParameters = innerContext.getParameters();
    for (auto& param : innerParameters)
        innerContext.setParameter(param.first, context.getParameter(param.first));
}

2506
void ReferenceCalcCustomCVForceKernel::copyParametersToContext(ContextImpl& context, const CustomCVForce& force) {
2507
    ixn->updateTabulatedFunctions(force);
2508
2509
}

2510
2511
void ReferenceCalcRMSDForceKernel::initialize(const System& system, const RMSDForce& force) {
    particles = force.getParticles();
peastman's avatar
peastman committed
2512
2513
2514
    if (particles.size() == 0)
        for (int i = 0; i < system.getNumParticles(); i++)
            particles.push_back(i);
2515
2516
2517
2518
2519
2520
2521
2522
2523
2524
2525
2526
2527
2528
2529
2530
2531
2532
2533
2534
    referencePos = force.getReferencePositions();
    Vec3 center;
    for (int i : particles)
        center += referencePos[i];
    center /= particles.size();
    for (Vec3& p : referencePos)
        p -= center;
}

double ReferenceCalcRMSDForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    ReferenceRMSDForce rmsd(referencePos, particles);
    return rmsd.calculateIxn(posData, forceData);
}

void ReferenceCalcRMSDForceKernel::copyParametersToContext(ContextImpl& context, const RMSDForce& force) {
    if (referencePos.size() != force.getReferencePositions().size())
        throw OpenMMException("updateParametersInContext: The number of reference positions has changed");
    particles = force.getParticles();
peastman's avatar
peastman committed
2535
2536
2537
    if (particles.size() == 0)
        for (int i = 0; i < referencePos.size(); i++)
            particles.push_back(i);
2538
2539
2540
2541
2542
2543
2544
2545
2546
    referencePos = force.getReferencePositions();
    Vec3 center;
    for (int i : particles)
        center += referencePos[i];
    center /= particles.size();
    for (Vec3& p : referencePos)
        p -= center;
}

2547
2548
2549
2550
2551
2552
2553
2554
2555
2556
2557
2558
2559
2560
void ReferenceCalcRGForceKernel::initialize(const System& system, const RGForce& force) {
    particles = force.getParticles();
    if (particles.size() == 0)
        for (int i = 0; i < system.getNumParticles(); i++)
            particles.push_back(i);
}

double ReferenceCalcRGForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    ReferenceRGForce rg(particles);
    return rg.calculateIxn(posData, forceData);
}

2561
2562
2563
2564
2565
2566
2567
2568
2569
2570
2571
2572
2573
2574
2575
2576
2577
2578
2579
2580
2581
2582
2583
2584
2585
2586
2587
2588
2589
2590
2591
2592
2593
2594
2595
2596
2597
2598
2599
void ReferenceCalcOrientationRestraintForceKernel::initialize(const System& system, const OrientationRestraintForce& force) {
    k = force.getK();
    particles = force.getParticles();
    if (particles.size() == 0)
        for (int i = 0; i < system.getNumParticles(); i++)
            particles.push_back(i);
    referencePos = force.getReferencePositions();
    Vec3 center;
    for (int i : particles)
        center += referencePos[i];
    center /= particles.size();
    for (Vec3& p : referencePos)
        p -= center;
}

double ReferenceCalcOrientationRestraintForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    ReferenceOrientationRestraintForce f(k, referencePos, particles);
    return f.calculateIxn(posData, forceData);
}

void ReferenceCalcOrientationRestraintForceKernel::copyParametersToContext(ContextImpl& context, const OrientationRestraintForce& force) {
    if (referencePos.size() != force.getReferencePositions().size())
        throw OpenMMException("updateParametersInContext: The number of reference positions has changed");
    k = force.getK();
    particles = force.getParticles();
    if (particles.size() == 0)
        for (int i = 0; i < referencePos.size(); i++)
            particles.push_back(i);
    referencePos = force.getReferencePositions();
    Vec3 center;
    for (int i : particles)
        center += referencePos[i];
    center /= particles.size();
    for (Vec3& p : referencePos)
        p -= center;
}

2600
2601
2602
2603
2604
ReferenceIntegrateVerletStepKernel::~ReferenceIntegrateVerletStepKernel() {
    if (dynamics)
        delete dynamics;
}

2605
void ReferenceIntegrateVerletStepKernel::initialize(const System& system, const VerletIntegrator& integrator) {
Peter Eastman's avatar
Peter Eastman committed
2606
    int numParticles = system.getNumParticles();
2607
    masses.resize(numParticles);
2608
    for (int i = 0; i < numParticles; ++i)
peastman's avatar
peastman committed
2609
        masses[i] = system.getParticleMass(i);
2610
2611
}

2612
void ReferenceIntegrateVerletStepKernel::execute(ContextImpl& context, const VerletIntegrator& integrator) {
2613
    double stepSize = integrator.getStepSize();
peastman's avatar
peastman committed
2614
2615
2616
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& velData = extractVelocities(context);
    vector<Vec3>& forceData = extractForces(context);
2617
2618
2619
    if (dynamics == 0 || stepSize != prevStepSize) {
        // Recreate the computation objects with the new parameters.
        
2620
        if (dynamics)
2621
            delete dynamics;
peastman's avatar
peastman committed
2622
        dynamics = new ReferenceVerletDynamics(context.getSystem().getNumParticles(), stepSize);
2623
        dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
2624
        dynamics->setVirtualSites(extractVirtualSites(context));
2625
2626
        prevStepSize = stepSize;
    }
2627
    dynamics->update(context.getSystem(), posData, velData, forceData, masses, integrator.getConstraintTolerance(), extractBoxVectors(context));
2628
    data.time += stepSize;
2629
    data.stepCount++;
2630
}
2631

2632
double ReferenceIntegrateVerletStepKernel::computeKineticEnergy(ContextImpl& context, const VerletIntegrator& integrator) {
2633
    return computeShiftedKineticEnergy(context, masses, 0.5*integrator.getStepSize());
2634
2635
}

2636
2637
2638
ReferenceIntegrateNoseHooverStepKernel::~ReferenceIntegrateNoseHooverStepKernel() {
    if (chainPropagator)
        delete chainPropagator;
2639
2640
2641
2642
    if (dynamics)
        delete dynamics;
}

2643
void ReferenceIntegrateNoseHooverStepKernel::initialize(const System& system, const NoseHooverIntegrator& integrator) {
2644
2645
2646
2647
    int numParticles = system.getNumParticles();
    masses.resize(numParticles);
    for (int i = 0; i < numParticles; ++i)
        masses[i] = system.getParticleMass(i);
2648
    this->chainPropagator = new ReferenceNoseHooverChain();
2649
2650
}

2651
void ReferenceIntegrateNoseHooverStepKernel::execute(ContextImpl& context, const NoseHooverIntegrator& integrator) {
2652
2653
2654
2655
2656
2657
2658
2659
    double stepSize = integrator.getStepSize();
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& velData = extractVelocities(context);
    vector<Vec3>& forceData = extractForces(context);
    if (dynamics == 0 || stepSize != prevStepSize) {
        // Recreate the computation objects with the new parameters.
        if (dynamics)
            delete dynamics;
2660
        dynamics = new ReferenceNoseHooverDynamics(context.getSystem().getNumParticles(), stepSize);
2661
        dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
2662
        dynamics->setVirtualSites(extractVirtualSites(context));
2663
2664
        prevStepSize = stepSize;
    }
2665
    dynamics->step1(context, context.getSystem(), posData, velData, forceData, masses, integrator.getConstraintTolerance(),
2666
2667
2668
2669
2670
2671
2672
2673
                     integrator.getAllThermostatedIndividualParticles(), integrator.getAllThermostatedPairs(), integrator.getMaximumPairDistance());
    int numChains = integrator.getNumThermostats();
    for(int chain = 0; chain < numChains; ++chain) {
        const auto &thermostatChain = integrator.getThermostat(chain);
        std::pair<double, double> KEs = computeMaskedKineticEnergy(context, thermostatChain, true);
        std::pair<double, double> scaleFactors = propagateChain(context, thermostatChain, KEs, stepSize);
        scaleVelocities(context, thermostatChain, scaleFactors);
    }
2674
    dynamics->step2(context, context.getSystem(), posData, velData, forceData, masses, integrator.getConstraintTolerance(),
2675
2676
                     integrator.getAllThermostatedIndividualParticles(), integrator.getAllThermostatedPairs(), integrator.getMaximumPairDistance(),
            extractBoxVectors(context));
2677
2678
2679
2680
    data.time += stepSize;
    data.stepCount++;
}

2681
double ReferenceIntegrateNoseHooverStepKernel::computeKineticEnergy(ContextImpl& context, const NoseHooverIntegrator& integrator) {
2682
2683
2684
    return computeShiftedKineticEnergy(context, masses, 0);
}

2685
2686
2687
2688
2689
2690
2691
2692
2693
2694
2695
2696
2697
2698
std::pair<double, double> ReferenceIntegrateNoseHooverStepKernel::propagateChain(ContextImpl& context, const NoseHooverChain &nhc,
                                                                                     std::pair<double, double> kineticEnergy, double timeStep) {
    double absKE = kineticEnergy.first;
    double relKE = kineticEnergy.second;
    if (absKE < 1e-8) return {1.0, 1.0};  // (catches the problem of zero velocities in the first dynamics step, where we have nothing to scale)
    // Get the variables describing the NHC
    int chainLength = nhc.getChainLength();
    int chainID = nhc.getChainID();
    int numDOFs = nhc.getNumDegreesOfFreedom();
    int numMTS = nhc.getNumMultiTimeSteps();

    int nAtoms = nhc.getThermostatedAtoms().size();
    double absScale = 0;
    if (nAtoms) {
2699
2700
        if (chainPositions.size() < 2*chainID+1){
            chainPositions.resize(2*chainID+1);
2701
        }
2702
2703
        if (chainVelocities.size() < 2*chainID+1){
            chainVelocities.resize(2*chainID+1);
2704
        }
2705
2706
2707
2708
        auto& positions = chainPositions.at(2*chainID);
        auto& velocities = chainVelocities.at(2*chainID);
        if (positions.size() < chainLength){
            positions.resize(chainLength, 0);
2709
        }
2710
2711
        if (velocities.size() < chainLength){
            velocities.resize(chainLength, 0);
2712
2713
2714
        }
        double temperature = nhc.getTemperature();
        double collisionFrequency = nhc.getCollisionFrequency();
2715
        absScale = chainPropagator->propagate(absKE, velocities, positions, numDOFs,
2716
2717
2718
2719
2720
2721
                                              temperature, collisionFrequency, timeStep,
                                              numMTS, nhc.getYoshidaSuzukiWeights());
    }
    double relScale = 0;
    int nPairs = nhc.getThermostatedPairs().size();
    if (nPairs) {
2722
2723
        if (chainPositions.size() < 2*chainID+2){
            chainPositions.resize(2*chainID+2);
2724
        }
2725
2726
        if (chainVelocities.size() < 2*chainID+2){
            chainVelocities.resize(2*chainID+2);
2727
        }
2728
2729
2730
2731
        auto& positions = chainPositions.at(2*chainID+1);
        auto& velocities = chainVelocities.at(2*chainID+1);
        if (positions.size() < chainLength){
            positions.resize(chainLength, 0);
2732
        }
2733
2734
        if (velocities.size() < chainLength){
            velocities.resize(chainLength, 0);
2735
2736
2737
        }
        double temperature = nhc.getRelativeTemperature();
        double collisionFrequency = nhc.getRelativeCollisionFrequency();
2738
        relScale = chainPropagator->propagate(relKE, velocities, positions, 3*nPairs,
2739
2740
2741
2742
2743
2744
2745
2746
2747
2748
2749
2750
2751
2752
2753
2754
2755
2756
2757
2758
2759
                                              temperature, collisionFrequency, timeStep,
                                              numMTS, nhc.getYoshidaSuzukiWeights());
    }
    return {absScale, relScale};
}

double ReferenceIntegrateNoseHooverStepKernel::computeHeatBathEnergy(ContextImpl& context, const NoseHooverChain &nhc) {
    double potentialEnergy = 0;
    double kineticEnergy = 0;
    int chainLength = nhc.getChainLength();
    int chainID = nhc.getChainID();
    int nAtoms = nhc.getThermostatedAtoms().size();
    int nPairs = nhc.getThermostatedPairs().size();
    if (nAtoms) {
        double temperature = nhc.getTemperature();
        double collisionFrequency = nhc.getCollisionFrequency();
        double kT = temperature * BOLTZ;
        int numDOFs = nhc.getNumDegreesOfFreedom();
        for(int i = 0; i < chainLength; ++i) {
            double prefac = i ? 1 : numDOFs;
            double mass = prefac * kT / (collisionFrequency * collisionFrequency);
2760
            double velocity = chainVelocities[2*chainID][i];
2761
2762
2763
            // The kinetic energy of this bead
            kineticEnergy += 0.5 * mass * velocity * velocity;
            // The potential energy of this bead
2764
            double position = chainPositions[2*chainID][i];
2765
2766
2767
2768
2769
2770
2771
2772
2773
2774
2775
            potentialEnergy += prefac * kT * position;
        }
    }
    if (nPairs) {
        double temperature = nhc.getRelativeTemperature();
        double collisionFrequency = nhc.getRelativeCollisionFrequency();
        double kT = temperature * BOLTZ;
        int numDOFs = 3 * nPairs;
        for(int i = 0; i < chainLength; ++i) {
            double prefac = i ? 1 : numDOFs;
            double mass = prefac * kT / (collisionFrequency * collisionFrequency);
2776
            double velocity = chainVelocities[2*chainID+1][i];
2777
2778
2779
            // The kinetic energy of this bead
            kineticEnergy += 0.5 * mass * velocity * velocity;
            // The potential energy of this bead
2780
            double position = chainPositions[2*chainID+1][i];
2781
2782
2783
2784
2785
2786
2787
2788
2789
2790
2791
2792
2793
2794
2795
2796
2797
2798
2799
2800
2801
2802
2803
2804
2805
2806
2807
2808
2809
2810
2811
2812
2813
2814
2815
2816
2817
2818
2819
2820
2821
2822
2823
2824
2825
2826
2827
2828
2829
2830
2831
2832
2833
2834
2835
2836
2837
2838
2839
2840
2841
2842
2843
2844
2845
2846
2847
2848
2849
2850
2851
2852
2853
2854
2855
2856
2857
2858
            potentialEnergy += prefac * kT * position;
        }
    }
    return kineticEnergy + potentialEnergy;
}

std::pair<double, double> ReferenceIntegrateNoseHooverStepKernel::computeMaskedKineticEnergy(ContextImpl& context,
                                                                const NoseHooverChain &noseHooverChain, bool downloadValue) {
    const std::vector<int>& atomsList = noseHooverChain.getThermostatedAtoms();
    const std::vector<std::pair<int,int>>& pairsList = noseHooverChain.getThermostatedPairs();
    std::vector<Vec3>& velocities = extractVelocities(context);
    const System& system = context.getSystem();
    int numParticles = system.getNumParticles();
    std::vector<double> masses(numParticles);
    for (int i = 0; i < numParticles; ++i)
        masses[i] = system.getParticleMass(i);

    double comKE = 0;
    double relKE = 0;
    // kinetic energy of individual atoms
    for (const auto &m: atomsList){
        comKE += 0.5 * masses[m] * velocities[m].dot(velocities[m]);
    }
    // center of mass kinetic energy of pairs
    for (const auto &p: pairsList){
        double m1 = masses[p.first];
        double m2 = masses[p.second];
        Vec3 v1 = velocities[p.first];
        Vec3 v2 = velocities[p.second];
        double invMass = 1.0 / (m1 + m2);
        double redMass = m1 * m2 * invMass;
        double fracM1 = m1 * invMass;
        double fracM2 = m2 * invMass;
        Vec3 comVelocity = fracM1 * v1 + fracM2 * v2;
        Vec3 relVelocity = v2 - v1;

        comKE += 0.5 * (m1 + m2) * comVelocity.dot(comVelocity);
        relKE += 0.5 * redMass * relVelocity.dot(relVelocity);
    }
    // We ignore the downloadValue argument here and always return the correct value
    return {comKE, relKE};
}


void ReferenceIntegrateNoseHooverStepKernel::scaleVelocities(ContextImpl& context, const NoseHooverChain &noseHooverChain, std::pair<double, double> scaleFactors) {
    const auto& atoms = noseHooverChain.getThermostatedAtoms();
    const auto& pairs = noseHooverChain.getThermostatedPairs();
    std::vector<Vec3>& velocities = extractVelocities(context);
    double absScale = scaleFactors.first;
    double relScale = scaleFactors.second;

    const System& system = context.getSystem();
    int numParticles = system.getNumParticles();
    std::vector<double> masses(numParticles);
    for (int i = 0; i < numParticles; ++i)
        masses[i] = system.getParticleMass(i);
    // scale absolute velocities
    for (const auto &a: atoms){
        velocities[a] *= absScale;
    }
    // scale relative velocities and absolute center of mass velocities for each pair
    for (const auto &p: pairs){
        int p1 = p.first;
        int p2 = p.second;
        double m1 = masses[p.first];
        double m2 = masses[p.second];
        Vec3 v1 = velocities[p.first];
        Vec3 v2 = velocities[p.second];
        double invMass = 1.0 / (m1 + m2);
        double fracM1 = m1 * invMass;
        double fracM2 = m2 * invMass;
        Vec3 comVelocity = fracM1 * v1 + fracM2 * v2;
        Vec3 relVelocity = v2 - v1;
        velocities[p1] = absScale * comVelocity - relScale * relVelocity * fracM2;
        velocities[p2] = absScale * comVelocity + relScale * relVelocity * fracM1;
    }
}

2859
2860
2861
2862
2863
2864
2865
2866
2867
2868
2869
2870
2871
2872
2873
2874
2875
2876
2877
2878
2879
2880
2881
2882
2883
2884
2885
2886
2887
2888
2889
void ReferenceIntegrateNoseHooverStepKernel::createCheckpoint(ContextImpl& context, ostream& stream) const {
    size_t numChains = chainPositions.size();
    assert(numChains == chainVelocities.size());
    stream.write((char*) &numChains, sizeof(size_t));
    for (size_t i=0; i<numChains; i++){
        auto & noseHooverPositions = chainPositions.at(i);
        auto & noseHooverVelocities = chainVelocities.at(i);
        size_t numBeads = noseHooverPositions.size();
        assert(numBeads == noseHooverVelocities.size());
        stream.write((char*) &numBeads, sizeof(size_t));
        stream.write((char*) noseHooverPositions.data(), sizeof(double)*numBeads);
        stream.write((char*) noseHooverVelocities.data(), sizeof(double)*numBeads);
    }
}

void ReferenceIntegrateNoseHooverStepKernel::loadCheckpoint(ContextImpl& context, istream& stream) {
    size_t numChains, numBeads;
    stream.read((char*) &numChains, sizeof(size_t));
    chainPositions.clear();
    chainVelocities.clear();
    for (size_t i=0; i<numChains; i++){
        stream.read((char*) &numBeads, sizeof(size_t));
        std::vector<double> noseHooverPositions(numBeads);
        std::vector<double> noseHooverVelocities(numBeads);
        stream.read((char*) &noseHooverPositions[0], sizeof(double)*numBeads);
        stream.read((char*) &noseHooverVelocities[0], sizeof(double)*numBeads);
        chainPositions.push_back(noseHooverPositions);
        chainVelocities.push_back(noseHooverVelocities);
    }
}

2890
2891
2892
2893
2894
2895
2896
2897
2898
2899
void ReferenceIntegrateNoseHooverStepKernel::getChainStates(ContextImpl& context, vector<vector<double> >& positions, vector<vector<double> >& velocities) const {
    positions = chainPositions;
    velocities = chainVelocities;
}

void ReferenceIntegrateNoseHooverStepKernel::setChainStates(ContextImpl& context, const vector<vector<double> >& positions, const vector<vector<double> >& velocities) {
    chainPositions = positions;
    chainVelocities = velocities;
}

2900
ReferenceIntegrateLangevinMiddleStepKernel::~ReferenceIntegrateLangevinMiddleStepKernel() {
2901
2902
2903
2904
    if (dynamics)
        delete dynamics;
}

2905
void ReferenceIntegrateLangevinMiddleStepKernel::initialize(const System& system, const LangevinMiddleIntegrator& integrator) {
2906
2907
2908
2909
2910
2911
2912
    int numParticles = system.getNumParticles();
    masses.resize(numParticles);
    for (int i = 0; i < numParticles; ++i)
        masses[i] = system.getParticleMass(i);
    SimTKOpenMMUtilities::setRandomNumberSeed((unsigned int) integrator.getRandomNumberSeed());
}

2913
void ReferenceIntegrateLangevinMiddleStepKernel::execute(ContextImpl& context, const LangevinMiddleIntegrator& integrator) {
2914
2915
2916
2917
2918
2919
2920
2921
2922
2923
    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;
2924
        dynamics = new ReferenceLangevinMiddleDynamics(
2925
2926
2927
2928
2929
                context.getSystem().getNumParticles(), 
                stepSize, 
                friction, 
                temperature);
        dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
2930
        dynamics->setVirtualSites(extractVirtualSites(context));
2931
2932
2933
2934
        prevTemp = temperature;
        prevFriction = friction;
        prevStepSize = stepSize;
    }
2935
    dynamics->update(context, posData, velData, masses, integrator.getConstraintTolerance(), extractBoxVectors(context));
2936
2937
2938
2939
    data.time += stepSize;
    data.stepCount++;
}

2940
double ReferenceIntegrateLangevinMiddleStepKernel::computeKineticEnergy(ContextImpl& context, const LangevinMiddleIntegrator& integrator) {
2941
2942
2943
    return computeShiftedKineticEnergy(context, masses, 0.0);
}

2944
2945
2946
2947
2948
ReferenceIntegrateBrownianStepKernel::~ReferenceIntegrateBrownianStepKernel() {
    if (dynamics)
        delete dynamics;
}

2949
void ReferenceIntegrateBrownianStepKernel::initialize(const System& system, const BrownianIntegrator& integrator) {
Peter Eastman's avatar
Peter Eastman committed
2950
    int numParticles = system.getNumParticles();
2951
    masses.resize(numParticles);
2952
    for (int i = 0; i < numParticles; ++i)
peastman's avatar
peastman committed
2953
        masses[i] = system.getParticleMass(i);
2954
    SimTKOpenMMUtilities::setRandomNumberSeed((unsigned int) integrator.getRandomNumberSeed());
2955
2956
}

2957
void ReferenceIntegrateBrownianStepKernel::execute(ContextImpl& context, const BrownianIntegrator& integrator) {
2958
2959
2960
    double temperature = integrator.getTemperature();
    double friction = integrator.getFriction();
    double stepSize = integrator.getStepSize();
peastman's avatar
peastman committed
2961
2962
2963
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& velData = extractVelocities(context);
    vector<Vec3>& forceData = extractForces(context);
2964
2965
2966
    if (dynamics == 0 || temperature != prevTemp || friction != prevFriction || stepSize != prevStepSize) {
        // Recreate the computation objects with the new parameters.
        
2967
        if (dynamics)
2968
            delete dynamics;
2969
        dynamics = new ReferenceBrownianDynamics(
2970
                context.getSystem().getNumParticles(), 
peastman's avatar
peastman committed
2971
2972
2973
                stepSize, 
                friction, 
                temperature);
2974
        dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
2975
        dynamics->setVirtualSites(extractVirtualSites(context));
2976
2977
2978
2979
        prevTemp = temperature;
        prevFriction = friction;
        prevStepSize = stepSize;
    }
2980
    dynamics->update(context.getSystem(), posData, velData, forceData, masses, integrator.getConstraintTolerance(), extractBoxVectors(context));
2981
    data.time += stepSize;
2982
    data.stepCount++;
2983
2984
}

2985
double ReferenceIntegrateBrownianStepKernel::computeKineticEnergy(ContextImpl& context, const BrownianIntegrator& integrator) {
2986
    return computeShiftedKineticEnergy(context, masses, 0);
2987
2988
}

2989
2990
2991
2992
2993
2994
2995
ReferenceIntegrateVariableLangevinStepKernel::~ReferenceIntegrateVariableLangevinStepKernel() {
    if (dynamics)
        delete dynamics;
}

void ReferenceIntegrateVariableLangevinStepKernel::initialize(const System& system, const VariableLangevinIntegrator& integrator) {
    int numParticles = system.getNumParticles();
2996
    masses.resize(numParticles);
2997
    for (int i = 0; i < numParticles; ++i)
peastman's avatar
peastman committed
2998
        masses[i] = system.getParticleMass(i);
2999
3000
3001
    SimTKOpenMMUtilities::setRandomNumberSeed((unsigned int) integrator.getRandomNumberSeed());
}

3002
double ReferenceIntegrateVariableLangevinStepKernel::execute(ContextImpl& context, const VariableLangevinIntegrator& integrator, double maxTime) {
3003
3004
3005
    double temperature = integrator.getTemperature();
    double friction = integrator.getFriction();
    double errorTol = integrator.getErrorTolerance();
peastman's avatar
peastman committed
3006
3007
3008
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& velData = extractVelocities(context);
    vector<Vec3>& forceData = extractForces(context);
3009
3010
3011
    if (dynamics == 0 || temperature != prevTemp || friction != prevFriction || errorTol != prevErrorTol) {
        // Recreate the computation objects with the new parameters.

3012
        if (dynamics)
3013
            delete dynamics;
peastman's avatar
peastman committed
3014
        dynamics = new ReferenceVariableStochasticDynamics(context.getSystem().getNumParticles(), friction, temperature, errorTol);
3015
        dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
3016
        dynamics->setVirtualSites(extractVirtualSites(context));
3017
3018
3019
3020
        prevTemp = temperature;
        prevFriction = friction;
        prevErrorTol = errorTol;
    }
peastman's avatar
peastman committed
3021
    double maxStepSize = maxTime-data.time;
3022
3023
    if (integrator.getMaximumStepSize() > 0)
        maxStepSize = min(integrator.getMaximumStepSize(), maxStepSize);
3024
    dynamics->update(context.getSystem(), posData, velData, forceData, masses, maxStepSize, integrator.getConstraintTolerance(), extractBoxVectors(context));
3025
3026
3027
3028
    data.time += dynamics->getDeltaT();
    if (dynamics->getDeltaT() == maxStepSize)
        data.time = maxTime; // Avoid round-off error
    data.stepCount++;
3029
    return dynamics->getDeltaT();
3030
3031
}

3032
double ReferenceIntegrateVariableLangevinStepKernel::computeKineticEnergy(ContextImpl& context, const VariableLangevinIntegrator& integrator) {
3033
    return computeShiftedKineticEnergy(context, masses, 0.5*integrator.getStepSize());
3034
3035
}

3036
3037
3038
3039
3040
3041
3042
ReferenceIntegrateVariableVerletStepKernel::~ReferenceIntegrateVariableVerletStepKernel() {
    if (dynamics)
        delete dynamics;
}

void ReferenceIntegrateVariableVerletStepKernel::initialize(const System& system, const VariableVerletIntegrator& integrator) {
    int numParticles = system.getNumParticles();
3043
    masses.resize(numParticles);
3044
    for (int i = 0; i < numParticles; ++i)
peastman's avatar
peastman committed
3045
        masses[i] = system.getParticleMass(i);
3046
3047
}

3048
double ReferenceIntegrateVariableVerletStepKernel::execute(ContextImpl& context, const VariableVerletIntegrator& integrator, double maxTime) {
3049
    double errorTol = integrator.getErrorTolerance();
peastman's avatar
peastman committed
3050
3051
3052
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& velData = extractVelocities(context);
    vector<Vec3>& forceData = extractForces(context);
3053
    if (dynamics == 0 || errorTol != prevErrorTol) {
3054
3055
        // Recreate the computation objects with the new parameters.

3056
        if (dynamics)
3057
            delete dynamics;
peastman's avatar
peastman committed
3058
        dynamics = new ReferenceVariableVerletDynamics(context.getSystem().getNumParticles(), errorTol);
3059
        dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
3060
        dynamics->setVirtualSites(extractVirtualSites(context));
3061
        prevErrorTol = errorTol;
3062
    }
peastman's avatar
peastman committed
3063
    double maxStepSize = maxTime-data.time;
3064
3065
    if (integrator.getMaximumStepSize() > 0)
        maxStepSize = min(integrator.getMaximumStepSize(), maxStepSize);
3066
    dynamics->update(context.getSystem(), posData, velData, forceData, masses, maxStepSize, integrator.getConstraintTolerance(), extractBoxVectors(context));
3067
    data.time += dynamics->getDeltaT();
3068
3069
3070
    if (dynamics->getDeltaT() == maxStepSize)
        data.time = maxTime; // Avoid round-off error
    data.stepCount++;
3071
    return dynamics->getDeltaT();
3072
3073
}

3074
double ReferenceIntegrateVariableVerletStepKernel::computeKineticEnergy(ContextImpl& context, const VariableVerletIntegrator& integrator) {
3075
    return computeShiftedKineticEnergy(context, masses, 0.5*integrator.getStepSize());
3076
3077
}

3078
3079
3080
3081
3082
3083
3084
3085
3086
ReferenceIntegrateCustomStepKernel::~ReferenceIntegrateCustomStepKernel() {
    if (dynamics)
        delete dynamics;
}

void ReferenceIntegrateCustomStepKernel::initialize(const System& system, const CustomIntegrator& integrator) {
    int numParticles = system.getNumParticles();
    masses.resize(numParticles);
    for (int i = 0; i < numParticles; ++i)
peastman's avatar
peastman committed
3087
        masses[i] = system.getParticleMass(i);
3088
    perDofValues.resize(integrator.getNumPerDofVariables());
peastman's avatar
peastman committed
3089
3090
    for (auto& values : perDofValues)
        values.resize(numParticles);
3091
3092
3093
3094

    // Create the computation objects.

    dynamics = new ReferenceCustomDynamics(system.getNumParticles(), integrator);
3095
    SimTKOpenMMUtilities::setRandomNumberSeed((unsigned int) integrator.getRandomNumberSeed());
3096
3097
3098
}

void ReferenceIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegrator& integrator, bool& forcesAreValid) {
peastman's avatar
peastman committed
3099
3100
3101
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& velData = extractVelocities(context);
    vector<Vec3>& forceData = extractForces(context);
3102
3103
3104
3105
3106
3107
3108
3109
3110
3111
    
    // Record global variables.
    
    map<string, double> globals;
    globals["dt"] = integrator.getStepSize();
    for (int i = 0; i < integrator.getNumGlobalVariables(); i++)
        globals[integrator.getGlobalVariableName(i)] = globalValues[i];
    
    // Execute the step.
    
3112
    dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
3113
    dynamics->setVirtualSites(extractVirtualSites(context));
3114
    dynamics->update(context, context.getSystem().getNumParticles(), posData, velData, forceData, masses, globals, perDofValues, forcesAreValid, integrator.getConstraintTolerance(), extractBoxVectors(context));
3115
3116
3117
3118
3119
3120
3121
3122
3123
3124
    
    // Record changed global variables.
    
    integrator.setStepSize(globals["dt"]);
    for (int i = 0; i < (int) globalValues.size(); i++)
        globalValues[i] = globals[integrator.getGlobalVariableName(i)];
    data.time += dynamics->getDeltaT();
    data.stepCount++;
}

3125
double ReferenceIntegrateCustomStepKernel::computeKineticEnergy(ContextImpl& context, CustomIntegrator& integrator, bool& forcesAreValid) {
peastman's avatar
peastman committed
3126
3127
3128
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& velData = extractVelocities(context);
    vector<Vec3>& forceData = extractForces(context);
3129
3130
3131
3132
3133
3134
3135
3136
3137
3138
3139
3140
3141
    
    // Record global variables.
    
    map<string, double> globals;
    globals["dt"] = integrator.getStepSize();
    for (int i = 0; i < integrator.getNumGlobalVariables(); i++)
        globals[integrator.getGlobalVariableName(i)] = globalValues[i];
    
    // Compute the kinetic energy.
    
    return dynamics->computeKineticEnergy(context, context.getSystem().getNumParticles(), posData, velData, forceData, masses, globals, perDofValues, forcesAreValid);
}

3142
3143
3144
3145
3146
3147
3148
3149
3150
3151
3152
3153
3154
3155
3156
3157
3158
3159
3160
3161
void ReferenceIntegrateCustomStepKernel::getGlobalVariables(ContextImpl& context, vector<double>& values) const {
    values = globalValues;
}

void ReferenceIntegrateCustomStepKernel::setGlobalVariables(ContextImpl& context, const vector<double>& values) {
    globalValues = values;
}

void ReferenceIntegrateCustomStepKernel::getPerDofVariable(ContextImpl& context, int variable, vector<Vec3>& values) const {
    values.resize(perDofValues[variable].size());
    for (int i = 0; i < (int) values.size(); i++)
        values[i] = perDofValues[variable][i];
}

void ReferenceIntegrateCustomStepKernel::setPerDofVariable(ContextImpl& context, int variable, const vector<Vec3>& values) {
    perDofValues[variable].resize(values.size());
    for (int i = 0; i < (int) values.size(); i++)
        perDofValues[variable][i] = values[i];
}

Peter Eastman's avatar
Peter Eastman committed
3162
3163
3164
3165
3166
3167
3168
3169
3170
3171
3172
3173
3174
3175
3176
3177
3178
3179
3180
3181
3182
ReferenceIntegrateDPDStepKernel::~ReferenceIntegrateDPDStepKernel() {
    if (dynamics != NULL)
        delete dynamics;
}

void ReferenceIntegrateDPDStepKernel::initialize(const System& system, const DPDIntegrator& integrator) {
    masses.resize(system.getNumParticles());
    for (int i = 0; i < system.getNumParticles(); ++i)
        masses[i] = system.getParticleMass(i);
    dynamics = new ReferenceDPDDynamics(system, integrator);
    SimTKOpenMMUtilities::setRandomNumberSeed((unsigned int) integrator.getRandomNumberSeed());
}

void ReferenceIntegrateDPDStepKernel::execute(ContextImpl& context, const DPDIntegrator& integrator) {
    dynamics->setTemperature(integrator.getTemperature());
    dynamics->setDeltaT(integrator.getStepSize());
    dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
    dynamics->setVirtualSites(extractVirtualSites(context));
    dynamics->setPeriodicBoxVectors(extractBoxVectors(context));
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& velData = extractVelocities(context);
3183
    dynamics->update(context, posData, velData, masses, integrator.getConstraintTolerance(), extractBoxVectors(context));
Peter Eastman's avatar
Peter Eastman committed
3184
3185
3186
3187
3188
3189
3190
3191
    data.time += integrator.getStepSize();
    data.stepCount++;
}

double ReferenceIntegrateDPDStepKernel::computeKineticEnergy(ContextImpl& context, const DPDIntegrator& integrator) {
    return computeShiftedKineticEnergy(context, masses, 0.0);
}

3192
3193
3194
3195
3196
3197
3198
3199
3200
3201
3202
3203
3204
3205
3206
3207
3208
3209
3210
3211
3212
3213
3214
3215
3216
3217
3218
3219
3220
3221
3222
3223
3224
3225
3226
3227
3228
3229
3230
3231
3232
3233
3234
3235
3236
3237
3238
3239
3240
ReferenceIntegrateQTBStepKernel::~ReferenceIntegrateQTBStepKernel() {
    if (dynamics != NULL)
        delete dynamics;
}

void ReferenceIntegrateQTBStepKernel::initialize(const System& system, const QTBIntegrator& integrator) {
    int numParticles = system.getNumParticles();
    masses.resize(numParticles);
    for (int i = 0; i < numParticles; ++i)
        masses[i] = system.getParticleMass(i);
    SimTKOpenMMUtilities::setRandomNumberSeed((unsigned int) integrator.getRandomNumberSeed());
    dynamics = new ReferenceQTBDynamics(system, integrator);
}

void ReferenceIntegrateQTBStepKernel::execute(ContextImpl& context, const QTBIntegrator& integrator) {
    double stepSize = integrator.getStepSize();
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& velData = extractVelocities(context);
    if (!hasInitialized) {
        hasInitialized = true;
        dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
        dynamics->setVirtualSites(extractVirtualSites(context));
    }
    dynamics->setTemperature(integrator.getTemperature());
    dynamics->update(context, posData, velData, masses, integrator.getConstraintTolerance(), extractBoxVectors(context), extractThreadPool(context));
    data.time += stepSize;
    data.stepCount++;
}

double ReferenceIntegrateQTBStepKernel::computeKineticEnergy(ContextImpl& context, const QTBIntegrator& integrator) {
    return computeShiftedKineticEnergy(context, masses, 0.0);
}

void ReferenceIntegrateQTBStepKernel::getAdaptedFriction(ContextImpl& context, int particle, std::vector<double>& friction) const {
    dynamics->getAdaptedFriction(particle, friction);
}

void ReferenceIntegrateQTBStepKernel::setAdaptedFriction(ContextImpl& context, int particle, const std::vector<double>& friction) {
    dynamics->setAdaptedFriction(particle, friction);
}

void ReferenceIntegrateQTBStepKernel::createCheckpoint(ContextImpl& context, ostream& stream) const {
    dynamics->createCheckpoint(stream);
}

void ReferenceIntegrateQTBStepKernel::loadCheckpoint(ContextImpl& context, istream& stream) {
    dynamics->loadCheckpoint(stream);
}

3241
3242
3243
3244
3245
ReferenceApplyAndersenThermostatKernel::~ReferenceApplyAndersenThermostatKernel() {
    if (thermostat)
        delete thermostat;
}

3246
void ReferenceApplyAndersenThermostatKernel::initialize(const System& system, const AndersenThermostat& thermostat) {
Peter Eastman's avatar
Peter Eastman committed
3247
    int numParticles = system.getNumParticles();
3248
    masses.resize(numParticles);
3249
    for (int i = 0; i < numParticles; ++i)
peastman's avatar
peastman committed
3250
        masses[i] = system.getParticleMass(i);
3251
    this->thermostat = new ReferenceAndersenThermostat();
3252
    SimTKOpenMMUtilities::setRandomNumberSeed((unsigned int) thermostat.getRandomNumberSeed());
3253
    particleGroups = AndersenThermostatImpl::calcParticleGroups(system);
3254
3255
}

3256
void ReferenceApplyAndersenThermostatKernel::execute(ContextImpl& context) {
peastman's avatar
peastman committed
3257
    vector<Vec3>& velData = extractVelocities(context);
3258
    thermostat->applyThermostat(particleGroups, velData, masses,
peastman's avatar
peastman committed
3259
3260
3261
        context.getParameter(AndersenThermostat::Temperature()),
        context.getParameter(AndersenThermostat::CollisionFrequency()),
        context.getIntegrator().getStepSize());
3262
3263
}

3264
3265
3266
3267
3268
ReferenceApplyMonteCarloBarostatKernel::~ReferenceApplyMonteCarloBarostatKernel() {
    if (barostat)
        delete barostat;
}

3269
3270
void ReferenceApplyMonteCarloBarostatKernel::initialize(const System& system, const Force& barostat, int components, bool rigidMolecules) {
    this->components = components;
3271
    this->rigidMolecules = rigidMolecules;
3272
3273
}

3274
void ReferenceApplyMonteCarloBarostatKernel::saveCoordinates(ContextImpl& context) {
3275
    if (barostat == NULL) {
3276
3277
3278
3279
        const System& system = context.getSystem();
        vector<double> masses;
        for (int i = 0; i < system.getNumParticles(); i++)
            masses.push_back(system.getParticleMass(i));
3280
        if (rigidMolecules)
3281
            barostat = new ReferenceMonteCarloBarostat(system.getNumParticles(), context.getMolecules(), masses);
3282
        else {
3283
            vector<vector<int> > molecules(system.getNumParticles());
3284
3285
            for (int i = 0; i < molecules.size(); i++)
                molecules[i].push_back(i);
3286
            barostat = new ReferenceMonteCarloBarostat(system.getNumParticles(), molecules, masses);
3287
3288
        }
    }
3289
3290
3291
3292
3293
    vector<Vec3>& posData = extractPositions(context);
    barostat->savePositions(posData);
}

void ReferenceApplyMonteCarloBarostatKernel::scaleCoordinates(ContextImpl& context, double scaleX, double scaleY, double scaleZ) {
peastman's avatar
peastman committed
3294
3295
    vector<Vec3>& posData = extractPositions(context);
    Vec3* boxVectors = extractBoxVectors(context);
3296
    barostat->applyBarostat(posData, boxVectors, scaleX, scaleY, scaleZ);
3297
3298
3299
}

void ReferenceApplyMonteCarloBarostatKernel::restoreCoordinates(ContextImpl& context) {
peastman's avatar
peastman committed
3300
    vector<Vec3>& posData = extractPositions(context);
3301
3302
3303
    barostat->restorePositions(posData);
}

3304
3305
3306
3307
void ReferenceApplyMonteCarloBarostatKernel::computeKineticEnergy(ContextImpl& context, vector<double>& ke) {
    barostat->computeMolecularKineticEnergy(extractVelocities(context), ke, components);
}

3308
3309
void ReferenceRemoveCMMotionKernel::initialize(const System& system, const CMMotionRemover& force) {
    frequency = force.getFrequency();
Peter Eastman's avatar
Peter Eastman committed
3310
    masses.resize(system.getNumParticles());
3311
    for (size_t i = 0; i < masses.size(); ++i)
Peter Eastman's avatar
Peter Eastman committed
3312
        masses[i] = system.getParticleMass(i);
3313
3314
}

3315
void ReferenceRemoveCMMotionKernel::execute(ContextImpl& context) {
3316
    if (data.stepCount%frequency != 0)
3317
        return;
peastman's avatar
peastman committed
3318
    vector<Vec3>& velData = extractVelocities(context);
3319
3320
3321
    
    // Calculate the center of mass momentum.
    
peastman's avatar
peastman committed
3322
3323
    double momentum[] = {0.0, 0.0, 0.0};
    double mass = 0.0;
3324
    for (size_t i = 0; i < masses.size(); ++i) {
peastman's avatar
peastman committed
3325
3326
3327
3328
        momentum[0] += masses[i]*velData[i][0];
        momentum[1] += masses[i]*velData[i][1];
        momentum[2] += masses[i]*velData[i][2];
        mass += masses[i];
3329
3330
    }
    
Peter Eastman's avatar
Peter Eastman committed
3331
    // Adjust the particle velocities.
3332
    
3333
3334
3335
    momentum[0] /= mass;
    momentum[1] /= mass;
    momentum[2] /= mass;
3336
    for (size_t i = 0; i < masses.size(); ++i) {
3337
3338
3339
3340
3341
        if (masses[i] != 0.0) {
            velData[i][0] -= momentum[0];
            velData[i][1] -= momentum[1];
            velData[i][2] -= momentum[2];
        }
3342
3343
    }
}
3344

3345
3346
3347
3348
3349
3350
3351
3352
3353
3354
3355
3356
3357
3358
3359
3360
3361
3362
3363
3364
3365
3366
3367
3368
3369
3370
3371
3372
3373
3374
3375
3376
3377
3378
void ReferenceCalcATMForceKernel::loadParams(int numParticles, const ATMForce& force) {
    //vector displacements
    displacement1.resize(numParticles);
    displacement0.resize(numParticles);
    //particle distance displacements
    pj1.resize(numParticles);
    pi1.resize(numParticles);
    pj0.resize(numParticles);
    pi0.resize(numParticles);
    for (int i = 0; i < numParticles; i++) {
        const ATMForce::CoordinateTransformation& transformation = force.getParticleTransformation(i);
        if (dynamic_cast<const ATMForce::FixedDisplacement*>(&transformation) != NULL) {
            const ATMForce::FixedDisplacement* fd = dynamic_cast<const ATMForce::FixedDisplacement*>(&transformation);
            const Vec3 d1 = fd->getFixedDisplacement1();
            const Vec3 d0 = fd->getFixedDisplacement0();
            displacement1[i] = d1;
            displacement0[i] = d0;
            pj1[i] = pi1[i] = pj0[i] = pi0[i] = -1;
        }
        else if (dynamic_cast<const ATMForce::ParticleOffsetDisplacement*>(&transformation) != NULL) {
          const ATMForce::ParticleOffsetDisplacement* vd = dynamic_cast<const ATMForce::ParticleOffsetDisplacement*>(&transformation);
            displacement1[i] = Vec3(0, 0, 0);
            displacement0[i] = Vec3(0, 0, 0);
            pj1[i] = vd->getDestinationParticle1();
            pi1[i] = vd->getOriginParticle1();
            pj0[i] = vd->getDestinationParticle0();
            pi0[i] = vd->getOriginParticle0();
        }
        else {
            throw OpenMMException("loadParams(): invalid particle Transformation");
        }
    }
}

3379
3380
3381
3382
3383
void ReferenceCalcATMForceKernel::initialize(const System& system, const ATMForce& force) {
    numParticles = force.getNumParticles();
    //displacement map
    displ1.resize(numParticles);
    displ0.resize(numParticles);
3384
3385
3386
3387
3388
3389
3390
3391
3392
3393
3394
3395
3396
3397
3398
3399
3400
3401
3402
3403
3404
3405
3406
3407
3408
3409
3410
3411
3412
3413
3414
3415
3416
3417
3418
3419
3420
3421
3422
3423
3424
3425
3426
3427
3428
    //load particle parameters from the force object
    loadParams(numParticles, force);
}

void ReferenceCalcATMForceKernel::setDisplacements(vector<Vec3>& pos){
  numParticles = pos.size();

  for (int i = 0; i < numParticles; i++) {
    if (pj1[i] >= 0 && pi1[i] >= 0){
        displ1[i] = pos[pj1[i]] - pos[pi1[i]];
        if (pi0[i] >= 0 && pj0[i] >= 0){
            displ0[i] = pos[pj0[i]] - pos[pi0[i]];
        }else{
            displ0[i] = Vec3();
        }
    }else{
        displ1[i] = displacement1[i];
        displ0[i] = displacement0[i];
    }
  }
}


//Add forces from variable displacements
void ReferenceCalcATMForceKernel::displForces(vector<Vec3>& force0, vector<Vec3>& force1){
    vector<Vec3> dforce1(numParticles), dforce0(numParticles);

    for (int i = 0; i < numParticles; i++){
        if (pj1[i] >= 0 && pi1[i] >= 0){
            dforce1[pj1[i]] += force1[i];
            dforce1[pi1[i]] -= force1[i];
        }
    }
    for (int i = 0; i < numParticles; i++){
        force1[i] += dforce1[i];
    }

    for (int i = 0; i < numParticles; i++){
        if (pj0[i] >= 0 && pi0[i] >= 0){
            dforce0[pj0[i]] += force0[i];
            dforce0[pi0[i]] -= force0[i];
        }
    }
    for (int i = 0; i < numParticles; i++){
        force0[i] += dforce0[i];
3429
3430
3431
3432
3433
3434
3435
3436
    }
}

void ReferenceCalcATMForceKernel::applyForces(ContextImpl& context, ContextImpl& innerContext0, ContextImpl& innerContext1,
        double dEdu0, double dEdu1, const map<string, double>& energyParamDerivs) {
    vector<Vec3>& force = extractForces(context);
    vector<Vec3>& force0 = extractForces(innerContext0);
    vector<Vec3>& force1 = extractForces(innerContext1);
3437

3438
3439
3440
    //add gradients from variable displacements
    displForces(force0, force1);

3441
3442
3443
3444
3445
    //protects from infinite forces when the hybrid potential does
    //not depend on u1 or u0, typically at the endpoints
    double epsi = std::numeric_limits<float>::min();
    for (int i = 0; i < force.size(); i++) {
        if (fabs(dEdu0) > epsi)
3446
3447
3448
            force[i] += dEdu0*force0[i];
        if (fabs(dEdu1) > epsi)
            force[i] += dEdu1*force1[i];
3449
3450
    }

3451
3452
3453
3454
3455
3456
3457
3458
    map<string, double>& derivs = extractEnergyParameterDerivatives(context);
    for (auto deriv : energyParamDerivs)
        derivs[deriv.first] += deriv.second;
}

void ReferenceCalcATMForceKernel::copyState(ContextImpl& context, ContextImpl& innerContext0, ContextImpl& innerContext1) {
    vector<Vec3>& pos = extractPositions(context);

3459
3460
3461
    //calculate displacement vectors
    setDisplacements(pos);

3462
3463
3464
3465
    //in the initial state, particles are displaced by displ0
    vector<Vec3> pos0(pos);
    for (int i = 0; i < pos0.size(); i++)
        pos0[i] += displ0[i];
3466
    innerContext0.setPositions(pos0);
3467
3468
3469
3470
3471

    //in the target state, particles are displaced by displ1
    vector<Vec3> pos1(pos);
    for (int i = 0; i < pos1.size(); i++)
        pos1[i] += displ1[i];
3472
    innerContext1.setPositions(pos1);
3473
3474
3475
3476
3477
3478
3479
3480
3481
3482
3483
3484
3485
3486
3487
3488
3489
3490
3491
3492
3493
3494
3495
3496
3497
3498

    Vec3 a, b, c;
    context.getPeriodicBoxVectors(a, b, c);
    innerContext0.setPeriodicBoxVectors(a, b, c);
    innerContext1.setPeriodicBoxVectors(a, b, c);

    innerContext0.setTime(context.getTime());
    innerContext1.setTime(context.getTime());

    map<string, double> innerParameters;

    innerParameters = innerContext0.getParameters();
    for (auto& param : innerParameters)
        innerContext0.setParameter(param.first, context.getParameter(param.first));

    innerParameters = innerContext1.getParameters();
    for (auto& param : innerParameters)
        innerContext1.setParameter(param.first, context.getParameter(param.first));

}

void ReferenceCalcATMForceKernel::copyParametersToContext(ContextImpl& context, const ATMForce& force) {
    if (force.getNumParticles() != numParticles)
          throw OpenMMException("copyParametersToContext: The number of ATMForce particles has changed");
    displ1.resize(numParticles);
    displ0.resize(numParticles);
3499
    loadParams(numParticles, force);
3500
}
3501
3502
3503
3504
3505
3506
3507
3508
3509
3510
3511
3512
3513
3514
3515

void ReferenceCalcCustomCPPForceKernel::initialize(const System& system, CustomCPPForceImpl& force) {
    this->force = &force;
    forces.resize(system.getNumParticles());
}

double ReferenceCalcCustomCPPForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = force->computeForce(context, posData, forces);
    if (includeForces)
        for (int i = 0; i < forces.size(); i++)
            forceData[i] += forces[i];
    return energy;
}