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
    if (includeForces) {
        int numParticles = context.getSystem().getNumParticles();
Peter Eastman's avatar
Peter Eastman committed
176
177
        for (int i = 0; i < numParticles; ++i)
            forceData[i] = Vec3();
178
    }
179
180
    else
        savedForces = forceData;
peastman's avatar
peastman committed
181
182
    for (auto& param : context.getParameters())
        extractEnergyParameterDerivatives(context)[param.first] = 0;
183
184
}

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

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

200
double ReferenceUpdateStateDataKernel::getTime(const ContextImpl& context) const {
201
202
203
    return data.time;
}

204
void ReferenceUpdateStateDataKernel::setTime(ContextImpl& context, double time) {
205
206
207
    data.time = time;
}

208
209
210
211
212
213
214
215
long long ReferenceUpdateStateDataKernel::getStepCount(const ContextImpl& context) const {
    return data.stepCount;
}

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

Peter Eastman's avatar
Peter Eastman committed
216
void ReferenceUpdateStateDataKernel::getPositions(ContextImpl& context, std::vector<Vec3>& positions, bool allowPeriodic) {
Peter Eastman's avatar
Peter Eastman committed
217
    positions = extractPositions(context);
218
219
220
}

void ReferenceUpdateStateDataKernel::setPositions(ContextImpl& context, const std::vector<Vec3>& positions) {
Peter Eastman's avatar
Peter Eastman committed
221
    extractPositions(context) = positions;
222
223
224
}

void ReferenceUpdateStateDataKernel::getVelocities(ContextImpl& context, std::vector<Vec3>& velocities) {
Peter Eastman's avatar
Peter Eastman committed
225
    velocities = extractVelocities(context);
226
227
228
}

void ReferenceUpdateStateDataKernel::setVelocities(ContextImpl& context, const std::vector<Vec3>& velocities) {
Peter Eastman's avatar
Peter Eastman committed
229
    extractVelocities(context) = velocities;
230
231
}

232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
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.
254
255
256

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

259
void ReferenceUpdateStateDataKernel::getForces(ContextImpl& context, std::vector<Vec3>& forces) {
Peter Eastman's avatar
Peter Eastman committed
260
    forces = extractForces(context);
261
262
}

263
264
265
266
void ReferenceUpdateStateDataKernel::getEnergyParameterDerivatives(ContextImpl& context, map<string, double>& derivs) {
    derivs = extractEnergyParameterDerivatives(context);
}

267
void ReferenceUpdateStateDataKernel::getPeriodicBoxVectors(ContextImpl& context, Vec3& a, Vec3& b, Vec3& c) const {
peastman's avatar
peastman committed
268
    Vec3* vectors = extractBoxVectors(context);
269
270
271
    a = vectors[0];
    b = vectors[1];
    c = vectors[2];
272
273
}

274
void ReferenceUpdateStateDataKernel::setPeriodicBoxVectors(ContextImpl& context, const Vec3& a, const Vec3& b, const Vec3& c) {
peastman's avatar
peastman committed
275
276
277
278
279
    Vec3& box = extractBoxSize(context);
    box[0] = a[0];
    box[1] = b[1];
    box[2] = c[2];
    Vec3* vectors = extractBoxVectors(context);
280
281
282
    vectors[0] = a;
    vectors[1] = b;
    vectors[2] = c;
283
284
}

Peter Eastman's avatar
Peter Eastman committed
285
void ReferenceUpdateStateDataKernel::createCheckpoint(ContextImpl& context, ostream& stream) {
286
    int version = 3;
Peter Eastman's avatar
Peter Eastman committed
287
288
    stream.write((char*) &version, sizeof(int));
    stream.write((char*) &data.time, sizeof(data.time));
289
    stream.write((char*) &data.stepCount, sizeof(long long));
peastman's avatar
peastman committed
290
291
292
293
294
295
    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
296
297
298
299
300
301
    SimTKOpenMMUtilities::createCheckpoint(stream);
}

void ReferenceUpdateStateDataKernel::loadCheckpoint(ContextImpl& context, istream& stream) {
    int version;
    stream.read((char*) &version, sizeof(int));
302
    if (version != 3)
Peter Eastman's avatar
Peter Eastman committed
303
304
        throw OpenMMException("Checkpoint was created with a different version of OpenMM");
    stream.read((char*) &data.time, sizeof(data.time));
305
    stream.read((char*) &data.stepCount, sizeof(long long));
peastman's avatar
peastman committed
306
307
308
309
310
311
    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
312
313
314
    SimTKOpenMMUtilities::loadCheckpoint(stream);
}

315
316
void ReferenceApplyConstraintsKernel::initialize(const System& system) {
    int numParticles = system.getNumParticles();
317
318
    masses.resize(numParticles);
    inverseMasses.resize(numParticles);
319
    for (int i = 0; i < numParticles; ++i) {
peastman's avatar
peastman committed
320
        masses[i] = system.getParticleMass(i);
321
322
323
324
325
326
327
328
        inverseMasses[i] = 1.0/masses[i];
    }
}

ReferenceApplyConstraintsKernel::~ReferenceApplyConstraintsKernel() {
}

void ReferenceApplyConstraintsKernel::apply(ContextImpl& context, double tol) {
peastman's avatar
peastman committed
329
    vector<Vec3>& positions = extractPositions(context);
330
    extractConstraints(context).apply(positions, positions, inverseMasses, tol);
331
    extractVirtualSites(context).computePositions(context.getSystem(), positions, extractBoxVectors(context));
332
333
}

334
void ReferenceApplyConstraintsKernel::applyToVelocities(ContextImpl& context, double tol) {
peastman's avatar
peastman committed
335
336
    vector<Vec3>& positions = extractPositions(context);
    vector<Vec3>& velocities = extractVelocities(context);
337
    extractConstraints(context).applyToVelocities(positions, velocities, inverseMasses, tol);
338
339
}

340
341
342
343
void ReferenceVirtualSitesKernel::initialize(const System& system) {
}

void ReferenceVirtualSitesKernel::computePositions(ContextImpl& context) {
peastman's avatar
peastman committed
344
    vector<Vec3>& positions = extractPositions(context);
345
    extractVirtualSites(context).computePositions(context.getSystem(), positions, extractBoxVectors(context));
346
347
}

348
void ReferenceCalcHarmonicBondForceKernel::initialize(const System& system, const HarmonicBondForce& force) {
349
    numBonds = force.getNumBonds();
350
351
    bondIndexArray.resize(numBonds, vector<int>(2));
    bondParamArray.resize(numBonds, vector<double>(2));
352
    for (int i = 0; i < numBonds; ++i) {
Peter Eastman's avatar
Peter Eastman committed
353
        int particle1, particle2;
354
        double length, k;
Peter Eastman's avatar
Peter Eastman committed
355
356
357
        force.getBondParameters(i, particle1, particle2, length, k);
        bondIndexArray[i][0] = particle1;
        bondIndexArray[i][1] = particle2;
peastman's avatar
peastman committed
358
359
        bondParamArray[i][0] = length;
        bondParamArray[i][1] = k;
360
    }
361
    usePeriodic = force.usesPeriodicBoundaryConditions();
362
363
}

364
double ReferenceCalcHarmonicBondForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
365
366
367
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
368
369
    ReferenceBondForce refBondForce;
    ReferenceHarmonicBondIxn harmonicBond;
370
371
    if (usePeriodic)
        harmonicBond.setPeriodic(extractBoxVectors(context));
372
    refBondForce.calculateForce(numBonds, bondIndexArray, posData, bondParamArray, forceData, includeEnergy ? &energy : NULL, harmonicBond);
373
374
375
    return energy;
}

376
void ReferenceCalcHarmonicBondForceKernel::copyParametersToContext(ContextImpl& context, const HarmonicBondForce& force, int firstBond, int lastBond) {
377
378
379
380
381
    if (numBonds != force.getNumBonds())
        throw OpenMMException("updateParametersInContext: The number of bonds has changed");

    // Record the values.

382
    for (int i = firstBond; i <= lastBond; ++i) {
383
384
385
386
387
388
389
        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
390
391
        bondParamArray[i][0] = length;
        bondParamArray[i][1] = k;
392
393
394
    }
}

395
396
397
398
399
ReferenceCalcCustomBondForceKernel::~ReferenceCalcCustomBondForceKernel() {
    if (ixn != NULL)
        delete ixn;
}

400
401
402
void ReferenceCalcCustomBondForceKernel::initialize(const System& system, const CustomBondForce& force) {
    numBonds = force.getNumBonds();
    int numParameters = force.getNumPerBondParameters();
403
    usePeriodic = force.usesPeriodicBoundaryConditions();
404
405
406

    // Build the arrays.

407
408
    bondIndexArray.resize(numBonds, vector<int>(2));
    bondParamArray.resize(numBonds, vector<double>(numParameters));
409
    vector<double> params;
410
    for (int i = 0; i < numBonds; ++i) {
411
412
413
414
415
        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
416
            bondParamArray[i][j] = params[j];
417
418
419
420
421
    }

    // Parse the expression used to calculate the force.

    Lepton::ParsedExpression expression = Lepton::Parser::parse(force.getEnergyFunction()).optimize();
422
423
    energyExpression = expression.createCompiledExpression();
    forceExpression = expression.differentiate("r").createCompiledExpression();
424
425
426
427
    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));
428
429
430
431
432
    for (int i = 0; i < force.getNumEnergyParameterDerivatives(); i++) {
        string param = force.getEnergyParameterDerivativeName(i);
        energyParamDerivNames.push_back(param);
        energyParamDerivExpressions.push_back(expression.differentiate(param).createCompiledExpression());
    }
433
434
435
436
437
    set<string> variables;
    variables.insert("r");
    variables.insert(parameterNames.begin(), parameterNames.end());
    variables.insert(globalParameterNames.begin(), globalParameterNames.end());
    validateVariables(expression.getRootNode(), variables);
438
    ixn = new ReferenceCustomBondIxn(energyExpression, forceExpression, parameterNames, energyParamDerivExpressions);
439
440
}

441
double ReferenceCalcCustomBondForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
442
443
444
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
445
    map<string, double> globalParameters;
peastman's avatar
peastman committed
446
447
    for (auto& name : globalParameterNames)
        globalParameters[name] = context.getParameter(name);
448
    ixn->setGlobalParameters(globalParameters);
449
    if (usePeriodic)
450
        ixn->setPeriodic(extractBoxVectors(context));
451
452
    vector<double> energyParamDerivValues(energyParamDerivNames.size()+1, 0.0);
    for (int i = 0; i < numBonds; i++)
453
        ixn->calculateBondIxn(bondIndexArray[i], posData, bondParamArray[i], forceData, includeEnergy ? &energy : NULL, &energyParamDerivValues[0]);
454
455
456
    map<string, double>& energyParamDerivs = extractEnergyParameterDerivatives(context);
    for (int i = 0; i < energyParamDerivNames.size(); i++)
        energyParamDerivs[energyParamDerivNames[i]] += energyParamDerivValues[i];
457
458
459
    return energy;
}

460
void ReferenceCalcCustomBondForceKernel::copyParametersToContext(ContextImpl& context, const CustomBondForce& force, int firstBond, int lastBond) {
461
462
463
464
465
466
467
    if (numBonds != force.getNumBonds())
        throw OpenMMException("updateParametersInContext: The number of bonds has changed");

    // Record the values.

    int numParameters = force.getNumPerBondParameters();
    vector<double> params;
468
    for (int i = firstBond; i <= lastBond; ++i) {
469
470
471
472
473
        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
474
            bondParamArray[i][j] = params[j];
475
476
477
    }
}

478
479
void ReferenceCalcHarmonicAngleForceKernel::initialize(const System& system, const HarmonicAngleForce& force) {
    numAngles = force.getNumAngles();
480
481
    angleIndexArray.resize(numAngles, vector<int>(3));
    angleParamArray.resize(numAngles, vector<double>(2));
482
    for (int i = 0; i < numAngles; ++i) {
Peter Eastman's avatar
Peter Eastman committed
483
        int particle1, particle2, particle3;
484
        double angle, k;
Peter Eastman's avatar
Peter Eastman committed
485
486
487
488
        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
489
490
        angleParamArray[i][0] = angle;
        angleParamArray[i][1] = k;
491
    }
492
    usePeriodic = force.usesPeriodicBoundaryConditions();
493
494
}

495
double ReferenceCalcHarmonicAngleForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
496
497
498
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
499
500
    ReferenceBondForce refBondForce;
    ReferenceAngleBondIxn angleBond;
501
502
    if (usePeriodic)
        angleBond.setPeriodic(extractBoxVectors(context));
503
    refBondForce.calculateForce(numAngles, angleIndexArray, posData, angleParamArray, forceData, includeEnergy ? &energy : NULL, angleBond);
504
505
506
    return energy;
}

507
void ReferenceCalcHarmonicAngleForceKernel::copyParametersToContext(ContextImpl& context, const HarmonicAngleForce& force, int firstAngle, int lastAngle) {
508
509
510
511
512
    if (numAngles != force.getNumAngles())
        throw OpenMMException("updateParametersInContext: The number of angles has changed");

    // Record the values.

513
    for (int i = firstAngle; i <= lastAngle; ++i) {
514
515
516
517
518
        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
519
520
        angleParamArray[i][0] = angle;
        angleParamArray[i][1] = k;
521
522
523
    }
}

524
525
526
527
528
ReferenceCalcCustomAngleForceKernel::~ReferenceCalcCustomAngleForceKernel() {
    if (ixn != NULL)
        delete ixn;
}

529
530
531
void ReferenceCalcCustomAngleForceKernel::initialize(const System& system, const CustomAngleForce& force) {
    numAngles = force.getNumAngles();
    int numParameters = force.getNumPerAngleParameters();
532
    usePeriodic = force.usesPeriodicBoundaryConditions();
533
534
535

    // Build the arrays.

536
537
    angleIndexArray.resize(numAngles, vector<int>(3));
    angleParamArray.resize(numAngles, vector<double>(numParameters));
538
    vector<double> params;
539
    for (int i = 0; i < numAngles; ++i) {
540
541
542
543
544
545
        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
546
            angleParamArray[i][j] = params[j];
547
548
549
550
551
    }

    // Parse the expression used to calculate the force.

    Lepton::ParsedExpression expression = Lepton::Parser::parse(force.getEnergyFunction()).optimize();
552
553
    energyExpression = expression.createCompiledExpression();
    forceExpression = expression.differentiate("theta").createCompiledExpression();
554
555
556
557
    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));
558
559
560
561
562
    for (int i = 0; i < force.getNumEnergyParameterDerivatives(); i++) {
        string param = force.getEnergyParameterDerivativeName(i);
        energyParamDerivNames.push_back(param);
        energyParamDerivExpressions.push_back(expression.differentiate(param).createCompiledExpression());
    }
563
564
565
566
567
    set<string> variables;
    variables.insert("theta");
    variables.insert(parameterNames.begin(), parameterNames.end());
    variables.insert(globalParameterNames.begin(), globalParameterNames.end());
    validateVariables(expression.getRootNode(), variables);
568
    ixn = new ReferenceCustomAngleIxn(energyExpression, forceExpression, parameterNames, energyParamDerivExpressions);
569
570
}

571
double ReferenceCalcCustomAngleForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
572
573
574
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
575
    map<string, double> globalParameters;
peastman's avatar
peastman committed
576
577
    for (auto& name : globalParameterNames)
        globalParameters[name] = context.getParameter(name);
578
    ixn->setGlobalParameters(globalParameters);
579
    if (usePeriodic)
580
        ixn->setPeriodic(extractBoxVectors(context));
581
582
    vector<double> energyParamDerivValues(energyParamDerivNames.size()+1, 0.0);
    for (int i = 0; i < numAngles; i++)
583
        ixn->calculateBondIxn(angleIndexArray[i], posData, angleParamArray[i], forceData, includeEnergy ? &energy : NULL, &energyParamDerivValues[0]);
584
585
586
    map<string, double>& energyParamDerivs = extractEnergyParameterDerivatives(context);
    for (int i = 0; i < energyParamDerivNames.size(); i++)
        energyParamDerivs[energyParamDerivNames[i]] += energyParamDerivValues[i];
587
588
589
    return energy;
}

590
void ReferenceCalcCustomAngleForceKernel::copyParametersToContext(ContextImpl& context, const CustomAngleForce& force, int firstAngle, int lastAngle) {
591
592
593
594
595
596
597
    if (numAngles != force.getNumAngles())
        throw OpenMMException("updateParametersInContext: The number of angles has changed");

    // Record the values.

    int numParameters = force.getNumPerAngleParameters();
    vector<double> params;
598
    for (int i = firstAngle; i <= lastAngle; ++i) {
599
600
601
602
603
        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
604
            angleParamArray[i][j] = params[j];
605
606
607
    }
}

608
609
void ReferenceCalcPeriodicTorsionForceKernel::initialize(const System& system, const PeriodicTorsionForce& force) {
    numTorsions = force.getNumTorsions();
610
611
    torsionIndexArray.resize(numTorsions, vector<int>(4));
    torsionParamArray.resize(numTorsions, vector<double>(3));
612
    for (int i = 0; i < numTorsions; ++i) {
Peter Eastman's avatar
Peter Eastman committed
613
        int particle1, particle2, particle3, particle4, periodicity;
614
        double phase, k;
Peter Eastman's avatar
Peter Eastman committed
615
616
617
618
619
        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
620
621
622
        torsionParamArray[i][0] = k;
        torsionParamArray[i][1] = phase;
        torsionParamArray[i][2] = periodicity;
623
    }
624
    usePeriodic = force.usesPeriodicBoundaryConditions();
625
626
}

627
double ReferenceCalcPeriodicTorsionForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
628
629
630
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
631
632
    ReferenceBondForce refBondForce;
    ReferenceProperDihedralBond periodicTorsionBond;
633
634
    if (usePeriodic)
        periodicTorsionBond.setPeriodic(extractBoxVectors(context));
635
    refBondForce.calculateForce(numTorsions, torsionIndexArray, posData, torsionParamArray, forceData, includeEnergy ? &energy : NULL, periodicTorsionBond);
636
637
638
    return energy;
}

639
void ReferenceCalcPeriodicTorsionForceKernel::copyParametersToContext(ContextImpl& context, const PeriodicTorsionForce& force, int firstTorsion, int lastTorsion) {
640
641
642
643
644
    if (numTorsions != force.getNumTorsions())
        throw OpenMMException("updateParametersInContext: The number of torsions has changed");

    // Record the values.

645
    for (int i = firstTorsion; i <= lastTorsion; ++i) {
646
647
648
649
650
        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
651
652
653
        torsionParamArray[i][0] = k;
        torsionParamArray[i][1] = phase;
        torsionParamArray[i][2] = periodicity;
654
655
656
    }
}

657
658
void ReferenceCalcRBTorsionForceKernel::initialize(const System& system, const RBTorsionForce& force) {
    numTorsions = force.getNumTorsions();
659
660
    torsionIndexArray.resize(numTorsions, vector<int>(4));
    torsionParamArray.resize(numTorsions, vector<double>(6));
661
    for (int i = 0; i < numTorsions; ++i) {
Peter Eastman's avatar
Peter Eastman committed
662
        int particle1, particle2, particle3, particle4;
663
        double c0, c1, c2, c3, c4, c5;
Peter Eastman's avatar
Peter Eastman committed
664
665
666
667
668
        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
669
670
671
672
673
674
        torsionParamArray[i][0] = c0;
        torsionParamArray[i][1] = c1;
        torsionParamArray[i][2] = c2;
        torsionParamArray[i][3] = c3;
        torsionParamArray[i][4] = c4;
        torsionParamArray[i][5] = c5;
675
    }
676
    usePeriodic = force.usesPeriodicBoundaryConditions();
677
678
}

679
double ReferenceCalcRBTorsionForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
680
681
682
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
683
684
    ReferenceBondForce refBondForce;
    ReferenceRbDihedralBond rbTorsionBond;
685
686
    if (usePeriodic)
        rbTorsionBond.setPeriodic(extractBoxVectors(context));
687
    refBondForce.calculateForce(numTorsions, torsionIndexArray, posData, torsionParamArray, forceData, includeEnergy ? &energy : NULL, rbTorsionBond);
688
689
690
    return energy;
}

691
692
693
694
695
696
697
698
699
700
701
702
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
703
704
705
706
707
708
        torsionParamArray[i][0] = c0;
        torsionParamArray[i][1] = c1;
        torsionParamArray[i][2] = c2;
        torsionParamArray[i][3] = c3;
        torsionParamArray[i][4] = c4;
        torsionParamArray[i][5] = c5;
709
710
711
    }
}

712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
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]);
    }
736
    usePeriodic = force.usesPeriodicBoundaryConditions();
737
738
}

739
double ReferenceCalcCMAPTorsionForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
740
741
742
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double totalEnergy = 0;
743
    ReferenceCMAPTorsionIxn torsion(coeff, torsionMaps, torsionIndices);
744
745
    if (usePeriodic)
        torsion.setPeriodic(extractBoxVectors(context));
746
747
748
749
    torsion.calculateIxn(posData, forceData, &totalEnergy);
    return totalEnergy;
}

750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
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");
    }
}

784
785
786
787
788
ReferenceCalcCustomTorsionForceKernel::~ReferenceCalcCustomTorsionForceKernel() {
    if (ixn != NULL)
        delete ixn;
}

789
790
791
void ReferenceCalcCustomTorsionForceKernel::initialize(const System& system, const CustomTorsionForce& force) {
    numTorsions = force.getNumTorsions();
    int numParameters = force.getNumPerTorsionParameters();
792
    usePeriodic = force.usesPeriodicBoundaryConditions();
793
794
795

    // Build the arrays.

796
797
    torsionIndexArray.resize(numTorsions, vector<int>(4));
    torsionParamArray.resize(numTorsions, vector<double>(numParameters));
798
    vector<double> params;
799
    for (int i = 0; i < numTorsions; ++i) {
800
801
802
803
804
805
806
        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
807
            torsionParamArray[i][j] = params[j];
808
809
810
811
812
    }

    // Parse the expression used to calculate the force.

    Lepton::ParsedExpression expression = Lepton::Parser::parse(force.getEnergyFunction()).optimize();
813
814
    energyExpression = expression.createCompiledExpression();
    forceExpression = expression.differentiate("theta").createCompiledExpression();
815
816
817
818
    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));
819
820
821
822
823
    for (int i = 0; i < force.getNumEnergyParameterDerivatives(); i++) {
        string param = force.getEnergyParameterDerivativeName(i);
        energyParamDerivNames.push_back(param);
        energyParamDerivExpressions.push_back(expression.differentiate(param).createCompiledExpression());
    }
824
825
826
827
828
    set<string> variables;
    variables.insert("theta");
    variables.insert(parameterNames.begin(), parameterNames.end());
    variables.insert(globalParameterNames.begin(), globalParameterNames.end());
    validateVariables(expression.getRootNode(), variables);
829
    ixn = new ReferenceCustomTorsionIxn(energyExpression, forceExpression, parameterNames, energyParamDerivExpressions);
830
831
}

832
double ReferenceCalcCustomTorsionForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
833
834
835
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
836
    map<string, double> globalParameters;
peastman's avatar
peastman committed
837
838
    for (auto& name : globalParameterNames)
        globalParameters[name] = context.getParameter(name);
839
    ixn->setGlobalParameters(globalParameters);
840
    if (usePeriodic)
841
        ixn->setPeriodic(extractBoxVectors(context));
842
843
    vector<double> energyParamDerivValues(energyParamDerivNames.size()+1, 0.0);
    for (int i = 0; i < numTorsions; i++)
844
        ixn->calculateBondIxn(torsionIndexArray[i], posData, torsionParamArray[i], forceData, includeEnergy ? &energy : NULL, &energyParamDerivValues[0]);
845
846
847
    map<string, double>& energyParamDerivs = extractEnergyParameterDerivatives(context);
    for (int i = 0; i < energyParamDerivNames.size(); i++)
        energyParamDerivs[energyParamDerivNames[i]] += energyParamDerivValues[i];
848
849
850
    return energy;
}

851
void ReferenceCalcCustomTorsionForceKernel::copyParametersToContext(ContextImpl& context, const CustomTorsionForce& force, int firstTorsion, int lastTorsion) {
852
853
854
855
856
857
858
    if (numTorsions != force.getNumTorsions())
        throw OpenMMException("updateParametersInContext: The number of torsions has changed");

    // Record the values.

    int numParameters = force.getNumPerTorsionParameters();
    vector<double> params;
859
    for (int i = firstTorsion; i <= lastTorsion; ++i) {
860
861
862
863
864
        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
865
            torsionParamArray[i][j] = params[j];
866
867
868
    }
}

869
870
871
872
873
ReferenceCalcNonbondedForceKernel::~ReferenceCalcNonbondedForceKernel() {
    if (neighborList != NULL)
        delete neighborList;
}

874
875
876
877
void ReferenceCalcNonbondedForceKernel::initialize(const System& system, const NonbondedForce& force) {

    // Identify which exceptions are 1-4 interactions.

878
879
880
881
882
883
884
885
    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
886
    numParticles = force.getNumParticles();
887
888
889
890
891
892
893
894
    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);
895
896
        if (chargeProd != 0.0 || epsilon != 0.0 || exceptionsWithOffsets.find(i) != exceptionsWithOffsets.end()) {
            nb14Index[i] = nb14s.size();
897
            nb14s.push_back(i);
898
        }
