Created
December 20, 2015 07:32
-
-
Save mabraham/5db12701291e3f5e8754 to your computer and use it in GitHub Desktop.
Sketch of new GROMACS integrator formalism
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
/* This code sketches a State object, some example integrator objects | |
* that own a State (and other pieces of machinery), and use those to | |
* implement a wide variety of simulation algorithms where each | |
* conforms to the IIntegrator interface. There's copious details | |
* missing, multiple TODOs noted, some notes about optimization | |
* opportunities, and ways we will later implement stuff with task | |
* parallelism. I haven't tried to implement Michael's collectives and | |
* elements model, but we can reconsider that if we identify something | |
* here that we don't like. | |
* | |
* The plan here is to have a very flat IIntegrator class hierarchy. | |
* See https://en.wikipedia.org/wiki/Composition_over_inheritance (and | |
* other places) why this is a good idea. If two integrator classes | |
* end up looking like they have enough code duplication to be a | |
* problem, then we extract that responsibility to an object of a new | |
* class, and each of the integrator classes owns an instance of that | |
* object. That is, we don't make sub-classes in order to re-use code | |
* from the super-class. | |
* | |
* Some details have been filled in for a subclass of IIntegrator that | |
* does Bussi-style velocity-Verlet unconstrained dynamics. For | |
* example, it contains Thermostats and IVelocityUpdater objects that | |
* are set up in the factory function to do the right thing according | |
* to the inputrec. Some of those objects will get re-used for | |
* e.g. Bussi-style leapfrog dynamics, which will also be an immediate | |
* subclass of IIntegrator. For example, adding support for | |
* constraints could be as straightforward as configuring the object | |
* with a NullConstrainer, ShakeConstrainer, or LincsConstrainer, | |
* according to the .tpr contents. If we then measured too much | |
* overhead for .tprs we care about, we can specialize the | |
* integrators accordingly. | |
* | |
* Once in the MD loop, the integrator just keeps on doing the right | |
* thing that it was built to do, rather than continually looking up | |
* in the inputrec what it should be doing. This requires a handful of | |
* virtual function calls, but passes fewer function parameters and | |
* checks fewer flags. (Once we are constructing task DAGs, such | |
* concerns move to an entirely different scope, anyway.) | |
* | |
* Names should be clear and descriptive. If a Chinese grad student | |
* can't tell the role of your identifier from the name, then you | |
* named it wrongly. :-) We need to write code to be conveniently | |
* read, not to be conveniently typed. | |
* | |
* I haven't sketched how one will implement a multi-stage or MC-MD | |
* integrator, but I think the required principles follow from what | |
* Michael has planned and I have sketched here. | |
*/ | |
#include "config.h" | |
#include "gromacs/math/vectypes.h" | |
namespace gmx | |
{ | |
//! Permit zero-cost run-time checks active only for debug builds | |
const bool isDebugMode = IS_DEBUG_BUILD_TYPE; | |
class StateVector | |
{ | |
public: | |
// Vector-like interface, including const and non-const | |
// data(), and size(). | |
void setValid(bool validity) { isValid_ = validity; } | |
bool valid() const { return isValid_; } | |
private: | |
bool isValid_; | |
// Later, we decide how to minimize any overhead here | |
RVec *velocities_; | |
rvec *rawVelocities_; // legacy | |
int rawVelocitiesAllocSize_; // legacy | |
}; | |
// === State === | |
class State | |
{ | |
public: | |
const StateVector *velocities() const | |
{ | |
if (isDebugMode) | |
{ | |
GMX_RELEASE_ASSERT(velocities.valid(), "Can't read invalid velocities (because ...)"); | |
} | |
return &velocities_; | |
} | |
// Maybe this should actually return a view object that | |
// e.g. can't be re-sized, or even has a destructor that calls | |
// validateVelocities()? That means other objects don't have | |
// to depend on the main interface of State, but only the view | |
// that they need to use. | |
StateVector *velocities() | |
{ | |
if (isDebugMode) | |
{ | |
GMX_RELEASE_ASSERT(velocities.valid(), "Can't modify invalid velocities (because ...)"); | |
velocitiesAreValid_ = false; | |
// This forces people to be const correct if they just | |
// want to write code that inspects | |
// velocities. They'll get a run-time error from the | |
// next velocity update if they've used the non-const | |
// getter where they intended read-only access. | |
} | |
return &velocities_; | |
} | |
void validateVelocities() | |
{ | |
if (isDebugMode) | |
{ | |
velocities.setValid(true); | |
} | |
} | |
/* Similar stuff for forces, positions and box (at | |
least). Probably virial. Probably not any thermostat or | |
barostat stuff - the need for those is particular to | |
certain integrators. */ | |
private: | |
// Later, we might extract the concept of validity to FancyRVec type | |
StateVector velocities_; | |
}; | |
// === Thermostats === | |
class IThermostat; | |
IThermostat *buildThermostatForCouplingGroup(t_inputrec ir); | |
//! This is the interface supported by all thermostats | |
class IThermostat | |
{ | |
friend buildThermostatForCouplingGroup; | |
public: | |
IThermostat(int indexOfCouplingGroup) : | |
indexOfCouplingGroup_(indexOfCouplingGroup) {} | |
virtual void updateVelocities(real timeIncrement, gmx_int64_t *step, gmx_ekindata_t *ekind, StateVector *velocities) = 0; | |
private: | |
int indexOfCouplingGroup_; // legacy - needed for using gmx_ekindata_t in current form | |
}; | |
//! Leaves velocity unmodified, e.g. NVE or testing code | |
class NullThermostat : public IThermostat | |
{ | |
// TODO assume single coupling group for now | |
friend buildThermostatForCouplingGroup; | |
public: | |
NullThermostat(int indexOfCouplingGroup) : | |
IThermostat(indexOfCouplingGroup) {} | |
virtual void updateVelocities(real timeIncrement, gmx_int64_t *step, gmx_ekindata_t *ekind, StateVector *velocities) | |
{ | |
// do nothing, not even re-validate velocities | |
// (Much) later, calling such functions will spawn tasks | |
// into a task DAG, and the job of this function call will | |
// be to connect the parent directly to the child, i.e | |
// zero overhead when the DAG is later executed. | |
} | |
}; | |
/*! \brief Specializes IThermostat for thermostats that do a Bussi-style resample of KE | |
* | |
* TODO Decide if this approach is reasonable. This extends the | |
* IThermostat interface to form a new interface for Bussi | |
* thermostats. I'm not sure we need more than one implementation of | |
* it, however, if we fix the way ekind stores KE. If we do need | |
* multiple implementations, then it would also be reasonable to | |
* instead have a KineticEnergyResampler class, and have the different | |
* Bussi thermostats subclass IThermostat directory, and call a method | |
* on the KineticEnergyResampler instance it contains. */ | |
class IBussiThermostat : public IThermostat | |
{ | |
friend buildThermostatForCouplingGroup; | |
public: | |
IBussiThermostat(int indexOfCouplingGroup /* and other settings from inputrec */) : | |
IThermostat(indexOfCouplingGroup) /* and other initializers */ {} | |
virtual void updateVelocities(real timeIncrement, gmx_int64_t *step, gmx_ekindata_t *ekind, StateVector *velocities) = 0; | |
private: | |
real resampleKineticEnergy(real currentKineticEnergy) | |
{ | |
// do actual work, calling RNG, etc. | |
} | |
real tau_; | |
int numDegreesOfFreedom_; | |
real targetAverageKineticEnergy_; | |
// perhaps other pre-computed quantities here | |
}; | |
//! Does Bussi v-rescale thermostat for integrators with on-step KE, e.g VV | |
class OnStepKineticEnergyBussiThermostat : public IBussiThermostat | |
{ | |
friend buildThermostatForCouplingGroup; | |
public: | |
OnStepKineticEnergyBussiThermostat(int indexOfCouplingGroup /* and other settings from inputrec */) : | |
IBussiThermostat(indexOfCouplingGroup /* and other initializers */) {} | |
virtual void updateVelocities(real timeIncrement, gmx_int64_t *step, gmx_ekindata_t *ekind, StateVector *velocities) | |
{ | |
// NB no conditional on the integrator or tau or anything | |
// - if we get here we just do the right work for this | |
// case. | |
real newKineticEnergy = IBussiThermostat::resampleKineticEnergy(whateverFunction(ekind->tcstat[indexOfCouplingGroup_].ekinf)); | |
// Here, do actual rescale of velocity components | |
} | |
private: | |
}; | |
//! Factory function | |
IThermostat *buildThermostatForCouplingGroup(t_inputrec *ir, int i) | |
{ | |
GMX_ASSERT(ir->etc != etcNO, "Only real thermostats get here"); | |
if (ir->opts->tau_t[i] == 0 || ir->opts->ndrf[i] == 0) | |
{ | |
// There will never be anything to do for this group, so don't do anything! | |
return new NullThermostat(); | |
} | |
if (ir->etc == etcVRESCALE) | |
{ | |
if (EI_VV(ir->eI)) | |
{ | |
return new OnStepKineticEnergyBussiThermostat(i /* and other stuff from inputrec */); | |
} | |
else | |
{ | |
// ... | |
} | |
// Also, bump a counter that will schedule a please_cite() | |
} | |
else | |
{ | |
// ... | |
} | |
} | |
// TODO use a container that knows it has to call delete on its | |
// contents (needs to contain pointers to IThermostat for dynamic | |
// dispatch to work) | |
typedef std::vector<IThermostat *> Thermostats; | |
//! Factory function | |
Thermostats buildThermostats(t_inputrec *ir) | |
{ | |
std::vector<Thermostat *> thermostats; | |
if (ir->etc != etcNO) | |
{ | |
for (int i = 0; i < ir->opts->ngtc; i++) | |
{ | |
thermostats.push_back(buildThermostatForCouplingGroup(ir, i)); | |
} | |
} | |
} | |
// === Velocity Updaters === | |
class IVelocityUpdater | |
{ | |
public: | |
IVelocityUpdater() {}; | |
virtual ~IVelocityUpdater() {}; | |
virtual void updater(real timeStep, const StateVector *forces, StateVector *velocities) = 0; | |
}; | |
class NormalVelocityUpdater : public IVelocityUpdater | |
{ | |
public: | |
NormalVelocityUpdater() {} | |
virtual ~NormalVelocityUpdater() {}; | |
virtual void updater(real timeStep, const StateVector *forces, StateVector *velocities) | |
{ | |
// loop over atoms i | |
{ | |
// loop over dimensions d | |
{ | |
velocities[i][d] += timeStep * forces[i][d] * massFactor_[i]; | |
} | |
} | |
} | |
private: | |
std::vector<real> massFactor_; // 1/(2 * mass) | |
}; | |
// Other specializations of IVelocityUpdater do freeze groups, | |
// acceleration, shells, vsites, etc. We also include a fallback "do | |
// everything supported" implementation, with all the conditionals | |
// that that requires, and we will unit test that the specializations | |
// do the same as the fallback. | |
// Position updaters will follow similar style, but are probably simpler. | |
// === Integrators === | |
class IIntegrator | |
{ | |
public: | |
IIntegrator(State *state, ForceCalculator *forceCalculator) : | |
state_(state), forceCalculator_(forceCalculator) {}; | |
virtual ~IIntegrator() {}; | |
virtual void doStep() = 0; | |
void doSimulation(); | |
protected: | |
State state_; | |
//! Configurable object that will compute forces (e.g. normal or shell) | |
ForceCalculator *forceCalculator_; | |
gmx_int64_t step_; | |
}; | |
/* TODO Later, we will probably specialize a few forms of doStep(), | |
* and have doSimulation construct a DAG of e.g. every task to do | |
* between successive pair-search steps. */ | |
void IIntegrator::doSimulation() | |
{ | |
// Any remaining setup? | |
while (true) | |
{ | |
if (/* time to renew pair list */) | |
{ | |
forceCalculator_.updatePairList(state_.positions()); | |
} | |
doStep(); | |
/* Michael's planning document doesn't address where in the | |
* loops we do Hamiltonian updates, but I think this is | |
* right. */ | |
if (/* should update Hamiltonian */) | |
{ | |
forceCalculator_.updateHamiltonian(); // e.g. do replica exchange | |
} | |
// Handle house-keeping, such as global signals and checkpoint | |
// checks | |
// ... | |
if (checkForLastStep()) | |
{ | |
break; | |
} | |
++step_; | |
} | |
writeFinalOutput(step_); | |
} | |
class GenericIntegrator : public IIntegrator | |
{ | |
public: | |
GenericIntegrator(State *state, ForceCalculator *forceCalculator, gmx_ekindata_t *ekind, real timeStep) : | |
IIntegrator(state, forceCalculator), ekind_(*ekind), | |
timeStep_(ir->delta_t), halfTimeStep_(0.5*timeStep_) | |
/* probably other stuff here */ {} | |
virtual ~GenericIntegrator() {}; | |
virtual void doStep() | |
{ | |
// This will look much like the current do_md loop, except | |
// for the bits handled by doSimulation(). We implement | |
// this first (so that we will have code that we can ship | |
// that does all the old things). In some cases, this code | |
// will be a reference for versions that are specialized | |
// by either the kind of integrator or the kind of MD step | |
// (e.g. the majority of steps do nothing fancy). | |
} | |
private: | |
// Probably stuff here like generic velocity and position | |
// updaters, and initially some of the crud from do_md. | |
}; | |
/*! \brief Implements e.g. Bussi-Parrinello dynamics with unconstrained velocity Verlet | |
* | |
* ie. B2A|B!X (where X might be the rescaling of KE as in | |
* Bussi-Parrinello). */ | |
class ThermostattedVelocityVerletIntegrator : public IIntegrator | |
{ | |
public: | |
friend buildIntegrator; | |
ThermostattedVelocityVerletIntegrator(State *state, ForceCalculator *forceCalculator, gmx_ekindata_t *ekind, real timeStep) : | |
IIntegrator(state, forceCalculator), ekind_(*ekind), | |
timeStep_(timeStep), halfTimeStep_(0.5*timeStep), | |
computeGlobalsFlags(CGLO_ENERGY), | |
positionUpdater_(nullptr), velocityUpdater_(nullptr), thermostats_() {} | |
virtual ~ThermostattedVelocityVerletIntegrator() {}; | |
virtual void doStep() | |
{ | |
/* Initial conditions: forces are current for x, v is at | |
the same time. */ | |
// TODO As an optimization, we might fuse these first two update stages | |
velocityUpdater_->update(halfTimeStep, state_.velocities()); | |
state_->validateVelocities(); | |
positionUpdater_->update(timeStep, state_.velocities(), state_.positions()); | |
state_->validatePositions(); | |
forceCalculator_->calculate(state_.positions(), state_.forces()); | |
state_->validateForces(); | |
velocityUpdater_->update(halfTimeStep, state_.velocities()); | |
state_->validateVelocities(); | |
if (/* should communicate */) | |
{ | |
compute_globals(/* all the arguments */, computeGlobalsFlags | flagsForThisStep); | |
} | |
if (/* should write trajectories */) | |
{ | |
writeTrajectory(/* pass required state variables as const */); // Like now, write x on step, v trailing by half a step | |
} | |
if (/* should write energies */) | |
{ | |
writeEnergies(); // including dhdl, pull and such | |
} | |
if (/* should update thermostats */) | |
{ | |
StateVector *velocities = state_.velocities(); | |
for(auto thermostat : Thermostats) | |
{ | |
thermostat->updateVelocities(step_, &ekind_, velocities); | |
} | |
state_->validateVelocities(); | |
} | |
} | |
private: | |
gmx_ekindata_t ekind_; | |
real timeStep_; | |
real halfTimeStep_; | |
int computeGlobalsFlags_; // plus any other junk to make compute_globals work. Or do we want a GlobalsComputer object? | |
//! Configurable object that will handle dynamical position updates | |
IPositionUpdater *positionUpdater_; | |
//! Configurable object that will handle dynamical velocity updates | |
IVelocityUpdater *velocityUpdater_; | |
//! Configurable object that will handle thermostatting | |
Thermostats thermostats_; | |
}; | |
//! Factory function for integrators. We pass in stuff read from .tpr and/or .cpt. | |
IIntegrator *buildIntegrator(t_inputrec *ir, State *state, gmx_ekindata_t *ekind, ForceCalculator *forceCalculator) | |
{ | |
IIntegrator *integrator; | |
/* Make sure that if we don't finish this for GROMACS 2016 | |
* release, users can't get broken things. This avoids needing to | |
* fork the code base and/or have it sit broken for ages. */ | |
if (!getenv("GMX_USE_NEW_INTEGRATOR_FORMALISM")) | |
{ | |
integrator = new GenericIntegrator(state, forceCalculator, ekind, ir); | |
} | |
/* This logic is going to get extremely long and complex, but the | |
* benefits are that | |
* - we currently don't have any good place for mdrun to do run-time | |
* sanity checks for supported algorithm combinations, | |
* - this code only runs at setup time, | |
* - we will probably have fewer things that can go wrong at run | |
* time, | |
* - it's harder to inadvertently support a broken combination of | |
* stuff, and | |
* - in our optimized code paths we end up calling a handful of | |
* virtual function calls per step, rather than evaluate scores | |
* of conditionals, pass lots of function arguments on the stack | |
* including many possibly unused ones, and have global | |
* dependencies on stuff like inputrec. */ | |
else if (EI_VV(ir->eI) && ir->etc != etcNO) | |
{ | |
// Pass in the standard components shared by all integrators | |
ThermostattedVelocityVerletIntegrator thisIntegrator = | |
new ThermostattedVelocityVerletIntegrator(state, forceCalculator, ekind, ir->delta_t); | |
// Build the components that are composed to form the integrator | |
thisIntegrator->positionUpdater_ = buildPositionUpdater(ir); | |
thisIntegrator->velocityUpdater_ = buildVelocityUpdater(ir); | |
thisIntegrator->thermostats_ = buildThermostats(ir); | |
integrator = thisIntegrator; | |
} | |
else if(/* whatever */) | |
{ | |
// ... | |
} | |
else /* fallback */ | |
{ | |
integrator = new GenericIntegrator(state, forceCalculator, ekind, ir); | |
} | |
return integrator; | |
} | |
// === High-level calling code === | |
void mdrunner() | |
{ | |
IIntegrator *integrator; | |
// Do hardware detection and thread launch and stuff wherever makes sense | |
// ... | |
{ | |
State *state, *globalState; | |
gmx_ekindata_t *ekind; | |
t_inputrec *ir; | |
// Read all that stuff from .tpr and/or .cpt files. | |
// Do any pre-processing such that we have a consistent state from all starting paths. | |
ForceCalculator forceCalculator(/* e.g. whatever init_forcerec gets now */); // Stuff like replica exchange coordinated here | |
// ... | |
state = DomainDecomposition.doPartition(globalState); | |
forceCalculator.searchForPairs(state.positions()); // need initial pair list | |
forceCalculator.calculate(state.positions(), state.forces()); // need initial forces | |
integrator = buildIntegrator(ir, state, ekind, forceCalculator); | |
// Various initialization-time stuff now goes out of scope? | |
} | |
if (/* rerun */) | |
{ | |
callRerunCode(); | |
} | |
else | |
{ | |
/* Call whatever "integrator," just like now. So we'll also | |
* wrap/adapt the EM and TPI code so they fit this | |
* formalism. No more crusty switch statement. */ | |
integrator->doSimulation(); | |
} | |
// Clean up | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment