Commit b69f0c15 authored by Peter Eastman's avatar Peter Eastman
Browse files

CustomIntegrator can depend on the energy of a force group

parent 0844f4fe
......@@ -91,6 +91,10 @@ namespace OpenMM {
* <li>dt: (global) This is the step size being used by the integrator.</li>
* <li>energy: (global, read-only) This is the current potential energy of the
* system.</li>
* <li>energy0, energy1, energy2, ...: (global, read-only) This is similar to
* energy, but includes only the contribution from forces in one force group.
* A single computation step may only depend on a single energy variable
* (energy, energy0, energy1, etc.).</li>
* <li>x: (per-DOF) This is the current value of the degree of freedom (the x,
* y, or z coordinate of a particle).</li>
* <li>v: (per-DOF) This is the current velocity associated with the degree of
......
......@@ -3606,7 +3606,7 @@ void OpenCLIntegrateCustomStepKernel::initialize(const System& system, const Cus
SimTKOpenMMUtilities::setRandomNumberSeed(integrator.getRandomNumberSeed());
}
string OpenCLIntegrateCustomStepKernel::createGlobalComputation(const string& variable, const Lepton::ParsedExpression& expr, CustomIntegrator& integrator) {
string OpenCLIntegrateCustomStepKernel::createGlobalComputation(const string& variable, const Lepton::ParsedExpression& expr, CustomIntegrator& integrator, const string& energyName) {
map<string, Lepton::ParsedExpression> expressions;
if (variable == "dt")
expressions["dt[0].y = "] = expr;
......@@ -3626,7 +3626,7 @@ string OpenCLIntegrateCustomStepKernel::createGlobalComputation(const string& va
variables["dt"] = "dt[0].y";
variables["uniform"] = "uniform";
variables["gaussian"] = "gaussian";
variables["energy"] = "energy[0]";
variables[energyName] = "energy[0]";
for (int i = 0; i < integrator.getNumGlobalVariables(); i++)
variables[integrator.getGlobalVariableName(i)] = "globals["+intToString(i)+"]";
for (int i = 0; i < (int) parameterNames.size(); i++)
......@@ -3635,7 +3635,7 @@ string OpenCLIntegrateCustomStepKernel::createGlobalComputation(const string& va
return OpenCLExpressionUtilities::createExpressions(expressions, variables, functions, "temp", "");
}
string OpenCLIntegrateCustomStepKernel::createPerDofComputation(const string& variable, const Lepton::ParsedExpression& expr, int component, CustomIntegrator& integrator, const string& forceName) {
string OpenCLIntegrateCustomStepKernel::createPerDofComputation(const string& variable, const Lepton::ParsedExpression& expr, int component, CustomIntegrator& integrator, const string& forceName, const string& energyName) {
const string suffixes[] = {".x", ".y", ".z"};
string suffix = suffixes[component];
map<string, Lepton::ParsedExpression> expressions;
......@@ -3660,7 +3660,7 @@ string OpenCLIntegrateCustomStepKernel::createPerDofComputation(const string& va
variables["uniform"] = "uniform"+suffix;
variables["m"] = "mass";
variables["dt"] = "stepSize";
variables["energy"] = "energy[0]";
variables[energyName] = "energy[0]";
for (int i = 0; i < integrator.getNumGlobalVariables(); i++)
variables[integrator.getGlobalVariableName(i)] = "globals["+intToString(i)+"]";
for (int i = 0; i < integrator.getNumPerDofVariables(); i++)
......@@ -3736,12 +3736,17 @@ void OpenCLIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegr
vector<string> variable(numSteps);
vector<Lepton::ParsedExpression> expression(numSteps);
vector<string> forceGroupName;
vector<string> energyGroupName;
for (int i = 0; i < 32; i++) {
stringstream str;
str << "f" << i;
forceGroupName.push_back(str.str());
stringstream fname;
fname << "f" << i;
forceGroupName.push_back(fname.str());
stringstream ename;
ename << "energy" << i;
energyGroupName.push_back(ename.str());
}
vector<string> forceName(numSteps, "f");
vector<string> energyName(numSteps, "energy");
for (int step = 0; step < numSteps; step++) {
string expr;
integrator.getComputationStep(step, stepType[step], variable[step], expr);
......@@ -3763,6 +3768,13 @@ void OpenCLIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegr
forceGroup[step] = 1<<i;
forceName[step] = forceGroupName[i];
}
if (usesVariable(expression[step], energyGroupName[i])) {
if (forceGroup[step] != -2)
throw OpenMMException("A single computation step cannot depend on multiple force groups");
needsEnergy[step] = true;
forceGroup[step] = 1<<i;
energyName[step] = energyGroupName[i];
}
}
}
invalidatesForces[step] = (stepType[step] == CustomIntegrator::ConstrainPositions || affectsForce.find(variable[step]) != affectsForce.end());
......@@ -3820,7 +3832,7 @@ void OpenCLIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegr
for (int j = step; j < numSteps && (j == step || merged[j]); j++) {
compute << "{\n";
for (int i = 0; i < 3; i++)
compute << createPerDofComputation(stepType[j] == CustomIntegrator::ComputePerDof ? variable[j] : "", expression[j], i, integrator, forceName[j]);
compute << createPerDofComputation(stepType[j] == CustomIntegrator::ComputePerDof ? variable[j] : "", expression[j], i, integrator, forceName[j], energyName[j]);
if (variable[j] == "x") {
if (storePosAsDelta[j]) {
if (cl.getSupportsDoublePrecision())
......@@ -3910,7 +3922,7 @@ void OpenCLIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegr
stringstream compute;
for (int i = step; i < numSteps && (i == step || merged[i]); i++)
compute << "{\n" << createGlobalComputation(variable[i], expression[i], integrator) << "}\n";
compute << "{\n" << createGlobalComputation(variable[i], expression[i], integrator, energyName[i]) << "}\n";
map<string, string> replacements;
replacements["COMPUTE_STEP"] = compute.str();
cl::Program program = cl.createProgram(cl.replaceStrings(OpenCLKernelSources::customIntegratorGlobal, replacements), defines);
......
......@@ -979,8 +979,8 @@ public:
void setPerDofVariable(ContextImpl& context, int variable, const std::vector<Vec3>& values);
private:
class ReorderListener;
std::string createGlobalComputation(const std::string& variable, const Lepton::ParsedExpression& expr, CustomIntegrator& integrator);
std::string createPerDofComputation(const std::string& variable, const Lepton::ParsedExpression& expr, int component, CustomIntegrator& integrator, const std::string& forceName);
std::string createGlobalComputation(const std::string& variable, const Lepton::ParsedExpression& expr, CustomIntegrator& integrator, const std::string& energyName);
std::string createPerDofComputation(const std::string& variable, const Lepton::ParsedExpression& expr, int component, CustomIntegrator& integrator, const std::string& forceName, const std::string& energyName);
void recordChangedParameters(ContextImpl& context);
OpenCLContext& cl;
double prevStepSize;
......
......@@ -539,9 +539,15 @@ void testForceGroups() {
integrator.addPerDofVariable("outf", 0);
integrator.addPerDofVariable("outf1", 0);
integrator.addPerDofVariable("outf2", 0);
integrator.addGlobalVariable("oute", 0);
integrator.addGlobalVariable("oute1", 0);
integrator.addGlobalVariable("oute2", 0);
integrator.addComputePerDof("outf", "f");
integrator.addComputePerDof("outf1", "f1");
integrator.addComputePerDof("outf2", "f2");
integrator.addComputeGlobal("oute", "energy");
integrator.addComputeGlobal("oute1", "energy1");
integrator.addComputeGlobal("oute2", "energy2");
HarmonicBondForce* bonds = new HarmonicBondForce();
bonds->addBond(0, 1, 1.5, 1.1);
bonds->setForceGroup(1);
......@@ -561,6 +567,8 @@ void testForceGroups() {
integrator.step(1);
vector<Vec3> f, f1, f2;
double e1 = 0.5*1.1*0.5*0.5;
double e2 = 138.935456*0.2*0.2/2.0;
integrator.getPerDofVariable(0, f);
integrator.getPerDofVariable(1, f1);
integrator.getPerDofVariable(2, f2);
......@@ -570,9 +578,15 @@ void testForceGroups() {
ASSERT_EQUAL_VEC(Vec3(138.935456*0.2*0.2/4.0, 0, 0), f2[1], 1e-5);
ASSERT_EQUAL_VEC(f1[0]+f2[0], f[0], 1e-5);
ASSERT_EQUAL_VEC(f1[1]+f2[1], f[1], 1e-5);
ASSERT_EQUAL_TOL(e1, integrator.getGlobalVariable(1), 1e-5);
ASSERT_EQUAL_TOL(e2, integrator.getGlobalVariable(2), 1e-5);
ASSERT_EQUAL_TOL(e1+e2, integrator.getGlobalVariable(0), 1e-5);
// Make sure they also match the values returned by the Context.
State s = context.getState(State::Forces | State::Energy, false);
State s1 = context.getState(State::Forces | State::Energy, false, 2);
State s2 = context.getState(State::Forces | State::Energy, false, 4);
vector<Vec3> c, c1, c2;
c = context.getState(State::Forces, false).getForces();
c1 = context.getState(State::Forces, false, 2).getForces();
......@@ -583,6 +597,9 @@ void testForceGroups() {
ASSERT_EQUAL_VEC(f1[1], c1[1], 1e-5);
ASSERT_EQUAL_VEC(f2[0], c2[0], 1e-5);
ASSERT_EQUAL_VEC(f2[1], c2[1], 1e-5);
ASSERT_EQUAL_TOL(s.getPotentialEnergy(), integrator.getGlobalVariable(0), 1e-5);
ASSERT_EQUAL_TOL(s1.getPotentialEnergy(), integrator.getGlobalVariable(1), 1e-5);
ASSERT_EQUAL_TOL(s2.getPotentialEnergy(), integrator.getGlobalVariable(2), 1e-5);
}
/**
......
......@@ -103,6 +103,7 @@ void ReferenceCustomDynamics::update(ContextImpl& context, int numberOfAtoms, ve
needsEnergy.resize(numSteps, false);
forceGroup.resize(numSteps, -2);
forceName.resize(numSteps, "f");
energyName.resize(numSteps, "energy");
set<string> affectsForce;
affectsForce.insert("x");
for (vector<ForceImpl*>::const_iterator iter = context.getForceImpls().begin(); iter != context.getForceImpls().end(); ++iter) {
......@@ -116,10 +117,14 @@ void ReferenceCustomDynamics::update(ContextImpl& context, int numberOfAtoms, ve
// Make a list of which steps require valid forces or energy to be known.
vector<string> forceGroupName;
vector<string> energyGroupName;
for (int i = 0; i < 32; i++) {
stringstream str;
str << "f" << i;
forceGroupName.push_back(str.str());
stringstream fname;
fname << "f" << i;
forceGroupName.push_back(fname.str());
stringstream ename;
ename << "energy" << i;
energyGroupName.push_back(ename.str());
}
for (int i = 0; i < numSteps; i++) {
if (stepType[i] == CustomIntegrator::ComputeGlobal || stepType[i] == CustomIntegrator::ComputePerDof || stepType[i] == CustomIntegrator::ComputeSum) {
......@@ -132,6 +137,17 @@ void ReferenceCustomDynamics::update(ContextImpl& context, int numberOfAtoms, ve
needsEnergy[i] = true;
forceGroup[i] = -1;
}
else if (op.getName().substr(0, 6) == "energy") {
for (int k = 0; k < (int) energyGroupName.size(); k++)
if (op.getName() == energyGroupName[k]) {
if (forceGroup[i] != -2)
throw OpenMMException("A single computation step cannot depend on multiple force groups");
needsForces[i] = true;
forceGroup[i] = 1<<k;
energyName[i] = energyGroupName[k];
break;
}
}
else if (op.getName() == "f") {
if (forceGroup[i] != -2)
throw OpenMMException("A single computation step cannot depend on multiple force groups");
......@@ -191,7 +207,7 @@ void ReferenceCustomDynamics::update(ContextImpl& context, int numberOfAtoms, ve
energy = e;
forcesAreValid = true;
}
globals["energy"] = energy;
globals[energyName[i]] = energy;
// Execute the step.
......
......@@ -43,7 +43,7 @@ private:
std::vector<RealOpenMM> inverseMasses;
std::vector<OpenMM::RealVec> sumBuffer, oldPos;
std::vector<OpenMM::CustomIntegrator::ComputationType> stepType;
std::vector<std::string> stepVariable, forceName;
std::vector<std::string> stepVariable, forceName, energyName;
std::vector<Lepton::ExpressionProgram> stepExpression;
std::vector<bool> invalidatesForces, needsForces, needsEnergy;
std::vector<int> forceGroup;
......
......@@ -524,9 +524,15 @@ void testForceGroups() {
integrator.addPerDofVariable("outf", 0);
integrator.addPerDofVariable("outf1", 0);
integrator.addPerDofVariable("outf2", 0);
integrator.addGlobalVariable("oute", 0);
integrator.addGlobalVariable("oute1", 0);
integrator.addGlobalVariable("oute2", 0);
integrator.addComputePerDof("outf", "f");
integrator.addComputePerDof("outf1", "f1");
integrator.addComputePerDof("outf2", "f2");
integrator.addComputeGlobal("oute", "energy");
integrator.addComputeGlobal("oute1", "energy1");
integrator.addComputeGlobal("oute2", "energy2");
HarmonicBondForce* bonds = new HarmonicBondForce();
bonds->addBond(0, 1, 1.5, 1.1);
bonds->setForceGroup(1);
......@@ -546,6 +552,8 @@ void testForceGroups() {
integrator.step(1);
vector<Vec3> f, f1, f2;
double e1 = 0.5*1.1*0.5*0.5;
double e2 = 138.935456*0.2*0.2/2.0;
integrator.getPerDofVariable(0, f);
integrator.getPerDofVariable(1, f1);
integrator.getPerDofVariable(2, f2);
......@@ -555,9 +563,15 @@ void testForceGroups() {
ASSERT_EQUAL_VEC(Vec3(138.935456*0.2*0.2/4.0, 0, 0), f2[1], 1e-5);
ASSERT_EQUAL_VEC(f1[0]+f2[0], f[0], 1e-5);
ASSERT_EQUAL_VEC(f1[1]+f2[1], f[1], 1e-5);
ASSERT_EQUAL_TOL(e1, integrator.getGlobalVariable(1), 1e-5);
ASSERT_EQUAL_TOL(e2, integrator.getGlobalVariable(2), 1e-5);
ASSERT_EQUAL_TOL(e1+e2, integrator.getGlobalVariable(0), 1e-5);
// Make sure they also match the values returned by the Context.
State s = context.getState(State::Forces | State::Energy, false);
State s1 = context.getState(State::Forces | State::Energy, false, 2);
State s2 = context.getState(State::Forces | State::Energy, false, 4);
vector<Vec3> c, c1, c2;
c = context.getState(State::Forces, false).getForces();
c1 = context.getState(State::Forces, false, 2).getForces();
......@@ -568,6 +582,9 @@ void testForceGroups() {
ASSERT_EQUAL_VEC(f1[1], c1[1], 1e-5);
ASSERT_EQUAL_VEC(f2[0], c2[0], 1e-5);
ASSERT_EQUAL_VEC(f2[1], c2[1], 1e-5);
ASSERT_EQUAL_TOL(s.getPotentialEnergy(), integrator.getGlobalVariable(0), 1e-5);
ASSERT_EQUAL_TOL(s1.getPotentialEnergy(), integrator.getGlobalVariable(1), 1e-5);
ASSERT_EQUAL_TOL(s2.getPotentialEnergy(), integrator.getGlobalVariable(2), 1e-5);
}
/**
......
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