899
900
901
902
903
    }

    // Build the arrays.

    num14 = nb14s.size();
904
905
906
    bonded14IndexArray.resize(num14, vector<int>(2));
    bonded14ParamArray.resize(num14, vector<double>(3));
    particleParamArray.resize(numParticles, vector<double>(3));
907
908
909
910
    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]);
911
    for (int i = 0; i < num14; ++i) {
Peter Eastman's avatar
Peter Eastman committed
912
        int particle1, particle2;
913
        force.getExceptionParameters(nb14s[i], particle1, particle2, baseExceptionParams[i][0], baseExceptionParams[i][1], baseExceptionParams[i][2]);
Peter Eastman's avatar
Peter Eastman committed
914
915
        bonded14IndexArray[i][0] = particle1;
        bonded14IndexArray[i][1] = particle2;
916
917
918
919
920
921
922
923
924
925
926
927
928
    }
    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);
929
        exceptionParamOffsets[make_pair(param, nb14Index[exception])] = {charge, sigma, epsilon};
930
    }
931
    nonbondedMethod = CalcNonbondedForceKernel::NonbondedMethod(force.getNonbondedMethod());
peastman's avatar
peastman committed
932
    nonbondedCutoff = force.getCutoffDistance();
933
    if (nonbondedMethod == NoCutoff) {
934
        neighborList = NULL;
935
936
937
        useSwitchingFunction = false;
    }
    else {
938
        neighborList = new NeighborList();
939
940
941
        useSwitchingFunction = force.getUseSwitchingFunction();
        switchingDistance = force.getSwitchingDistance();
    }
942
943
944
    if (nonbondedMethod == Ewald) {
        double alpha;
        NonbondedForceImpl::calcEwaldParameters(system, force, alpha, kmax[0], kmax[1], kmax[2]);
peastman's avatar
peastman committed
945
        ewaldAlpha = alpha;
946
947
948
    }
    else if (nonbondedMethod == PME) {
        double alpha;
949
        NonbondedForceImpl::calcPMEParameters(system, force, alpha, gridSize[0], gridSize[1], gridSize[2], false);
peastman's avatar
peastman committed
950
        ewaldAlpha = alpha;
951
    }
952
953
954
    else if (nonbondedMethod == LJPME) {
        double alpha;
        NonbondedForceImpl::calcPMEParameters(system, force, alpha, gridSize[0], gridSize[1], gridSize[2], false);
peastman's avatar
peastman committed
955
        ewaldAlpha = alpha;
956
        NonbondedForceImpl::calcPMEParameters(system, force, alpha, dispersionGridSize[0], dispersionGridSize[1], dispersionGridSize[2], true);
peastman's avatar
peastman committed
957
        ewaldDispersionAlpha = alpha;
958
959
        useSwitchingFunction = false;
    }
960
961
962
963
    if (nonbondedMethod == NoCutoff || nonbondedMethod == CutoffNonPeriodic)
        exceptionsArePeriodic = false;
    else
        exceptionsArePeriodic = force.getExceptionsUsePeriodicBoundaryConditions();
peastman's avatar
peastman committed
964
    rfDielectric = force.getReactionFieldDielectric();
965
966
967
968
    if (force.getUseDispersionCorrection())
        dispersionCoefficient = NonbondedForceImpl::calcDispersionCorrection(system, force);
    else
        dispersionCoefficient = 0.0;
969
970
}

971
double ReferenceCalcNonbondedForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy, bool includeDirect, bool includeReciprocal) {
972
    computeParameters(context);
peastman's avatar
peastman committed
973
974
975
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
976
    ReferenceLJCoulombIxn clj;
977
    bool periodic = (nonbondedMethod == CutoffPeriodic);
978
    bool ewald  = (nonbondedMethod == Ewald);
979
    bool pme  = (nonbondedMethod == PME);
980
    bool ljpme = (nonbondedMethod == LJPME);
981
    if (nonbondedMethod != NoCutoff) {
982
        computeNeighborListVoxelHash(*neighborList, numParticles, posData, exclusions, extractBoxVectors(context), periodic || ewald || pme || ljpme, nonbondedCutoff, 0.0);
983
        clj.setUseCutoff(nonbondedCutoff, *neighborList, rfDielectric);
984
    }
985
    if (periodic || ewald || pme || ljpme) {
peastman's avatar
peastman committed
986
        Vec3* boxVectors = extractBoxVectors(context);
987
        double minAllowedSize = 1.999999*nonbondedCutoff;
988
        if (boxVectors[0][0] < minAllowedSize || boxVectors[1][1] < minAllowedSize || boxVectors[2][2] < minAllowedSize)
989
            throw OpenMMException("The periodic box size has decreased to less than twice the nonbonded cutoff.");
990
        clj.setPeriodic(boxVectors);
991
        clj.setPeriodicExceptions(exceptionsArePeriodic);
992
    }
993
994
    if (ewald)
        clj.setUseEwald(ewaldAlpha, kmax[0], kmax[1], kmax[2]);
995
    if (pme)
996
        clj.setUsePME(ewaldAlpha, gridSize);
997
998
999
1000
    if (ljpme){
        clj.setUsePME(ewaldAlpha, gridSize);
        clj.setUseLJPME(ewaldDispersionAlpha, dispersionGridSize);
    }
1001
1002
    if (useSwitchingFunction)
        clj.setUseSwitchingFunction(switchingDistance);
1003
    clj.calculatePairIxn(numParticles, posData, particleParamArray, exclusions, forceData, includeEnergy ? &energy : NULL, includeDirect, includeReciprocal);
1004
1005
1006
    if (includeDirect) {
        ReferenceBondForce refBondForce;
        ReferenceLJCoulomb14 nonbonded14;
1007
1008
1009
1010
        if (exceptionsArePeriodic) {
            Vec3* boxVectors = extractBoxVectors(context);
            nonbonded14.setPeriodic(boxVectors);
        }
1011
        refBondForce.calculateForce(num14, bonded14IndexArray, posData, bonded14ParamArray, forceData, includeEnergy ? &energy : NULL, nonbonded14);
1012
        if (periodic || ewald || pme) {
peastman's avatar
peastman committed
1013
            Vec3* boxVectors = extractBoxVectors(context);
1014
            energy += dispersionCoefficient/(boxVectors[0][0]*boxVectors[1][1]*boxVectors[2][2]);
1015
        }
1016
    }
1017
1018
1019
    return energy;
}

1020
void ReferenceCalcNonbondedForceKernel::copyParametersToContext(ContextImpl& context, const NonbondedForce& force, int firstParticle, int lastParticle, int firstException, int lastException) {
1021
1022
    if (force.getNumParticles() != numParticles)
        throw OpenMMException("updateParametersInContext: The number of particles has changed");
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
    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
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044

    // 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);
1045
1046
        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
1047
    }
1048
1049
1050
1051
1052
    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);
1053
1054
1055
1056
1057
        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
1058
1059
1060
1061
1062
1063
1064
            nb14s.push_back(i);
    }
    if (nb14s.size() != num14)
        throw OpenMMException("updateParametersInContext: The number of non-excluded exceptions has changed");

    // Record the values.

1065
    for (int i = firstParticle; i <= lastParticle; ++i)
1066
        force.getParticleParameters(i, baseParticleParams[i][0], baseParticleParams[i][1], baseParticleParams[i][2]);
1067
1068
    for (int i = 0; i < num14; ++i) {
        int particle1, particle2;
1069
        force.getExceptionParameters(nb14s[i], particle1, particle2, baseExceptionParams[i][0], baseExceptionParams[i][1], baseExceptionParams[i][2]);
1070
1071
1072
        bonded14IndexArray[i][0] = particle1;
        bonded14IndexArray[i][1] = particle2;
    }
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
    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};
    }
1089
1090
1091
1092
1093
1094
1095
1096
    
    // 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);
}

1097
void ReferenceCalcNonbondedForceKernel::getPMEParameters(double& alpha, int& nx, int& ny, int& nz) const {
1098
1099
    if (nonbondedMethod != PME && nonbondedMethod != LJPME)
        throw OpenMMException("getPMEParametersInContext: This Context is not using PME or LJPME");
1100
1101
1102
1103
1104
1105
    alpha = ewaldAlpha;
    nx = gridSize[0];
    ny = gridSize[1];
    nz = gridSize[2];
}

1106
void ReferenceCalcNonbondedForceKernel::getLJPMEParameters(double& alpha, int& nx, int& ny, int& nz) const {
1107
1108
    if (nonbondedMethod != LJPME)
        throw OpenMMException("getPMEParametersInContext: This Context is not using LJPME");
1109
1110
1111
1112
    alpha = ewaldDispersionAlpha;
    nx = dispersionGridSize[0];
    ny = dispersionGridSize[1];
    nz = dispersionGridSize[2];
1113
1114
}

1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
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];
    }
}

1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
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
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);
}

1394
1395
1396
ReferenceCalcCustomNonbondedForceKernel::~ReferenceCalcCustomNonbondedForceKernel() {
    if (neighborList != NULL)
        delete neighborList;
1397
1398
    if (forceCopy != NULL)
        delete forceCopy;
1399
1400
1401
1402
}

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

1403
    // Record the exclusions.
1404
1405
1406

    numParticles = force.getNumParticles();
    exclusions.resize(numParticles);
1407
    for (int i = 0; i < force.getNumExclusions(); i++) {
1408
        int particle1, particle2;
1409
        force.getExclusionParticles(i, particle1, particle2);
1410
1411
1412
1413
1414
1415
        exclusions[particle1].insert(particle2);
        exclusions[particle2].insert(particle1);
    }

    // Build the arrays.

1416
1417
1418
    particleParamArray.resize(numParticles);
    for (int i = 0; i < numParticles; ++i)
        force.getParticleParameters(i, particleParamArray[i]);
1419
    nonbondedMethod = CalcCustomNonbondedForceKernel::NonbondedMethod(force.getNonbondedMethod());
peastman's avatar
peastman committed
1420
    nonbondedCutoff = force.getCutoffDistance();
1421
    if (nonbondedMethod == NoCutoff) {
1422
        neighborList = NULL;
1423
1424
1425
        useSwitchingFunction = false;
    }
    else {
1426
        neighborList = new NeighborList();
1427
1428
1429
        useSwitchingFunction = force.getUseSwitchingFunction();
        switchingDistance = force.getSwitchingDistance();
    }
1430

1431
    // Record the tabulated function update counts for future reference.
1432
1433

    for (int i = 0; i < force.getNumTabulatedFunctions(); i++)
1434
        tabulatedFunctionUpdateCount[force.getTabulatedFunctionName(i)] = force.getTabulatedFunction(i).getUpdateCount();
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460

    // 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) {
1461
1462
1463
    // Create custom functions for the tabulated functions.

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

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

1469
    Lepton::ParsedExpression expression = Lepton::Parser::parse(force.getEnergyFunction(), functions).optimize();
1470
1471
    energyExpression = expression.createCompiledExpression();
    forceExpression = expression.differentiate("r").createCompiledExpression();
1472
1473
1474
    parameterNames.clear();
    globalParameterNames.clear();
    globalParamValues.clear();
1475
1476
    computedValueNames.clear();
    computedValueExpressions.clear();
1477
1478
1479
    energyParamDerivNames.clear();
    energyParamDerivExpressions.clear();
    for (int i = 0; i < force.getNumPerParticleParameters(); i++)
1480
        parameterNames.push_back(force.getPerParticleParameterName(i));
1481
    for (int i = 0; i < force.getNumGlobalParameters(); i++) {
1482
        globalParameterNames.push_back(force.getGlobalParameterName(i));
1483
1484
        globalParamValues[force.getGlobalParameterName(i)] = force.getGlobalParameterDefaultValue(i);
    }
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
    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());
    }
1503
1504
1505
1506
1507
    for (int i = 0; i < force.getNumEnergyParameterDerivatives(); i++) {
        string param = force.getEnergyParameterDerivativeName(i);
        energyParamDerivNames.push_back(param);
        energyParamDerivExpressions.push_back(expression.differentiate(param).createCompiledExpression());
    }
1508
1509
1510
    for (auto& name : computedValueNames) {
        pairVariables.insert(name+"1");
        pairVariables.insert(name+"2");
1511
    }
1512
    validateVariables(expression.getRootNode(), pairVariables);
1513
1514
1515

    // Delete the custom functions.

peastman's avatar
peastman committed
1516
1517
    for (auto& function : functions)
        delete function.second;
1518
1519
}

1520
double ReferenceCalcCustomNonbondedForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
1521
1522
1523
1524
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    Vec3* boxVectors = extractBoxVectors(context);
    double energy = 0;
1525
    ReferenceCustomNonbondedIxn ixn(energyExpression, forceExpression, parameterNames, energyParamDerivExpressions, computedValueNames, computedValueExpressions);
1526
1527
    bool periodic = (nonbondedMethod == CutoffPeriodic);
    if (nonbondedMethod != NoCutoff) {
1528
        computeNeighborListVoxelHash(*neighborList, numParticles, posData, exclusions, extractBoxVectors(context), periodic, nonbondedCutoff, 0.0);
1529
1530
        ixn.setUseCutoff(nonbondedCutoff, *neighborList);
    }
1531
1532
    if (periodic) {
        double minAllowedSize = 2*nonbondedCutoff;
1533
        if (boxVectors[0][0] < minAllowedSize || boxVectors[1][1] < minAllowedSize || boxVectors[2][2] < minAllowedSize)
1534
            throw OpenMMException("The periodic box size has decreased to less than twice the nonbonded cutoff.");
1535
        ixn.setPeriodic(boxVectors);
1536
    }
1537
1538
    if (interactionGroups.size() > 0)
        ixn.setInteractionGroups(interactionGroups);
1539
    bool globalParamsChanged = false;
peastman's avatar
peastman committed
1540
1541
1542
    for (auto& name : globalParameterNames) {
        double value = context.getParameter(name);
        if (globalParamValues[name] != value)
1543
            globalParamsChanged = true;
peastman's avatar
peastman committed
1544
        globalParamValues[name] = value;
1545
    }
1546
1547
    if (useSwitchingFunction)
        ixn.setUseSwitchingFunction(switchingDistance);
1548
    vector<double> energyParamDerivValues(energyParamDerivNames.size()+1, 0.0);
1549
    ixn.calculatePairIxn(numParticles, posData, particleParamArray, exclusions, globalParamValues, forceData, includeEnergy ? &energy : NULL, &energyParamDerivValues[0]);
1550
1551
1552
    map<string, double>& energyParamDerivs = extractEnergyParameterDerivatives(context);
    for (int i = 0; i < energyParamDerivNames.size(); i++)
        energyParamDerivs[energyParamDerivNames[i]] += energyParamDerivValues[i];
1553
1554
1555
    
    // Add in the long range correction.
    
1556
    if (!hasInitializedLongRangeCorrection) {
1557
        ThreadPool& threads = extractThreadPool(context);
1558
        longRangeCorrectionData = CustomNonbondedForceImpl::prepareLongRangeCorrection(*forceCopy, threads.getNumThreads());
1559
        CustomNonbondedForceImpl::calcLongRangeCorrection(*forceCopy, longRangeCorrectionData, context.getOwner(), longRangeCoefficient, longRangeCoefficientDerivs, threads);
1560
1561
        hasInitializedLongRangeCorrection = true;
    }
1562
    else if (globalParamsChanged && forceCopy != NULL) {
1563
        ThreadPool& threads = extractThreadPool(context);
1564
1565
        CustomNonbondedForceImpl::calcLongRangeCorrection(*forceCopy, longRangeCorrectionData, context.getOwner(), longRangeCoefficient, longRangeCoefficientDerivs, threads);
    }
1566
1567
1568
1569
    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;
1570
1571
1572
    return energy;
}

1573
void ReferenceCalcCustomNonbondedForceKernel::copyParametersToContext(ContextImpl& context, const CustomNonbondedForce& force, int firstParticle, int lastParticle) {
1574
1575
1576
1577
1578
1579
    if (numParticles != force.getNumParticles())
        throw OpenMMException("updateParametersInContext: The number of particles has changed");

    // Record the values.

    int numParameters = force.getNumPerParticleParameters();
1580
1581
    vector<double> parameters;
    for (int i = firstParticle; i <= lastParticle; ++i) {
1582
1583
        force.getParticleParameters(i, parameters);
        for (int j = 0; j < numParameters; j++)
peastman's avatar
peastman committed
1584
            particleParamArray[i][j] = parameters[j];
1585
    }
1586
1587
1588
1589
    
    // If necessary, recompute the long range correction.
    
    if (forceCopy != NULL) {
1590
        ThreadPool& threads = extractThreadPool(context);
1591
        longRangeCorrectionData = CustomNonbondedForceImpl::prepareLongRangeCorrection(force, threads.getNumThreads());
1592
        CustomNonbondedForceImpl::calcLongRangeCorrection(force, longRangeCorrectionData, context.getOwner(), longRangeCoefficient, longRangeCoefficientDerivs, threads);
1593
1594
1595
        hasInitializedLongRangeCorrection = true;
        *forceCopy = force;
    }
1596
1597
1598
1599
1600
1601

    // See if any tabulated functions have changed.

    bool changed = false;
    for (int i = 0; i < force.getNumTabulatedFunctions(); i++) {
        string name = force.getTabulatedFunctionName(i);
1602
1603
        if (force.getTabulatedFunction(i).getUpdateCount() != tabulatedFunctionUpdateCount[name]) {
            tabulatedFunctionUpdateCount[name] = force.getTabulatedFunction(i).getUpdateCount();
1604
1605
1606
1607
1608
            changed = true;
        }
    }
    if (changed)
        createExpressions(force);
1609
1610
}

1611
ReferenceCalcGBSAOBCForceKernel::~ReferenceCalcGBSAOBCForceKernel() {
1612
    if (obc) {
Peter Eastman's avatar
Peter Eastman committed
1613
        delete obc->getObcParameters();
1614
1615
1616
1617
        delete obc;
    }
}

1618
void ReferenceCalcGBSAOBCForceKernel::initialize(const System& system, const GBSAOBCForce& force) {
Peter Eastman's avatar
Peter Eastman committed
1619
1620
    int numParticles = system.getNumParticles();
    charges.resize(numParticles);
peastman's avatar
peastman committed
1621
1622
    vector<double> atomicRadii(numParticles);
    vector<double> scaleFactors(numParticles);
Peter Eastman's avatar
Peter Eastman committed
1623
    for (int i = 0; i < numParticles; ++i) {
1624
        double charge, radius, scalingFactor;
Peter Eastman's avatar
Peter Eastman committed
1625
        force.getParticleParameters(i, charge, radius, scalingFactor);
peastman's avatar
peastman committed
1626
1627
1628
        charges[i] = charge;
        atomicRadii[i] = radius;
        scaleFactors[i] = scalingFactor;
1629
    }
1630
    ObcParameters* obcParameters = new ObcParameters(numParticles, ObcParameters::ObcTypeII);
1631
    obcParameters->setAtomicRadii(atomicRadii);
1632
    obcParameters->setScaledRadiusFactors(scaleFactors);
peastman's avatar
peastman committed
1633
1634
    obcParameters->setSolventDielectric(force.getSolventDielectric());
    obcParameters->setSoluteDielectric(force.getSoluteDielectric());
1635
    obcParameters->setPi4Asolv(4*M_PI*force.getSurfaceAreaEnergy());
1636
    if (force.getNonbondedMethod() != GBSAOBCForce::NoCutoff)
peastman's avatar
peastman committed
1637
        obcParameters->setUseCutoff(force.getCutoffDistance());
1638
    isPeriodic = (force.getNonbondedMethod() == GBSAOBCForce::CutoffPeriodic);
1639
    obc = new ReferenceObc(obcParameters);
1640
    obc->setIncludeAceApproximation(true);
1641
1642
}

1643
double ReferenceCalcGBSAOBCForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
1644
1645
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
1646
    if (isPeriodic)
1647
        obc->getObcParameters()->setPeriodic(extractBoxVectors(context));
Mark Friedrichs's avatar
Mark Friedrichs committed
1648
    return obc->computeBornEnergyForces(posData, charges, forceData);
1649
1650
}

1651
1652
1653
1654
1655
1656
1657
1658
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
1659
1660
    vector<double> atomicRadii(numParticles);
    vector<double> scaleFactors(numParticles);
1661
1662
1663
    for (int i = 0; i < numParticles; ++i) {
        double charge, radius, scalingFactor;
        force.getParticleParameters(i, charge, radius, scalingFactor);
peastman's avatar
peastman committed
1664
1665
1666
        charges[i] = charge;
        atomicRadii[i] = radius;
        scaleFactors[i] = scalingFactor;
1667
1668
1669
1670
1671
    }
    obcParameters->setAtomicRadii(atomicRadii);
    obcParameters->setScaledRadiusFactors(scaleFactors);
}

1672
1673
1674
1675
1676
1677
ReferenceCalcCustomGBForceKernel::~ReferenceCalcCustomGBForceKernel() {
    if (neighborList != NULL)
        delete neighborList;
}

void ReferenceCalcCustomGBForceKernel::initialize(const System& system, const CustomGBForce& force) {
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
    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.");
        }
    }
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703

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

1704
1705
1706
    particleParamArray.resize(numParticles);
    for (int i = 0; i < numParticles; ++i)
        force.getParticleParameters(i, particleParamArray[i]);
1707
    for (int i = 0; i < force.getNumPerParticleParameters(); i++)
1708
1709
1710
1711
        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
1712
    nonbondedCutoff = force.getCutoffDistance();
1713
1714
1715
1716
1717
    if (nonbondedMethod == NoCutoff)
        neighborList = NULL;
    else
        neighborList = new NeighborList();

1718
    // Record the tabulated function update counts for future reference.
1719
1720

    for (int i = 0; i < force.getNumTabulatedFunctions(); i++)
1721
        tabulatedFunctionUpdateCount[force.getTabulatedFunctionName(i)] = force.getTabulatedFunction(i).getUpdateCount();
1722
1723
1724
1725
1726
1727
1728

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

void ReferenceCalcCustomGBForceKernel::createExpressions(const CustomGBForce& force) {
1729
1730
1731
    // Create custom functions for the tabulated functions.

    map<string, Lepton::CustomFunction*> functions;
1732
    for (int i = 0; i < force.getNumTabulatedFunctions(); i++)
1733
        functions[force.getTabulatedFunctionName(i)] = createReferenceTabulatedFunction(force.getTabulatedFunction(i));
1734
1735
1736

    // Parse the expressions for computed values.

1737
1738
1739
1740
1741
1742
1743
    valueExpressions.clear();
    valueTypes.clear();
    valueNames.clear();
    energyParamDerivNames.clear();
    valueDerivExpressions.clear();
    valueGradientExpressions.clear();
    valueParamDerivExpressions.clear();
1744
    valueDerivExpressions.resize(force.getNumComputedValues());
1745
    valueGradientExpressions.resize(force.getNumComputedValues());
1746
    valueParamDerivExpressions.resize(force.getNumComputedValues());
1747
1748
1749
1750
1751
    set<string> particleVariables, pairVariables;
    pairVariables.insert("r");
    particleVariables.insert("x");
    particleVariables.insert("y");
    particleVariables.insert("z");
1752
    for (int i = 0; i < force.getNumPerParticleParameters(); i++) {
1753
1754
1755
1756
1757
1758
        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());
1759
1760
1761
1762
1763
    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();
1764
        valueExpressions.push_back(ex.createCompiledExpression());
1765
1766
        valueTypes.push_back(type);
        valueNames.push_back(name);
1767
        if (i == 0) {
1768
            valueDerivExpressions[i].push_back(ex.differentiate("r").createCompiledExpression());
1769
1770
            validateVariables(ex.getRootNode(), pairVariables);
        }
1771
        else {
1772
1773
1774
            valueGradientExpressions[i].push_back(ex.differentiate("x").createCompiledExpression());
            valueGradientExpressions[i].push_back(ex.differentiate("y").createCompiledExpression());
            valueGradientExpressions[i].push_back(ex.differentiate("z").createCompiledExpression());
1775
            for (int j = 0; j < i; j++)
1776
                valueDerivExpressions[i].push_back(ex.differentiate(valueNames[j]).createCompiledExpression());
1777
            validateVariables(ex.getRootNode(), particleVariables);
1778
        }
1779
1780
1781
1782
1783
        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());
        }
1784
1785
1786
        particleVariables.insert(name);
        pairVariables.insert(name+"1");
        pairVariables.insert(name+"2");
1787
1788
    }

1789
    // Parse the expressions for energy terms.
1790

1791
1792
1793
1794
1795
    energyExpressions.clear();
    energyTypes.clear();
    energyDerivExpressions.clear();
    energyGradientExpressions.clear();
    energyParamDerivExpressions.clear();
1796
    energyDerivExpressions.resize(force.getNumEnergyTerms());
1797
    energyGradientExpressions.resize(force.getNumEnergyTerms());
1798
    energyParamDerivExpressions.resize(force.getNumEnergyTerms());
1799
1800
1801
1802
1803
    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();
1804
        energyExpressions.push_back(ex.createCompiledExpression());
1805
1806
        energyTypes.push_back(type);
        if (type != CustomGBForce::SingleParticle)
1807
            energyDerivExpressions[i].push_back(ex.differentiate("r").createCompiledExpression());
1808
        for (int j = 0; j < force.getNumComputedValues(); j++) {
1809
            if (type == CustomGBForce::SingleParticle) {
1810
1811
1812
1813
                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());
1814
                validateVariables(ex.getRootNode(), particleVariables);
1815
            }
