ReferenceKernels.cpp 157 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"
Evan Pretti's avatar
Evan Pretti committed
56
#include "ReferenceLCPOIxn.h"
57
58
59
#include "ReferenceLJCoulomb14.h"
#include "ReferenceLJCoulombIxn.h"
#include "ReferenceMonteCarloBarostat.h"
60
#include "ReferenceNoseHooverChain.h"
61
#include "ReferenceNoseHooverDynamics.h"
62
#include "ReferenceOrientationRestraintForce.h"
63
#include "ReferencePointFunctions.h"
64
#include "ReferenceProperDihedralBond.h"
65
#include "ReferenceQTBDynamics.h"
66
#include "ReferenceRbDihedralBond.h"
67
#include "ReferenceRGForce.h"
68
#include "ReferenceRMSDForce.h"
69
#include "ReferenceTabulatedFunction.h"
70
71
72
73
#include "ReferenceVariableStochasticDynamics.h"
#include "ReferenceVariableVerletDynamics.h"
#include "ReferenceVerletDynamics.h"
#include "ReferenceVirtualSites.h"
74
#include "openmm/CMMotionRemover.h"
75
#include "openmm/Context.h"
76
#include "openmm/System.h"
77
#include "openmm/internal/AndersenThermostatImpl.h"
78
#include "openmm/internal/ContextImpl.h"
79
#include "openmm/internal/CustomCentroidBondForceImpl.h"
80
#include "openmm/internal/CustomCompoundBondForceImpl.h"
81
#include "openmm/internal/CustomHbondForceImpl.h"
82
#include "openmm/internal/CMAPTorsionForceImpl.h"
83
#include "openmm/internal/NonbondedForceImpl.h"
84
#include "openmm/internal/ConstantPotentialForceImpl.h"
85
#include "openmm/Integrator.h"
86
#include "openmm/OpenMMException.h"
87
#include "SimTKOpenMMUtilities.h"
88
#include "lepton/CustomFunction.h"
89
#include "lepton/Operation.h"
90
91
#include "lepton/Parser.h"
#include "lepton/ParsedExpression.h"
92
#include <cmath>
Peter Eastman's avatar
Peter Eastman committed
93
#include <iostream>
94
#include <limits>
95
96
97
98

using namespace OpenMM;
using namespace std;

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

209
210
211
212
213
214
215
216
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
217
void ReferenceUpdateStateDataKernel::getPositions(ContextImpl& context, std::vector<Vec3>& positions, bool allowPeriodic) {
Peter Eastman's avatar
Peter Eastman committed
218
    positions = extractPositions(context);
219
220
221
}

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

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

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

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

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

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

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

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

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

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

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

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

ReferenceApplyConstraintsKernel::~ReferenceApplyConstraintsKernel() {
}

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

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

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

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

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

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

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

    // Record the values.

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

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

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

    // Build the arrays.

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

    // Parse the expression used to calculate the force.

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

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

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

    // Record the values.

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

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

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

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

    // Record the values.

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

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

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

    // Build the arrays.

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

    // Parse the expression used to calculate the force.

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

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

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

    // Record the values.

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

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

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

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

    // Record the values.

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

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

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

692
693
694
695
696
697
698
699
700
701
702
703
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
704
705
706
707
708
709
        torsionParamArray[i][0] = c0;
        torsionParamArray[i][1] = c1;
        torsionParamArray[i][2] = c2;
        torsionParamArray[i][3] = c3;
        torsionParamArray[i][4] = c4;
        torsionParamArray[i][5] = c5;
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
736
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]);
    }
737
    usePeriodic = force.usesPeriodicBoundaryConditions();
738
739
}

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

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
784
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");
    }
}

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

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

    // Build the arrays.

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

    // Parse the expression used to calculate the force.

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

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

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

    // Record the values.

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

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

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

    // Identify which exceptions are 1-4 interactions.

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

    // Build the arrays.

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

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

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

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

    // Record the values.

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

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

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

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

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
1394
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);
}

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

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

1404
    // Record the exclusions.
1405
1406
1407

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

    // Build the arrays.

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

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

    for (int i = 0; i < force.getNumTabulatedFunctions(); i++)
1435
        tabulatedFunctionUpdateCount[force.getTabulatedFunctionName(i)] = force.getTabulatedFunction(i).getUpdateCount();
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
1461

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

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

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

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

    // Delete the custom functions.

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

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

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

    // Record the values.

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

    // See if any tabulated functions have changed.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    // Parse the expressions for computed values.

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

1790
    // Parse the expressions for energy terms.
1791

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

    // Delete the custom functions.

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

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

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

    // See if any tabulated functions have changed.

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

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

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

    // Build the arrays.

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

    // Parse the expression used to calculate the force.

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

1926
1927
}

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

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

    // Record the values.

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

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

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.

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

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

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

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

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

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

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

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

    // Delete the custom functions.

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

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

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

    // See if any tabulated functions have changed.

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

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

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

    // Build the arrays.

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

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

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

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

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

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

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

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

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

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

    // Delete the custom functions.

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

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

    // See if any tabulated functions have changed.

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

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

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

    // Build the arrays.

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

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

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

    // Create the interaction.

    createInteraction(force);
}

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

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

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

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

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

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

    // Delete the custom functions.

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

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

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

    // See if any tabulated functions have changed.

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

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

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

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

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

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

    // Create the interaction.

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

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

    // See if any tabulated functions have changed.

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

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

Evan Pretti's avatar
Evan Pretti committed
2442
2443
2444
2445
2446
2447
2448
2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
2461
2462
2463
2464
2465
2466
2467
2468
2469
2470
2471
2472
2473
2474
2475
2476
2477
2478
2479
2480
2481
2482
2483
2484
2485
2486
2487
2488
2489
2490
2491
2492
2493
2494
2495
void ReferenceCalcLCPOForceKernel::initialize(const System& system, const LCPOForce& force) {
    oneBodyEnergy = 0.0;
    double maxRadius = 0.0;

    double surfaceTension = force.getSurfaceTension();
    for (int i = 0; i < force.getNumParticles(); i++) {
        double radius, p1, p2, p3, p4;
        force.getParticleParameters(i, radius, p1, p2, p3, p4);
        p1 *= surfaceTension;
        p2 *= surfaceTension;
        p3 *= surfaceTension;
        p4 *= surfaceTension;
        oneBodyEnergy += 4.0 * PI_M * p1 * radius * radius;

        if (radius != 0.0) {
            activeParticlesInv.push_back(activeParticles.size());
            activeParticles.push_back(i);
            parameters.push_back({radius, p2, p3, p4});
            maxRadius = max(maxRadius, radius);
        }
        else {
            activeParticlesInv.push_back(-1);
        }
    }

    cutoff = 2.0 * maxRadius;
    usePeriodic = force.usesPeriodicBoundaryConditions();
}

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

    if (usePeriodic) {
        double minAllowedSize = 1.999999 * cutoff;
        if (boxVectors[0][0] < minAllowedSize || boxVectors[1][1] < minAllowedSize || boxVectors[2][2] < minAllowedSize) {
            throw OpenMMException("The periodic box size is less than twice the required cutoff for LCPO.");
        }
    }

    ReferenceLCPOIxn lcpo(activeParticles, activeParticlesInv, parameters, cutoff, usePeriodic);
    return oneBodyEnergy + lcpo.execute(boxVectors, posData, forceData, includeForces, includeEnergy);
}

void ReferenceCalcLCPOForceKernel::copyParametersToContext(ContextImpl& context, const LCPOForce& force) {
    // For the reference implementation, just reinitialize everything.

    activeParticles.clear();
    activeParticlesInv.clear();
    parameters.clear();
    initialize(context.getSystem(), force);
}

2496
2497
2498
2499
2500
ReferenceCalcCustomCVForceKernel::~ReferenceCalcCustomCVForceKernel() {
    if (ixn != NULL)
        delete ixn;
}

2501
void ReferenceCalcCustomCVForceKernel::initialize(const System& system, const CustomCVForce& force, ContextImpl& innerContext) {
2502
2503
2504
2505
2506
2507
2508
2509
2510
2511
2512
2513
2514
2515
2516
2517
2518
2519
2520
2521
2522
2523
2524
    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);
2525
2526
2527
2528
    Vec3 a, b, c;
    context.getPeriodicBoxVectors(a, b, c);
    innerContext.setPeriodicBoxVectors(a, b, c);
    innerContext.setTime(context.getTime());
2529
2530
2531
2532
2533
    map<string, double> innerParameters = innerContext.getParameters();
    for (auto& param : innerParameters)
        innerContext.setParameter(param.first, context.getParameter(param.first));
}

2534
void ReferenceCalcCustomCVForceKernel::copyParametersToContext(ContextImpl& context, const CustomCVForce& force) {
2535
    ixn->updateTabulatedFunctions(force);
2536
2537
}

2538
2539
void ReferenceCalcRMSDForceKernel::initialize(const System& system, const RMSDForce& force) {
    particles = force.getParticles();
peastman's avatar
peastman committed
2540
2541
2542
    if (particles.size() == 0)
        for (int i = 0; i < system.getNumParticles(); i++)
            particles.push_back(i);
2543
2544
2545
2546
2547
2548
2549
2550
2551
2552
2553
2554
2555
2556
2557
2558
2559
2560
2561
2562
    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
2563
2564
2565
    if (particles.size() == 0)
        for (int i = 0; i < referencePos.size(); i++)
            particles.push_back(i);
2566
2567
2568
2569
2570
2571
2572
2573
2574
    referencePos = force.getReferencePositions();
    Vec3 center;
    for (int i : particles)
        center += referencePos[i];
    center /= particles.size();
    for (Vec3& p : referencePos)
        p -= center;
}

