CpuCustomGBForce.cpp 27.3 KB
Newer Older
1

2
/* Portions copyright (c) 2009-2016 Stanford University and Simbios.
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
 * Contributors: Peter Eastman
 *
 * 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 <string.h>
#include <sstream>

#include "SimTKOpenMMUtilities.h"
#include "ReferenceForce.h"
#include "CpuCustomGBForce.h"
31
#include "openmm/internal/gmx_atomic.h"
32

33
34
using namespace OpenMM;
using namespace std;
35

peastman's avatar
peastman committed
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
class CpuCustomGBForce::ComputeForceTask : public ThreadPool::Task {
public:
    ComputeForceTask(CpuCustomGBForce& owner) : owner(owner) {
    }
    void execute(ThreadPool& threads, int threadIndex) {
        owner.threadComputeForce(threads, threadIndex);
    }
    CpuCustomGBForce& owner;
};

CpuCustomGBForce::ThreadData::ThreadData(int numAtoms, int numThreads, int threadIndex,
                      const vector<Lepton::CompiledExpression>& valueExpressions,
                      const vector<vector<Lepton::CompiledExpression> >& valueDerivExpressions,
                      const vector<vector<Lepton::CompiledExpression> >& valueGradientExpressions,
                      const vector<string>& valueNames,
                      const vector<Lepton::CompiledExpression>& energyExpressions,
                      const vector<vector<Lepton::CompiledExpression> >& energyDerivExpressions,
                      const vector<vector<Lepton::CompiledExpression> >& energyGradientExpressions,
                      const vector<string>& parameterNames) :
            valueExpressions(valueExpressions), valueDerivExpressions(valueDerivExpressions), valueGradientExpressions(valueGradientExpressions),
            energyExpressions(energyExpressions), energyDerivExpressions(energyDerivExpressions), energyGradientExpressions(energyGradientExpressions) {
    firstAtom = (threadIndex*(long long) numAtoms)/numThreads;
    lastAtom = ((threadIndex+1)*(long long) numAtoms)/numThreads;
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
    for (int i = 0; i < (int) valueExpressions.size(); i++)
        expressionSet.registerExpression(this->valueExpressions[i]);
    for (int i = 0; i < (int) valueDerivExpressions.size(); i++)
        for (int j = 0; j < (int) valueDerivExpressions[i].size(); j++)
            expressionSet.registerExpression(this->valueDerivExpressions[i][j]);
    for (int i = 0; i < (int) valueGradientExpressions.size(); i++)
        for (int j = 0; j < (int) valueGradientExpressions[i].size(); j++)
            expressionSet.registerExpression(this->valueGradientExpressions[i][j]);
    for (int i = 0; i < (int) energyExpressions.size(); i++)
        expressionSet.registerExpression(this->energyExpressions[i]);
    for (int i = 0; i < (int) energyDerivExpressions.size(); i++)
        for (int j = 0; j < (int) energyDerivExpressions[i].size(); j++)
            expressionSet.registerExpression(this->energyDerivExpressions[i][j]);
    for (int i = 0; i < (int) energyGradientExpressions.size(); i++)
        for (int j = 0; j < (int) energyGradientExpressions[i].size(); j++)
            expressionSet.registerExpression(this->energyGradientExpressions[i][j]);
    xindex = expressionSet.getVariableIndex("x");
    yindex = expressionSet.getVariableIndex("y");
    zindex = expressionSet.getVariableIndex("z");
    rindex = expressionSet.getVariableIndex("r");
peastman's avatar
peastman committed
79
80
    for (int i = 0; i < (int) parameterNames.size(); i++) {
        paramIndex.push_back(expressionSet.getVariableIndex(parameterNames[i]));
81
82
        for (int j = 1; j < 3; j++) {
            stringstream name;
peastman's avatar
peastman committed
83
            name << parameterNames[i] << j;
84
85
86
87
88
89
90
91
92
93
94
            particleParamIndex.push_back(expressionSet.getVariableIndex(name.str()));
        }
    }
    for (int i = 0; i < (int) valueNames.size(); i++) {
        valueIndex.push_back(expressionSet.getVariableIndex(valueNames[i]));
        for (int j = 1; j < 3; j++) {
            stringstream name;
            name << valueNames[i] << j;
            particleValueIndex.push_back(expressionSet.getVariableIndex(name.str()));
        }
    }
peastman's avatar
peastman committed
95
96
97
    value0.resize(numAtoms);
    dEdV.resize(valueNames.size());
    for (int i = 0; i < (int) dEdV.size(); i++)
98
99
100
101
102
103
        dEdV[i].resize(numAtoms);
    dVdX.resize(valueDerivExpressions.size());
    dVdY.resize(valueDerivExpressions.size());
    dVdZ.resize(valueDerivExpressions.size());
    dVdR1.resize(valueDerivExpressions.size());
    dVdR2.resize(valueDerivExpressions.size());
104
105
}

peastman's avatar
peastman committed
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
CpuCustomGBForce::CpuCustomGBForce(int numAtoms, const std::vector<std::set<int> >& exclusions,
                     const vector<Lepton::CompiledExpression>& valueExpressions,
                     const vector<vector<Lepton::CompiledExpression> >& valueDerivExpressions,
                     const vector<vector<Lepton::CompiledExpression> >& valueGradientExpressions,
                     const vector<string>& valueNames,
                     const vector<CustomGBForce::ComputationType>& valueTypes,
                     const vector<Lepton::CompiledExpression>& energyExpressions,
                     const vector<vector<Lepton::CompiledExpression> >& energyDerivExpressions,
                     const vector<vector<Lepton::CompiledExpression> >& energyGradientExpressions,
                     const vector<CustomGBForce::ComputationType>& energyTypes,
                     const vector<string>& parameterNames, ThreadPool& threads) :
            exclusions(exclusions), cutoff(false), periodic(false), valueNames(valueNames), valueTypes(valueTypes),
            energyTypes(energyTypes), paramNames(parameterNames), threads(threads) {
    for (int i = 0; i < threads.getNumThreads(); i++)
        threadData.push_back(new ThreadData(numAtoms, threads.getNumThreads(), i, valueExpressions, valueDerivExpressions, valueGradientExpressions, valueNames,
                      energyExpressions, energyDerivExpressions, energyGradientExpressions, parameterNames));
    values.resize(valueNames.size());
    dEdV.resize(valueNames.size());
    for (int i = 0; i < (int) values.size(); i++) {
        values[i].resize(numAtoms);
        dEdV[i].resize(numAtoms);
    }
}

130
CpuCustomGBForce::~CpuCustomGBForce() {
peastman's avatar
peastman committed
131
132
    for (int i = 0; i < (int) threadData.size(); i++)
        delete threadData[i];
133
134
}

135
void CpuCustomGBForce::setUseCutoff(float distance, const CpuNeighborList& neighbors) {
136
137
    cutoff = true;
    cutoffDistance = distance;
138
    cutoffDistance2 = distance*distance;
139
    neighborList = &neighbors;
140
  }
141
142
143
144
145
146
147
148
149
150
151
152
153

void CpuCustomGBForce::setPeriodic(RealVec& boxSize) {
    if (cutoff) {
        assert(boxSize[0] >= 2.0*cutoffDistance);
        assert(boxSize[1] >= 2.0*cutoffDistance);
        assert(boxSize[2] >= 2.0*cutoffDistance);
    }
    periodic = true;
    periodicBoxSize[0] = boxSize[0];
    periodicBoxSize[1] = boxSize[1];
    periodicBoxSize[2] = boxSize[2];
  }

154
void CpuCustomGBForce::calculateIxn(int numberOfAtoms, float* posq, RealOpenMM** atomParameters,
peastman's avatar
peastman committed
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
                                           map<string, double>& globalParameters, vector<AlignedArray<float> >& threadForce,
                                           bool includeForce, bool includeEnergy, double& totalEnergy) {
    // Record the parameters for the threads.
    
    this->numberOfAtoms = numberOfAtoms;
    this->posq = posq;
    this->atomParameters = atomParameters;
    this->globalParameters = &globalParameters;
    this->threadForce = &threadForce;
    this->includeForce = includeForce;
    this->includeEnergy = includeEnergy;
    threadEnergy.resize(threads.getNumThreads());
    gmx_atomic_t counter;
    this->atomicCounter = &counter;

    // Calculate the first computed value.

    ComputeForceTask task(*this);
    gmx_atomic_set(&counter, 0);
    threads.execute(task);
    threads.waitForThreads();
    
    // Calculate the remaining computed values.
    
    threads.resumeThreads();
    threads.waitForThreads();

    // Calculate the energy terms.

    for (int i = 0; i < (int) threadData[0]->energyExpressions.size(); i++) {
        gmx_atomic_set(&counter, 0);
        threads.execute(task);
        threads.waitForThreads();
    }

    // Sum the energy derivatives.

    threads.resumeThreads();
    threads.waitForThreads();
    
    // Apply the chain rule to evaluate forces.

    gmx_atomic_set(&counter, 0);
    threads.resumeThreads();
    threads.waitForThreads();

    // Combine the energies from all the threads.
    
    if (includeEnergy) {
        int numThreads = threads.getNumThreads();
        for (int i = 0; i < numThreads; i++)
            totalEnergy += threadEnergy[i];
    }
}

void CpuCustomGBForce::threadComputeForce(ThreadPool& threads, int threadIndex) {
    // Compute this thread's subset of interactions.

    int numThreads = threads.getNumThreads();
    threadEnergy[threadIndex] = 0;
    double& energy = threadEnergy[threadIndex];
    float* forces = &(*threadForce)[threadIndex][0];
    ThreadData& data = *threadData[threadIndex];
218
219
    fvec4 boxSize(periodicBoxSize[0], periodicBoxSize[1], periodicBoxSize[2], 0);
    fvec4 invBoxSize((1/periodicBoxSize[0]), (1/periodicBoxSize[1]), (1/periodicBoxSize[2]), 0);
peastman's avatar
peastman committed
220
221
222
223
224
225
226
227
228
229
230
231
232
    for (map<string, double>::const_iterator iter = globalParameters->begin(); iter != globalParameters->end(); ++iter)
        data.expressionSet.setVariable(data.expressionSet.getVariableIndex(iter->first), iter->second);

    // Calculate the first computed value.

    for (int i = 0; i < (int) data.value0.size(); i++)
        data.value0[i] = 0.0f;
    if (valueTypes[0] == CustomGBForce::ParticlePair)
        calculateParticlePairValue(0, data, numberOfAtoms, posq, atomParameters, true, boxSize, invBoxSize);
    else
        calculateParticlePairValue(0, data, numberOfAtoms, posq, atomParameters, false, boxSize, invBoxSize);
    threads.syncThreads();

233
    // Sum the first computed value and calculate the remaining ones.
peastman's avatar
peastman committed
234

235
    int numValues = valueTypes.size();
peastman's avatar
peastman committed
236
237
238
239
240
    for (int atom = data.firstAtom; atom < data.lastAtom; atom++) {
        float sum = 0.0f;
        for (int j = 0; j < (int) threadData.size(); j++)
            sum += threadData[j]->value0[atom];
        values[0][atom] = sum;
241
242
243
244
245
246
247
248
249
        data.expressionSet.setVariable(data.xindex, posq[4*atom]);
        data.expressionSet.setVariable(data.yindex, posq[4*atom+1]);
        data.expressionSet.setVariable(data.zindex, posq[4*atom+2]);
        for (int j = 0; j < (int) paramNames.size(); j++)
            data.expressionSet.setVariable(data.paramIndex[j], atomParameters[atom][j]);
        for (int i = 1; i < numValues; i++) {
            data.expressionSet.setVariable(data.valueIndex[i-1], values[i-1][atom]);
            values[i][atom] = (float) data.valueExpressions[i].evaluate();
        }
peastman's avatar
peastman committed
250
251
    }
    threads.syncThreads();
252
253
254

    // Now calculate the energy and its derivatives.

peastman's avatar
peastman committed
255
256
257
258
    for (int i = 0; i < (int) data.dEdV.size(); i++)
        for (int j = 0; j < (int) data.dEdV[i].size(); j++)
            data.dEdV[i][j] = 0.0;
    for (int termIndex = 0; termIndex < (int) data.energyExpressions.size(); termIndex++) {
259
        if (energyTypes[termIndex] == CustomGBForce::SingleParticle)
peastman's avatar
peastman committed
260
            calculateSingleParticleEnergyTerm(termIndex, data, numberOfAtoms, posq, atomParameters, forces, energy);
261
        else if (energyTypes[termIndex] == CustomGBForce::ParticlePair)
peastman's avatar
peastman committed
262
            calculateParticlePairEnergyTerm(termIndex, data, numberOfAtoms, posq, atomParameters, true, forces, energy, boxSize, invBoxSize);
263
        else
peastman's avatar
peastman committed
264
265
266
267
268
269
270
271
272
273
274
275
276
            calculateParticlePairEnergyTerm(termIndex, data, numberOfAtoms, posq, atomParameters, false, forces, energy, boxSize, invBoxSize);
        threads.syncThreads();
    }

    // Sum the energy derivatives.

    for (int atom = data.firstAtom; atom < data.lastAtom; atom++) {
        for (int i = 0; i < (int) dEdV.size(); i++) {
            float sum = 0.0f;
            for (int j = 0; j < (int) threadData.size(); j++)
                sum += threadData[j]->dEdV[i][atom];
            dEdV[i][atom] = sum;
        }
277
    }
peastman's avatar
peastman committed
278
    threads.syncThreads();
279
280
281

    // Apply the chain rule to evaluate forces.

peastman's avatar
peastman committed
282
    calculateChainRuleForces(data, numberOfAtoms, posq, atomParameters, forces, boxSize, invBoxSize);
283
284
}

peastman's avatar
peastman committed
285
286
void CpuCustomGBForce::calculateParticlePairValue(int index, ThreadData& data, int numAtoms, float* posq, RealOpenMM** atomParameters,
        bool useExclusions, const fvec4& boxSize, const fvec4& invBoxSize) {
287
    for (int i = 0; i < numAtoms; i++)
288
        values[index][i] = 0.0f;
peastman's avatar
peastman committed
289
    vector<float>& valueArray = (index == 0 ? data.value0 : values[index]);
290
291
292
    if (cutoff) {
        // Loop over all pairs in the neighbor list.

peastman's avatar
peastman committed
293
294
295
296
        while (true) {
            int blockIndex = gmx_atomic_fetch_add(reinterpret_cast<gmx_atomic_t*>(atomicCounter), 1);
            if (blockIndex >= neighborList->getNumBlocks())
                break;
297
298
            const int blockSize = neighborList->getBlockSize();
            const int* blockAtom = &neighborList->getSortedAtoms()[blockSize*blockIndex];
299
300
301
302
            const vector<int>& neighbors = neighborList->getBlockNeighbors(blockIndex);
            const vector<char>& blockExclusions = neighborList->getBlockExclusions(blockIndex);
            for (int i = 0; i < (int) neighbors.size(); i++) {
                int first = neighbors[i];
303
                for (int k = 0; k < blockSize; k++) {
304
305
306
307
                    if ((blockExclusions[i] & (1<<k)) == 0) {
                        int second = blockAtom[k];
                        if (useExclusions && exclusions[first].find(second) != exclusions[first].end())
                            continue;
peastman's avatar
peastman committed
308
309
                        calculateOnePairValue(index, first, second, data, posq, atomParameters, valueArray, boxSize, invBoxSize);
                        calculateOnePairValue(index, second, first, data, posq, atomParameters, valueArray, boxSize, invBoxSize);
310
311
312
                    }
                }
            }
313
314
315
316
317
        }
    }
    else {
        // Perform an O(N^2) loop over all atom pairs.

peastman's avatar
peastman committed
318
319
320
321
        while (true) {
            int i = gmx_atomic_fetch_add(reinterpret_cast<gmx_atomic_t*>(atomicCounter), 1);
            if (i >= numAtoms)
                break;
322
323
324
            for (int j = i+1; j < numAtoms; j++) {
                if (useExclusions && exclusions[i].find(j) != exclusions[i].end())
                    continue;
peastman's avatar
peastman committed
325
326
                calculateOnePairValue(index, i, j, data, posq, atomParameters, valueArray, boxSize, invBoxSize);
                calculateOnePairValue(index, j, i, data, posq, atomParameters, valueArray, boxSize, invBoxSize);
327
328
329
330
331
           }
        }
    }
}

peastman's avatar
peastman committed
332
333
void CpuCustomGBForce::calculateOnePairValue(int index, int atom1, int atom2, ThreadData& data, float* posq, RealOpenMM** atomParameters,
        vector<float>& valueArray, const fvec4& boxSize, const fvec4& invBoxSize) {
334
335
336
337
338
    fvec4 deltaR;
    fvec4 pos1(posq+4*atom1);
    fvec4 pos2(posq+4*atom2);
    float r2;
    getDeltaR(pos2, pos1, deltaR, r2, periodic, boxSize, invBoxSize);
339
    if (cutoff && r2 >= cutoffDistance2)
340
        return;
341
    float r = sqrtf(r2);
342
    for (int i = 0; i < (int) paramNames.size(); i++) {
peastman's avatar
peastman committed
343
344
        data.expressionSet.setVariable(data.particleParamIndex[i*2], atomParameters[atom1][i]);
        data.expressionSet.setVariable(data.particleParamIndex[i*2+1], atomParameters[atom2][i]);
345
    }
peastman's avatar
peastman committed
346
    data.expressionSet.setVariable(data.rindex, r);
347
    for (int i = 0; i < index; i++) {
peastman's avatar
peastman committed
348
349
        data.expressionSet.setVariable(data.particleValueIndex[i*2], values[i][atom1]);
        data.expressionSet.setVariable(data.particleValueIndex[i*2+1], values[i][atom2]);
350
    }
peastman's avatar
peastman committed
351
    valueArray[atom1] += (float) data.valueExpressions[index].evaluate();
352
353
}

peastman's avatar
peastman committed
354
355
356
357
358
359
void CpuCustomGBForce::calculateSingleParticleEnergyTerm(int index, ThreadData& data, int numAtoms, float* posq,
        RealOpenMM** atomParameters, float* forces, double& totalEnergy) {
    for (int i = data.firstAtom; i < data.lastAtom; i++) {
        data.expressionSet.setVariable(data.xindex, posq[4*i]);
        data.expressionSet.setVariable(data.yindex, posq[4*i+1]);
        data.expressionSet.setVariable(data.zindex, posq[4*i+2]);
360
        for (int j = 0; j < (int) paramNames.size(); j++)
peastman's avatar
peastman committed
361
            data.expressionSet.setVariable(data.paramIndex[j], atomParameters[i][j]);
362
        for (int j = 0; j < (int) valueNames.size(); j++)
peastman's avatar
peastman committed
363
364
365
            data.expressionSet.setVariable(data.valueIndex[j], values[j][i]);
        if (includeEnergy)
            totalEnergy += (float) data.energyExpressions[index].evaluate();
366
        for (int j = 0; j < (int) valueNames.size(); j++)
peastman's avatar
peastman committed
367
368
369
370
            data.dEdV[j][i] += (float) data.energyDerivExpressions[index][j].evaluate();
        forces[4*i+0] -= (float) data.energyGradientExpressions[index][0].evaluate();
        forces[4*i+1] -= (float) data.energyGradientExpressions[index][1].evaluate();
        forces[4*i+2] -= (float) data.energyGradientExpressions[index][2].evaluate();
371
372
373
    }
}

peastman's avatar
peastman committed
374
375
void CpuCustomGBForce::calculateParticlePairEnergyTerm(int index, ThreadData& data, int numAtoms, float* posq, RealOpenMM** atomParameters,
        bool useExclusions, float* forces, double& totalEnergy, const fvec4& boxSize, const fvec4& invBoxSize) {
376
377
378
    if (cutoff) {
        // Loop over all pairs in the neighbor list.

peastman's avatar
peastman committed
379
380
381
382
        while (true) {
            int blockIndex = gmx_atomic_fetch_add(reinterpret_cast<gmx_atomic_t*>(atomicCounter), 1);
            if (blockIndex >= neighborList->getNumBlocks())
                break;
383
384
            const int blockSize = neighborList->getBlockSize();
            const int* blockAtom = &neighborList->getSortedAtoms()[blockSize*blockIndex];
385
386
387
388
            const vector<int>& neighbors = neighborList->getBlockNeighbors(blockIndex);
            const vector<char>& blockExclusions = neighborList->getBlockExclusions(blockIndex);
            for (int i = 0; i < (int) neighbors.size(); i++) {
                int first = neighbors[i];
389
                for (int k = 0; k < blockSize; k++) {
390
391
392
393
                    if ((blockExclusions[i] & (1<<k)) == 0) {
                        int second = blockAtom[k];
                        if (useExclusions && exclusions[first].find(second) != exclusions[first].end())
                            continue;
peastman's avatar
peastman committed
394
                        calculateOnePairEnergyTerm(index, first, second, data, posq, atomParameters, forces, totalEnergy, boxSize, invBoxSize);
395
396
397
                    }
                }
            }
398
399
400
401
402
        }
    }
    else {
        // Perform an O(N^2) loop over all atom pairs.

peastman's avatar
peastman committed
403
404
405
406
        while (true) {
            int i = gmx_atomic_fetch_add(reinterpret_cast<gmx_atomic_t*>(atomicCounter), 1);
            if (i >= numAtoms)
                break;
407
408
409
            for (int j = i+1; j < numAtoms; j++) {
                if (useExclusions && exclusions[i].find(j) != exclusions[i].end())
                    continue;
peastman's avatar
peastman committed
410
                calculateOnePairEnergyTerm(index, i, j, data, posq, atomParameters, forces, totalEnergy, boxSize, invBoxSize);
411
412
413
414
415
           }
        }
    }
}

peastman's avatar
peastman committed
416
417
void CpuCustomGBForce::calculateOnePairEnergyTerm(int index, int atom1, int atom2, ThreadData& data, float* posq, RealOpenMM** atomParameters,
        float* forces, double& totalEnergy, const fvec4& boxSize, const fvec4& invBoxSize) {
418
419
    // Compute the displacement.

420
421
422
423
424
    fvec4 deltaR;
    fvec4 pos1(posq+4*atom1);
    fvec4 pos2(posq+4*atom2);
    float r2;
    getDeltaR(pos2, pos1, deltaR, r2, periodic, boxSize, invBoxSize);
425
    if (cutoff && r2 >= cutoffDistance2)
426
        return;
427
    float r = sqrtf(r2);
428
429
430
431

    // Record variables for evaluating expressions.

    for (int i = 0; i < (int) paramNames.size(); i++) {
peastman's avatar
peastman committed
432
433
        data.expressionSet.setVariable(data.particleParamIndex[i*2], atomParameters[atom1][i]);
        data.expressionSet.setVariable(data.particleParamIndex[i*2+1], atomParameters[atom2][i]);
434
    }
peastman's avatar
peastman committed
435
    data.expressionSet.setVariable(data.rindex, r);
436
    for (int i = 0; i < (int) valueNames.size(); i++) {
peastman's avatar
peastman committed
437
438
        data.expressionSet.setVariable(data.particleValueIndex[i*2], values[i][atom1]);
        data.expressionSet.setVariable(data.particleValueIndex[i*2+1], values[i][atom2]);
439
440
441
442
    }

    // Evaluate the energy and its derivatives.

peastman's avatar
peastman committed
443
444
445
    if (includeEnergy)
        totalEnergy += (float) data.energyExpressions[index].evaluate();
    float dEdR = (float) data.energyDerivExpressions[index][0].evaluate();
446
    dEdR *= 1/r;
447
448
449
    fvec4 result = deltaR*dEdR;
    (fvec4(forces+4*atom1)-result).store(forces+4*atom1);
    (fvec4(forces+4*atom2)+result).store(forces+4*atom2);
450
    for (int i = 0; i < (int) valueNames.size(); i++) {
peastman's avatar
peastman committed
451
452
        data.dEdV[i][atom1] += (float) data.energyDerivExpressions[index][2*i+1].evaluate();
        data.dEdV[i][atom2] += (float) data.energyDerivExpressions[index][2*i+2].evaluate();
453
454
455
    }
}

peastman's avatar
peastman committed
456
457
void CpuCustomGBForce::calculateChainRuleForces(ThreadData& data, int numAtoms, float* posq, RealOpenMM** atomParameters,
        float* forces, const fvec4& boxSize, const fvec4& invBoxSize) {
458
459
460
    if (cutoff) {
        // Loop over all pairs in the neighbor list.

peastman's avatar
peastman committed
461
462
463
464
        while (true) {
            int blockIndex = gmx_atomic_fetch_add(reinterpret_cast<gmx_atomic_t*>(atomicCounter), 1);
            if (blockIndex >= neighborList->getNumBlocks())
                break;
465
466
            const int blockSize = neighborList->getBlockSize();
            const int* blockAtom = &neighborList->getSortedAtoms()[blockSize*blockIndex];
467
468
469
470
            const vector<int>& neighbors = neighborList->getBlockNeighbors(blockIndex);
            const vector<char>& blockExclusions = neighborList->getBlockExclusions(blockIndex);
            for (int i = 0; i < (int) neighbors.size(); i++) {
                int first = neighbors[i];
471
                for (int k = 0; k < blockSize; k++) {
472
473
474
                    if ((blockExclusions[i] & (1<<k)) == 0) {
                        int second = blockAtom[k];
                        bool isExcluded = (exclusions[first].find(second) != exclusions[first].end());
peastman's avatar
peastman committed
475
476
                        calculateOnePairChainRule(first, second, data, posq, atomParameters, forces, isExcluded, boxSize, invBoxSize);
                        calculateOnePairChainRule(second, first, data, posq, atomParameters, forces, isExcluded, boxSize, invBoxSize);
477
478
479
                    }
                }
            }
480
481
482
483
484
        }
    }
    else {
        // Perform an O(N^2) loop over all atom pairs.

peastman's avatar
peastman committed
485
486
487
488
        while (true) {
            int i = gmx_atomic_fetch_add(reinterpret_cast<gmx_atomic_t*>(atomicCounter), 1);
            if (i >= numAtoms)
                break;
489
490
            for (int j = i+1; j < numAtoms; j++) {
                bool isExcluded = (exclusions[i].find(j) != exclusions[i].end());
peastman's avatar
peastman committed
491
492
                calculateOnePairChainRule(i, j, data, posq, atomParameters, forces, isExcluded, boxSize, invBoxSize);
                calculateOnePairChainRule(j, i, data, posq, atomParameters, forces, isExcluded, boxSize, invBoxSize);
493
494
495
496
497
498
           }
        }
    }

    // Compute chain rule terms for computed values that depend explicitly on particle coordinates.

peastman's avatar
peastman committed
499
500
501
502
    for (int i = data.firstAtom; i < data.lastAtom; i++) {
        data.expressionSet.setVariable(data.xindex, posq[4*i]);
        data.expressionSet.setVariable(data.yindex, posq[4*i+1]);
        data.expressionSet.setVariable(data.zindex, posq[4*i+2]);
503
        for (int j = 0; j < (int) paramNames.size(); j++)
peastman's avatar
peastman committed
504
            data.expressionSet.setVariable(data.paramIndex[j], atomParameters[i][j]);
505
        for (int j = 1; j < (int) valueNames.size(); j++) {
peastman's avatar
peastman committed
506
507
508
509
            data.expressionSet.setVariable(data.valueIndex[j-1], values[j-1][i]);
            data.dVdX[j] = 0.0;
            data.dVdY[j] = 0.0;
            data.dVdZ[j] = 0.0;
510
            for (int k = 1; k < j; k++) {
peastman's avatar
peastman committed
511
512
513
514
                float dVdV = (float) data.valueDerivExpressions[j][k].evaluate();
                data.dVdX[j] += dVdV*data.dVdX[k];
                data.dVdY[j] += dVdV*data.dVdY[k];
                data.dVdZ[j] += dVdV*data.dVdZ[k];
515
            }
peastman's avatar
peastman committed
516
517
518
519
520
521
            data.dVdX[j] += (float) data.valueGradientExpressions[j][0].evaluate();
            data.dVdY[j] += (float) data.valueGradientExpressions[j][1].evaluate();
            data.dVdZ[j] += (float) data.valueGradientExpressions[j][2].evaluate();
            forces[4*i+0] -= dEdV[j][i]*data.dVdX[j];
            forces[4*i+1] -= dEdV[j][i]*data.dVdY[j];
            forces[4*i+2] -= dEdV[j][i]*data.dVdZ[j];
522
523
524
525
        }
    }
}

peastman's avatar
peastman committed
526
527
void CpuCustomGBForce::calculateOnePairChainRule(int atom1, int atom2, ThreadData& data, float* posq, RealOpenMM** atomParameters,
        float* forces, bool isExcluded, const fvec4& boxSize, const fvec4& invBoxSize) {
528
529
    // Compute the displacement.

530
531
532
533
534
    fvec4 deltaR;
    fvec4 pos1(posq+4*atom1);
    fvec4 pos2(posq+4*atom2);
    float r2;
    getDeltaR(pos2, pos1, deltaR, r2, periodic, boxSize, invBoxSize);
535
    if (cutoff && r2 >= cutoffDistance2)
536
        return;
537
    float r = sqrtf(r2);
538
539
540
541

    // Record variables for evaluating expressions.

    for (int i = 0; i < (int) paramNames.size(); i++) {
peastman's avatar
peastman committed
542
543
        data.expressionSet.setVariable(data.particleParamIndex[i*2], atomParameters[atom1][i]);
        data.expressionSet.setVariable(data.particleParamIndex[i*2+1], atomParameters[atom2][i]);
544
        data.expressionSet.setVariable(data.paramIndex[i], atomParameters[atom1][i]);
545
    }
546
547
548
549
    data.expressionSet.setVariable(data.valueIndex[0], values[0][atom1]);
    data.expressionSet.setVariable(data.xindex, posq[4*atom1]);
    data.expressionSet.setVariable(data.yindex, posq[4*atom1+1]);
    data.expressionSet.setVariable(data.zindex, posq[4*atom1+2]);
peastman's avatar
peastman committed
550
551
552
    data.expressionSet.setVariable(data.rindex, r);
    data.expressionSet.setVariable(data.particleValueIndex[0], values[0][atom1]);
    data.expressionSet.setVariable(data.particleValueIndex[1], values[0][atom2]);
553
554
555

    // Evaluate the derivative of each parameter with respect to position and apply forces.

556
557
    float rinv = 1/r;
    deltaR *= rinv;
558
    fvec4 f1(0.0f), f2(0.0f);
559
    if (!isExcluded || valueTypes[0] != CustomGBForce::ParticlePair) {
peastman's avatar
peastman committed
560
561
        data.dVdR1[0] = (float) data.valueDerivExpressions[0][0].evaluate();
        data.dVdR2[0] = -data.dVdR1[0];
562
563
        f1 -= deltaR*(dEdV[0][atom1]*data.dVdR1[0]);
        f2 -= deltaR*(dEdV[0][atom1]*data.dVdR2[0]);
564
565
    }
    for (int i = 1; i < (int) valueNames.size(); i++) {
peastman's avatar
peastman committed
566
567
568
        data.expressionSet.setVariable(data.valueIndex[i], values[i][atom1]);
        data.dVdR1[i] = 0.0;
        data.dVdR2[i] = 0.0;
569
        for (int j = 0; j < i; j++) {
peastman's avatar
peastman committed
570
571
572
            float dVdV = (float) data.valueDerivExpressions[i][j].evaluate();
            data.dVdR1[i] += dVdV*data.dVdR1[j];
            data.dVdR2[i] += dVdV*data.dVdR2[j];
573
        }
574
575
        f1 -= deltaR*(dEdV[i][atom1]*data.dVdR1[i]);
        f2 -= deltaR*(dEdV[i][atom1]*data.dVdR2[i]);
576
    }
577
578
    (fvec4(forces+4*atom1)+f1).store(forces+4*atom1);
    (fvec4(forces+4*atom2)+f2).store(forces+4*atom2);
579
580
581
582
583
584
585
}

void CpuCustomGBForce::getDeltaR(const fvec4& posI, const fvec4& posJ, fvec4& deltaR, float& r2, bool periodic, const fvec4& boxSize, const fvec4& invBoxSize) const {
    deltaR = posJ-posI;
    if (periodic) {
        fvec4 base = round(deltaR*invBoxSize)*boxSize;
        deltaR = deltaR-base;
586
    }
587
    r2 = dot3(deltaR, deltaR);
588
}