CpuCustomGBForce.cpp 27.5 KB
Newer Older
1
2
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
31
32

/* Portions copyright (c) 2009-2014 Stanford University and Simbios.
 * 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 "SimTKOpenMMCommon.h"
#include "SimTKOpenMMLog.h"
#include "SimTKOpenMMUtilities.h"
#include "ReferenceForce.h"
#include "CpuCustomGBForce.h"
peastman's avatar
peastman committed
33
#include "gmx_atomic.h"
34

35
36
using namespace OpenMM;
using namespace std;
37

peastman's avatar
peastman committed
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
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;
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
    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
81
82
    for (int i = 0; i < (int) parameterNames.size(); i++) {
        paramIndex.push_back(expressionSet.getVariableIndex(parameterNames[i]));
83
84
        for (int j = 1; j < 3; j++) {
            stringstream name;
peastman's avatar
peastman committed
85
            name << parameterNames[i] << j;
86
87
88
89
90
91
92
93
94
95
96
            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
97
98
99
    value0.resize(numAtoms);
    dEdV.resize(valueNames.size());
    for (int i = 0; i < (int) dEdV.size(); i++)
100
101
102
103
104
105
        dEdV[i].resize(numAtoms);
    dVdX.resize(valueDerivExpressions.size());
    dVdY.resize(valueDerivExpressions.size());
    dVdZ.resize(valueDerivExpressions.size());
    dVdR1.resize(valueDerivExpressions.size());
    dVdR2.resize(valueDerivExpressions.size());
106
107
}

peastman's avatar
peastman committed
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
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);
    }
}

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

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

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

155
void CpuCustomGBForce::calculateIxn(int numberOfAtoms, float* posq, RealOpenMM** atomParameters,
peastman's avatar
peastman committed
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
218
                                           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];
219
220
    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
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
    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();

    // Sum the first computed value.

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

peastman's avatar
peastman committed
243
    // Calculate the remaining computed values.
244
245

    int numValues = valueTypes.size();
peastman's avatar
peastman committed
246
247
248
    for (int i = 1; i < numValues; i++)
        calculateSingleParticleValue(i, data, numberOfAtoms, posq, atomParameters);
    threads.syncThreads();
249
250
251

    // Now calculate the energy and its derivatives.

peastman's avatar
peastman committed
252
253
254
255
    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++) {
256
        if (energyTypes[termIndex] == CustomGBForce::SingleParticle)
peastman's avatar
peastman committed
257
            calculateSingleParticleEnergyTerm(termIndex, data, numberOfAtoms, posq, atomParameters, forces, energy);
258
        else if (energyTypes[termIndex] == CustomGBForce::ParticlePair)
peastman's avatar
peastman committed
259
            calculateParticlePairEnergyTerm(termIndex, data, numberOfAtoms, posq, atomParameters, true, forces, energy, boxSize, invBoxSize);
260
        else
peastman's avatar
peastman committed
261
262
263
264
265
266
267
268
269
270
271
272
273
            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;
        }
274
    }
peastman's avatar
peastman committed
275
    threads.syncThreads();
276
277
278

    // Apply the chain rule to evaluate forces.

peastman's avatar
peastman committed
279
    calculateChainRuleForces(data, numberOfAtoms, posq, atomParameters, forces, boxSize, invBoxSize);
280
281
}

peastman's avatar
peastman committed
282
283
284
285
286
void CpuCustomGBForce::calculateSingleParticleValue(int index, ThreadData& data, int numAtoms, float* posq, RealOpenMM** atomParameters) {
    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]);
287
        for (int j = 0; j < (int) paramNames.size(); j++)
peastman's avatar
peastman committed
288
            data.expressionSet.setVariable(data.paramIndex[j], atomParameters[i][j]);
289
        for (int j = 0; j < index; j++)
peastman's avatar
peastman committed
290
291
            data.expressionSet.setVariable(data.valueIndex[j], values[j][i]);
        values[index][i] = (float) data.valueExpressions[index].evaluate();
292
293
294
    }
}

peastman's avatar
peastman committed
295
296
void CpuCustomGBForce::calculateParticlePairValue(int index, ThreadData& data, int numAtoms, float* posq, RealOpenMM** atomParameters,
        bool useExclusions, const fvec4& boxSize, const fvec4& invBoxSize) {
297
    for (int i = 0; i < numAtoms; i++)
298
        values[index][i] = 0.0f;
peastman's avatar
peastman committed
299
    vector<float>& valueArray = (index == 0 ? data.value0 : values[index]);
300
301
302
    if (cutoff) {
        // Loop over all pairs in the neighbor list.

peastman's avatar
peastman committed
303
304
305
306
        while (true) {
            int blockIndex = gmx_atomic_fetch_add(reinterpret_cast<gmx_atomic_t*>(atomicCounter), 1);
            if (blockIndex >= neighborList->getNumBlocks())
                break;
307
308
309
310
311
312
313
314
315
316
            const int* blockAtom = &neighborList->getSortedAtoms()[4*blockIndex];
            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];
                for (int k = 0; k < 4; k++) {
                    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
317
318
                        calculateOnePairValue(index, first, second, data, posq, atomParameters, valueArray, boxSize, invBoxSize);
                        calculateOnePairValue(index, second, first, data, posq, atomParameters, valueArray, boxSize, invBoxSize);
319
320
321
                    }
                }
            }
322
323
324
325
326
        }
    }
    else {
        // Perform an O(N^2) loop over all atom pairs.

peastman's avatar
peastman committed
327
328
329
330
        while (true) {
            int i = gmx_atomic_fetch_add(reinterpret_cast<gmx_atomic_t*>(atomicCounter), 1);
            if (i >= numAtoms)
                break;
331
332
333
            for (int j = i+1; j < numAtoms; j++) {
                if (useExclusions && exclusions[i].find(j) != exclusions[i].end())
                    continue;
peastman's avatar
peastman committed
334
335
                calculateOnePairValue(index, i, j, data, posq, atomParameters, valueArray, boxSize, invBoxSize);
                calculateOnePairValue(index, j, i, data, posq, atomParameters, valueArray, boxSize, invBoxSize);
336
337
338
339
340
           }
        }
    }
}

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

peastman's avatar
peastman committed
363
364
365
366
367
368
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]);
369
        for (int j = 0; j < (int) paramNames.size(); j++)
peastman's avatar
peastman committed
370
            data.expressionSet.setVariable(data.paramIndex[j], atomParameters[i][j]);
371
        for (int j = 0; j < (int) valueNames.size(); j++)
peastman's avatar
peastman committed
372
373
374
            data.expressionSet.setVariable(data.valueIndex[j], values[j][i]);
        if (includeEnergy)
            totalEnergy += (float) data.energyExpressions[index].evaluate();
375
        for (int j = 0; j < (int) valueNames.size(); j++)
peastman's avatar
peastman committed
376
377
378
379
            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();
380
381
382
    }
}

peastman's avatar
peastman committed
383
384
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) {
385
386
387
    if (cutoff) {
        // Loop over all pairs in the neighbor list.

peastman's avatar
peastman committed
388
389
390
391
        while (true) {
            int blockIndex = gmx_atomic_fetch_add(reinterpret_cast<gmx_atomic_t*>(atomicCounter), 1);
            if (blockIndex >= neighborList->getNumBlocks())
                break;
392
393
394
395
396
397
398
399
400
401
            const int* blockAtom = &neighborList->getSortedAtoms()[4*blockIndex];
            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];
                for (int k = 0; k < 4; k++) {
                    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
402
                        calculateOnePairEnergyTerm(index, first, second, data, posq, atomParameters, forces, totalEnergy, boxSize, invBoxSize);
403
404
405
                    }
                }
            }
406
407
408
409
410
        }
    }
    else {
        // Perform an O(N^2) loop over all atom pairs.

peastman's avatar
peastman committed
411
412
413
414
        while (true) {
            int i = gmx_atomic_fetch_add(reinterpret_cast<gmx_atomic_t*>(atomicCounter), 1);
            if (i >= numAtoms)
                break;
415
416
417
            for (int j = i+1; j < numAtoms; j++) {
                if (useExclusions && exclusions[i].find(j) != exclusions[i].end())
                    continue;
peastman's avatar
peastman committed
418
                calculateOnePairEnergyTerm(index, i, j, data, posq, atomParameters, forces, totalEnergy, boxSize, invBoxSize);
419
420
421
422
423
           }
        }
    }
}

peastman's avatar
peastman committed
424
425
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) {
426
427
    // Compute the displacement.

428
429
430
431
432
433
    fvec4 deltaR;
    fvec4 pos1(posq+4*atom1);
    fvec4 pos2(posq+4*atom2);
    float r2;
    getDeltaR(pos2, pos1, deltaR, r2, periodic, boxSize, invBoxSize);
    float r = sqrtf(r2);
434
435
436
437
438
439
    if (cutoff && r >= cutoffDistance)
        return;

    // Record variables for evaluating expressions.

    for (int i = 0; i < (int) paramNames.size(); i++) {
peastman's avatar
peastman committed
440
441
        data.expressionSet.setVariable(data.particleParamIndex[i*2], atomParameters[atom1][i]);
        data.expressionSet.setVariable(data.particleParamIndex[i*2+1], atomParameters[atom2][i]);
442
    }
peastman's avatar
peastman committed
443
    data.expressionSet.setVariable(data.rindex, r);
444
    for (int i = 0; i < (int) valueNames.size(); i++) {
peastman's avatar
peastman committed
445
446
        data.expressionSet.setVariable(data.particleValueIndex[i*2], values[i][atom1]);
        data.expressionSet.setVariable(data.particleValueIndex[i*2+1], values[i][atom2]);
447
448
449
450
    }

    // Evaluate the energy and its derivatives.

peastman's avatar
peastman committed
451
452
453
    if (includeEnergy)
        totalEnergy += (float) data.energyExpressions[index].evaluate();
    float dEdR = (float) data.energyDerivExpressions[index][0].evaluate();
454
    dEdR *= 1/r;
455
456
457
    fvec4 result = deltaR*dEdR;
    (fvec4(forces+4*atom1)-result).store(forces+4*atom1);
    (fvec4(forces+4*atom2)+result).store(forces+4*atom2);
458
    for (int i = 0; i < (int) valueNames.size(); i++) {
peastman's avatar
peastman committed
459
460
        data.dEdV[i][atom1] += (float) data.energyDerivExpressions[index][2*i+1].evaluate();
        data.dEdV[i][atom2] += (float) data.energyDerivExpressions[index][2*i+2].evaluate();
461
462
463
    }
}

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

peastman's avatar
peastman committed
469
470
471
472
        while (true) {
            int blockIndex = gmx_atomic_fetch_add(reinterpret_cast<gmx_atomic_t*>(atomicCounter), 1);
            if (blockIndex >= neighborList->getNumBlocks())
                break;
473
474
475
476
477
478
479
480
481
            const int* blockAtom = &neighborList->getSortedAtoms()[4*blockIndex];
            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];
                for (int k = 0; k < 4; k++) {
                    if ((blockExclusions[i] & (1<<k)) == 0) {
                        int second = blockAtom[k];
                        bool isExcluded = (exclusions[first].find(second) != exclusions[first].end());
peastman's avatar
peastman committed
482
483
                        calculateOnePairChainRule(first, second, data, posq, atomParameters, forces, isExcluded, boxSize, invBoxSize);
                        calculateOnePairChainRule(second, first, data, posq, atomParameters, forces, isExcluded, boxSize, invBoxSize);
484
485
486
                    }
                }
            }
487
488
489
490
491
        }
    }
    else {
        // Perform an O(N^2) loop over all atom pairs.

peastman's avatar
peastman committed
492
493
494
495
        while (true) {
            int i = gmx_atomic_fetch_add(reinterpret_cast<gmx_atomic_t*>(atomicCounter), 1);
            if (i >= numAtoms)
                break;
496
497
            for (int j = i+1; j < numAtoms; j++) {
                bool isExcluded = (exclusions[i].find(j) != exclusions[i].end());
peastman's avatar
peastman committed
498
499
                calculateOnePairChainRule(i, j, data, posq, atomParameters, forces, isExcluded, boxSize, invBoxSize);
                calculateOnePairChainRule(j, i, data, posq, atomParameters, forces, isExcluded, boxSize, invBoxSize);
500
501
502
503
504
505
           }
        }
    }

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

peastman's avatar
peastman committed
506
507
508
509
    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]);
510
        for (int j = 0; j < (int) paramNames.size(); j++)
peastman's avatar
peastman committed
511
            data.expressionSet.setVariable(data.paramIndex[j], atomParameters[i][j]);
512
        for (int j = 1; j < (int) valueNames.size(); j++) {
peastman's avatar
peastman committed
513
514
515
516
            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;
517
            for (int k = 1; k < j; k++) {
peastman's avatar
peastman committed
518
519
520
521
                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];
522
            }
peastman's avatar
peastman committed
523
524
525
526
527
528
            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];
529
530
531
532
        }
    }
}

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

537
538
539
540
541
542
    fvec4 deltaR;
    fvec4 pos1(posq+4*atom1);
    fvec4 pos2(posq+4*atom2);
    float r2;
    getDeltaR(pos2, pos1, deltaR, r2, periodic, boxSize, invBoxSize);
    float r = sqrtf(r2);
543
544
545
546
547
548
    if (cutoff && r >= cutoffDistance)
        return;

    // Record variables for evaluating expressions.

    for (int i = 0; i < (int) paramNames.size(); i++) {
peastman's avatar
peastman committed
549
550
        data.expressionSet.setVariable(data.particleParamIndex[i*2], atomParameters[atom1][i]);
        data.expressionSet.setVariable(data.particleParamIndex[i*2+1], atomParameters[atom2][i]);
551
    }
peastman's avatar
peastman committed
552
553
554
    data.expressionSet.setVariable(data.rindex, r);
    data.expressionSet.setVariable(data.particleValueIndex[0], values[0][atom1]);
    data.expressionSet.setVariable(data.particleValueIndex[1], values[0][atom2]);
555
556
557

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

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

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;
591
    }
592
    r2 = dot3(deltaR, deltaR);
593
}