1816
            else {
1817
1818
                energyDerivExpressions[i].push_back(ex.differentiate(valueNames[j]+"1").createCompiledExpression());
                energyDerivExpressions[i].push_back(ex.differentiate(valueNames[j]+"2").createCompiledExpression());
1819
                validateVariables(ex.getRootNode(), pairVariables);
1820
1821
            }
        }
1822
1823
        for (int j = 0; j < force.getNumEnergyParameterDerivatives(); j++)
            energyParamDerivExpressions[i].push_back(ex.differentiate(force.getEnergyParameterDerivativeName(j)).createCompiledExpression());
1824
1825
1826
1827
    }

    // Delete the custom functions.

peastman's avatar
peastman committed
1828
1829
    for (auto& function : functions)
        delete function.second;
1830
1831
}

1832
double ReferenceCalcCustomGBForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
1833
1834
1835
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
1836
1837
    ReferenceCustomGBIxn ixn(valueExpressions, valueDerivExpressions, valueGradientExpressions, valueParamDerivExpressions, valueNames, valueTypes,
        energyExpressions, energyDerivExpressions, energyGradientExpressions, energyParamDerivExpressions, energyTypes, particleParameterNames);
1838
    bool periodic = (nonbondedMethod == CutoffPeriodic);
1839
    if (periodic)
1840
        ixn.setPeriodic(extractBoxVectors(context));
1841
    if (nonbondedMethod != NoCutoff) {
Peter Eastman's avatar
Peter Eastman committed
1842
1843
        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);
1844
1845
1846
        ixn.setUseCutoff(nonbondedCutoff, *neighborList);
    }
    map<string, double> globalParameters;
peastman's avatar
peastman committed
1847
1848
    for (auto& name : globalParameterNames)
        globalParameters[name] = context.getParameter(name);
1849
1850
1851
1852
1853
    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];
1854
1855
1856
    return energy;
}

1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
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
1869
            particleParamArray[i][j] = parameters[j];
1870
    }
1871
1872
1873
1874
1875
1876

    // See if any tabulated functions have changed.

    bool changed = false;
    for (int i = 0; i < force.getNumTabulatedFunctions(); i++) {
        string name = force.getTabulatedFunctionName(i);
1877
1878
        if (force.getTabulatedFunction(i).getUpdateCount() != tabulatedFunctionUpdateCount[name]) {
            tabulatedFunctionUpdateCount[name] = force.getTabulatedFunction(i).getUpdateCount();
1879
1880
1881
1882
1883
            changed = true;
        }
    }
    if (changed)
        createExpressions(force);
1884
1885
}

1886
1887
1888
1889
1890
ReferenceCalcCustomExternalForceKernel::~ReferenceCalcCustomExternalForceKernel() {
    if (ixn != NULL)
        delete ixn;
}

1891
1892
1893
1894
1895
1896
1897
void ReferenceCalcCustomExternalForceKernel::initialize(const System& system, const CustomExternalForce& force) {
    numParticles = force.getNumParticles();
    int numParameters = force.getNumPerParticleParameters();

    // Build the arrays.

    particles.resize(numParticles);
1898
1899
1900
    particleParamArray.resize(numParticles);
    for (int i = 0; i < numParticles; ++i)
        force.getParticleParameters(i, particles[i], particleParamArray[i]);
1901
1902
1903

    // Parse the expression used to calculate the force.

1904
    map<string, Lepton::CustomFunction*> functions;
1905
    ReferencePointDistanceFunction periodicDistance(true, &boxVectors);
1906
1907
    functions["periodicdistance"] = &periodicDistance;
    Lepton::ParsedExpression expression = Lepton::Parser::parse(force.getEnergyFunction(), functions).optimize();
1908
1909
1910
1911
    energyExpression = expression.createCompiledExpression();
    forceExpressionX = expression.differentiate("x").createCompiledExpression();
    forceExpressionY = expression.differentiate("y").createCompiledExpression();
    forceExpressionZ = expression.differentiate("z").createCompiledExpression();
1912
1913
1914
1915
    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));
1916
1917
1918
1919
1920
1921
1922
    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);
1923
1924
    ixn = new ReferenceCustomExternalIxn(energyExpression, forceExpressionX, forceExpressionY, forceExpressionZ, parameterNames);

1925
1926
}

1927
double ReferenceCalcCustomExternalForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
1928
1929
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
1930
    boxVectors = extractBoxVectors(context);
peastman's avatar
peastman committed
1931
    double energy = 0;
1932
    map<string, double> globalParameters;
peastman's avatar
peastman committed
1933
1934
    for (auto& name : globalParameterNames)
        globalParameters[name] = context.getParameter(name);
1935
    ixn->setGlobalParameters(globalParameters);
1936
    for (int i = 0; i < numParticles; ++i)
1937
        ixn->calculateForce(particles[i], posData, particleParamArray[i], forceData, includeEnergy ? &energy : NULL);
1938
1939
1940
    return energy;
}

1941
void ReferenceCalcCustomExternalForceKernel::copyParametersToContext(ContextImpl& context, const CustomExternalForce& force, int firstParticle, int lastParticle) {
1942
1943
1944
1945
1946
1947
1948
    if (numParticles != force.getNumParticles())
        throw OpenMMException("updateParametersInContext: The number of particles has changed");

    // Record the values.

    int numParameters = force.getNumPerParticleParameters();
    vector<double> params;
1949
    for (int i = firstParticle; i <= lastParticle; ++i) {
1950
        int particle;
1951
        force.getParticleParameters(i, particle, params);
1952
1953
1954
        if (particle != particles[i])
            throw OpenMMException("updateParametersInContext: A particle index has changed");
        for (int j = 0; j < numParameters; j++)
1955
            particleParamArray[i][j] = params[j];
1956
1957
1958
    }
}

1959
ReferenceCalcCustomHbondForceKernel::~ReferenceCalcCustomHbondForceKernel() {
1960
1961
    if (ixn != NULL)
        delete ixn;
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
}

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.

1980
    donorParticles.resize(numDonors);
1981
    donorParamArray.resize(numDonors);
1982
    for (int i = 0; i < numDonors; ++i) {
1983
        int d1, d2, d3;
1984
        force.getDonorParameters(i, d1, d2, d3, donorParamArray[i]);
1985
1986
1987
        donorParticles[i].push_back(d1);
        donorParticles[i].push_back(d2);
        donorParticles[i].push_back(d3);
1988
    }
1989
    acceptorParticles.resize(numAcceptors);
1990
    acceptorParamArray.resize(numAcceptors);
1991
    for (int i = 0; i < numAcceptors; ++i) {
1992
        int a1, a2, a3;
1993
        force.getAcceptorParameters(i, a1, a2, a3, acceptorParamArray[i]);
1994
1995
1996
        acceptorParticles[i].push_back(a1);
        acceptorParticles[i].push_back(a2);
        acceptorParticles[i].push_back(a3);
1997
    }
1998
1999
    for (int i = 0; i < force.getNumGlobalParameters(); i++)
        globalParameterNames.push_back(force.getGlobalParameterName(i));
peastman's avatar
peastman committed
2000
    nonbondedCutoff = force.getCutoffDistance();
2001

2002
    // Record the tabulated function update counts for future reference.
2003
2004

    for (int i = 0; i < force.getNumTabulatedFunctions(); i++)
2005
        tabulatedFunctionUpdateCount[force.getTabulatedFunctionName(i)] = force.getTabulatedFunction(i).getUpdateCount();
2006
2007
2008
2009
2010
2011
2012

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

void ReferenceCalcCustomHbondForceKernel::createInteraction(const CustomHbondForce& force) {
2013
2014
2015
    // Create custom functions for the tabulated functions.

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

2019
    // Parse the expression and create the object used to calculate the interaction.
2020

2021
2022
2023
    map<string, vector<int> > distances;
    map<string, vector<int> > angles;
    map<string, vector<int> > dihedrals;
2024
    Lepton::ParsedExpression energyExpression = CustomHbondForceImpl::prepareExpression(force, functions, distances, angles, dihedrals);
2025
2026
    vector<string> donorParameterNames;
    vector<string> acceptorParameterNames;
2027
    for (int i = 0; i < force.getNumPerDonorParameters(); i++)
2028
        donorParameterNames.push_back(force.getPerDonorParameterName(i));
2029
    for (int i = 0; i < force.getNumPerAcceptorParameters(); i++)
2030
        acceptorParameterNames.push_back(force.getPerAcceptorParameterName(i));
2031
    ixn = new ReferenceCustomHbondIxn(donorParticles, acceptorParticles, energyExpression, donorParameterNames, acceptorParameterNames, distances, angles, dihedrals);
2032
    NonbondedMethod nonbondedMethod = CalcCustomHbondForceKernel::NonbondedMethod(force.getNonbondedMethod());
2033
    isPeriodic = (nonbondedMethod == CutoffPeriodic);
2034
2035
    if (nonbondedMethod != NoCutoff)
        ixn->setUseCutoff(nonbondedCutoff);
2036
2037
2038

    // Delete the custom functions.

peastman's avatar
peastman committed
2039
2040
    for (auto& function : functions)
        delete function.second;
2041
2042
}

2043
double ReferenceCalcCustomHbondForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
2044
2045
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
2046
    if (isPeriodic)
2047
        ixn->setPeriodic(extractBoxVectors(context));
peastman's avatar
peastman committed
2048
    double energy = 0;
2049
    map<string, double> globalParameters;
peastman's avatar
peastman committed
2050
2051
    for (auto& name : globalParameterNames)
        globalParameters[name] = context.getParameter(name);
2052
    ixn->calculatePairIxn(posData, donorParamArray, acceptorParamArray, exclusions, globalParameters, forceData, includeEnergy ? &energy : NULL);
2053
2054
2055
    return energy;
}

2056
2057
2058
2059
2060
2061
2062
2063
2064
2065
2066
2067
2068
2069
2070
2071
2072
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
2073
            donorParamArray[i][j] = parameters[j];
2074
2075
2076
2077
2078
2079
2080
2081
2082
    }
    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
2083
            acceptorParamArray[i][j] = parameters[j];
2084
    }
2085
2086
2087
2088
2089
2090

    // See if any tabulated functions have changed.

    bool changed = false;
    for (int i = 0; i < force.getNumTabulatedFunctions(); i++) {
        string name = force.getTabulatedFunctionName(i);
2091
2092
        if (force.getTabulatedFunction(i).getUpdateCount() != tabulatedFunctionUpdateCount[name]) {
            tabulatedFunctionUpdateCount[name] = force.getTabulatedFunction(i).getUpdateCount();
2093
2094
2095
2096
2097
2098
2099
2100
            changed = true;
        }
    }
    if (changed) {
        delete ixn;
        ixn = NULL;
        createInteraction(force);
    }
2101
2102
}

2103
2104
2105
2106
2107
2108
ReferenceCalcCustomCentroidBondForceKernel::~ReferenceCalcCustomCentroidBondForceKernel() {
    if (ixn != NULL)
        delete ixn;
}

void ReferenceCalcCustomCentroidBondForceKernel::initialize(const System& system, const CustomCentroidBondForce& force) {
2109
    usePeriodic = force.usesPeriodicBoundaryConditions();
2110
2111
2112
2113

    // Build the arrays.

    int numGroups = force.getNumGroups();
2114
    groupAtoms.resize(numGroups);
2115
2116
2117
2118
2119
    vector<double> ignored;
    for (int i = 0; i < numGroups; i++)
        force.getGroupParameters(i, groupAtoms[i], ignored);
    CustomCentroidBondForceImpl::computeNormalizedWeights(force, system, normalizedWeights);
    numBonds = force.getNumBonds();
2120
    bondGroups.resize(numBonds);
2121
2122
2123
    bondParamArray.resize(numBonds);
    for (int i = 0; i < numBonds; ++i)
        force.getBondParameters(i, bondGroups[i], bondParamArray[i]);
2124

2125
    // Record the tabulated function update counts for future reference.
2126
2127

    for (int i = 0; i < force.getNumTabulatedFunctions(); i++)
2128
        tabulatedFunctionUpdateCount[force.getTabulatedFunctionName(i)] = force.getTabulatedFunction(i).getUpdateCount();
2129
2130
2131
2132
2133
2134
2135

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

void ReferenceCalcCustomCentroidBondForceKernel::createInteraction(const CustomCentroidBondForce& force) {
2136
2137
2138
    // Create custom functions for the tabulated functions.

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

2142
2143
2144
2145
2146
2147
    // Create implementations of point functions.

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

2148
2149
    // Parse the expression and create the object used to calculate the interaction.

2150
    int numGroups = force.getNumGroups();
2151
    Lepton::ParsedExpression energyExpression = CustomCentroidBondForceImpl::prepareExpression(force, functions);
2152
    vector<string> bondParameterNames;
2153
    for (int i = 0; i < force.getNumPerBondParameters(); i++)
2154
2155
2156
        bondParameterNames.push_back(force.getPerBondParameterName(i));
    for (int i = 0; i < force.getNumGlobalParameters(); i++)
        globalParameterNames.push_back(force.getGlobalParameterName(i));
2157
2158
2159
2160
2161
2162
    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());
    }
2163
    ixn = new ReferenceCustomCentroidBondIxn(force.getNumGroupsPerBond(), groupAtoms, normalizedWeights, bondGroups, energyExpression, bondParameterNames, energyParamDerivExpressions);
2164
2165
2166

    // Delete the custom functions.

peastman's avatar
peastman committed
2167
2168
    for (auto& function : functions)
        delete function.second;
2169
2170
2171
}

double ReferenceCalcCustomCentroidBondForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
2172
2173
2174
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
2175
    map<string, double> globalParameters;
peastman's avatar
peastman committed
2176
2177
    for (auto& name : globalParameterNames)
        globalParameters[name] = context.getParameter(name);
2178
2179
2180
2181
    if (usePeriodic) {
        boxVectors = extractBoxVectors(context);
        ixn->setPeriodic(boxVectors);
    }
2182
2183
2184
2185
2186
    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];
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
    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
2206
            bondParamArray[i][j] = params[j];
2207
    }
2208
2209
2210
2211
2212
2213

    // See if any tabulated functions have changed.

    bool changed = false;
    for (int i = 0; i < force.getNumTabulatedFunctions(); i++) {
        string name = force.getTabulatedFunctionName(i);
2214
2215
        if (force.getTabulatedFunction(i).getUpdateCount() != tabulatedFunctionUpdateCount[name]) {
            tabulatedFunctionUpdateCount[name] = force.getTabulatedFunction(i).getUpdateCount();
2216
2217
2218
2219
2220
2221
2222
2223
            changed = true;
        }
    }
    if (changed) {
        delete ixn;
        ixn = NULL;
        createInteraction(force);
    }
2224
2225
}

2226
2227
2228
2229
2230
2231
ReferenceCalcCustomCompoundBondForceKernel::~ReferenceCalcCustomCompoundBondForceKernel() {
    if (ixn != NULL)
        delete ixn;
}

void ReferenceCalcCustomCompoundBondForceKernel::initialize(const System& system, const CustomCompoundBondForce& force) {
2232
    usePeriodic = force.usesPeriodicBoundaryConditions();
2233
2234
2235
2236

    // Build the arrays.

    numBonds = force.getNumBonds();
2237
    bondParticles.resize(numBonds);
2238
2239
2240
    bondParamArray.resize(numBonds);
    for (int i = 0; i < numBonds; ++i)
        force.getBondParameters(i, bondParticles[i], bondParamArray[i]);
2241

2242
    // Record the tabulated function update counts for future reference.
2243
2244

    for (int i = 0; i < force.getNumTabulatedFunctions(); i++)
2245
        tabulatedFunctionUpdateCount[force.getTabulatedFunctionName(i)] = force.getTabulatedFunction(i).getUpdateCount();
2246
2247
2248
2249
2250
2251
2252

    // Create the interaction.

    createInteraction(force);
}

void ReferenceCalcCustomCompoundBondForceKernel::createInteraction(const CustomCompoundBondForce& force) {
2253
2254
2255
    // Create custom functions for the tabulated functions.

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

2259
2260
2261
2262
2263
2264
    // Create implementations of point functions.

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

2265
2266
    // Parse the expression and create the object used to calculate the interaction.

2267
    Lepton::ParsedExpression energyExpression = CustomCompoundBondForceImpl::prepareExpression(force, functions);
2268
    vector<string> bondParameterNames;
2269
    for (int i = 0; i < force.getNumPerBondParameters(); i++)
2270
2271
2272
        bondParameterNames.push_back(force.getPerBondParameterName(i));
    for (int i = 0; i < force.getNumGlobalParameters(); i++)
        globalParameterNames.push_back(force.getGlobalParameterName(i));
2273
2274
2275
2276
2277
2278
    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());
    }
2279
    ixn = new ReferenceCustomCompoundBondIxn(force.getNumParticlesPerBond(), bondParticles, energyExpression, bondParameterNames, energyParamDerivExpressions);
2280
2281
2282

    // Delete the custom functions.

peastman's avatar
peastman committed
2283
2284
    for (auto& function : functions)
        delete function.second;
2285
2286
2287
}

double ReferenceCalcCustomCompoundBondForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
2288
2289
2290
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
2291
    map<string, double> globalParameters;
peastman's avatar
peastman committed
2292
2293
    for (auto& name : globalParameterNames)
        globalParameters[name] = context.getParameter(name);
2294
2295
2296
2297
    if (usePeriodic) {
        boxVectors = extractBoxVectors(context);
        ixn->setPeriodic(boxVectors);
    }
2298
2299
2300
2301
2302
    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];
2303
2304
2305
    return energy;
}

2306
2307
2308
2309
2310
2311
2312
2313
2314
2315
2316
2317
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);
2318
        for (int j = 0; j < particles.size(); j++)
2319
2320
2321
            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
2322
            bondParamArray[i][j] = params[j];
2323
    }
2324
2325
2326
2327
2328
2329

    // See if any tabulated functions have changed.

    bool changed = false;
    for (int i = 0; i < force.getNumTabulatedFunctions(); i++) {
        string name = force.getTabulatedFunctionName(i);
2330
2331
        if (force.getTabulatedFunction(i).getUpdateCount() != tabulatedFunctionUpdateCount[name]) {
            tabulatedFunctionUpdateCount[name] = force.getTabulatedFunction(i).getUpdateCount();
2332
2333
2334
2335
2336
2337
2338
2339
            changed = true;
        }
    }
    if (changed) {
        delete ixn;
        ixn = NULL;
        createInteraction(force);
    }
2340
2341
}

2342
2343
2344
2345
2346
2347
2348
2349
2350
ReferenceCalcCustomManyParticleForceKernel::~ReferenceCalcCustomManyParticleForceKernel() {
    if (ixn != NULL)
        delete ixn;
}

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

    numParticles = system.getNumParticles();
2351
    particleParamArray.resize(numParticles);
2352
2353
    for (int i = 0; i < numParticles; ++i) {
        int type;
2354
        force.getParticleParameters(i, particleParamArray[i], type);
2355
2356
2357
    }
    for (int i = 0; i < force.getNumGlobalParameters(); i++)
        globalParameterNames.push_back(force.getGlobalParameterName(i));
2358

2359
    // Record the tabulated function update counts for future reference.
2360
2361

    for (int i = 0; i < force.getNumTabulatedFunctions(); i++)
2362
        tabulatedFunctionUpdateCount[force.getTabulatedFunctionName(i)] = force.getTabulatedFunction(i).getUpdateCount();
2363
2364
2365

    // Create the interaction.

2366
    ixn = new ReferenceCustomManyParticleIxn(force);
2367
2368
2369
2370
2371
    nonbondedMethod = CalcCustomManyParticleForceKernel::NonbondedMethod(force.getNonbondedMethod());
    cutoffDistance = force.getCutoffDistance();
}

double ReferenceCalcCustomManyParticleForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
peastman's avatar
peastman committed
2372
2373
2374
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    double energy = 0;
2375
    map<string, double> globalParameters;
peastman's avatar
peastman committed
2376
2377
    for (auto& name : globalParameterNames)
        globalParameters[name] = context.getParameter(name);
2378
    if (nonbondedMethod == CutoffPeriodic) {
peastman's avatar
peastman committed
2379
        Vec3* boxVectors = extractBoxVectors(context);
2380
        double minAllowedSize = 2*cutoffDistance;
2381
        if (boxVectors[0][0] < minAllowedSize || boxVectors[1][1] < minAllowedSize || boxVectors[2][2] < minAllowedSize)
2382
            throw OpenMMException("The periodic box size has decreased to less than twice the nonbonded cutoff.");
2383
        ixn->setPeriodic(boxVectors);
2384
2385
2386
2387
2388
2389
2390
2391
2392
2393
2394
2395
2396
2397
2398
2399
2400
2401
    }
    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
2402
            particleParamArray[i][j] = parameters[j];
2403
    }
2404
2405
2406
2407
2408
2409

    // See if any tabulated functions have changed.

    bool changed = false;
    for (int i = 0; i < force.getNumTabulatedFunctions(); i++) {
        string name = force.getTabulatedFunctionName(i);
2410
2411
        if (force.getTabulatedFunction(i).getUpdateCount() != tabulatedFunctionUpdateCount[name]) {
            tabulatedFunctionUpdateCount[name] = force.getTabulatedFunction(i).getUpdateCount();
2412
2413
2414
2415
2416
2417
2418
2419
            changed = true;
        }
    }
    if (changed) {
        delete ixn;
        ixn = NULL;
        ixn = new ReferenceCustomManyParticleIxn(force);
    }
2420
2421
}

2422
2423
2424
2425
2426
2427
2428
2429
2430
2431
2432
2433
2434
2435
2436
2437
2438
2439
2440
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);
}

2441
2442
2443
2444
2445
ReferenceCalcCustomCVForceKernel::~ReferenceCalcCustomCVForceKernel() {
    if (ixn != NULL)
        delete ixn;
}

2446
void ReferenceCalcCustomCVForceKernel::initialize(const System& system, const CustomCVForce& force, ContextImpl& innerContext) {
2447
2448
2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
2461
2462
2463
2464
2465
2466
2467
2468
2469
    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);
2470
2471
2472
2473
    Vec3 a, b, c;
    context.getPeriodicBoxVectors(a, b, c);
    innerContext.setPeriodicBoxVectors(a, b, c);
    innerContext.setTime(context.getTime());
2474
2475
2476
2477
2478
    map<string, double> innerParameters = innerContext.getParameters();
    for (auto& param : innerParameters)
        innerContext.setParameter(param.first, context.getParameter(param.first));
}

2479
void ReferenceCalcCustomCVForceKernel::copyParametersToContext(ContextImpl& context, const CustomCVForce& force) {
2480
    ixn->updateTabulatedFunctions(force);
2481
2482
}

2483
2484
void ReferenceCalcRMSDForceKernel::initialize(const System& system, const RMSDForce& force) {
    particles = force.getParticles();
peastman's avatar
peastman committed
2485
2486
2487
    if (particles.size() == 0)
        for (int i = 0; i < system.getNumParticles(); i++)
            particles.push_back(i);
2488
2489
2490
2491
2492
2493
2494
2495
2496
2497
2498
2499
2500
2501
2502
2503
2504
2505
2506
2507
    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
2508
2509
2510
    if (particles.size() == 0)
        for (int i = 0; i < referencePos.size(); i++)
            particles.push_back(i);
2511
2512
2513
2514
2515
2516
2517
2518
2519
    referencePos = force.getReferencePositions();
    Vec3 center;
    for (int i : particles)
        center += referencePos[i];
    center /= particles.size();
    for (Vec3& p : referencePos)
        p -= center;
}

2520
2521
2522
2523
2524
2525
2526
2527
2528
2529
2530
2531
2532
2533
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);
}

2534
2535
2536
2537
2538
2539
2540
2541
2542
2543
2544
2545
2546
2547
2548
2549
2550
2551
2552
2553
2554
2555
2556
2557
2558
2559
2560
2561
2562
2563
2564
2565
2566
2567
2568
2569
2570
2571
2572
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;
}

2573
2574
2575
2576
2577
ReferenceIntegrateVerletStepKernel::~ReferenceIntegrateVerletStepKernel() {
    if (dynamics)
        delete dynamics;
}

2578
void ReferenceIntegrateVerletStepKernel::initialize(const System& system, const VerletIntegrator& integrator) {
Peter Eastman's avatar
Peter Eastman committed
2579
    int numParticles = system.getNumParticles();
2580
    masses.resize(numParticles);
2581
    for (int i = 0; i < numParticles; ++i)
peastman's avatar
peastman committed
2582
        masses[i] = system.getParticleMass(i);
2583
2584
}

2585
void ReferenceIntegrateVerletStepKernel::execute(ContextImpl& context, const VerletIntegrator& integrator) {
2586
    double stepSize = integrator.getStepSize();
peastman's avatar
peastman committed
2587
2588
2589
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& velData = extractVelocities(context);
    vector<Vec3>& forceData = extractForces(context);
2590
2591
2592
    if (dynamics == 0 || stepSize != prevStepSize) {
        // Recreate the computation objects with the new parameters.
        
2593
        if (dynamics)
2594
            delete dynamics;
peastman's avatar
peastman committed
2595
        dynamics = new ReferenceVerletDynamics(context.getSystem().getNumParticles(), stepSize);
2596
        dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
2597
        dynamics->setVirtualSites(extractVirtualSites(context));
2598
2599
        prevStepSize = stepSize;
    }
2600
    dynamics->update(context.getSystem(), posData, velData, forceData, masses, integrator.getConstraintTolerance(), extractBoxVectors(context));
2601
    data.time += stepSize;
2602
    data.stepCount++;
2603
}
2604

2605
double ReferenceIntegrateVerletStepKernel::computeKineticEnergy(ContextImpl& context, const VerletIntegrator& integrator) {
2606
    return computeShiftedKineticEnergy(context, masses, 0.5*integrator.getStepSize());
2607
2608
}

2609
2610
2611
ReferenceIntegrateNoseHooverStepKernel::~ReferenceIntegrateNoseHooverStepKernel() {
    if (chainPropagator)
        delete chainPropagator;
2612
2613
2614
2615
    if (dynamics)
        delete dynamics;
}

2616
void ReferenceIntegrateNoseHooverStepKernel::initialize(const System& system, const NoseHooverIntegrator& integrator) {
2617
2618
2619
2620
    int numParticles = system.getNumParticles();
    masses.resize(numParticles);
    for (int i = 0; i < numParticles; ++i)
        masses[i] = system.getParticleMass(i);
2621
    this->chainPropagator = new ReferenceNoseHooverChain();
2622
2623
}

2624
void ReferenceIntegrateNoseHooverStepKernel::execute(ContextImpl& context, const NoseHooverIntegrator& integrator) {
2625
2626
2627
2628
2629
2630
2631
2632
    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;
2633
        dynamics = new ReferenceNoseHooverDynamics(context.getSystem().getNumParticles(), stepSize);
2634
        dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
2635
        dynamics->setVirtualSites(extractVirtualSites(context));
2636
2637
        prevStepSize = stepSize;
    }
2638
    dynamics->step1(context, context.getSystem(), posData, velData, forceData, masses, integrator.getConstraintTolerance(),
2639
2640
2641
2642
2643
2644
2645
2646
                     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);
    }
2647
    dynamics->step2(context, context.getSystem(), posData, velData, forceData, masses, integrator.getConstraintTolerance(),
2648
2649
                     integrator.getAllThermostatedIndividualParticles(), integrator.getAllThermostatedPairs(), integrator.getMaximumPairDistance(),
            extractBoxVectors(context));
2650
2651
2652
2653
    data.time += stepSize;
    data.stepCount++;
}

2654
double ReferenceIntegrateNoseHooverStepKernel::computeKineticEnergy(ContextImpl& context, const NoseHooverIntegrator& integrator) {
2655
2656
2657
    return computeShiftedKineticEnergy(context, masses, 0);
}

2658
2659
2660
2661
2662
2663
2664
2665
2666
2667
2668
2669
2670
2671
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) {
2672
2673
        if (chainPositions.size() < 2*chainID+1){
            chainPositions.resize(2*chainID+1);
2674
        }
2675
2676
        if (chainVelocities.size() < 2*chainID+1){
            chainVelocities.resize(2*chainID+1);
2677
        }
2678
2679
2680
2681
        auto& positions = chainPositions.at(2*chainID);
        auto& velocities = chainVelocities.at(2*chainID);
        if (positions.size() < chainLength){
            positions.resize(chainLength, 0);
2682
        }
2683
2684
        if (velocities.size() < chainLength){
            velocities.resize(chainLength, 0);
2685
2686
2687
        }
        double temperature = nhc.getTemperature();
        double collisionFrequency = nhc.getCollisionFrequency();
2688
        absScale = chainPropagator->propagate(absKE, velocities, positions, numDOFs,
2689
2690
2691
2692
2693
2694
                                              temperature, collisionFrequency, timeStep,
                                              numMTS, nhc.getYoshidaSuzukiWeights());
    }
    double relScale = 0;
    int nPairs = nhc.getThermostatedPairs().size();
    if (nPairs) {
2695
2696
        if (chainPositions.size() < 2*chainID+2){
            chainPositions.resize(2*chainID+2);
2697
        }
2698
2699
        if (chainVelocities.size() < 2*chainID+2){
            chainVelocities.resize(2*chainID+2);
2700
        }
2701
2702
2703
2704
        auto& positions = chainPositions.at(2*chainID+1);
        auto& velocities = chainVelocities.at(2*chainID+1);
        if (positions.size() < chainLength){
            positions.resize(chainLength, 0);
2705
        }
2706
2707
        if (velocities.size() < chainLength){
            velocities.resize(chainLength, 0);
2708
2709
2710
        }
        double temperature = nhc.getRelativeTemperature();
        double collisionFrequency = nhc.getRelativeCollisionFrequency();
2711
        relScale = chainPropagator->propagate(relKE, velocities, positions, 3*nPairs,
2712
2713
2714
2715
2716
2717
2718
2719
2720
2721
2722
2723
2724
2725
2726
2727
2728
2729
2730
2731
2732
                                              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);
2733
            double velocity = chainVelocities[2*chainID][i];
2734
2735
2736
            // The kinetic energy of this bead
            kineticEnergy += 0.5 * mass * velocity * velocity;
            // The potential energy of this bead
2737
            double position = chainPositions[2*chainID][i];
2738
2739
2740
2741
2742
2743
2744
2745
2746
2747
2748
            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);
2749
            double velocity = chainVelocities[2*chainID+1][i];
2750
2751
2752
            // The kinetic energy of this bead
            kineticEnergy += 0.5 * mass * velocity * velocity;
            // The potential energy of this bead
2753
            double position = chainPositions[2*chainID+1][i];
2754
2755
2756
2757
2758
2759
2760
2761
2762
2763
2764
2765
2766
2767
2768
2769
2770
2771
2772
2773
2774
2775
2776
2777
2778
2779
2780
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
            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;
    }
}

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
2859
2860
2861
2862
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);
    }
}

2863
2864
2865
2866
2867
2868
2869
2870
2871
2872
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;
}

2873
ReferenceIntegrateLangevinMiddleStepKernel::~ReferenceIntegrateLangevinMiddleStepKernel() {
2874
2875
2876
2877
    if (dynamics)
        delete dynamics;
}

2878
void ReferenceIntegrateLangevinMiddleStepKernel::initialize(const System& system, const LangevinMiddleIntegrator& integrator) {
2879
2880
2881
2882
2883
2884
2885
    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());
}

2886
void ReferenceIntegrateLangevinMiddleStepKernel::execute(ContextImpl& context, const LangevinMiddleIntegrator& integrator) {
2887
2888
2889
2890
2891
2892
2893
2894
2895
2896
    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;
2897
        dynamics = new ReferenceLangevinMiddleDynamics(
2898
2899
2900
2901
2902
                context.getSystem().getNumParticles(), 
                stepSize, 
                friction, 
                temperature);
        dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
2903
        dynamics->setVirtualSites(extractVirtualSites(context));
2904
2905
2906
2907
        prevTemp = temperature;
        prevFriction = friction;
        prevStepSize = stepSize;
    }
2908
    dynamics->update(context, posData, velData, masses, integrator.getConstraintTolerance(), extractBoxVectors(context));
2909
2910
2911
2912
    data.time += stepSize;
    data.stepCount++;
}

2913
double ReferenceIntegrateLangevinMiddleStepKernel::computeKineticEnergy(ContextImpl& context, const LangevinMiddleIntegrator& integrator) {
2914
2915
2916
    return computeShiftedKineticEnergy(context, masses, 0.0);
}

2917
2918
2919
2920
2921
ReferenceIntegrateBrownianStepKernel::~ReferenceIntegrateBrownianStepKernel() {
    if (dynamics)
        delete dynamics;
}

2922
void ReferenceIntegrateBrownianStepKernel::initialize(const System& system, const BrownianIntegrator& integrator) {
Peter Eastman's avatar
Peter Eastman committed
2923
    int numParticles = system.getNumParticles();
2924
    masses.resize(numParticles);
2925
    for (int i = 0; i < numParticles; ++i)
peastman's avatar
peastman committed
2926
        masses[i] = system.getParticleMass(i);
2927
    SimTKOpenMMUtilities::setRandomNumberSeed((unsigned int) integrator.getRandomNumberSeed());
2928
2929
}

2930
void ReferenceIntegrateBrownianStepKernel::execute(ContextImpl& context, const BrownianIntegrator& integrator) {
2931
2932
2933
    double temperature = integrator.getTemperature();
    double friction = integrator.getFriction();
    double stepSize = integrator.getStepSize();
peastman's avatar
peastman committed
2934
2935
2936
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& velData = extractVelocities(context);
    vector<Vec3>& forceData = extractForces(context);
2937
2938
2939
    if (dynamics == 0 || temperature != prevTemp || friction != prevFriction || stepSize != prevStepSize) {
        // Recreate the computation objects with the new parameters.
        
2940
        if (dynamics)
2941
            delete dynamics;
2942
        dynamics = new ReferenceBrownianDynamics(
2943
                context.getSystem().getNumParticles(), 
peastman's avatar
peastman committed
2944
2945
2946
                stepSize, 
                friction, 
                temperature);
2947
        dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
2948
        dynamics->setVirtualSites(extractVirtualSites(context));
2949
2950
2951
2952
        prevTemp = temperature;
        prevFriction = friction;
        prevStepSize = stepSize;
    }
2953
    dynamics->update(context.getSystem(), posData, velData, forceData, masses, integrator.getConstraintTolerance(), extractBoxVectors(context));
2954
    data.time += stepSize;
2955
    data.stepCount++;
2956
2957
}

2958
double ReferenceIntegrateBrownianStepKernel::computeKineticEnergy(ContextImpl& context, const BrownianIntegrator& integrator) {
2959
    return computeShiftedKineticEnergy(context, masses, 0);
2960
2961
}

2962
2963
2964
2965
2966
2967
2968
ReferenceIntegrateVariableLangevinStepKernel::~ReferenceIntegrateVariableLangevinStepKernel() {
    if (dynamics)
        delete dynamics;
}

void ReferenceIntegrateVariableLangevinStepKernel::initialize(const System& system, const VariableLangevinIntegrator& integrator) {
    int numParticles = system.getNumParticles();
2969
    masses.resize(numParticles);
2970
    for (int i = 0; i < numParticles; ++i)
peastman's avatar
peastman committed
2971
        masses[i] = system.getParticleMass(i);
2972
2973
2974
    SimTKOpenMMUtilities::setRandomNumberSeed((unsigned int) integrator.getRandomNumberSeed());
}

2975
double ReferenceIntegrateVariableLangevinStepKernel::execute(ContextImpl& context, const VariableLangevinIntegrator& integrator, double maxTime) {
2976
2977
2978
    double temperature = integrator.getTemperature();
    double friction = integrator.getFriction();
    double errorTol = integrator.getErrorTolerance();
peastman's avatar
peastman committed
2979
2980
2981
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& velData = extractVelocities(context);
    vector<Vec3>& forceData = extractForces(context);
2982
2983
2984
    if (dynamics == 0 || temperature != prevTemp || friction != prevFriction || errorTol != prevErrorTol) {
        // Recreate the computation objects with the new parameters.

2985
        if (dynamics)
2986
            delete dynamics;
peastman's avatar
peastman committed
2987
        dynamics = new ReferenceVariableStochasticDynamics(context.getSystem().getNumParticles(), friction, temperature, errorTol);
2988
        dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
2989
        dynamics->setVirtualSites(extractVirtualSites(context));
2990
2991
2992
2993
        prevTemp = temperature;
        prevFriction = friction;
        prevErrorTol = errorTol;
    }
peastman's avatar
peastman committed
2994
    double maxStepSize = maxTime-data.time;
2995
2996
    if (integrator.getMaximumStepSize() > 0)
        maxStepSize = min(integrator.getMaximumStepSize(), maxStepSize);
2997
    dynamics->update(context.getSystem(), posData, velData, forceData, masses, maxStepSize, integrator.getConstraintTolerance(), extractBoxVectors(context));
2998
2999
3000
3001
    data.time += dynamics->getDeltaT();
    if (dynamics->getDeltaT() == maxStepSize)
        data.time = maxTime; // Avoid round-off error
    data.stepCount++;
3002
    return dynamics->getDeltaT();
3003
3004
}

3005
double ReferenceIntegrateVariableLangevinStepKernel::computeKineticEnergy(ContextImpl& context, const VariableLangevinIntegrator& integrator) {
3006
    return computeShiftedKineticEnergy(context, masses, 0.5*integrator.getStepSize());
3007
3008
}

3009
3010
3011
3012
3013
3014
3015
ReferenceIntegrateVariableVerletStepKernel::~ReferenceIntegrateVariableVerletStepKernel() {
    if (dynamics)
        delete dynamics;
}

void ReferenceIntegrateVariableVerletStepKernel::initialize(const System& system, const VariableVerletIntegrator& integrator) {
    int numParticles = system.getNumParticles();
3016
    masses.resize(numParticles);
3017
    for (int i = 0; i < numParticles; ++i)
peastman's avatar
peastman committed
3018
        masses[i] = system.getParticleMass(i);
3019
3020
}

3021
double ReferenceIntegrateVariableVerletStepKernel::execute(ContextImpl& context, const VariableVerletIntegrator& integrator, double maxTime) {
3022
    double errorTol = integrator.getErrorTolerance();
peastman's avatar
peastman committed
3023
3024
3025
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& velData = extractVelocities(context);
    vector<Vec3>& forceData = extractForces(context);
3026
    if (dynamics == 0 || errorTol != prevErrorTol) {
3027
3028
        // Recreate the computation objects with the new parameters.

3029
        if (dynamics)
3030
            delete dynamics;
peastman's avatar
peastman committed
3031
        dynamics = new ReferenceVariableVerletDynamics(context.getSystem().getNumParticles(), errorTol);
3032
        dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
3033
        dynamics->setVirtualSites(extractVirtualSites(context));
3034
        prevErrorTol = errorTol;
3035
    }
peastman's avatar
peastman committed
3036
    double maxStepSize = maxTime-data.time;
3037
3038
    if (integrator.getMaximumStepSize() > 0)
        maxStepSize = min(integrator.getMaximumStepSize(), maxStepSize);
3039
    dynamics->update(context.getSystem(), posData, velData, forceData, masses, maxStepSize, integrator.getConstraintTolerance(), extractBoxVectors(context));
3040
    data.time += dynamics->getDeltaT();
3041
3042
3043
    if (dynamics->getDeltaT() == maxStepSize)
        data.time = maxTime; // Avoid round-off error
    data.stepCount++;
3044
    return dynamics->getDeltaT();
3045
3046
}

3047
double ReferenceIntegrateVariableVerletStepKernel::computeKineticEnergy(ContextImpl& context, const VariableVerletIntegrator& integrator) {
3048
    return computeShiftedKineticEnergy(context, masses, 0.5*integrator.getStepSize());
3049
3050
}

3051
3052
3053
3054
3055
3056
3057
3058
3059
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
3060
        masses[i] = system.getParticleMass(i);
3061
    perDofValues.resize(integrator.getNumPerDofVariables());
peastman's avatar
peastman committed
3062
3063
    for (auto& values : perDofValues)
        values.resize(numParticles);
3064
3065
3066
3067

    // Create the computation objects.

    dynamics = new ReferenceCustomDynamics(system.getNumParticles(), integrator);
3068
    SimTKOpenMMUtilities::setRandomNumberSeed((unsigned int) integrator.getRandomNumberSeed());
3069
3070
3071
}

void ReferenceIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegrator& integrator, bool& forcesAreValid) {
peastman's avatar
peastman committed
3072
3073
3074
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& velData = extractVelocities(context);
    vector<Vec3>& forceData = extractForces(context);
3075
3076
3077
3078
3079
3080
3081
3082
3083
3084
    
    // 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.
    
3085
    dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
3086
    dynamics->setVirtualSites(extractVirtualSites(context));
3087
    dynamics->update(context, context.getSystem().getNumParticles(), posData, velData, forceData, masses, globals, perDofValues, forcesAreValid, integrator.getConstraintTolerance(), extractBoxVectors(context));
3088
3089
3090
3091
3092
3093
3094
3095
3096
3097
    
    // 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++;
}

3098
double ReferenceIntegrateCustomStepKernel::computeKineticEnergy(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
3112
3113
3114
    
    // 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);
}

3115
3116
3117
3118
3119
3120
3121
3122
3123
3124
3125
3126
3127
3128
3129
3130
3131
3132
3133
3134
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
3135
3136
3137
3138
3139
3140
3141
3142
3143
3144
3145
3146
3147
3148
3149
3150
3151
3152
3153
3154
3155
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);
3156
    dynamics->update(context, posData, velData, masses, integrator.getConstraintTolerance(), extractBoxVectors(context));
Peter Eastman's avatar
Peter Eastman committed
3157
3158
3159
3160
3161
3162
3163
3164
    data.time += integrator.getStepSize();
    data.stepCount++;
}

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

3165
3166
3167
3168
3169
3170
3171
3172
3173
3174
3175
3176
3177
3178
3179
3180
3181
3182
3183
3184
3185
3186
3187
3188
3189
3190
3191
3192
3193
3194
3195
3196
3197
3198
3199
3200
3201
3202
3203
3204
3205
3206
3207
3208
3209
3210
3211
3212
3213
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);
}

3214
3215
3216
3217
3218
ReferenceApplyAndersenThermostatKernel::~ReferenceApplyAndersenThermostatKernel() {
    if (thermostat)
        delete thermostat;
}

3219
void ReferenceApplyAndersenThermostatKernel::initialize(const System& system, const AndersenThermostat& thermostat) {
Peter Eastman's avatar
Peter Eastman committed
3220
    int numParticles = system.getNumParticles();
3221
    masses.resize(numParticles);
3222
    for (int i = 0; i < numParticles; ++i)
peastman's avatar
peastman committed
3223
        masses[i] = system.getParticleMass(i);
3224
    this->thermostat = new ReferenceAndersenThermostat();
3225
    SimTKOpenMMUtilities::setRandomNumberSeed((unsigned int) thermostat.getRandomNumberSeed());
3226
    particleGroups = AndersenThermostatImpl::calcParticleGroups(system);
3227
3228
}

3229
void ReferenceApplyAndersenThermostatKernel::execute(ContextImpl& context) {
peastman's avatar
peastman committed
3230
    vector<Vec3>& velData = extractVelocities(context);
3231
    thermostat->applyThermostat(particleGroups, velData, masses,
peastman's avatar
peastman committed
3232
3233
3234
        context.getParameter(AndersenThermostat::Temperature()),
        context.getParameter(AndersenThermostat::CollisionFrequency()),
        context.getIntegrator().getStepSize());
3235
3236
}

3237
3238
3239
3240
3241
ReferenceApplyMonteCarloBarostatKernel::~ReferenceApplyMonteCarloBarostatKernel() {
    if (barostat)
        delete barostat;
}

3242
3243
void ReferenceApplyMonteCarloBarostatKernel::initialize(const System& system, const Force& barostat, int components, bool rigidMolecules) {
    this->components = components;
3244
    this->rigidMolecules = rigidMolecules;
3245
3246
}

3247
void ReferenceApplyMonteCarloBarostatKernel::saveCoordinates(ContextImpl& context) {
3248
    if (barostat == NULL) {
3249
3250
3251
3252
        const System& system = context.getSystem();
        vector<double> masses;
        for (int i = 0; i < system.getNumParticles(); i++)
            masses.push_back(system.getParticleMass(i));
3253
        if (rigidMolecules)
3254
            barostat = new ReferenceMonteCarloBarostat(system.getNumParticles(), context.getMolecules(), masses);
3255
        else {
3256
            vector<vector<int> > molecules(system.getNumParticles());
3257
3258
            for (int i = 0; i < molecules.size(); i++)
                molecules[i].push_back(i);
3259
            barostat = new ReferenceMonteCarloBarostat(system.getNumParticles(), molecules, masses);
3260
3261
        }
    }
3262
3263
3264
3265
3266
    vector<Vec3>& posData = extractPositions(context);
    barostat->savePositions(posData);
}

void ReferenceApplyMonteCarloBarostatKernel::scaleCoordinates(ContextImpl& context, double scaleX, double scaleY, double scaleZ) {
peastman's avatar
peastman committed
3267
3268
    vector<Vec3>& posData = extractPositions(context);
    Vec3* boxVectors = extractBoxVectors(context);
3269
    barostat->applyBarostat(posData, boxVectors, scaleX, scaleY, scaleZ);
3270
3271
3272
}

void ReferenceApplyMonteCarloBarostatKernel::restoreCoordinates(ContextImpl& context) {
peastman's avatar
peastman committed
3273
    vector<Vec3>& posData = extractPositions(context);
3274
3275
3276
    barostat->restorePositions(posData);
}

3277
3278
3279
3280
void ReferenceApplyMonteCarloBarostatKernel::computeKineticEnergy(ContextImpl& context, vector<double>& ke) {
    barostat->computeMolecularKineticEnergy(extractVelocities(context), ke, components);
}

3281
3282
void ReferenceRemoveCMMotionKernel::initialize(const System& system, const CMMotionRemover& force) {
    frequency = force.getFrequency();
Peter Eastman's avatar
Peter Eastman committed
3283
    masses.resize(system.getNumParticles());
3284
    for (size_t i = 0; i < masses.size(); ++i)
Peter Eastman's avatar
Peter Eastman committed
3285
        masses[i] = system.getParticleMass(i);
3286
3287
}

3288
void ReferenceRemoveCMMotionKernel::execute(ContextImpl& context) {
3289
    if (data.stepCount%frequency != 0)
3290
        return;
peastman's avatar
peastman committed
3291
    vector<Vec3>& velData = extractVelocities(context);
3292
3293
3294
    
    // Calculate the center of mass momentum.
    
peastman's avatar
peastman committed
3295
3296
    double momentum[] = {0.0, 0.0, 0.0};
    double mass = 0.0;
3297
    for (size_t i = 0; i < masses.size(); ++i) {
peastman's avatar
peastman committed
3298
3299
3300
3301
        momentum[0] += masses[i]*velData[i][0];
        momentum[1] += masses[i]*velData[i][1];
        momentum[2] += masses[i]*velData[i][2];
        mass += masses[i];
3302
3303
    }
    
Peter Eastman's avatar
Peter Eastman committed
3304
    // Adjust the particle velocities.
3305
    
3306
3307
3308
    momentum[0] /= mass;
    momentum[1] /= mass;
    momentum[2] /= mass;
3309
    for (size_t i = 0; i < masses.size(); ++i) {
3310
3311
3312
3313
3314
        if (masses[i] != 0.0) {
            velData[i][0] -= momentum[0];
            velData[i][1] -= momentum[1];
            velData[i][2] -= momentum[2];
        }
3315
3316
    }
}
3317

3318
3319
3320
3321
3322
3323
3324
3325
3326
3327
3328
3329
3330
3331
3332
3333
3334
3335
3336
3337
3338
3339
3340
3341
3342
3343
3344
3345
3346
3347
3348
3349
3350
3351
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");
        }
    }
}

3352
3353
3354
3355
3356
void ReferenceCalcATMForceKernel::initialize(const System& system, const ATMForce& force) {
    numParticles = force.getNumParticles();
    //displacement map
    displ1.resize(numParticles);
    displ0.resize(numParticles);
3357
3358
3359
3360
3361
3362
3363
3364
3365
3366
3367
3368
3369
3370
3371
3372
3373
3374
3375
3376
3377
3378
3379
3380
3381
3382
3383
3384
3385
3386
3387
3388
3389
3390
3391
3392
3393
3394
3395
3396
3397
3398
3399
3400
3401
    //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];
3402
3403
3404
3405
3406
3407
3408
3409
    }
}

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);
3410

3411
3412
3413
    //add gradients from variable displacements
    displForces(force0, force1);

3414
3415
3416
3417
3418
    //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)
3419
3420
3421
            force[i] += dEdu0*force0[i];
        if (fabs(dEdu1) > epsi)
            force[i] += dEdu1*force1[i];
3422
3423
    }

3424
3425
3426
3427
3428
3429
3430
3431
    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);

3432
3433
3434
    //calculate displacement vectors
    setDisplacements(pos);

3435
3436
3437
3438
    //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];
3439
    innerContext0.setPositions(pos0);
3440
3441
3442
3443
3444

    //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];
3445
    innerContext1.setPositions(pos1);
3446
3447
3448
3449
3450
3451
3452
3453
3454
3455
3456
3457
3458
3459
3460
3461
3462
3463
3464
3465
3466
3467
3468
3469
3470
3471

    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);
3472
    loadParams(numParticles, force);
3473
}
3474
3475
3476
3477
3478
3479
3480
3481
3482
3483
3484
3485
3486
3487
3488

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;
}
Peter Eastman's avatar
Peter Eastman committed
3489
3490
3491
3492
3493
3494
3495
3496
3497
3498
3499
3500
3501
3502
3503
3504
3505
3506
3507
3508
3509
3510
3511
3512
3513
3514

void ReferenceCalcPythonForceKernel::initialize(const System& system, const PythonForce& force) {
    computation = &force.getComputation();
    forces.resize(system.getNumParticles());
    usePeriodic = force.usesPeriodicBoundaryConditions();
}

double ReferenceCalcPythonForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& forceData = extractForces(context);
    State::StateBuilder builder(context.getTime(), context.getStepCount());
    builder.setPositions(posData);
    builder.setParameters(context.getParameters());
    if (usePeriodic) {
        Vec3 a, b, c;
        context.getPeriodicBoxVectors(a, b, c);
        builder.setPeriodicBoxVectors(a, b, c);
    }
    double energy;
    State state = builder.getState();
    computation->compute(state, energy, forces.data(), true);
    if (includeForces)
        for (int i = 0; i < forces.size(); i++)
            forceData[i] += forces[i];
    return energy;
}