Commit bb70341a authored by peastman's avatar peastman
Browse files

Further optimizations to CustomGBForce

parent 83a0b613
...@@ -29,6 +29,7 @@ ...@@ -29,6 +29,7 @@
#include "CpuNeighborList.h" #include "CpuNeighborList.h"
#include "lepton/CompiledExpression.h" #include "lepton/CompiledExpression.h"
#include "openmm/CustomGBForce.h" #include "openmm/CustomGBForce.h"
#include "openmm/internal/vectorize.h"
#include <map> #include <map>
#include <set> #include <set>
#include <vector> #include <vector>
...@@ -40,8 +41,8 @@ private: ...@@ -40,8 +41,8 @@ private:
bool cutoff; bool cutoff;
bool periodic; bool periodic;
const CpuNeighborList* neighborList; const CpuNeighborList* neighborList;
RealOpenMM periodicBoxSize[3]; float periodicBoxSize[3];
RealOpenMM cutoffDistance; float cutoffDistance;
CompiledExpressionSet expressionSet; CompiledExpressionSet expressionSet;
std::vector<Lepton::CompiledExpression> valueExpressions; std::vector<Lepton::CompiledExpression> valueExpressions;
std::vector<std::vector<Lepton::CompiledExpression> > valueDerivExpressions; std::vector<std::vector<Lepton::CompiledExpression> > valueDerivExpressions;
...@@ -59,20 +60,20 @@ private: ...@@ -59,20 +60,20 @@ private:
std::vector<int> particleValueIndex; std::vector<int> particleValueIndex;
int xindex, yindex, zindex, rindex; int xindex, yindex, zindex, rindex;
// Workspace vectors // Workspace vectors
std::vector<std::vector<RealOpenMM> > values, dEdV; std::vector<std::vector<float> > values, dEdV;
std::vector<RealOpenMM> dVdR1, dVdR2, dVdX, dVdY, dVdZ; std::vector<float> dVdR1, dVdR2, dVdX, dVdY, dVdZ;
/** /**
* Calculate a computed value of type SingleParticle * Calculate a computed value of type SingleParticle
* *
* @param index the index of the value to compute * @param index the index of the value to compute
* @param numAtoms number of atoms * @param numAtoms number of atoms
* @param atomCoordinates atom coordinates * @param posq atom coordinates
* @param values the vector to store computed values into * @param values the vector to store computed values into
* @param atomParameters atomParameters[atomIndex][paramterIndex] * @param atomParameters atomParameters[atomIndex][paramterIndex]
*/ */
void calculateSingleParticleValue(int index, int numAtoms, std::vector<RealVec>& atomCoordinates, std::vector<std::vector<RealOpenMM> >& values, void calculateSingleParticleValue(int index, int numAtoms, float* posq, std::vector<std::vector<float> >& values,
RealOpenMM** atomParameters); RealOpenMM** atomParameters);
/** /**
...@@ -80,16 +81,16 @@ private: ...@@ -80,16 +81,16 @@ private:
* *
* @param index the index of the value to compute * @param index the index of the value to compute
* @param numAtoms number of atoms * @param numAtoms number of atoms
* @param atomCoordinates atom coordinates * @param posq atom coordinates
* @param atomParameters atomParameters[atomIndex][paramterIndex] * @param atomParameters atomParameters[atomIndex][paramterIndex]
* @param values the vector to store computed values into * @param values the vector to store computed values into
* @param exclusions exclusions[i] is the set of excluded indices for atom i * @param exclusions exclusions[i] is the set of excluded indices for atom i
* @param useExclusions specifies whether to use exclusions * @param useExclusions specifies whether to use exclusions
*/ */
void calculateParticlePairValue(int index, int numAtoms, std::vector<RealVec>& atomCoordinates, RealOpenMM** atomParameters, void calculateParticlePairValue(int index, int numAtoms, float* posq, RealOpenMM** atomParameters,
std::vector<std::vector<RealOpenMM> >& values, std::vector<std::vector<float> >& values,
const std::vector<std::set<int> >& exclusions, bool useExclusions); const std::vector<std::set<int> >& exclusions, bool useExclusions, const fvec4& boxSize, const fvec4& invBoxSize);
/** /**
* Evaluate a single atom pair as part of calculating a computed value * Evaluate a single atom pair as part of calculating a computed value
...@@ -97,20 +98,20 @@ private: ...@@ -97,20 +98,20 @@ private:
* @param index the index of the value to compute * @param index the index of the value to compute
* @param atom1 the index of the first atom in the pair * @param atom1 the index of the first atom in the pair
* @param atom2 the index of the second atom in the pair * @param atom2 the index of the second atom in the pair
* @param atomCoordinates atom coordinates * @param posq atom coordinates
* @param atomParameters atomParameters[atomIndex][paramterIndex] * @param atomParameters atomParameters[atomIndex][paramterIndex]
* @param values the vector to store computed values into * @param values the vector to store computed values into
*/ */
void calculateOnePairValue(int index, int atom1, int atom2, std::vector<RealVec>& atomCoordinates, RealOpenMM** atomParameters, void calculateOnePairValue(int index, int atom1, int atom2, float* posq, RealOpenMM** atomParameters,
std::vector<std::vector<RealOpenMM> >& values); std::vector<std::vector<float> >& values, const fvec4& boxSize, const fvec4& invBoxSize);
/** /**
* Calculate an energy term of type SingleParticle * Calculate an energy term of type SingleParticle
* *
* @param index the index of the value to compute * @param index the index of the value to compute
* @param numAtoms number of atoms * @param numAtoms number of atoms
* @param atomCoordinates atom coordinates * @param posq atom coordinates
* @param values the vector containing computed values * @param values the vector containing computed values
* @param atomParameters atomParameters[atomIndex][paramterIndex] * @param atomParameters atomParameters[atomIndex][paramterIndex]
* @param forces forces on atoms are added to this * @param forces forces on atoms are added to this
...@@ -118,16 +119,16 @@ private: ...@@ -118,16 +119,16 @@ private:
* @param dEdV the derivative of energy with respect to computed values is stored in this * @param dEdV the derivative of energy with respect to computed values is stored in this
*/ */
void calculateSingleParticleEnergyTerm(int index, int numAtoms, std::vector<RealVec>& atomCoordinates, const std::vector<std::vector<RealOpenMM> >& values, void calculateSingleParticleEnergyTerm(int index, int numAtoms, float* posq, const std::vector<std::vector<float> >& values,
RealOpenMM** atomParameters, std::vector<RealVec>& forces, RealOpenMM** atomParameters, float* forces,
RealOpenMM* totalEnergy, std::vector<std::vector<RealOpenMM> >& dEdV); double* totalEnergy, std::vector<std::vector<float> >& dEdV);
/** /**
* Calculate an energy term that is based on particle pairs * Calculate an energy term that is based on particle pairs
* *
* @param index the index of the term to compute * @param index the index of the term to compute
* @param numAtoms number of atoms * @param numAtoms number of atoms
* @param atomCoordinates atom coordinates * @param posq atom coordinates
* @param atomParameters atomParameters[atomIndex][paramterIndex] * @param atomParameters atomParameters[atomIndex][paramterIndex]
* @param values the vector containing computed values * @param values the vector containing computed values
* @param exclusions exclusions[i] is the set of excluded indices for atom i * @param exclusions exclusions[i] is the set of excluded indices for atom i
...@@ -137,10 +138,11 @@ private: ...@@ -137,10 +138,11 @@ private:
* @param dEdV the derivative of energy with respect to computed values is stored in this * @param dEdV the derivative of energy with respect to computed values is stored in this
*/ */
void calculateParticlePairEnergyTerm(int index, int numAtoms, std::vector<RealVec>& atomCoordinates, RealOpenMM** atomParameters, void calculateParticlePairEnergyTerm(int index, int numAtoms, float* posq, RealOpenMM** atomParameters,
const std::vector<std::vector<RealOpenMM> >& values, const std::vector<std::vector<float> >& values,
const std::vector<std::set<int> >& exclusions, bool useExclusions, const std::vector<std::set<int> >& exclusions, bool useExclusions,
std::vector<RealVec>& forces, RealOpenMM* totalEnergy, std::vector<std::vector<RealOpenMM> >& dEdV); float* forces, double* totalEnergy, std::vector<std::vector<float> >& dEdV,
const fvec4& boxSize, const fvec4& invBoxSize);
/** /**
* Evaluate a single atom pair as part of calculating an energy term * Evaluate a single atom pair as part of calculating an energy term
...@@ -148,7 +150,7 @@ private: ...@@ -148,7 +150,7 @@ private:
* @param index the index of the term to compute * @param index the index of the term to compute
* @param atom1 the index of the first atom in the pair * @param atom1 the index of the first atom in the pair
* @param atom2 the index of the second atom in the pair * @param atom2 the index of the second atom in the pair
* @param atomCoordinates atom coordinates * @param posq atom coordinates
* @param atomParameters atomParameters[atomIndex][paramterIndex] * @param atomParameters atomParameters[atomIndex][paramterIndex]
* @param values the vector containing computed values * @param values the vector containing computed values
* @param forces forces on atoms are added to this * @param forces forces on atoms are added to this
...@@ -156,15 +158,16 @@ private: ...@@ -156,15 +158,16 @@ private:
* @param dEdV the derivative of energy with respect to computed values is stored in this * @param dEdV the derivative of energy with respect to computed values is stored in this
*/ */
void calculateOnePairEnergyTerm(int index, int atom1, int atom2, std::vector<RealVec>& atomCoordinates, RealOpenMM** atomParameters, void calculateOnePairEnergyTerm(int index, int atom1, int atom2, float* posq, RealOpenMM** atomParameters,
const std::vector<std::vector<RealOpenMM> >& values, const std::vector<std::vector<float> >& values,
std::vector<RealVec>& forces, RealOpenMM* totalEnergy, std::vector<std::vector<RealOpenMM> >& dEdV); float* forces, double* totalEnergy, std::vector<std::vector<float> >& dEdV,
const fvec4& boxSize, const fvec4& invBoxSize);
/** /**
* Apply the chain rule to compute forces on atoms * Apply the chain rule to compute forces on atoms
* *
* @param numAtoms number of atoms * @param numAtoms number of atoms
* @param atomCoordinates atom coordinates * @param posq atom coordinates
* @param atomParameters atomParameters[atomIndex][paramterIndex] * @param atomParameters atomParameters[atomIndex][paramterIndex]
* @param values the vector containing computed values * @param values the vector containing computed values
* @param exclusions exclusions[i] is the set of excluded indices for atom i * @param exclusions exclusions[i] is the set of excluded indices for atom i
...@@ -172,17 +175,18 @@ private: ...@@ -172,17 +175,18 @@ private:
* @param dEdV the derivative of energy with respect to computed values is stored in this * @param dEdV the derivative of energy with respect to computed values is stored in this
*/ */
void calculateChainRuleForces(int numAtoms, std::vector<RealVec>& atomCoordinates, RealOpenMM** atomParameters, void calculateChainRuleForces(int numAtoms, float* posq, RealOpenMM** atomParameters,
const std::vector<std::vector<RealOpenMM> >& values, const std::vector<std::vector<float> >& values,
const std::vector<std::set<int> >& exclusions, const std::vector<std::set<int> >& exclusions,
std::vector<RealVec>& forces, std::vector<std::vector<RealOpenMM> >& dEdV); float* forces, std::vector<std::vector<float> >& dEdV,
const fvec4& boxSize, const fvec4& invBoxSize);
/** /**
* Evaluate a single atom pair as part of applying the chain rule * Evaluate a single atom pair as part of applying the chain rule
* *
* @param atom1 the index of the first atom in the pair * @param atom1 the index of the first atom in the pair
* @param atom2 the index of the second atom in the pair * @param atom2 the index of the second atom in the pair
* @param atomCoordinates atom coordinates * @param posq atom coordinates
* @param atomParameters atomParameters[atomIndex][paramterIndex] * @param atomParameters atomParameters[atomIndex][paramterIndex]
* @param values the vector containing computed values * @param values the vector containing computed values
* @param forces forces on atoms are added to this * @param forces forces on atoms are added to this
...@@ -190,10 +194,16 @@ private: ...@@ -190,10 +194,16 @@ private:
* @param isExcluded specifies whether this is an excluded pair * @param isExcluded specifies whether this is an excluded pair
*/ */
void calculateOnePairChainRule(int atom1, int atom2, std::vector<RealVec>& atomCoordinates, RealOpenMM** atomParameters, void calculateOnePairChainRule(int atom1, int atom2, float* posq, RealOpenMM** atomParameters,
const std::vector<std::vector<RealOpenMM> >& values, const std::vector<std::vector<float> >& values,
std::vector<RealVec>& forces, std::vector<std::vector<RealOpenMM> >& dEdV, float* forces, std::vector<std::vector<float> >& dEdV,
bool isExcluded); bool isExcluded, const fvec4& boxSize, const fvec4& invBoxSize);
/**
* Compute the displacement and squared distance between two points, optionally using
* periodic boundary conditions.
*/
void getDeltaR(const fvec4& posI, const fvec4& posJ, fvec4& deltaR, float& r2, bool periodic, const fvec4& boxSize, const fvec4& invBoxSize) const;
public: public:
...@@ -221,7 +231,7 @@ public: ...@@ -221,7 +231,7 @@ public:
* @param neighbors the neighbor list to use * @param neighbors the neighbor list to use
*/ */
void setUseCutoff(RealOpenMM distance, const CpuNeighborList& neighbors); void setUseCutoff(float distance, const CpuNeighborList& neighbors);
/** /**
* Set the force to use periodic boundary conditions. This requires that a cutoff has * Set the force to use periodic boundary conditions. This requires that a cutoff has
...@@ -237,7 +247,7 @@ public: ...@@ -237,7 +247,7 @@ public:
* Calculate custom GB ixn * Calculate custom GB ixn
* *
* @param numberOfAtoms number of atoms * @param numberOfAtoms number of atoms
* @param atomCoordinates atom coordinates * @param posq atom coordinates
* @param atomParameters atomParameters[atomIndex][paramterIndex] * @param atomParameters atomParameters[atomIndex][paramterIndex]
* @param exclusions exclusions[i] is the set of excluded indices for atom i * @param exclusions exclusions[i] is the set of excluded indices for atom i
* @param globalParameters the values of global parameters * @param globalParameters the values of global parameters
...@@ -245,8 +255,8 @@ public: ...@@ -245,8 +255,8 @@ public:
* @param totalEnergy total energy * @param totalEnergy total energy
*/ */
void calculateIxn(int numberOfAtoms, std::vector<RealVec>& atomCoordinates, RealOpenMM** atomParameters, const std::vector<std::set<int> >& exclusions, void calculateIxn(int numberOfAtoms, float* posq, RealOpenMM** atomParameters, const std::vector<std::set<int> >& exclusions,
std::map<std::string, double>& globalParameters, std::vector<RealVec>& forces, RealOpenMM* totalEnergy); std::map<std::string, double>& globalParameters, float* forces, double* totalEnergy);
}; };
} // namespace OpenMM } // namespace OpenMM
......
...@@ -99,7 +99,7 @@ CpuCustomGBForce::CpuCustomGBForce(int numAtoms, const vector<Lepton::CompiledEx ...@@ -99,7 +99,7 @@ CpuCustomGBForce::CpuCustomGBForce(int numAtoms, const vector<Lepton::CompiledEx
CpuCustomGBForce::~CpuCustomGBForce() { CpuCustomGBForce::~CpuCustomGBForce() {
} }
void CpuCustomGBForce::setUseCutoff(RealOpenMM distance, const CpuNeighborList& neighbors) { void CpuCustomGBForce::setUseCutoff(float distance, const CpuNeighborList& neighbors) {
cutoff = true; cutoff = true;
cutoffDistance = distance; cutoffDistance = distance;
neighborList = &neighbors; neighborList = &neighbors;
...@@ -118,9 +118,11 @@ void CpuCustomGBForce::setPeriodic(RealVec& boxSize) { ...@@ -118,9 +118,11 @@ void CpuCustomGBForce::setPeriodic(RealVec& boxSize) {
periodicBoxSize[2] = boxSize[2]; periodicBoxSize[2] = boxSize[2];
} }
void CpuCustomGBForce::calculateIxn(int numberOfAtoms, vector<RealVec>& atomCoordinates, RealOpenMM** atomParameters, void CpuCustomGBForce::calculateIxn(int numberOfAtoms, float* posq, RealOpenMM** atomParameters,
const vector<set<int> >& exclusions, map<string, double>& globalParameters, vector<RealVec>& forces, const vector<set<int> >& exclusions, map<string, double>& globalParameters, float* forces,
RealOpenMM* totalEnergy) { double* totalEnergy) {
fvec4 boxSize(periodicBoxSize[0], periodicBoxSize[1], periodicBoxSize[2], 0);
fvec4 invBoxSize((1/periodicBoxSize[0]), (1/periodicBoxSize[1]), (1/periodicBoxSize[2]), 0);
for (map<string, double>::const_iterator iter = globalParameters.begin(); iter != globalParameters.end(); ++iter) for (map<string, double>::const_iterator iter = globalParameters.begin(); iter != globalParameters.end(); ++iter)
expressionSet.setVariable(expressionSet.getVariableIndex(iter->first), iter->second); expressionSet.setVariable(expressionSet.getVariableIndex(iter->first), iter->second);
...@@ -129,11 +131,11 @@ void CpuCustomGBForce::calculateIxn(int numberOfAtoms, vector<RealVec>& atomCoor ...@@ -129,11 +131,11 @@ void CpuCustomGBForce::calculateIxn(int numberOfAtoms, vector<RealVec>& atomCoor
int numValues = valueTypes.size(); int numValues = valueTypes.size();
for (int valueIndex = 0; valueIndex < numValues; valueIndex++) { for (int valueIndex = 0; valueIndex < numValues; valueIndex++) {
if (valueTypes[valueIndex] == CustomGBForce::SingleParticle) if (valueTypes[valueIndex] == CustomGBForce::SingleParticle)
calculateSingleParticleValue(valueIndex, numberOfAtoms, atomCoordinates, values, atomParameters); calculateSingleParticleValue(valueIndex, numberOfAtoms, posq, values, atomParameters);
else if (valueTypes[valueIndex] == CustomGBForce::ParticlePair) else if (valueTypes[valueIndex] == CustomGBForce::ParticlePair)
calculateParticlePairValue(valueIndex, numberOfAtoms, atomCoordinates, atomParameters, values, exclusions, true); calculateParticlePairValue(valueIndex, numberOfAtoms, posq, atomParameters, values, exclusions, true, boxSize, invBoxSize);
else else
calculateParticlePairValue(valueIndex, numberOfAtoms, atomCoordinates, atomParameters, values, exclusions, false); calculateParticlePairValue(valueIndex, numberOfAtoms, posq, atomParameters, values, exclusions, false, boxSize, invBoxSize);
} }
// Now calculate the energy and its derivatives. // Now calculate the energy and its derivatives.
...@@ -143,38 +145,38 @@ void CpuCustomGBForce::calculateIxn(int numberOfAtoms, vector<RealVec>& atomCoor ...@@ -143,38 +145,38 @@ void CpuCustomGBForce::calculateIxn(int numberOfAtoms, vector<RealVec>& atomCoor
dEdV[i][j] = 0.0; dEdV[i][j] = 0.0;
for (int termIndex = 0; termIndex < (int) energyExpressions.size(); termIndex++) { for (int termIndex = 0; termIndex < (int) energyExpressions.size(); termIndex++) {
if (energyTypes[termIndex] == CustomGBForce::SingleParticle) if (energyTypes[termIndex] == CustomGBForce::SingleParticle)
calculateSingleParticleEnergyTerm(termIndex, numberOfAtoms, atomCoordinates, values, atomParameters, forces, totalEnergy, dEdV); calculateSingleParticleEnergyTerm(termIndex, numberOfAtoms, posq, values, atomParameters, forces, totalEnergy, dEdV);
else if (energyTypes[termIndex] == CustomGBForce::ParticlePair) else if (energyTypes[termIndex] == CustomGBForce::ParticlePair)
calculateParticlePairEnergyTerm(termIndex, numberOfAtoms, atomCoordinates, atomParameters, values, exclusions, true, forces, totalEnergy, dEdV); calculateParticlePairEnergyTerm(termIndex, numberOfAtoms, posq, atomParameters, values, exclusions, true, forces, totalEnergy, dEdV, boxSize, invBoxSize);
else else
calculateParticlePairEnergyTerm(termIndex, numberOfAtoms, atomCoordinates, atomParameters, values, exclusions, false, forces, totalEnergy, dEdV); calculateParticlePairEnergyTerm(termIndex, numberOfAtoms, posq, atomParameters, values, exclusions, false, forces, totalEnergy, dEdV, boxSize, invBoxSize);
} }
// Apply the chain rule to evaluate forces. // Apply the chain rule to evaluate forces.
calculateChainRuleForces(numberOfAtoms, atomCoordinates, atomParameters, values, exclusions, forces, dEdV); calculateChainRuleForces(numberOfAtoms, posq, atomParameters, values, exclusions, forces, dEdV, boxSize, invBoxSize);
} }
void CpuCustomGBForce::calculateSingleParticleValue(int index, int numAtoms, vector<RealVec>& atomCoordinates, vector<vector<RealOpenMM> >& values, void CpuCustomGBForce::calculateSingleParticleValue(int index, int numAtoms, float* posq, vector<vector<float> >& values,
RealOpenMM** atomParameters) { RealOpenMM** atomParameters) {
values[index].resize(numAtoms); values[index].resize(numAtoms);
for (int i = 0; i < numAtoms; i++) { for (int i = 0; i < numAtoms; i++) {
expressionSet.setVariable(xindex, atomCoordinates[i][0]); expressionSet.setVariable(xindex, posq[4*i]);
expressionSet.setVariable(yindex, atomCoordinates[i][1]); expressionSet.setVariable(yindex, posq[4*i+1]);
expressionSet.setVariable(zindex, atomCoordinates[i][2]); expressionSet.setVariable(zindex, posq[4*i+2]);
for (int j = 0; j < (int) paramNames.size(); j++) for (int j = 0; j < (int) paramNames.size(); j++)
expressionSet.setVariable(paramIndex[j], atomParameters[i][j]); expressionSet.setVariable(paramIndex[j], atomParameters[i][j]);
for (int j = 0; j < index; j++) for (int j = 0; j < index; j++)
expressionSet.setVariable(valueIndex[j], values[j][i]); expressionSet.setVariable(valueIndex[j], values[j][i]);
values[index][i] = (RealOpenMM) valueExpressions[index].evaluate(); values[index][i] = (float) valueExpressions[index].evaluate();
} }
} }
void CpuCustomGBForce::calculateParticlePairValue(int index, int numAtoms, vector<RealVec>& atomCoordinates, RealOpenMM** atomParameters, void CpuCustomGBForce::calculateParticlePairValue(int index, int numAtoms, float* posq, RealOpenMM** atomParameters,
vector<vector<RealOpenMM> >& values, const vector<set<int> >& exclusions, bool useExclusions) { vector<vector<float> >& values, const vector<set<int> >& exclusions, bool useExclusions, const fvec4& boxSize, const fvec4& invBoxSize) {
values[index].resize(numAtoms); values[index].resize(numAtoms);
for (int i = 0; i < numAtoms; i++) for (int i = 0; i < numAtoms; i++)
values[index][i] = (RealOpenMM) 0.0; values[index][i] = 0.0f;
if (cutoff) { if (cutoff) {
// Loop over all pairs in the neighbor list. // Loop over all pairs in the neighbor list.
...@@ -189,8 +191,8 @@ void CpuCustomGBForce::calculateParticlePairValue(int index, int numAtoms, vecto ...@@ -189,8 +191,8 @@ void CpuCustomGBForce::calculateParticlePairValue(int index, int numAtoms, vecto
int second = blockAtom[k]; int second = blockAtom[k];
if (useExclusions && exclusions[first].find(second) != exclusions[first].end()) if (useExclusions && exclusions[first].find(second) != exclusions[first].end())
continue; continue;
calculateOnePairValue(index, first, second, atomCoordinates, atomParameters, values); calculateOnePairValue(index, first, second, posq, atomParameters, values, boxSize, invBoxSize);
calculateOnePairValue(index, second, first, atomCoordinates, atomParameters, values); calculateOnePairValue(index, second, first, posq, atomParameters, values, boxSize, invBoxSize);
} }
} }
} }
...@@ -203,21 +205,21 @@ void CpuCustomGBForce::calculateParticlePairValue(int index, int numAtoms, vecto ...@@ -203,21 +205,21 @@ void CpuCustomGBForce::calculateParticlePairValue(int index, int numAtoms, vecto
for (int j = i+1; j < numAtoms; j++) { for (int j = i+1; j < numAtoms; j++) {
if (useExclusions && exclusions[i].find(j) != exclusions[i].end()) if (useExclusions && exclusions[i].find(j) != exclusions[i].end())
continue; continue;
calculateOnePairValue(index, i, j, atomCoordinates, atomParameters, values); calculateOnePairValue(index, i, j, posq, atomParameters, values, boxSize, invBoxSize);
calculateOnePairValue(index, j, i, atomCoordinates, atomParameters, values); calculateOnePairValue(index, j, i, posq, atomParameters, values, boxSize, invBoxSize);
} }
} }
} }
} }
void CpuCustomGBForce::calculateOnePairValue(int index, int atom1, int atom2, vector<RealVec>& atomCoordinates, RealOpenMM** atomParameters, void CpuCustomGBForce::calculateOnePairValue(int index, int atom1, int atom2, float* posq, RealOpenMM** atomParameters,
vector<vector<RealOpenMM> >& values) { vector<vector<float> >& values, const fvec4& boxSize, const fvec4& invBoxSize) {
RealOpenMM deltaR[ReferenceForce::LastDeltaRIndex]; fvec4 deltaR;
if (periodic) fvec4 pos1(posq+4*atom1);
ReferenceForce::getDeltaRPeriodic(atomCoordinates[atom2], atomCoordinates[atom1], periodicBoxSize, deltaR); fvec4 pos2(posq+4*atom2);
else float r2;
ReferenceForce::getDeltaR(atomCoordinates[atom2], atomCoordinates[atom1], deltaR); getDeltaR(pos2, pos1, deltaR, r2, periodic, boxSize, invBoxSize);
RealOpenMM r = deltaR[ReferenceForce::RIndex]; float r = sqrtf(r2);
if (cutoff && r >= cutoffDistance) if (cutoff && r >= cutoffDistance)
return; return;
for (int i = 0; i < (int) paramNames.size(); i++) { for (int i = 0; i < (int) paramNames.size(); i++) {
...@@ -229,33 +231,33 @@ void CpuCustomGBForce::calculateOnePairValue(int index, int atom1, int atom2, ve ...@@ -229,33 +231,33 @@ void CpuCustomGBForce::calculateOnePairValue(int index, int atom1, int atom2, ve
expressionSet.setVariable(particleValueIndex[i*2], values[i][atom1]); expressionSet.setVariable(particleValueIndex[i*2], values[i][atom1]);
expressionSet.setVariable(particleValueIndex[i*2+1], values[i][atom2]); expressionSet.setVariable(particleValueIndex[i*2+1], values[i][atom2]);
} }
values[index][atom1] += (RealOpenMM) valueExpressions[index].evaluate(); values[index][atom1] += (float) valueExpressions[index].evaluate();
} }
void CpuCustomGBForce::calculateSingleParticleEnergyTerm(int index, int numAtoms, vector<RealVec>& atomCoordinates, const vector<vector<RealOpenMM> >& values, void CpuCustomGBForce::calculateSingleParticleEnergyTerm(int index, int numAtoms, float* posq, const vector<vector<float> >& values,
RealOpenMM** atomParameters, vector<RealVec>& forces, RealOpenMM* totalEnergy, RealOpenMM** atomParameters, float* forces, double* totalEnergy,
vector<vector<RealOpenMM> >& dEdV) { vector<vector<float> >& dEdV) {
for (int i = 0; i < numAtoms; i++) { for (int i = 0; i < numAtoms; i++) {
expressionSet.setVariable(xindex, atomCoordinates[i][0]); expressionSet.setVariable(xindex, posq[4*i]);
expressionSet.setVariable(yindex, atomCoordinates[i][1]); expressionSet.setVariable(yindex, posq[4*i+1]);
expressionSet.setVariable(zindex, atomCoordinates[i][2]); expressionSet.setVariable(zindex, posq[4*i+2]);
for (int j = 0; j < (int) paramNames.size(); j++) for (int j = 0; j < (int) paramNames.size(); j++)
expressionSet.setVariable(paramIndex[j], atomParameters[i][j]); expressionSet.setVariable(paramIndex[j], atomParameters[i][j]);
for (int j = 0; j < (int) valueNames.size(); j++) for (int j = 0; j < (int) valueNames.size(); j++)
expressionSet.setVariable(valueIndex[j], values[j][i]); expressionSet.setVariable(valueIndex[j], values[j][i]);
if (totalEnergy != NULL) if (totalEnergy != NULL)
*totalEnergy += (RealOpenMM) energyExpressions[index].evaluate(); *totalEnergy += (float) energyExpressions[index].evaluate();
for (int j = 0; j < (int) valueNames.size(); j++) for (int j = 0; j < (int) valueNames.size(); j++)
dEdV[j][i] += (RealOpenMM) energyDerivExpressions[index][j].evaluate(); dEdV[j][i] += (float) energyDerivExpressions[index][j].evaluate();
forces[i][0] -= (RealOpenMM) energyGradientExpressions[index][0].evaluate(); forces[4*i+0] -= (float) energyGradientExpressions[index][0].evaluate();
forces[i][1] -= (RealOpenMM) energyGradientExpressions[index][1].evaluate(); forces[4*i+1] -= (float) energyGradientExpressions[index][1].evaluate();
forces[i][2] -= (RealOpenMM) energyGradientExpressions[index][2].evaluate(); forces[4*i+2] -= (float) energyGradientExpressions[index][2].evaluate();
} }
} }
void CpuCustomGBForce::calculateParticlePairEnergyTerm(int index, int numAtoms, vector<RealVec>& atomCoordinates, RealOpenMM** atomParameters, void CpuCustomGBForce::calculateParticlePairEnergyTerm(int index, int numAtoms, float* posq, RealOpenMM** atomParameters,
const vector<vector<RealOpenMM> >& values, const vector<set<int> >& exclusions, bool useExclusions, const vector<vector<float> >& values, const vector<set<int> >& exclusions, bool useExclusions,
vector<RealVec>& forces, RealOpenMM* totalEnergy, vector<vector<RealOpenMM> >& dEdV) { float* forces, double* totalEnergy, vector<vector<float> >& dEdV, const fvec4& boxSize, const fvec4& invBoxSize) {
if (cutoff) { if (cutoff) {
// Loop over all pairs in the neighbor list. // Loop over all pairs in the neighbor list.
...@@ -270,7 +272,7 @@ void CpuCustomGBForce::calculateParticlePairEnergyTerm(int index, int numAtoms, ...@@ -270,7 +272,7 @@ void CpuCustomGBForce::calculateParticlePairEnergyTerm(int index, int numAtoms,
int second = blockAtom[k]; int second = blockAtom[k];
if (useExclusions && exclusions[first].find(second) != exclusions[first].end()) if (useExclusions && exclusions[first].find(second) != exclusions[first].end())
continue; continue;
calculateOnePairEnergyTerm(index, first, second, atomCoordinates, atomParameters, values, forces, totalEnergy, dEdV); calculateOnePairEnergyTerm(index, first, second, posq, atomParameters, values, forces, totalEnergy, dEdV, boxSize, invBoxSize);
} }
} }
} }
...@@ -283,23 +285,23 @@ void CpuCustomGBForce::calculateParticlePairEnergyTerm(int index, int numAtoms, ...@@ -283,23 +285,23 @@ void CpuCustomGBForce::calculateParticlePairEnergyTerm(int index, int numAtoms,
for (int j = i+1; j < numAtoms; j++) { for (int j = i+1; j < numAtoms; j++) {
if (useExclusions && exclusions[i].find(j) != exclusions[i].end()) if (useExclusions && exclusions[i].find(j) != exclusions[i].end())
continue; continue;
calculateOnePairEnergyTerm(index, i, j, atomCoordinates, atomParameters, values, forces, totalEnergy, dEdV); calculateOnePairEnergyTerm(index, i, j, posq, atomParameters, values, forces, totalEnergy, dEdV, boxSize, invBoxSize);
} }
} }
} }
} }
void CpuCustomGBForce::calculateOnePairEnergyTerm(int index, int atom1, int atom2, vector<RealVec>& atomCoordinates, RealOpenMM** atomParameters, void CpuCustomGBForce::calculateOnePairEnergyTerm(int index, int atom1, int atom2, float* posq, RealOpenMM** atomParameters,
const vector<vector<RealOpenMM> >& values, vector<RealVec>& forces, RealOpenMM* totalEnergy, const vector<vector<float> >& values, float* forces, double* totalEnergy,
vector<vector<RealOpenMM> >& dEdV) { vector<vector<float> >& dEdV, const fvec4& boxSize, const fvec4& invBoxSize) {
// Compute the displacement. // Compute the displacement.
RealOpenMM deltaR[ReferenceForce::LastDeltaRIndex]; fvec4 deltaR;
if (periodic) fvec4 pos1(posq+4*atom1);
ReferenceForce::getDeltaRPeriodic(atomCoordinates[atom2], atomCoordinates[atom1], periodicBoxSize, deltaR); fvec4 pos2(posq+4*atom2);
else float r2;
ReferenceForce::getDeltaR(atomCoordinates[atom2], atomCoordinates[atom1], deltaR); getDeltaR(pos2, pos1, deltaR, r2, periodic, boxSize, invBoxSize);
RealOpenMM r = deltaR[ReferenceForce::RIndex]; float r = sqrtf(r2);
if (cutoff && r >= cutoffDistance) if (cutoff && r >= cutoffDistance)
return; return;
...@@ -318,21 +320,21 @@ void CpuCustomGBForce::calculateOnePairEnergyTerm(int index, int atom1, int atom ...@@ -318,21 +320,21 @@ void CpuCustomGBForce::calculateOnePairEnergyTerm(int index, int atom1, int atom
// Evaluate the energy and its derivatives. // Evaluate the energy and its derivatives.
if (totalEnergy != NULL) if (totalEnergy != NULL)
*totalEnergy += (RealOpenMM) energyExpressions[index].evaluate(); *totalEnergy += (float) energyExpressions[index].evaluate();
RealOpenMM dEdR = (RealOpenMM) energyDerivExpressions[index][0].evaluate(); float dEdR = (float) energyDerivExpressions[index][0].evaluate();
dEdR *= 1/r; dEdR *= 1/r;
for (int i = 0; i < 3; i++) { fvec4 result = deltaR*dEdR;
forces[atom1][i] -= dEdR*deltaR[i]; (fvec4(forces+4*atom1)-result).store(forces+4*atom1);
forces[atom2][i] += dEdR*deltaR[i]; (fvec4(forces+4*atom2)+result).store(forces+4*atom2);
}
for (int i = 0; i < (int) valueNames.size(); i++) { for (int i = 0; i < (int) valueNames.size(); i++) {
dEdV[i][atom1] += (RealOpenMM) energyDerivExpressions[index][2*i+1].evaluate(); dEdV[i][atom1] += (float) energyDerivExpressions[index][2*i+1].evaluate();
dEdV[i][atom2] += (RealOpenMM) energyDerivExpressions[index][2*i+2].evaluate(); dEdV[i][atom2] += (float) energyDerivExpressions[index][2*i+2].evaluate();
} }
} }
void CpuCustomGBForce::calculateChainRuleForces(int numAtoms, vector<RealVec>& atomCoordinates, RealOpenMM** atomParameters, void CpuCustomGBForce::calculateChainRuleForces(int numAtoms, float* posq, RealOpenMM** atomParameters,
const vector<vector<RealOpenMM> >& values, const vector<set<int> >& exclusions, vector<RealVec>& forces, vector<vector<RealOpenMM> >& dEdV) { const vector<vector<float> >& values, const vector<set<int> >& exclusions, float* forces, vector<vector<float> >& dEdV,
const fvec4& boxSize, const fvec4& invBoxSize) {
if (cutoff) { if (cutoff) {
// Loop over all pairs in the neighbor list. // Loop over all pairs in the neighbor list.
...@@ -346,8 +348,8 @@ void CpuCustomGBForce::calculateChainRuleForces(int numAtoms, vector<RealVec>& a ...@@ -346,8 +348,8 @@ void CpuCustomGBForce::calculateChainRuleForces(int numAtoms, vector<RealVec>& a
if ((blockExclusions[i] & (1<<k)) == 0) { if ((blockExclusions[i] & (1<<k)) == 0) {
int second = blockAtom[k]; int second = blockAtom[k];
bool isExcluded = (exclusions[first].find(second) != exclusions[first].end()); bool isExcluded = (exclusions[first].find(second) != exclusions[first].end());
calculateOnePairChainRule(first, second, atomCoordinates, atomParameters, values, forces, dEdV, isExcluded); calculateOnePairChainRule(first, second, posq, atomParameters, values, forces, dEdV, isExcluded, boxSize, invBoxSize);
calculateOnePairChainRule(second, first, atomCoordinates, atomParameters, values, forces, dEdV, isExcluded); calculateOnePairChainRule(second, first, posq, atomParameters, values, forces, dEdV, isExcluded, boxSize, invBoxSize);
} }
} }
} }
...@@ -359,8 +361,8 @@ void CpuCustomGBForce::calculateChainRuleForces(int numAtoms, vector<RealVec>& a ...@@ -359,8 +361,8 @@ void CpuCustomGBForce::calculateChainRuleForces(int numAtoms, vector<RealVec>& a
for (int i = 0; i < numAtoms; i++) { for (int i = 0; i < numAtoms; i++) {
for (int j = i+1; j < numAtoms; j++) { for (int j = i+1; j < numAtoms; j++) {
bool isExcluded = (exclusions[i].find(j) != exclusions[i].end()); bool isExcluded = (exclusions[i].find(j) != exclusions[i].end());
calculateOnePairChainRule(i, j, atomCoordinates, atomParameters, values, forces, dEdV, isExcluded); calculateOnePairChainRule(i, j, posq, atomParameters, values, forces, dEdV, isExcluded, boxSize, invBoxSize);
calculateOnePairChainRule(j, i, atomCoordinates, atomParameters, values, forces, dEdV, isExcluded); calculateOnePairChainRule(j, i, posq, atomParameters, values, forces, dEdV, isExcluded, boxSize, invBoxSize);
} }
} }
} }
...@@ -368,9 +370,9 @@ void CpuCustomGBForce::calculateChainRuleForces(int numAtoms, vector<RealVec>& a ...@@ -368,9 +370,9 @@ void CpuCustomGBForce::calculateChainRuleForces(int numAtoms, vector<RealVec>& a
// Compute chain rule terms for computed values that depend explicitly on particle coordinates. // Compute chain rule terms for computed values that depend explicitly on particle coordinates.
for (int i = 0; i < numAtoms; i++) { for (int i = 0; i < numAtoms; i++) {
expressionSet.setVariable(xindex, atomCoordinates[i][0]); expressionSet.setVariable(xindex, posq[4*i]);
expressionSet.setVariable(yindex, atomCoordinates[i][1]); expressionSet.setVariable(yindex, posq[4*i+1]);
expressionSet.setVariable(zindex, atomCoordinates[i][2]); expressionSet.setVariable(zindex, posq[4*i+2]);
for (int j = 0; j < (int) paramNames.size(); j++) for (int j = 0; j < (int) paramNames.size(); j++)
expressionSet.setVariable(paramIndex[j], atomParameters[i][j]); expressionSet.setVariable(paramIndex[j], atomParameters[i][j]);
for (int j = 1; j < (int) valueNames.size(); j++) { for (int j = 1; j < (int) valueNames.size(); j++) {
...@@ -379,31 +381,31 @@ void CpuCustomGBForce::calculateChainRuleForces(int numAtoms, vector<RealVec>& a ...@@ -379,31 +381,31 @@ void CpuCustomGBForce::calculateChainRuleForces(int numAtoms, vector<RealVec>& a
dVdY[j] = 0.0; dVdY[j] = 0.0;
dVdZ[j] = 0.0; dVdZ[j] = 0.0;
for (int k = 1; k < j; k++) { for (int k = 1; k < j; k++) {
RealOpenMM dVdV = (RealOpenMM) valueDerivExpressions[j][k].evaluate(); float dVdV = (float) valueDerivExpressions[j][k].evaluate();
dVdX[j] += dVdV*dVdX[k]; dVdX[j] += dVdV*dVdX[k];
dVdY[j] += dVdV*dVdY[k]; dVdY[j] += dVdV*dVdY[k];
dVdZ[j] += dVdV*dVdZ[k]; dVdZ[j] += dVdV*dVdZ[k];
} }
dVdX[j] += (RealOpenMM) valueGradientExpressions[j][0].evaluate(); dVdX[j] += (float) valueGradientExpressions[j][0].evaluate();
dVdY[j] += (RealOpenMM) valueGradientExpressions[j][1].evaluate(); dVdY[j] += (float) valueGradientExpressions[j][1].evaluate();
dVdZ[j] += (RealOpenMM) valueGradientExpressions[j][2].evaluate(); dVdZ[j] += (float) valueGradientExpressions[j][2].evaluate();
forces[i][0] -= dEdV[j][i]*dVdX[j]; forces[4*i+0] -= dEdV[j][i]*dVdX[j];
forces[i][1] -= dEdV[j][i]*dVdY[j]; forces[4*i+1] -= dEdV[j][i]*dVdY[j];
forces[i][2] -= dEdV[j][i]*dVdZ[j]; forces[4*i+2] -= dEdV[j][i]*dVdZ[j];
} }
} }
} }
void CpuCustomGBForce::calculateOnePairChainRule(int atom1, int atom2, vector<RealVec>& atomCoordinates, RealOpenMM** atomParameters, void CpuCustomGBForce::calculateOnePairChainRule(int atom1, int atom2, float* posq, RealOpenMM** atomParameters,
const vector<vector<RealOpenMM> >& values, vector<RealVec>& forces, vector<vector<RealOpenMM> >& dEdV, bool isExcluded) { const vector<vector<float> >& values, float* forces, vector<vector<float> >& dEdV, bool isExcluded, const fvec4& boxSize, const fvec4& invBoxSize) {
// Compute the displacement. // Compute the displacement.
RealOpenMM deltaR[ReferenceForce::LastDeltaRIndex]; fvec4 deltaR;
if (periodic) fvec4 pos1(posq+4*atom1);
ReferenceForce::getDeltaRPeriodic(atomCoordinates[atom2], atomCoordinates[atom1], periodicBoxSize, deltaR); fvec4 pos2(posq+4*atom2);
else float r2;
ReferenceForce::getDeltaR(atomCoordinates[atom2], atomCoordinates[atom1], deltaR); getDeltaR(pos2, pos1, deltaR, r2, periodic, boxSize, invBoxSize);
RealOpenMM r = deltaR[ReferenceForce::RIndex]; float r = sqrtf(r2);
if (cutoff && r >= cutoffDistance) if (cutoff && r >= cutoffDistance)
return; return;
...@@ -419,36 +421,39 @@ void CpuCustomGBForce::calculateOnePairChainRule(int atom1, int atom2, vector<Re ...@@ -419,36 +421,39 @@ void CpuCustomGBForce::calculateOnePairChainRule(int atom1, int atom2, vector<Re
// Evaluate the derivative of each parameter with respect to position and apply forces. // Evaluate the derivative of each parameter with respect to position and apply forces.
RealOpenMM rinv = 1/r; float rinv = 1/r;
deltaR[0] *= rinv; deltaR *= rinv;
deltaR[1] *= rinv;
deltaR[2] *= rinv;
if (!isExcluded || valueTypes[0] != CustomGBForce::ParticlePair) { if (!isExcluded || valueTypes[0] != CustomGBForce::ParticlePair) {
dVdR1[0] = (RealOpenMM) valueDerivExpressions[0][0].evaluate(); dVdR1[0] = (float) valueDerivExpressions[0][0].evaluate();
dVdR2[0] = -dVdR1[0]; dVdR2[0] = -dVdR1[0];
for (int i = 0; i < 3; i++) { (fvec4(forces+4*atom1)-deltaR*(dEdV[0][atom1]*dVdR1[0])).store(forces+4*atom1);
forces[atom1][i] -= dEdV[0][atom1]*dVdR1[0]*deltaR[i]; (fvec4(forces+4*atom2)-deltaR*(dEdV[0][atom1]*dVdR2[0])).store(forces+4*atom2);
forces[atom2][i] -= dEdV[0][atom1]*dVdR2[0]*deltaR[i];
}
} }
for (int i = 0; i < (int) paramNames.size(); i++) for (int i = 0; i < (int) paramNames.size(); i++)
expressionSet.setVariable(paramIndex[i], atomParameters[atom1][i]); expressionSet.setVariable(paramIndex[i], atomParameters[atom1][i]);
expressionSet.setVariable(valueIndex[0], values[0][atom1]); expressionSet.setVariable(valueIndex[0], values[0][atom1]);
for (int i = 1; i < (int) valueNames.size(); i++) { for (int i = 1; i < (int) valueNames.size(); i++) {
expressionSet.setVariable(valueIndex[i], values[i][atom1]); expressionSet.setVariable(valueIndex[i], values[i][atom1]);
expressionSet.setVariable(xindex, atomCoordinates[atom1][0]); expressionSet.setVariable(xindex, posq[4*atom1]);
expressionSet.setVariable(yindex, atomCoordinates[atom1][1]); expressionSet.setVariable(yindex, posq[4*atom1+1]);
expressionSet.setVariable(zindex, atomCoordinates[atom1][2]); expressionSet.setVariable(zindex, posq[4*atom1+2]);
dVdR1[i] = 0.0; dVdR1[i] = 0.0;
dVdR2[i] = 0.0; dVdR2[i] = 0.0;
for (int j = 0; j < i; j++) { for (int j = 0; j < i; j++) {
RealOpenMM dVdV = (RealOpenMM) valueDerivExpressions[i][j].evaluate(); float dVdV = (float) valueDerivExpressions[i][j].evaluate();
dVdR1[i] += dVdV*dVdR1[j]; dVdR1[i] += dVdV*dVdR1[j];
dVdR2[i] += dVdV*dVdR2[j]; dVdR2[i] += dVdV*dVdR2[j];
} }
for (int k = 0; k < 3; k++) { (fvec4(forces+4*atom1)-deltaR*(dEdV[i][atom1]*dVdR1[i])).store(forces+4*atom1);
forces[atom1][k] -= dEdV[i][atom1]*dVdR1[i]*deltaR[k]; (fvec4(forces+4*atom2)-deltaR*(dEdV[i][atom1]*dVdR2[i])).store(forces+4*atom2);
forces[atom2][k] -= dEdV[i][atom1]*dVdR2[i]*deltaR[k];
} }
}
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;
} }
r2 = dot3(deltaR, deltaR);
} }
...@@ -963,7 +963,6 @@ void CpuCalcCustomGBForceKernel::initialize(const System& system, const CustomGB ...@@ -963,7 +963,6 @@ void CpuCalcCustomGBForceKernel::initialize(const System& system, const CustomGB
} }
double CpuCalcCustomGBForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) { double CpuCalcCustomGBForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
vector<RealVec>& posData = extractPositions(context);
vector<RealVec>& forceData = extractForces(context); vector<RealVec>& forceData = extractForces(context);
RealOpenMM energy = 0; RealOpenMM energy = 0;
RealVec& box = extractBoxSize(context); RealVec& box = extractBoxSize(context);
...@@ -978,7 +977,7 @@ double CpuCalcCustomGBForceKernel::execute(ContextImpl& context, bool includeFor ...@@ -978,7 +977,7 @@ double CpuCalcCustomGBForceKernel::execute(ContextImpl& context, bool includeFor
map<string, double> globalParameters; map<string, double> globalParameters;
for (int i = 0; i < (int) globalParameterNames.size(); i++) for (int i = 0; i < (int) globalParameterNames.size(); i++)
globalParameters[globalParameterNames[i]] = context.getParameter(globalParameterNames[i]); globalParameters[globalParameterNames[i]] = context.getParameter(globalParameterNames[i]);
ixn->calculateIxn(numParticles, posData, particleParamArray, exclusions, globalParameters, forceData, includeEnergy ? &energy : NULL); ixn->calculateIxn(numParticles, &data.posq[0], particleParamArray, exclusions, globalParameters, &data.threadForce[0][0], includeEnergy ? &energy : NULL);
return energy; return energy;
} }
......
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