CpuCustomNonbondedForce.cpp 17.7 KB
Newer Older
1

2
/* Portions copyright (c) 2009-2022 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
 * 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 "CpuCustomNonbondedForce.h"
26
27
#include "openmm/internal/hardware.h"
#include <cmath>
28
29

using namespace OpenMM;
30
using namespace Lepton;
31
32
using namespace std;

33
34
35
36
CpuCustomNonbondedForce::ThreadData::ThreadData(const CompiledExpression& energyExpression, const CompiledVectorExpression& energyVecExpression,
            const CompiledExpression& forceExpression, const CompiledVectorExpression& forceVecExpression,
            const vector<string>& parameterNames, const std::vector<CompiledExpression> energyParamDerivExpressions,
            const vector<string>& computedValueNames, const vector<CompiledExpression> computedValueExpressions,
37
            vector<vector<double> >& atomComputedValues) :
38
39
            energyExpression(energyExpression), energyVecExpression(energyVecExpression), forceExpression(forceExpression),
            forceVecExpression(forceVecExpression), energyParamDerivExpressions(energyParamDerivExpressions),
40
41
42
            computedValueExpressions(computedValueExpressions), atomComputedValues(atomComputedValues) {
    // Prepare for passing variables to expressions.

43
44
45
    map<string, double*> variableLocations;
    variableLocations["r"] = &r;
    particleParam.resize(2*parameterNames.size());
46
    computedValues.resize(2*computedValueNames.size());
47
48
49
    for (int i = 0; i < parameterNames.size(); i++) {
        variableLocations[parameterNames[i]+"1"] = &particleParam[i*2];
        variableLocations[parameterNames[i]+"2"] = &particleParam[i*2+1];
50
    }
51
52
53
    for (int i = 0; i < computedValueNames.size(); i++) {
        variableLocations[computedValueNames[i]+"1"] = &computedValues[i*2];
        variableLocations[computedValueNames[i]+"2"] = &computedValues[i*2+1];
54
    }
55
    energyParamDerivs.resize(energyParamDerivExpressions.size());
56
57
58
59
    this->energyExpression.setVariableLocations(variableLocations);
    this->forceExpression.setVariableLocations(variableLocations);
    expressionSet.registerExpression(this->energyExpression);
    expressionSet.registerExpression(this->forceExpression);
peastman's avatar
peastman committed
60
61
62
    for (auto& expression : this->energyParamDerivExpressions) {
        expression.setVariableLocations(variableLocations);
        expressionSet.registerExpression(expression);
63
    }
64

65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
    // Prepare for passing variables to vectorized expressions.

    map<string, float*> vecVariableLocations;
    int blockSize = getVectorWidth();
    rvec.resize(blockSize);
    vecParticle1Params.resize(blockSize*parameterNames.size());
    vecParticle2Params.resize(blockSize*parameterNames.size());
    vecParticle1Values.resize(blockSize*computedValueNames.size());
    vecParticle2Values.resize(blockSize*computedValueNames.size());
    vecVariableLocations["r"] = rvec.data();
    for (int i = 0; i < parameterNames.size(); i++) {
        vecVariableLocations[parameterNames[i]+"1"] = &vecParticle1Params[i*blockSize];
        vecVariableLocations[parameterNames[i]+"2"] = &vecParticle2Params[i*blockSize];
    }
    for (int i = 0; i < computedValueNames.size(); i++) {
        vecVariableLocations[computedValueNames[i]+"1"] = &vecParticle1Values[i*blockSize];
        vecVariableLocations[computedValueNames[i]+"2"] = &vecParticle2Values[i*blockSize];
    }
    this->energyVecExpression.setVariableLocations(vecVariableLocations);
    this->forceVecExpression.setVariableLocations(vecVariableLocations);

86
87
88
89
90
91
92
93
94
    // Prepare for passing variables to the computed value expressions.

    map<string, double*> valueVariableLocations;
    for (int i = 0; i < parameterNames.size(); i++)
        valueVariableLocations[parameterNames[i]] = &particleParam[i];
    for (auto& expression : this->computedValueExpressions) {
        expression.setVariableLocations(valueVariableLocations);
        expressionSet.registerExpression(expression);
    }
95
96
}

97
98
CpuCustomNonbondedForce::CpuCustomNonbondedForce(ThreadPool& threads, const CpuNeighborList& neighbors) : cutoff(false), useSwitch(false),
        periodic(false), useInteractionGroups(false), threads(threads), neighborList(&neighbors) {
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
}

void CpuCustomNonbondedForce::initialize(const ParsedExpression& energyExpression,
            const ParsedExpression& forceExpression, const vector<string>& parameterNames, const vector<set<int> >& exclusions,
            const vector<ParsedExpression> energyParamDerivExpressions, const vector<string>& computedValueNames,
            const vector<ParsedExpression> computedValueExpressions) {
    this->paramNames = parameterNames;
    this->exclusions = exclusions;
    this->computedValueNames = computedValueNames;
    CompiledExpression compiledEnergyExpression = energyExpression.createCompiledExpression();
    CompiledExpression compiledForceExpression = forceExpression.createCompiledExpression();
    CompiledVectorExpression energyVecExpression = energyExpression.createCompiledVectorExpression(getVectorWidth());
    CompiledVectorExpression forceVecExpression = forceExpression.createCompiledVectorExpression(getVectorWidth());
    vector<CompiledExpression> compiledDerivExpressions, compiledValueExpressions;
    for (auto& exp : energyParamDerivExpressions)
        compiledDerivExpressions.push_back(exp.createCompiledExpression());
    for (auto& exp : computedValueExpressions)
        compiledValueExpressions.push_back(exp.createCompiledExpression());
117
    for (int i = 0; i < threads.getNumThreads(); i++)
118
119
        threadData.push_back(new ThreadData(compiledEnergyExpression, energyVecExpression, compiledForceExpression, forceVecExpression, parameterNames,
                compiledDerivExpressions, computedValueNames, compiledValueExpressions, atomComputedValues));
120
121
}

122
CpuCustomNonbondedForce::~CpuCustomNonbondedForce() {
peastman's avatar
peastman committed
123
124
    for (auto data : threadData)
        delete data;
125
}
126

127
void CpuCustomNonbondedForce::setUseCutoff(double distance) {
128
129
130
131
132
    cutoff = true;
    cutoffDistance = distance;
  }

void CpuCustomNonbondedForce::setInteractionGroups(const vector<pair<set<int>, set<int> > >& groups) {
133
    useInteractionGroups = true;
peastman's avatar
peastman committed
134
135
136
    for (auto& group : groups) {
        const set<int>& set1 = group.first;
        const set<int>& set2 = group.second;
137
138
139
140
141
142
143
144
145
146
        for (set<int>::const_iterator atom1 = set1.begin(); atom1 != set1.end(); ++atom1) {
            for (set<int>::const_iterator atom2 = set2.begin(); atom2 != set2.end(); ++atom2) {
                if (*atom1 == *atom2 || exclusions[*atom1].find(*atom2) != exclusions[*atom1].end())
                    continue; // This is an excluded interaction.
                if (*atom1 > *atom2 && set1.find(*atom2) != set1.end() && set2.find(*atom1) != set2.end())
                    continue; // Both atoms are in both sets, so skip duplicate interactions.
                groupInteractions.push_back(make_pair(*atom1, *atom2));
            }
        }
    }
147
148
}

peastman's avatar
peastman committed
149
void CpuCustomNonbondedForce::setUseSwitchingFunction(double distance) {
150
151
152
153
    useSwitch = true;
    switchingDistance = distance;
}

peastman's avatar
peastman committed
154
void CpuCustomNonbondedForce::setPeriodic(Vec3* periodicBoxVectors) {
155
    assert(cutoff);
156
157
158
    assert(periodicBoxVectors[0][0] >= 2.0*cutoffDistance);
    assert(periodicBoxVectors[1][1] >= 2.0*cutoffDistance);
    assert(periodicBoxVectors[2][2] >= 2.0*cutoffDistance);
159
    periodic = true;
160
161
162
163
164
165
166
167
168
169
170
171
172
    this->periodicBoxVectors[0] = periodicBoxVectors[0];
    this->periodicBoxVectors[1] = periodicBoxVectors[1];
    this->periodicBoxVectors[2] = periodicBoxVectors[2];
    recipBoxSize[0] = (float) (1.0/periodicBoxVectors[0][0]);
    recipBoxSize[1] = (float) (1.0/periodicBoxVectors[1][1]);
    recipBoxSize[2] = (float) (1.0/periodicBoxVectors[2][2]);
    periodicBoxVec4.resize(3);
    periodicBoxVec4[0] = fvec4(periodicBoxVectors[0][0], periodicBoxVectors[0][1], periodicBoxVectors[0][2], 0);
    periodicBoxVec4[1] = fvec4(periodicBoxVectors[1][0], periodicBoxVectors[1][1], periodicBoxVectors[1][2], 0);
    periodicBoxVec4[2] = fvec4(periodicBoxVectors[2][0], periodicBoxVectors[2][1], periodicBoxVectors[2][2], 0);
    triclinic = (periodicBoxVectors[0][1] != 0.0 || periodicBoxVectors[0][2] != 0.0 ||
                 periodicBoxVectors[1][0] != 0.0 || periodicBoxVectors[1][2] != 0.0 ||
                 periodicBoxVectors[2][0] != 0.0 || periodicBoxVectors[2][1] != 0.0);
173
}
174
175


176
177
178
void CpuCustomNonbondedForce::calculatePairIxn(int numberOfAtoms, float* posq, vector<Vec3>& atomCoordinates, vector<vector<double> >& atomParameters,
                                               const map<string, double>& globalParameters, vector<AlignedArray<float> >& threadForce,
                                               bool includeForce, bool includeEnergy, double& totalEnergy, double* energyParamDerivs) {
179
180
181
182
183
    // Record the parameters for the threads.
    
    this->numberOfAtoms = numberOfAtoms;
    this->posq = posq;
    this->atomCoordinates = &atomCoordinates[0];
184
    this->atomParameters = &atomParameters[0];
185
    this->globalParameters = &globalParameters;
186
    this->threadForce = &threadForce;
187
188
    this->includeForce = includeForce;
    this->includeEnergy = includeEnergy;
189
    threadEnergy.resize(threads.getNumThreads());
190
    atomComputedValues.resize(computedValueNames.size(), vector<double>(numberOfAtoms));
peastman's avatar
peastman committed
191
    atomicCounter = 0;
192
    
193
    // Signal the threads to start running and wait for them to finish.
194
    
peastman's avatar
peastman committed
195
    threads.execute([&] (ThreadPool& threads, int threadIndex) { threadComputeForce(threads, threadIndex); });
196
197
198
    threads.waitForThreads(); // Computed values
    threads.resumeThreads();
    threads.waitForThreads(); // Interactions
199
    
200
201
    // Combine the energies from all the threads.
    
202
    int numThreads = threads.getNumThreads();
203
204
205
206
    if (includeEnergy) {
        for (int i = 0; i < numThreads; i++)
            totalEnergy += threadEnergy[i];
    }
207
208
209
210
211
212
213

    // Combine the energy derivatives from all threads.
    
    int numDerivs = threadData[0]->energyParamDerivs.size();
    for (int i = 0; i < numThreads; i++)
        for (int j = 0; j < numDerivs; j++)
            energyParamDerivs[j] += threadData[i]->energyParamDerivs[j];
214
215
216
}

void CpuCustomNonbondedForce::threadComputeForce(ThreadPool& threads, int threadIndex) {
217
    int numThreads = threads.getNumThreads();
218
    int blockSize = getVectorWidth();
219
    ThreadData& data = *threadData[threadIndex];
220
    for (auto& param : *globalParameters) {
221
        data.expressionSet.setVariable(data.expressionSet.getVariableIndex(param.first), param.second);
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
        try {
            float* p = data.energyVecExpression.getVariablePointer(param.first);
            for (int i = 0; i < blockSize; i++)
                p[i] = param.second;
        }
        catch (...) {
            // The expression doesn't use this parameter.
        }
        try {
            float* p = data.forceVecExpression.getVariablePointer(param.first);
            for (int i = 0; i < blockSize; i++)
                p[i] = param.second;
        }
        catch (...) {
            // The expression doesn't use this parameter.
        }
    }
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254

    // Process computed values for this thread's subset of interactions.

    int numComputedValues = atomComputedValues.size();
    if (numComputedValues > 0) {
        int start = threadIndex*numberOfAtoms/numThreads;
        int end = (threadIndex+1)*numberOfAtoms/numThreads;
        for (int i = start; i < end; i++) {
            for (int j = 0; j < paramNames.size(); j++)
                data.particleParam[j] = atomParameters[i][j];
            for (int j = 0; j < numComputedValues; j++)
                atomComputedValues[j][i] = data.computedValueExpressions[j].evaluate();
        }
    }
    threads.syncThreads();

255
256
257
258
259
    // Compute this thread's subset of interactions.

    threadEnergy[threadIndex] = 0;
    double& energy = threadEnergy[threadIndex];
    float* forces = &(*threadForce)[threadIndex][0];
peastman's avatar
peastman committed
260
261
    for (auto& deriv : data.energyParamDerivs)
        deriv = 0.0;
262
263
    fvec4 boxSize(periodicBoxVectors[0][0], periodicBoxVectors[1][1], periodicBoxVectors[2][2], 0);
    fvec4 invBoxSize(recipBoxSize[0], recipBoxSize[1], recipBoxSize[2], 0);
264
    if (useInteractionGroups) {
265
266
        // The user has specified interaction groups, so compute only the requested interactions.
        
267
268
269
        int start = threadIndex*groupInteractions.size()/numThreads;
        int end = (threadIndex+1)*groupInteractions.size()/numThreads;
        for (int i = start; i < end; i++) {
270
271
            int atom1 = groupInteractions[i].first;
            int atom2 = groupInteractions[i].second;
272
            for (int j = 0; j < paramNames.size(); j++) {
273
274
                data.particleParam[j*2] = atomParameters[atom1][j];
                data.particleParam[j*2+1] = atomParameters[atom2][j];
275
            }
276
277
278
279
            for (int j = 0; j < computedValueNames.size(); j++) {
                data.computedValues[j*2] = atomComputedValues[j][atom1];
                data.computedValues[j*2+1] = atomComputedValues[j][atom2];
            }
280
            calculateOneIxn(atom1, atom2, data, forces, energy, boxSize, invBoxSize);
281
282
        }
    }
283
284
    else {
        // Get the interactions from the neighbor list.
285
286

        while (true) {
peastman's avatar
peastman committed
287
            int blockIndex = atomicCounter++;
288
289
            if (blockIndex >= neighborList->getNumBlocks())
                break;
290
291
292
            if (data.energyParamDerivs.size() == 0)
                calculateBlockIxn(data, blockIndex, forces, energy, boxSize, invBoxSize);
            else {
293
294
295
296
297
                const int blockSize = neighborList->getBlockSize();
                const int32_t* blockAtom = &neighborList->getSortedAtoms()[blockSize*blockIndex];
                CpuNeighborList::NeighborIterator neighbors = neighborList->getNeighborIterator(blockIndex);
                while (neighbors.next()) {
                    int first = neighbors.getNeighbor();
298
299
300
301
                    for (int j = 0; j < paramNames.size(); j++)
                        data.particleParam[j*2] = atomParameters[first][j];
                    for (int j = 0; j < computedValueNames.size(); j++)
                        data.computedValues[j*2] = atomComputedValues[j][first];
302
                    auto exclusions = neighbors.getExclusions();
303
                    for (int k = 0; k < blockSize; k++) {
304
                        if ((exclusions & (1<<k)) == 0) {
305
306
307
308
309
310
311
                            int second = blockAtom[k];
                            for (int j = 0; j < paramNames.size(); j++)
                                data.particleParam[j*2+1] = atomParameters[second][j];
                            for (int j = 0; j < computedValueNames.size(); j++)
                                data.computedValues[j*2+1] = atomComputedValues[j][second];
                            calculateOneIxn(first, second, data, forces, energy, boxSize, invBoxSize);
                        }
312
313
314
315
316
317
318
                    }
                }
            }
        }
    }
}

319
320
321
void CpuCustomNonbondedForce::calculateOneIxn(int ii, int jj, ThreadData& data, 
        float* forces, double& totalEnergy, const fvec4& boxSize, const fvec4& invBoxSize) {
    // Get deltaR, R2, and R between 2 atoms
322
323
324
325
326
327
328
329
330

    fvec4 deltaR;
    fvec4 posI(posq+4*ii);
    fvec4 posJ(posq+4*jj);
    float r2;
    getDeltaR(posI, posJ, deltaR, r2, boxSize, invBoxSize);
    if (cutoff && r2 >= cutoffDistance*cutoffDistance)
        return;
    float r = sqrtf(r2);
331
    data.r = r;
332
333
334

    // accumulate forces

335
    double dEdR = (includeForce ? data.forceExpression.evaluate()/r : 0.0);
336
337
338
    double energy = 0.0;
    if (includeEnergy || (useSwitch && r > switchingDistance))
        energy = data.energyExpression.evaluate();
339
    double switchValue = 1.0;
340
341
    if (useSwitch) {
        if (r > switchingDistance) {
342
343
344
            double t = (r-switchingDistance)/(cutoffDistance-switchingDistance);
            switchValue = 1+t*t*t*(-10+t*(15-t*6));
            double switchDeriv = t*t*(-30+t*(60-t*30))/(cutoffDistance-switchingDistance);
345
346
347
348
349
350
351
352
353
354
            dEdR = switchValue*dEdR + energy*switchDeriv/r;
            energy *= switchValue;
        }
    }
    fvec4 result = deltaR*dEdR;
    (fvec4(forces+4*ii)+result).store(forces+4*ii);
    (fvec4(forces+4*jj)-result).store(forces+4*jj);

    // accumulate energies

355
    totalEnergy += energy;
356
357
358
359
360
    
    // Accumulate energy derivatives.

    for (int i = 0; i < data.energyParamDerivExpressions.size(); i++)
        data.energyParamDerivs[i] += switchValue*data.energyParamDerivExpressions[i].evaluate();
361
362
363
364
365
}

void CpuCustomNonbondedForce::getDeltaR(const fvec4& posI, const fvec4& posJ, fvec4& deltaR, float& r2, const fvec4& boxSize, const fvec4& invBoxSize) const {
    deltaR = posJ-posI;
    if (periodic) {
366
367
368
369
370
371
372
373
374
        if (triclinic) {
            deltaR -= periodicBoxVec4[2]*floorf(deltaR[2]*recipBoxSize[2]+0.5f);
            deltaR -= periodicBoxVec4[1]*floorf(deltaR[1]*recipBoxSize[1]+0.5f);
            deltaR -= periodicBoxVec4[0]*floorf(deltaR[0]*recipBoxSize[0]+0.5f);
        }
        else {
            fvec4 base = round(deltaR*invBoxSize)*boxSize;
            deltaR = deltaR-base;
        }
375
376
377
    }
    r2 = dot3(deltaR, deltaR);
}