Sketch of new GROMACS integrator formalism

integratorsketch.cpp

/* 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
}

}