2575
2576
2577
2578
2579
2580
2581
2582
2583
2584
2585
2586
2587
2588
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);
}

2589
2590
2591
2592
2593
2594
2595
2596
2597
2598
2599
2600
2601
2602
2603
2604
2605
2606
2607
2608
2609
2610
2611
2612
2613
2614
2615
2616
2617
2618
2619
2620
2621
2622
2623
2624
2625
2626
2627
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;
}

2628
2629
2630
2631
2632
ReferenceIntegrateVerletStepKernel::~ReferenceIntegrateVerletStepKernel() {
    if (dynamics)
        delete dynamics;
}

2633
void ReferenceIntegrateVerletStepKernel::initialize(const System& system, const VerletIntegrator& integrator) {
Peter Eastman's avatar
Peter Eastman committed
2634
    int numParticles = system.getNumParticles();
2635
    masses.resize(numParticles);
2636
    for (int i = 0; i < numParticles; ++i)
peastman's avatar
peastman committed
2637
        masses[i] = system.getParticleMass(i);
2638
2639
}

2640
void ReferenceIntegrateVerletStepKernel::execute(ContextImpl& context, const VerletIntegrator& integrator) {
2641
    double stepSize = integrator.getStepSize();
peastman's avatar
peastman committed
2642
2643
2644
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& velData = extractVelocities(context);
    vector<Vec3>& forceData = extractForces(context);
2645
2646
2647
    if (dynamics == 0 || stepSize != prevStepSize) {
        // Recreate the computation objects with the new parameters.
        
2648
        if (dynamics)
2649
            delete dynamics;
peastman's avatar
peastman committed
2650
        dynamics = new ReferenceVerletDynamics(context.getSystem().getNumParticles(), stepSize);
2651
        dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
2652
        dynamics->setVirtualSites(extractVirtualSites(context));
2653
2654
        prevStepSize = stepSize;
    }
2655
    dynamics->update(context.getSystem(), posData, velData, forceData, masses, integrator.getConstraintTolerance(), extractBoxVectors(context));
2656
    data.time += stepSize;
2657
    data.stepCount++;
2658
}
2659

2660
double ReferenceIntegrateVerletStepKernel::computeKineticEnergy(ContextImpl& context, const VerletIntegrator& integrator) {
2661
    return computeShiftedKineticEnergy(context, masses, 0.5*integrator.getStepSize());
2662
2663
}

2664
2665
2666
ReferenceIntegrateNoseHooverStepKernel::~ReferenceIntegrateNoseHooverStepKernel() {
    if (chainPropagator)
        delete chainPropagator;
2667
2668
2669
2670
    if (dynamics)
        delete dynamics;
}

2671
void ReferenceIntegrateNoseHooverStepKernel::initialize(const System& system, const NoseHooverIntegrator& integrator) {
2672
2673
2674
2675
    int numParticles = system.getNumParticles();
    masses.resize(numParticles);
    for (int i = 0; i < numParticles; ++i)
        masses[i] = system.getParticleMass(i);
2676
    this->chainPropagator = new ReferenceNoseHooverChain();
2677
2678
}

2679
void ReferenceIntegrateNoseHooverStepKernel::execute(ContextImpl& context, const NoseHooverIntegrator& integrator) {
2680
2681
2682
2683
2684
2685
2686
2687
    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;
2688
        dynamics = new ReferenceNoseHooverDynamics(context.getSystem().getNumParticles(), stepSize);
2689
        dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
2690
        dynamics->setVirtualSites(extractVirtualSites(context));
2691
2692
        prevStepSize = stepSize;
    }
2693
    dynamics->step1(context, context.getSystem(), posData, velData, forceData, masses, integrator.getConstraintTolerance(),
2694
2695
2696
2697
2698
2699
2700
2701
                     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);
    }
2702
    dynamics->step2(context, context.getSystem(), posData, velData, forceData, masses, integrator.getConstraintTolerance(),
2703
2704
                     integrator.getAllThermostatedIndividualParticles(), integrator.getAllThermostatedPairs(), integrator.getMaximumPairDistance(),
            extractBoxVectors(context));
2705
2706
2707
2708
    data.time += stepSize;
    data.stepCount++;
}

2709
double ReferenceIntegrateNoseHooverStepKernel::computeKineticEnergy(ContextImpl& context, const NoseHooverIntegrator& integrator) {
2710
2711
2712
    return computeShiftedKineticEnergy(context, masses, 0);
}

2713
2714
2715
2716
2717
2718
2719
2720
2721
2722
2723
2724
2725
2726
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) {
2727
2728
        if (chainPositions.size() < 2*chainID+1){
            chainPositions.resize(2*chainID+1);
2729
        }
2730
2731
        if (chainVelocities.size() < 2*chainID+1){
            chainVelocities.resize(2*chainID+1);
2732
        }
2733
2734
2735
2736
        auto& positions = chainPositions.at(2*chainID);
        auto& velocities = chainVelocities.at(2*chainID);
        if (positions.size() < chainLength){
            positions.resize(chainLength, 0);
2737
        }
2738
2739
        if (velocities.size() < chainLength){
            velocities.resize(chainLength, 0);
2740
2741
2742
        }
        double temperature = nhc.getTemperature();
        double collisionFrequency = nhc.getCollisionFrequency();
2743
        absScale = chainPropagator->propagate(absKE, velocities, positions, numDOFs,
2744
2745
2746
2747
2748
2749
                                              temperature, collisionFrequency, timeStep,
                                              numMTS, nhc.getYoshidaSuzukiWeights());
    }
    double relScale = 0;
    int nPairs = nhc.getThermostatedPairs().size();
    if (nPairs) {
2750
2751
        if (chainPositions.size() < 2*chainID+2){
            chainPositions.resize(2*chainID+2);
2752
        }
2753
2754
        if (chainVelocities.size() < 2*chainID+2){
            chainVelocities.resize(2*chainID+2);
2755
        }
2756
2757
2758
2759
        auto& positions = chainPositions.at(2*chainID+1);
        auto& velocities = chainVelocities.at(2*chainID+1);
        if (positions.size() < chainLength){
            positions.resize(chainLength, 0);
2760
        }
2761
2762
        if (velocities.size() < chainLength){
            velocities.resize(chainLength, 0);
2763
2764
2765
        }
        double temperature = nhc.getRelativeTemperature();
        double collisionFrequency = nhc.getRelativeCollisionFrequency();
2766
        relScale = chainPropagator->propagate(relKE, velocities, positions, 3*nPairs,
2767
2768
2769
2770
2771
2772
2773
2774
2775
2776
2777
2778
2779
2780
2781
2782
2783
2784
2785
2786
2787
                                              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);
2788
            double velocity = chainVelocities[2*chainID][i];
2789
2790
2791
            // The kinetic energy of this bead
            kineticEnergy += 0.5 * mass * velocity * velocity;
            // The potential energy of this bead
2792
            double position = chainPositions[2*chainID][i];
2793
2794
2795
2796
2797
2798
2799
2800
2801
2802
2803
            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);
2804
            double velocity = chainVelocities[2*chainID+1][i];
2805
2806
2807
            // The kinetic energy of this bead
            kineticEnergy += 0.5 * mass * velocity * velocity;
            // The potential energy of this bead
2808
            double position = chainPositions[2*chainID+1][i];
2809
2810
2811
2812
2813
2814
2815
2816
2817
2818
2819
2820
2821
2822
2823
2824
2825
2826
2827
2828
2829
2830
2831
2832
2833
2834
2835
2836
2837
2838
2839
2840
2841
2842
2843
2844
2845
2846
2847
2848
2849
2850
2851
2852
2853
2854
2855
2856
2857
2858
2859
2860
2861
2862
2863
2864
2865
2866
2867
2868
2869
2870
2871
2872
2873
2874
2875
2876
2877
2878
2879
2880
2881
2882
2883
2884
2885
2886
            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;
    }
}

2887
2888
2889
2890
2891
2892
2893
2894
2895
2896
2897
2898
2899
2900
2901
2902
2903
2904
2905
2906
2907
2908
2909
2910
2911
2912
2913
2914
2915
2916
2917
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);
    }
}

2918
2919
2920
2921
2922
2923
2924
2925
2926
2927
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;
}

2928
ReferenceIntegrateLangevinMiddleStepKernel::~ReferenceIntegrateLangevinMiddleStepKernel() {
2929
2930
2931
2932
    if (dynamics)
        delete dynamics;
}

2933
void ReferenceIntegrateLangevinMiddleStepKernel::initialize(const System& system, const LangevinMiddleIntegrator& integrator) {
2934
2935
2936
2937
2938
2939
2940
    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());
}

2941
void ReferenceIntegrateLangevinMiddleStepKernel::execute(ContextImpl& context, const LangevinMiddleIntegrator& integrator) {
2942
2943
2944
2945
2946
2947
2948
2949
2950
2951
    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;
2952
        dynamics = new ReferenceLangevinMiddleDynamics(
2953
2954
2955
2956
2957
                context.getSystem().getNumParticles(), 
                stepSize, 
                friction, 
                temperature);
        dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
2958
        dynamics->setVirtualSites(extractVirtualSites(context));
2959
2960
2961
2962
        prevTemp = temperature;
        prevFriction = friction;
        prevStepSize = stepSize;
    }
