Skip to content

GitLab

Commit ff8cac54 authored by peastman's avatar peastman
Browse files

CPU implementation of parameter derivatives for CustomGBForce

parent aeaa6f2e
Hide whitespace changes
Inline Side-by-side
Showing with 153 additions and 45 deletions
platforms/cpu/include/CpuCustomGBForce.h
View file @ ff8cac54
......@@ -47,14 +47,14 @@ private:
const CpuNeighborList* neighborList;
float periodicBoxSize[3];
float cutoffDistance, cutoffDistance2;
int numValues, numParams;
const std::vector<std::set<int> > exclusions;
std::vector<std::string> valueNames;
std::vector<CustomGBForce::ComputationType> valueTypes;
std::vector<std::string> paramNames;
std::vector<CustomGBForce::ComputationType> energyTypes;
ThreadPool& threads;
std::vector<ThreadData*> threadData;
std::vector<double> threadEnergy;
std::vector<std::vector<std::vector<float> > > dValuedParam;
// Workspace vectors
std::vector<std::vector<float> > values, dEdV;
// The following variables are used to make information accessible to the individual threads.
......@@ -189,11 +189,13 @@ public:
const std::vector<Lepton::CompiledExpression>& valueExpressions,
const std::vector<std::vector<Lepton::CompiledExpression> >& valueDerivExpressions,
const std::vector<std::vector<Lepton::CompiledExpression> >& valueGradientExpressions,
const std::vector<std::vector<Lepton::CompiledExpression> >& valueParamDerivExpressions,
const std::vector<std::string>& valueNames,
const std::vector<CustomGBForce::ComputationType>& valueTypes,
const std::vector<Lepton::CompiledExpression>& energyExpressions,
const std::vector<std::vector<Lepton::CompiledExpression> >& energyDerivExpressions,
const std::vector<std::vector<Lepton::CompiledExpression> >& energyGradientExpressions,
const std::vector<std::vector<Lepton::CompiledExpression> >& energyParamDerivExpressions,
const std::vector<CustomGBForce::ComputationType>& energyTypes,
const std::vector<std::string>& parameterNames, ThreadPool& threads);
......@@ -221,16 +223,17 @@ public:
/**
* Calculate custom GB ixn
*
* @param numberOfAtoms number of atoms
* @param posq atom coordinates
* @param atomParameters atomParameters[atomIndex][paramterIndex]
* @param globalParameters the values of global parameters
* @param forces force array (forces added)
* @param totalEnergy total energy
* @param numberOfAtoms number of atoms
* @param posq atom coordinates
* @param atomParameters atomParameters[atomIndex][paramterIndex]
* @param globalParameters the values of global parameters
* @param forces force array (forces added)
* @param totalEnergy total energy
* @param energyParamDerivs derivatives of the energy with respect to global parameters
*/
void calculateIxn(int numberOfAtoms, float* posq, RealOpenMM** atomParameters,
std::map<std::string, double>& globalParameters, std::vector<AlignedArray<float> >& threadForce, bool includeForce, bool includeEnergy, double& totalEnergy);
void calculateIxn(int numberOfAtoms, float* posq, RealOpenMM** atomParameters, std::map<std::string, double>& globalParameters,
std::vector<AlignedArray<float> >& threadForce, bool includeForce, bool includeEnergy, double& totalEnergy, double* energyParamDerivs);
};
class CpuCustomGBForce::ThreadData {
......@@ -239,19 +242,23 @@ public:
const std::vector<Lepton::CompiledExpression>& valueExpressions,
const std::vector<std::vector<Lepton::CompiledExpression> >& valueDerivExpressions,
const std::vector<std::vector<Lepton::CompiledExpression> >& valueGradientExpressions,
const std::vector<std::vector<Lepton::CompiledExpression> >& valueParamDerivExpressions,
const std::vector<std::string>& valueNames,
const std::vector<Lepton::CompiledExpression>& energyExpressions,
const std::vector<std::vector<Lepton::CompiledExpression> >& energyDerivExpressions,
const std::vector<std::vector<Lepton::CompiledExpression> >& energyGradientExpressions,
const std::vector<std::vector<Lepton::CompiledExpression> >& energyParamDerivExpressions,
const std::vector<std::string>& parameterNames);
CompiledExpressionSet expressionSet;
std::vector<Lepton::CompiledExpression> valueExpressions;
std::vector<std::vector<Lepton::CompiledExpression> > valueDerivExpressions;
std::vector<std::vector<Lepton::CompiledExpression> > valueGradientExpressions;
std::vector<std::vector<Lepton::CompiledExpression> > valueParamDerivExpressions;
std::vector<int> valueIndex;
std::vector<Lepton::CompiledExpression> energyExpressions;
std::vector<std::vector<Lepton::CompiledExpression> > energyDerivExpressions;
std::vector<std::vector<Lepton::CompiledExpression> > energyGradientExpressions;
std::vector<std::vector<Lepton::CompiledExpression> > energyParamDerivExpressions;
std::vector<int> paramIndex;
std::vector<int> particleParamIndex;
std::vector<int> particleValueIndex;
......@@ -260,6 +267,8 @@ public:
// Workspace vectors
std::vector<float> value0, dVdR1, dVdR2, dVdX, dVdY, dVdZ;
std::vector<std::vector<float> > dEdV;
std::vector<std::vector<float> > dValue0dParam;
std::vector<float> energyParamDerivs;
};
} // namespace OpenMM
......
platforms/cpu/include/CpuKernels.h
View file @ ff8cac54
......@@ -398,7 +398,7 @@ private:
RealOpenMM nonbondedCutoff;
CpuCustomGBForce* ixn;
std::vector<std::set<int> > exclusions;
std::vector<std::string> particleParameterNames, globalParameterNames, valueNames;
std::vector<std::string> particleParameterNames, globalParameterNames, energyParamDerivNames, valueNames;
std::vector<OpenMM::CustomGBForce::ComputationType> valueTypes;
std::vector<OpenMM::CustomGBForce::ComputationType> energyTypes;
NonbondedMethod nonbondedMethod;
......
platforms/cpu/src/CpuCustomGBForce.cpp
View file @ ff8cac54
......@@ -47,13 +47,16 @@ CpuCustomGBForce::ThreadData::ThreadData(int numAtoms, int numThreads, int threa
const vector<Lepton::CompiledExpression>& valueExpressions,
const vector<vector<Lepton::CompiledExpression> >& valueDerivExpressions,
const vector<vector<Lepton::CompiledExpression> >& valueGradientExpressions,
const vector<vector<Lepton::CompiledExpression> >& valueParamDerivExpressions,
const vector<string>& valueNames,
const vector<Lepton::CompiledExpression>& energyExpressions,
const vector<vector<Lepton::CompiledExpression> >& energyDerivExpressions,
const vector<vector<Lepton::CompiledExpression> >& energyGradientExpressions,
const vector<vector<Lepton::CompiledExpression> >& energyParamDerivExpressions,
const vector<string>& parameterNames) :
valueExpressions(valueExpressions), valueDerivExpressions(valueDerivExpressions), valueGradientExpressions(valueGradientExpressions),
energyExpressions(energyExpressions), energyDerivExpressions(energyDerivExpressions), energyGradientExpressions(energyGradientExpressions) {
valueParamDerivExpressions(valueParamDerivExpressions), energyExpressions(energyExpressions), energyDerivExpressions(energyDerivExpressions),
energyGradientExpressions(energyGradientExpressions), energyParamDerivExpressions(energyParamDerivExpressions) {
firstAtom = (threadIndex*(long long) numAtoms)/numThreads;
lastAtom = ((threadIndex+1)*(long long) numAtoms)/numThreads;
for (int i = 0; i < (int) valueExpressions.size(); i++)
......@@ -64,6 +67,9 @@ CpuCustomGBForce::ThreadData::ThreadData(int numAtoms, int numThreads, int threa
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) valueParamDerivExpressions.size(); i++)
for (int j = 0; j < (int) valueParamDerivExpressions[i].size(); j++)
expressionSet.registerExpression(this->valueParamDerivExpressions[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++)
......@@ -72,6 +78,9 @@ CpuCustomGBForce::ThreadData::ThreadData(int numAtoms, int numThreads, int threa
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]);
for (int i = 0; i < (int) energyParamDerivExpressions.size(); i++)
for (int j = 0; j < (int) energyParamDerivExpressions[i].size(); j++)
expressionSet.registerExpression(this->energyParamDerivExpressions[i][j]);
xindex = expressionSet.getVariableIndex("x");
yindex = expressionSet.getVariableIndex("y");
zindex = expressionSet.getVariableIndex("z");
......@@ -101,30 +110,37 @@ CpuCustomGBForce::ThreadData::ThreadData(int numAtoms, int numThreads, int threa
dVdZ.resize(valueDerivExpressions.size());
dVdR1.resize(valueDerivExpressions.size());
dVdR2.resize(valueDerivExpressions.size());
dValue0dParam.resize(valueParamDerivExpressions[0].size(), vector<float>(numAtoms));
energyParamDerivs.resize(valueParamDerivExpressions[0].size());
}
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<vector<Lepton::CompiledExpression> >& valueParamDerivExpressions,
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<vector<Lepton::CompiledExpression> >& energyParamDerivExpressions,
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) {
exclusions(exclusions), cutoff(false), periodic(false), valueTypes(valueTypes), energyTypes(energyTypes), numValues(valueNames.size()),
numParams(parameterNames.size()), 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());
threadData.push_back(new ThreadData(numAtoms, threads.getNumThreads(), i, valueExpressions, valueDerivExpressions, valueGradientExpressions,
valueParamDerivExpressions, valueNames, energyExpressions, energyDerivExpressions, energyGradientExpressions, energyParamDerivExpressions, parameterNames));
values.resize(numValues);
dEdV.resize(numValues);
for (int i = 0; i < (int) values.size(); i++) {
values[i].resize(numAtoms);
dEdV[i].resize(numAtoms);
}
dValuedParam.resize(numValues);
for (int i = 0; i < numValues; i++)
dValuedParam[i].resize(valueParamDerivExpressions[0].size(), vector<float>(numAtoms));
}
CpuCustomGBForce::~CpuCustomGBForce() {
......@@ -153,7 +169,7 @@ void CpuCustomGBForce::setPeriodic(RealVec& boxSize) {
void CpuCustomGBForce::calculateIxn(int numberOfAtoms, float* posq, RealOpenMM** atomParameters,
map<string, double>& globalParameters, vector<AlignedArray<float> >& threadForce,
bool includeForce, bool includeEnergy, double& totalEnergy) {
bool includeForce, bool includeEnergy, double& totalEnergy, double* energyParamDerivs) {
// Record the parameters for the threads.
this->numberOfAtoms = numberOfAtoms;
......@@ -173,12 +189,20 @@ void CpuCustomGBForce::calculateIxn(int numberOfAtoms, float* posq, RealOpenMM**
gmx_atomic_set(&counter, 0);
threads.execute(task);
threads.waitForThreads();
// Sum derivatives of the first computed value with respect to global parameters.
bool hasParamDerivs = (threadData[0]->dValue0dParam.size() > 0);
if (hasParamDerivs) {
threads.resumeThreads();
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++) {
......@@ -205,6 +229,10 @@ void CpuCustomGBForce::calculateIxn(int numberOfAtoms, float* posq, RealOpenMM**
for (int i = 0; i < numThreads; i++)
totalEnergy += threadEnergy[i];
}
if (hasParamDerivs)
for (int i = 0; i < threads.getNumThreads(); i++)
for (int j = 0; j < threadData[j]->energyParamDerivs.size(); j++)
energyParamDerivs[j] += threadData[i]->energyParamDerivs[j];
}
void CpuCustomGBForce::threadComputeForce(ThreadPool& threads, int threadIndex) {
......@@ -224,11 +252,28 @@ void CpuCustomGBForce::threadComputeForce(ThreadPool& threads, int threadIndex)
for (int i = 0; i < (int) data.value0.size(); i++)
data.value0[i] = 0.0f;
for (int i = 0; i < (int) data.dValue0dParam.size(); i++)
for (int j = 0; j < (int) data.dValue0dParam[i].size(); j++)
data.dValue0dParam[i][j] = 0.0;
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 derivatives of the first computed value with respect to global parameters.
bool hasParamDerivs = (data.dValue0dParam.size() > 0);
if (hasParamDerivs) {
for (int j = 0; j < data.dValue0dParam.size(); j++)
for (int k = data.firstAtom; k < data.lastAtom; k++) {
float sum = 0.0f;
for (int m = 0; m < threadData.size(); m++)
sum += threadData[m]->dValue0dParam[j][k];
dValuedParam[0][j][k] = sum;
}
threads.syncThreads();
}
// Sum the first computed value and calculate the remaining ones.
......@@ -241,11 +286,23 @@ void CpuCustomGBForce::threadComputeForce(ThreadPool& threads, int threadIndex)
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++)
for (int j = 0; j < numParams; 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();
// Calculate derivatives with respect to parameters.
if (hasParamDerivs) {
for (int j = 0; j < data.valueParamDerivExpressions[i].size(); j++)
dValuedParam[i][j][atom] = data.valueParamDerivExpressions[i][j].evaluate();
for (int j = 0; j < i; j++) {
float dVdV = data.valueDerivExpressions[i][j].evaluate();
for (int k = 0; k < data.valueParamDerivExpressions[i].size(); k++)
dValuedParam[i][k][atom] += dVdV*dValuedParam[j][k][atom];
}
}
}
}
threads.syncThreads();
......@@ -254,7 +311,9 @@ void CpuCustomGBForce::threadComputeForce(ThreadPool& threads, int threadIndex)
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;
data.dEdV[i][j] = 0.0f;
for (int i = 0; i < (int) data.energyParamDerivs.size(); i++)
data.energyParamDerivs[i] = 0.0f;
for (int termIndex = 0; termIndex < (int) data.energyExpressions.size(); termIndex++) {
if (energyTypes[termIndex] == CustomGBForce::SingleParticle)
calculateSingleParticleEnergyTerm(termIndex, data, numberOfAtoms, posq, atomParameters, forces, energy);
......@@ -339,7 +398,7 @@ void CpuCustomGBForce::calculateOnePairValue(int index, int atom1, int atom2, Th
if (cutoff && r2 >= cutoffDistance2)
return;
float r = sqrtf(r2);
for (int i = 0; i < (int) paramNames.size(); i++) {
for (int i = 0; i < numParams; i++) {
data.expressionSet.setVariable(data.particleParamIndex[i*2], atomParameters[atom1][i]);
data.expressionSet.setVariable(data.particleParamIndex[i*2+1], atomParameters[atom2][i]);
}
......@@ -349,6 +408,11 @@ void CpuCustomGBForce::calculateOnePairValue(int index, int atom1, int atom2, Th
data.expressionSet.setVariable(data.particleValueIndex[i*2+1], values[i][atom2]);
}
valueArray[atom1] += (float) data.valueExpressions[index].evaluate();
// Calculate derivatives with respect to parameters.
for (int i = 0; i < data.valueParamDerivExpressions[index].size(); i++)
data.dValue0dParam[i][atom1] += data.valueParamDerivExpressions[index][i].evaluate();
}
void CpuCustomGBForce::calculateSingleParticleEnergyTerm(int index, ThreadData& data, int numAtoms, float* posq,
......@@ -357,17 +421,22 @@ void CpuCustomGBForce::calculateSingleParticleEnergyTerm(int index, ThreadData&
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]);
for (int j = 0; j < (int) paramNames.size(); j++)
for (int j = 0; j < numParams; j++)
data.expressionSet.setVariable(data.paramIndex[j], atomParameters[i][j]);
for (int j = 0; j < (int) valueNames.size(); j++)
for (int j = 0; j < (int) values.size(); j++)
data.expressionSet.setVariable(data.valueIndex[j], values[j][i]);
if (includeEnergy)
totalEnergy += (float) data.energyExpressions[index].evaluate();
for (int j = 0; j < (int) valueNames.size(); j++)
for (int j = 0; j < (int) values.size(); j++)
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();
// Compute derivatives with respect to parameters.
for (int k = 0; k < data.energyParamDerivExpressions[index].size(); k++)
data.energyParamDerivs[k] += data.energyParamDerivExpressions[index][k].evaluate();
}
}
......@@ -428,12 +497,12 @@ void CpuCustomGBForce::calculateOnePairEnergyTerm(int index, int atom1, int atom
// Record variables for evaluating expressions.
for (int i = 0; i < (int) paramNames.size(); i++) {
for (int i = 0; i < numParams; i++) {
data.expressionSet.setVariable(data.particleParamIndex[i*2], atomParameters[atom1][i]);
data.expressionSet.setVariable(data.particleParamIndex[i*2+1], atomParameters[atom2][i]);
}
data.expressionSet.setVariable(data.rindex, r);
for (int i = 0; i < (int) valueNames.size(); i++) {
for (int i = 0; i < (int) values.size(); i++) {
data.expressionSet.setVariable(data.particleValueIndex[i*2], values[i][atom1]);
data.expressionSet.setVariable(data.particleValueIndex[i*2+1], values[i][atom2]);
}
......@@ -447,10 +516,15 @@ void CpuCustomGBForce::calculateOnePairEnergyTerm(int index, int atom1, int atom
fvec4 result = deltaR*dEdR;
(fvec4(forces+4*atom1)-result).store(forces+4*atom1);
(fvec4(forces+4*atom2)+result).store(forces+4*atom2);
for (int i = 0; i < (int) valueNames.size(); i++) {
for (int i = 0; i < (int) values.size(); i++) {
data.dEdV[i][atom1] += (float) data.energyDerivExpressions[index][2*i+1].evaluate();
data.dEdV[i][atom2] += (float) data.energyDerivExpressions[index][2*i+2].evaluate();
}
// Compute derivatives with respect to parameters.
for (int i = 0; i < data.energyParamDerivExpressions[index].size(); i++)
data.energyParamDerivs[i] += data.energyParamDerivExpressions[index][i].evaluate();
}
void CpuCustomGBForce::calculateChainRuleForces(ThreadData& data, int numAtoms, float* posq, RealOpenMM** atomParameters,
......@@ -500,9 +574,9 @@ void CpuCustomGBForce::calculateChainRuleForces(ThreadData& data, int numAtoms,
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]);
for (int j = 0; j < (int) paramNames.size(); j++)
for (int j = 0; j < numParams; j++)
data.expressionSet.setVariable(data.paramIndex[j], atomParameters[i][j]);
for (int j = 1; j < (int) valueNames.size(); j++) {
for (int j = 1; j < (int) values.size(); j++) {
data.expressionSet.setVariable(data.valueIndex[j-1], values[j-1][i]);
data.dVdX[j] = 0.0;
data.dVdY[j] = 0.0;
......@@ -521,6 +595,13 @@ void CpuCustomGBForce::calculateChainRuleForces(ThreadData& data, int numAtoms,
forces[4*i+2] -= dEdV[j][i]*data.dVdZ[j];
}
}
// Compute chain rule terms for derivatives with respect to parameters.
for (int i = data.firstAtom; i < data.lastAtom; i++)
for (int j = 0; j < data.valueIndex.size(); j++)
for (int k = 0; k < dValuedParam[j].size(); k++)
data.energyParamDerivs[k] += dEdV[j][i]*dValuedParam[j][k][i];
}
void CpuCustomGBForce::calculateOnePairChainRule(int atom1, int atom2, ThreadData& data, float* posq, RealOpenMM** atomParameters,
......@@ -538,7 +619,7 @@ void CpuCustomGBForce::calculateOnePairChainRule(int atom1, int atom2, ThreadDat
// Record variables for evaluating expressions.
for (int i = 0; i < (int) paramNames.size(); i++) {
for (int i = 0; i < numParams; i++) {
data.expressionSet.setVariable(data.particleParamIndex[i*2], atomParameters[atom1][i]);
data.expressionSet.setVariable(data.particleParamIndex[i*2+1], atomParameters[atom2][i]);
data.expressionSet.setVariable(data.paramIndex[i], atomParameters[atom1][i]);
......@@ -562,7 +643,7 @@ void CpuCustomGBForce::calculateOnePairChainRule(int atom1, int atom2, ThreadDat
f1 -= deltaR*(dEdV[0][atom1]*data.dVdR1[0]);
f2 -= deltaR*(dEdV[0][atom1]*data.dVdR2[0]);
}
for (int i = 1; i < (int) valueNames.size(); i++) {
for (int i = 1; i < (int) values.size(); i++) {
data.expressionSet.setVariable(data.valueIndex[i], values[i][atom1]);
data.dVdR1[i] = 0.0;
data.dVdR2[i] = 0.0;
......
platforms/cpu/src/CpuKernels.cpp
View file @ ff8cac54
......@@ -1045,6 +1045,7 @@ void CpuCalcCustomGBForceKernel::initialize(const System& system, const CustomGB
vector<vector<Lepton::CompiledExpression> > valueDerivExpressions(force.getNumComputedValues());
vector<vector<Lepton::CompiledExpression> > valueGradientExpressions(force.getNumComputedValues());
vector<vector<Lepton::CompiledExpression> > valueParamDerivExpressions(force.getNumComputedValues());
vector<Lepton::CompiledExpression> valueExpressions;
vector<Lepton::CompiledExpression> energyExpressions;
set<string> particleVariables, pairVariables;
......@@ -1079,6 +1080,11 @@ void CpuCalcCustomGBForceKernel::initialize(const System& system, const CustomGB
valueDerivExpressions[i].push_back(ex.differentiate(valueNames[j]).createCompiledExpression());
validateVariables(ex.getRootNode(), particleVariables);
}
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());
}
particleVariables.insert(name);
pairVariables.insert(name+"1");
pairVariables.insert(name+"2");
......@@ -1088,6 +1094,7 @@ void CpuCalcCustomGBForceKernel::initialize(const System& system, const CustomGB
vector<vector<Lepton::CompiledExpression> > energyDerivExpressions(force.getNumEnergyTerms());
vector<vector<Lepton::CompiledExpression> > energyGradientExpressions(force.getNumEnergyTerms());
vector<vector<Lepton::CompiledExpression> > energyParamDerivExpressions(force.getNumEnergyTerms());
for (int i = 0; i < force.getNumEnergyTerms(); i++) {
string expression;
CustomGBForce::ComputationType type;
......@@ -1111,14 +1118,17 @@ void CpuCalcCustomGBForceKernel::initialize(const System& system, const CustomGB
validateVariables(ex.getRootNode(), pairVariables);
}
}
for (int j = 0; j < force.getNumEnergyParameterDerivatives(); j++)
energyParamDerivExpressions[i].push_back(ex.differentiate(force.getEnergyParameterDerivativeName(j)).createCompiledExpression());
}
// Delete the custom functions.
for (map<string, Lepton::CustomFunction*>::iterator iter = functions.begin(); iter != functions.end(); iter++)
delete iter->second;
ixn = new CpuCustomGBForce(numParticles, exclusions, valueExpressions, valueDerivExpressions, valueGradientExpressions, valueNames, valueTypes, energyExpressions,
energyDerivExpressions, energyGradientExpressions, energyTypes, particleParameterNames, data.threads);
ixn = new CpuCustomGBForce(numParticles, exclusions, valueExpressions, valueDerivExpressions, valueGradientExpressions, valueParamDerivExpressions,
valueNames, valueTypes, energyExpressions, energyDerivExpressions, energyGradientExpressions, energyParamDerivExpressions, energyTypes,
particleParameterNames, data.threads);
data.isPeriodic = (force.getNonbondedMethod() == CustomGBForce::CutoffPeriodic);
}
......@@ -1135,7 +1145,11 @@ double CpuCalcCustomGBForceKernel::execute(ContextImpl& context, bool includeFor
map<string, double> globalParameters;
for (int i = 0; i < (int) globalParameterNames.size(); i++)
globalParameters[globalParameterNames[i]] = context.getParameter(globalParameterNames[i]);
ixn->calculateIxn(numParticles, &data.posq[0], particleParamArray, globalParameters, data.threadForce, includeForces, includeEnergy, energy);
vector<double> energyParamDerivValues(energyParamDerivNames.size()+1, 0.0);
ixn->calculateIxn(numParticles, &data.posq[0], particleParamArray, globalParameters, data.threadForce, includeForces, includeEnergy, energy, &energyParamDerivValues[0]);
map<string, double>& energyParamDerivs = extractEnergyParameterDerivatives(context);
for (int i = 0; i < energyParamDerivNames.size(); i++)
energyParamDerivs[energyParamDerivNames[i]] += energyParamDerivValues[i];
return energy;
}
......
platforms/reference/include/ReferenceCustomGBIxn.h
View file @ ff8cac54
......@@ -170,7 +170,7 @@ class ReferenceCustomGBIxn {
--------------------------------------------------------------------------------------- */
void calculateChainRuleForces(int numAtoms, std::vector<OpenMM::RealVec>& atomCoordinates, RealOpenMM** atomParameters,
const std::vector<std::set<int> >& exclusions, std::vector<OpenMM::RealVec>& forces);
const std::vector<std::set<int> >& exclusions, std::vector<OpenMM::RealVec>& forces, double* energyParamDerivs);
/**---------------------------------------------------------------------------------------
......
platforms/reference/src/SimTKReference/ReferenceCustomGBIxn.cpp
View file @ ff8cac54
......@@ -190,7 +190,7 @@ void ReferenceCustomGBIxn::calculateIxn(int numberOfAtoms, vector<RealVec>& atom
// Apply the chain rule to evaluate forces.
calculateChainRuleForces(numberOfAtoms, atomCoordinates, atomParameters, exclusions, forces);
calculateChainRuleForces(numberOfAtoms, atomCoordinates, atomParameters, exclusions, forces, energyParamDerivs);
}
void ReferenceCustomGBIxn::calculateSingleParticleValue(int index, int numAtoms, vector<RealVec>& atomCoordinates, RealOpenMM** atomParameters) {
......@@ -204,11 +204,11 @@ void ReferenceCustomGBIxn::calculateSingleParticleValue(int index, int numAtoms,
for (int j = 0; j < index; j++)
expressionSet.setVariable(valueIndex[j], values[j][i]);
values[index][i] = (RealOpenMM) valueExpressions[index].evaluate();
for (int j = 0; j < valueParamDerivExpressions[index].size(); j++)
dValuedParam[index][j][i] += valueParamDerivExpressions[index][j].evaluate();
// Calculate derivatives with respect to parameters.
for (int j = 0; j < valueParamDerivExpressions[index].size(); j++)
dValuedParam[index][j][i] += valueParamDerivExpressions[index][j].evaluate();
for (int j = 0; j < index; j++) {
RealOpenMM dVdV = valueDerivExpressions[index][j].evaluate();
for (int k = 0; k < valueParamDerivExpressions[index].size(); k++)
......@@ -296,11 +296,8 @@ void ReferenceCustomGBIxn::calculateSingleParticleEnergyTerm(int index, int numA
// Compute derivatives with respect to parameters.
for (int k = 0; k < energyParamDerivExpressions[index].size(); k++) {
for (int k = 0; k < energyParamDerivExpressions[index].size(); k++)
energyParamDerivs[k] += energyParamDerivExpressions[index][k].evaluate();
for (int j = 0; j < (int) valueIndex.size(); j++)
energyParamDerivs[k] += dEdV[j][i]*dValuedParam[j][k][i];
}
}
}
......@@ -376,7 +373,7 @@ void ReferenceCustomGBIxn::calculateOnePairEnergyTerm(int index, int atom1, int
}
void ReferenceCustomGBIxn::calculateChainRuleForces(int numAtoms, vector<RealVec>& atomCoordinates, RealOpenMM** atomParameters,
const vector<set<int> >& exclusions, vector<RealVec>& forces) {
const vector<set<int> >& exclusions, vector<RealVec>& forces, double* energyParamDerivs) {
if (cutoff) {
// Loop over all pairs in the neighbor list.
......@@ -426,6 +423,13 @@ void ReferenceCustomGBIxn::calculateChainRuleForces(int numAtoms, vector<RealVec
forces[i][2] -= dEdV[j][i]*dVdZ[j];
}
}
// Compute chain rule terms for derivatives with respect to parameters.
for (int i = 0; i < numAtoms; i++)
for (int j = 0; j < (int) valueIndex.size(); j++)
for (int k = 0; k < dValuedParam[j].size(); k++)
energyParamDerivs[k] += dEdV[j][i]*dValuedParam[j][k][i];
}
void ReferenceCustomGBIxn::calculateOnePairChainRule(int atom1, int atom2, vector<RealVec>& atomCoordinates, RealOpenMM** atomParameters,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment