Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Sketch of new GROMACS integrator formalism
/* 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
You can’t perform that action at this time.