2963
    dynamics->update(context, posData, velData, masses, integrator.getConstraintTolerance(), extractBoxVectors(context));
2964
2965
2966
2967
    data.time += stepSize;
    data.stepCount++;
}

2968
double ReferenceIntegrateLangevinMiddleStepKernel::computeKineticEnergy(ContextImpl& context, const LangevinMiddleIntegrator& integrator) {
2969
2970
2971
    return computeShiftedKineticEnergy(context, masses, 0.0);
}

2972
2973
2974
2975
2976
ReferenceIntegrateBrownianStepKernel::~ReferenceIntegrateBrownianStepKernel() {
    if (dynamics)
        delete dynamics;
}

2977
void ReferenceIntegrateBrownianStepKernel::initialize(const System& system, const BrownianIntegrator& integrator) {
Peter Eastman's avatar
Peter Eastman committed
2978
    int numParticles = system.getNumParticles();
2979
    masses.resize(numParticles);
2980
    for (int i = 0; i < numParticles; ++i)
peastman's avatar
peastman committed
2981
        masses[i] = system.getParticleMass(i);
2982
    SimTKOpenMMUtilities::setRandomNumberSeed((unsigned int) integrator.getRandomNumberSeed());
2983
2984
}

2985
void ReferenceIntegrateBrownianStepKernel::execute(ContextImpl& context, const BrownianIntegrator& integrator) {
2986
2987
2988
    double temperature = integrator.getTemperature();
    double friction = integrator.getFriction();
    double stepSize = integrator.getStepSize();
peastman's avatar
peastman committed
2989
2990
2991
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& velData = extractVelocities(context);
    vector<Vec3>& forceData = extractForces(context);
2992
2993
2994
    if (dynamics == 0 || temperature != prevTemp || friction != prevFriction || stepSize != prevStepSize) {
        // Recreate the computation objects with the new parameters.
        
2995
        if (dynamics)
2996
            delete dynamics;
2997
        dynamics = new ReferenceBrownianDynamics(
2998
                context.getSystem().getNumParticles(), 
peastman's avatar
peastman committed
2999
3000
3001
                stepSize, 
                friction, 
                temperature);
3002
        dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
3003
        dynamics->setVirtualSites(extractVirtualSites(context));
3004
3005
3006
3007
        prevTemp = temperature;
        prevFriction = friction;
        prevStepSize = stepSize;
    }
3008
    dynamics->update(context.getSystem(), posData, velData, forceData, masses, integrator.getConstraintTolerance(), extractBoxVectors(context));
3009
    data.time += stepSize;
3010
    data.stepCount++;
3011
3012
}

3013
double ReferenceIntegrateBrownianStepKernel::computeKineticEnergy(ContextImpl& context, const BrownianIntegrator& integrator) {
3014
    return computeShiftedKineticEnergy(context, masses, 0);
3015
3016
}

3017
3018
3019
3020
3021
3022
3023
ReferenceIntegrateVariableLangevinStepKernel::~ReferenceIntegrateVariableLangevinStepKernel() {
    if (dynamics)
        delete dynamics;
}

void ReferenceIntegrateVariableLangevinStepKernel::initialize(const System& system, const VariableLangevinIntegrator& integrator) {
    int numParticles = system.getNumParticles();
3024
    masses.resize(numParticles);
3025
    for (int i = 0; i < numParticles; ++i)
peastman's avatar
peastman committed
3026
        masses[i] = system.getParticleMass(i);
3027
3028
3029
    SimTKOpenMMUtilities::setRandomNumberSeed((unsigned int) integrator.getRandomNumberSeed());
}

3030
double ReferenceIntegrateVariableLangevinStepKernel::execute(ContextImpl& context, const VariableLangevinIntegrator& integrator, double maxTime) {
3031
3032
3033
    double temperature = integrator.getTemperature();
    double friction = integrator.getFriction();
    double errorTol = integrator.getErrorTolerance();
peastman's avatar
peastman committed
3034
3035
3036
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& velData = extractVelocities(context);
    vector<Vec3>& forceData = extractForces(context);
3037
3038
3039
    if (dynamics == 0 || temperature != prevTemp || friction != prevFriction || errorTol != prevErrorTol) {
        // Recreate the computation objects with the new parameters.

3040
        if (dynamics)
3041
            delete dynamics;
peastman's avatar
peastman committed
3042
        dynamics = new ReferenceVariableStochasticDynamics(context.getSystem().getNumParticles(), friction, temperature, errorTol);
3043
        dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
3044
        dynamics->setVirtualSites(extractVirtualSites(context));
3045
3046
3047
3048
        prevTemp = temperature;
        prevFriction = friction;
        prevErrorTol = errorTol;
    }
peastman's avatar
peastman committed
3049
    double maxStepSize = maxTime-data.time;
3050
3051
    if (integrator.getMaximumStepSize() > 0)
        maxStepSize = min(integrator.getMaximumStepSize(), maxStepSize);
3052
    dynamics->update(context.getSystem(), posData, velData, forceData, masses, maxStepSize, integrator.getConstraintTolerance(), extractBoxVectors(context));
3053
3054
3055
3056
    data.time += dynamics->getDeltaT();
    if (dynamics->getDeltaT() == maxStepSize)
        data.time = maxTime; // Avoid round-off error
    data.stepCount++;
3057
    return dynamics->getDeltaT();
3058
3059
}

3060
double ReferenceIntegrateVariableLangevinStepKernel::computeKineticEnergy(ContextImpl& context, const VariableLangevinIntegrator& integrator) {
3061
    return computeShiftedKineticEnergy(context, masses, 0.5*integrator.getStepSize());
3062
3063
}

3064
3065
3066
3067
3068
3069
3070
ReferenceIntegrateVariableVerletStepKernel::~ReferenceIntegrateVariableVerletStepKernel() {
    if (dynamics)
        delete dynamics;
}

void ReferenceIntegrateVariableVerletStepKernel::initialize(const System& system, const VariableVerletIntegrator& integrator) {
    int numParticles = system.getNumParticles();
3071
    masses.resize(numParticles);
3072
    for (int i = 0; i < numParticles; ++i)
peastman's avatar
peastman committed
3073
        masses[i] = system.getParticleMass(i);
3074
3075
}

3076
double ReferenceIntegrateVariableVerletStepKernel::execute(ContextImpl& context, const VariableVerletIntegrator& integrator, double maxTime) {
3077
    double errorTol = integrator.getErrorTolerance();
peastman's avatar
peastman committed
3078
3079
3080
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& velData = extractVelocities(context);
    vector<Vec3>& forceData = extractForces(context);
3081
    if (dynamics == 0 || errorTol != prevErrorTol) {
3082
3083
        // Recreate the computation objects with the new parameters.

3084
        if (dynamics)
3085
            delete dynamics;
peastman's avatar
peastman committed
3086
        dynamics = new ReferenceVariableVerletDynamics(context.getSystem().getNumParticles(), errorTol);
3087
        dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
3088
        dynamics->setVirtualSites(extractVirtualSites(context));
3089
        prevErrorTol = errorTol;
3090
    }
peastman's avatar
peastman committed
3091
    double maxStepSize = maxTime-data.time;
3092
3093
    if (integrator.getMaximumStepSize() > 0)
        maxStepSize = min(integrator.getMaximumStepSize(), maxStepSize);
3094
    dynamics->update(context.getSystem(), posData, velData, forceData, masses, maxStepSize, integrator.getConstraintTolerance(), extractBoxVectors(context));
3095
    data.time += dynamics->getDeltaT();
3096
3097
3098
    if (dynamics->getDeltaT() == maxStepSize)
        data.time = maxTime; // Avoid round-off error
    data.stepCount++;
3099
    return dynamics->getDeltaT();
3100
3101
}

3102
double ReferenceIntegrateVariableVerletStepKernel::computeKineticEnergy(ContextImpl& context, const VariableVerletIntegrator& integrator) {
3103
    return computeShiftedKineticEnergy(context, masses, 0.5*integrator.getStepSize());
3104
3105
}

3106
3107
3108
3109
3110
3111
3112
3113
3114
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
3115
        masses[i] = system.getParticleMass(i);
3116
    perDofValues.resize(integrator.getNumPerDofVariables());
peastman's avatar
peastman committed
3117
3118
    for (auto& values : perDofValues)
        values.resize(numParticles);
3119
3120
3121
3122

    // Create the computation objects.

    dynamics = new ReferenceCustomDynamics(system.getNumParticles(), integrator);
3123
    SimTKOpenMMUtilities::setRandomNumberSeed((unsigned int) integrator.getRandomNumberSeed());
3124
3125
3126
}

void ReferenceIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegrator& integrator, bool& forcesAreValid) {
peastman's avatar
peastman committed
3127
3128
3129
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& velData = extractVelocities(context);
    vector<Vec3>& forceData = extractForces(context);
3130
3131
3132
3133
3134
3135
3136
3137
3138
3139
    
    // 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.
    
3140
    dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
3141
    dynamics->setVirtualSites(extractVirtualSites(context));
3142
    dynamics->update(context, context.getSystem().getNumParticles(), posData, velData, forceData, masses, globals, perDofValues, forcesAreValid, integrator.getConstraintTolerance(), extractBoxVectors(context));
3143
3144
3145
3146
3147
3148
3149
3150
3151
3152
    
    // 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++;
}

3153
double ReferenceIntegrateCustomStepKernel::computeKineticEnergy(ContextImpl& context, CustomIntegrator& integrator, bool& forcesAreValid) {
peastman's avatar
peastman committed
3154
3155
3156
    vector<Vec3>& posData = extractPositions(context);
    vector<Vec3>& velData = extractVelocities(context);
    vector<Vec3>& forceData = extractForces(context);
3157
3158
3159
3160
3161
3162
3163
3164
3165
3166
3167
3168
3169
    
    // 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);
}

3170
3171
3172
3173
3174
3175
3176
3177
3178
3179
3180
3181
3182
3183
3184
3185
3186
3187
3188
3189
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
3190
3191
3192
3193
3194
3195
3196
3197
3198
3199
3200
3201
3202
3203
3204
3205
3206
3207
3208
3209
3210
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);
3211
    dynamics->update(context, posData, velData, masses, integrator.getConstraintTolerance(), extractBoxVectors(context));
Peter Eastman's avatar
Peter Eastman committed
3212
3213
3214
3215
3216
3217
3218
3219
    data.time += integrator.getStepSize();
    data.stepCount++;
}

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

3220
3221
3222
3223
3224
3225
3226
3227
3228
3229
3230
3231
3232
3233
3234
3235
3236
3237
3238
3239
3240
3241
3242
3243
3244
3245
3246
3247
3248
3249
3250
3251
3252
3253
3254
3255
3256
3257
3258
3259
3260
3261
3262
3263
3264
3265
3266
3267
3268
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);
}

3269
3270
3271
3272
3273
ReferenceApplyAndersenThermostatKernel::~ReferenceApplyAndersenThermostatKernel() {
    if (thermostat)
        delete thermostat;
}

3274
void ReferenceApplyAndersenThermostatKernel::initialize(const System& system, const AndersenThermostat& thermostat) {
Peter Eastman's avatar
Peter Eastman committed
3275
    int numParticles = system.getNumParticles();
3276
    masses.resize(numParticles);
3277
    for (int i = 0; i < numParticles; ++i)
peastman's avatar
peastman committed
3278
        masses[i] = system.getParticleMass(i);
3279
    this->thermostat = new ReferenceAndersenThermostat();
3280
    SimTKOpenMMUtilities::setRandomNumberSeed((unsigned int) thermostat.getRandomNumberSeed());
3281
    particleGroups = AndersenThermostatImpl::calcParticleGroups(system);
3282
3283
}

3284
void ReferenceApplyAndersenThermostatKernel::execute(ContextImpl& context) {
peastman's avatar
peastman committed
3285
    vector<Vec3>& velData = extractVelocities(context);
3286
    thermostat->applyThermostat(particleGroups, velData, masses,
peastman's avatar
peastman committed
3287
3288
3289
        context.getParameter(AndersenThermostat::Temperature()),
        context.getParameter(AndersenThermostat::CollisionFrequency()),
        context.getIntegrator().getStepSize());
3290
3291
}

3292
3293
3294
3295
3296
ReferenceApplyMonteCarloBarostatKernel::~ReferenceApplyMonteCarloBarostatKernel() {
    if (barostat)
        delete barostat;
}

3297
3298
void ReferenceApplyMonteCarloBarostatKernel::initialize(const System& system, const Force& barostat, int components, bool rigidMolecules) {
    this->components = components;
3299
    this->rigidMolecules = rigidMolecules;
3300
3301
}

3302
void ReferenceApplyMonteCarloBarostatKernel::saveCoordinates(ContextImpl& context) {
3303
    if (barostat == NULL) {
3304
3305
3306
3307
        const System& system = context.getSystem();
        vector<double> masses;
        for (int i = 0; i < system.getNumParticles(); i++)
            masses.push_back(system.getParticleMass(i));
3308
        if (rigidMolecules)
3309
            barostat = new ReferenceMonteCarloBarostat(system.getNumParticles(), context.getMolecules(), masses);
3310
        else {
3311
            vector<vector<int> > molecules(system.getNumParticles());
3312
3313
            for (int i = 0; i < molecules.size(); i++)
                molecules[i].push_back(i);
3314
            barostat = new ReferenceMonteCarloBarostat(system.getNumParticles(), molecules, masses);
3315
3316
        }
    }
3317
3318
3319
3320
3321
    vector<Vec3>& posData = extractPositions(context);
    barostat->savePositions(posData);
}

void ReferenceApplyMonteCarloBarostatKernel::scaleCoordinates(ContextImpl& context, double scaleX, double scaleY, double scaleZ) {
peastman's avatar
peastman committed
3322
3323
    vector<Vec3>& posData = extractPositions(context);
    Vec3* boxVectors = extractBoxVectors(context);
3324
    barostat->applyBarostat(posData, boxVectors, scaleX, scaleY, scaleZ);
3325
3326
3327
}

void ReferenceApplyMonteCarloBarostatKernel::restoreCoordinates(ContextImpl& context) {
peastman's avatar
peastman committed
3328
    vector<Vec3>& posData = extractPositions(context);
3329
3330
3331
    barostat->restorePositions(posData);
}

3332
3333
3334
3335
void ReferenceApplyMonteCarloBarostatKernel::computeKineticEnergy(ContextImpl& context, vector<double>& ke) {
    barostat->computeMolecularKineticEnergy(extractVelocities(context), ke, components);
}

3336
3337
void ReferenceRemoveCMMotionKernel::initialize(const System& system, const CMMotionRemover& force) {
    frequency = force.getFrequency();
Peter Eastman's avatar
Peter Eastman committed
3338
    masses.resize(system.getNumParticles());
3339
    for (size_t i = 0; i < masses.size(); ++i)
Peter Eastman's avatar
Peter Eastman committed
3340
        masses[i] = system.getParticleMass(i);
3341
3342
}

3343
void ReferenceRemoveCMMotionKernel::execute(ContextImpl& context) {
3344
    if (data.stepCount%frequency != 0)
3345
        return;
peastman's avatar
peastman committed
3346
    vector<Vec3>& velData = extractVelocities(context);
3347
3348
3349
    
    // Calculate the center of mass momentum.
    
peastman's avatar
peastman committed
3350
3351
    double momentum[] = {0.0, 0.0, 0.0};
    double mass = 0.0;
3352
    for (size_t i = 0; i < masses.size(); ++i) {
peastman's avatar
peastman committed
3353
3354
3355
3356
        momentum[0] += masses[i]*velData[i][0];
        momentum[1] += masses[i]*velData[i][1];
        momentum[2] += masses[i]*velData[i][2];
        mass += masses[i];
3357
3358
    }
    
Peter Eastman's avatar
Peter Eastman committed
3359
    // Adjust the particle velocities.
3360
    
3361
3362
3363
    momentum[0] /= mass;
    momentum[1] /= mass;
    momentum[2] /= mass;
3364
    for (size_t i = 0; i < masses.size(); ++i) {
3365
3366
3367
3368
3369
        if (masses[i] != 0.0) {
            velData[i][0] -= momentum[0];
            velData[i][1] -= momentum[1];
            velData[i][2] -= momentum[2];
        }
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
3402
3403
3404
3405
3406
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");
        }
    }
}

3407
3408
3409
3410
3411
void ReferenceCalcATMForceKernel::initialize(const System& system, const ATMForce& force) {
    numParticles = force.getNumParticles();
    //displacement map
    displ1.resize(numParticles);
    displ0.resize(numParticles);
3412
3413
3414
3415
3416
3417
3418
3419
3420
3421
3422
3423
3424
3425
3426
3427
3428
3429
3430
3431
3432
3433
3434
3435
3436
3437
3438
3439
3440
3441
3442
3443
3444
3445
3446
3447
3448
3449
3450
3451
3452
3453
3454
3455
3456
    //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];
3457
3458
3459
3460
3461
3462
3463
3464
    }
}

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

3466
3467
3468
    //add gradients from variable displacements
    displForces(force0, force1);

3469
3470
3471
3472
3473
    //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)
3474
3475
3476
            force[i] += dEdu0*force0[i];
        if (fabs(dEdu1) > epsi)
            force[i] += dEdu1*force1[i];
3477
3478
    }

3479
3480
3481
3482
3483
3484
3485
3486
    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);

3487
3488
3489
    //calculate displacement vectors
    setDisplacements(pos);

3490
3491
3492
3493
    //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];
3494
    innerContext0.setPositions(pos0);
3495
3496
3497
3498
3499

    //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];
3500
    innerContext1.setPositions(pos1);
3501
3502
3503
3504
3505
3506
3507
3508
3509
3510
3511
3512
3513
3514
3515
3516
3517
3518
3519
3520
3521
3522
3523
3524
3525
3526

    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);
3527
    loadParams(numParticles, force);
3528
}
3529
3530
3531
3532
3533
3534
3535
3536
3537
3538
3539
3540
3541
3542
3543

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
3544
3545
3546
3547
3548
3549
3550
3551
3552
3553
3554
3555
3556
3557
3558
3559
3560
3561
3562
3563
3564
3565
3566
3567
3568
3569

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;
}