Barebones LangevinSplittingIntegrator test

barebones_langevin_splitting_test.py

import numpy
import simtk.unit
import simtk.openmm as mm
print('OpenMM version: ', mm.version.full_version)
from openmmtools.constants import kB
import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt
from simtk.openmm import app

class LangevinSplittingIntegrator(mm.CustomIntegrator):
    '''
    Integrates Langevin dynamics with a prescribed splitting.
    One way to divide the Langevin system is into three parts which can each be solved "exactly:"
        - R: Linear "drift" -- Deterministic update of positions
        - V: Linear "kick" -- Deterministic update of momenta
        - O: Ornstein-Uhlenbeck -- Stochastic update of momenta
    We can then construct integrators by solving each part for a certain timestep in sequence.


    Examples:
    '''

    def __init__(self,
                 splitting="VRORV",
                 temperature=298.0 * simtk.unit.kelvin,
                 collision_rate=91.0 / simtk.unit.picoseconds,
                 timestep=1.0 * simtk.unit.femtoseconds,
                 constraint_tolerance=1e-8
                 ):
        '''
        Always monitor heat, since it costs essentially nothing...

        Accumulate the heat exchanged with the bath in each step, in the global `heat`

        Parameters
        ----------
        splitting : string
            Sequence of R, V, O steps executed each timestep
        temperature : numpy.unit.Quantity compatible with kelvin, default: 298.0*simtk.unit.kelvin
           Fictitious "bath" temperature
        collision_rate : numpy.unit.Quantity compatible with 1/picoseconds, default: 91.0/simtk.unit.picoseconds
           Collision rate
        timestep : numpy.unit.Quantity compatible with femtoseconds, default: 1.0*simtk.unit.femtoseconds
           Integration timestep
        '''
        # Compute constants
        kT = kB * temperature
        gamma = collision_rate

        splitting = splitting.upper()
        assert (set(splitting) == set("RVO"))

        # Count how many times each step appears, so we know how big each R/V/O substep will be
        n_R = sum([letter == "R" for letter in splitting])
        n_V = sum([letter == "V" for letter in splitting])
        n_O = sum([letter == "O" for letter in splitting])

        # Define substep functions
        def R_step():
            self.addComputePerDof("x", "x + (dt / n_R) * v")
            self.addComputePerDof("x1", "x")  # save pre-constraint positions in x1
            self.addConstrainPositions()  # x is now constrained
            self.addComputePerDof("v", "v + (x - x1) / (dt / n_R)")
            self.addConstrainVelocities()

        def V_step():
            self.addComputePerDof("v", "v + ((dt / n_V) * f / m)")
            self.addConstrainVelocities()

        kinetic_energy = "0.5 * m * v * v"
        def O_step():
            self.addComputeSum("old_ke", kinetic_energy)

            self.addComputePerDof("v", "(b * v) + sqrt((1 - b*b) * (kT / m)) * gaussian")
            self.addConstrainVelocities()

            self.addComputeSum("new_ke", kinetic_energy)

            self.addComputeGlobal("heat", "heat + (new_ke - old_ke)")

        substep_functions = {"O": O_step, "R": R_step, "V": V_step }

        # Create a new CustomIntegrator
        super(LangevinSplittingIntegrator, self).__init__(timestep)
        # Initialize
        self.addGlobalVariable("kT", kT)  # thermal energy
        self.addGlobalVariable("b", numpy.exp(-gamma * timestep / n_O))  # velocity mixing parameter
        self.addPerDofVariable('x1', 0) # positions before application of position constraints
        self.setConstraintTolerance(constraint_tolerance)

        self.addGlobalVariable("n_R", n_R)
        self.addGlobalVariable("n_V", n_V)
        self.addGlobalVariable("n_O", n_O)

        # Add bookkeeping variables
        self.addGlobalVariable("old_ke", 0)
        self.addGlobalVariable("new_ke", 0)
        self.addGlobalVariable("heat", 0)

        # Integrate, applying constraints or bookkeeping as necessary
        for step in splitting: substep_functions[step]()

if __name__=="__main__":
    test_system = "alanine"
    #test_system = "waterbox"

    def get_test_system(constrained=True):
        if test_system == 'alanine':
            from openmmtools.testsystems import AlanineDipeptideVacuum
            if constrained:
                testsystem = AlanineDipeptideVacuum()
            else:
                testsystem = AlanineDipeptideVacuum(constraints=None)
        elif test_system == 'waterbox':
            from openmmtools.testsystems import WaterBox
            testsystem = WaterBox(constrained=constrained)

        return testsystem

    testsystem = get_test_system()

    (system, positions) = testsystem.system, testsystem.positions
    print('Dimensionality: {}'.format(len(positions) * 3))

    from simtk import openmm, unit
    temperature = 300 * unit.kelvin

    def test(starting_positions,
             scheme='RVOVR',
             timestep=1.0*unit.femtoseconds,
             constrained=True,
             n_steps=5000
             ):
        testsystem = get_test_system(constrained)
        (system, positions) = testsystem.system, testsystem.positions
        integrator = LangevinSplittingIntegrator(scheme,
                                                 temperature=temperature,
                                                 timestep=timestep,
                                                 )

        platform = mm.Platform.getPlatformByName('Reference')
        simulation = app.Simulation(testsystem.topology, system, integrator, platform)

        #platform = mm.Platform.getPlatformByName('OpenCL')
        #platform.setPropertyDefaultValue('OpenCLPrecision', 'mixed')
        #simulation = app.Simulation(testsystem.topology, system, integrator, platform)

        context = simulation.context
        context.setPositions(starting_positions)
        context.setVelocitiesToTemperature(temperature)
        energy_unit = unit.kilojoule_per_mole

        def total_energy():
            state = context.getState(getEnergy=True)
            ke = state.getKineticEnergy()
            pe = state.getPotentialEnergy()
            return (pe + ke).value_in_unit(energy_unit)

        heats = [0]
        shadow_works = [0]
        energies = [total_energy()]

        for i in range(n_steps):
            integrator.step(1)
            energies.append(total_energy())
            heats.append(integrator.getGlobalVariableByName('heat'))
            shadow_works.append((energies[-1] - energies[0]) - heats[-1])
        return context, heats, shadow_works

    # non-equilibrated initial conditions
    n_replicates = 2
    plt.figure()
    for _ in range(n_replicates):
        context, heats, shadow_works = test(positions)
        plt.plot(shadow_works, c='blue')
    plt.xlabel('# steps')
    plt.ylabel('W_shad')
    plt.title('Shadow work (non-equilibrated initial conditions)')
    plt.savefig('shadow_work_plot.jpg')
    plt.close()


    # equilibrated initial conditions
    plt.figure()
    equilibrated = context.getState(getPositions=True).getPositions()
    for _ in range(n_replicates):
        context, heats, shadow_works = test(equilibrated)
        plt.plot(shadow_works, c='blue')

    plt.xlabel('# steps')
    plt.ylabel('W_shad')
    plt.title('Shadow work (equilibrated initial conditions)')
    plt.savefig('shadow_work_plot_after_equilibration.jpg')
    plt.close()

    # many choices now
    def do_comparison(constrained=True, transient=1000):

        colors = {'RVOVR': 'blue', 'VRORV':'red', 'OVRVO':'green'}
        plt.figure()

        timesteps = numpy.array([0.5, 1, 2, 3])
        if constrained: timesteps = timesteps * 2

        for scheme in colors.keys():
            for timestep in timesteps:
                label = '{} ({} fs)'.format(scheme, timestep)

                context, heats, shadow_works = test(positions, scheme, timestep * unit.femtoseconds, constrained)
                equilibrated = context.getState(getPositions=True).getPositions()

                context, heats, shadow_works = test(equilibrated,
                                             scheme=scheme,
                                             timestep=timestep*unit.femtoseconds,
                                             constrained=constrained
                                             )
                print('{}: {:.3f}'.format(label, shadow_works[-1]))
                plt.plot(numpy.array(shadow_works[transient:]) - shadow_works[transient], c=colors[scheme], label=label, linewidth=timestep)
        plt.legend(loc=(1,0))
        plt.tight_layout()
        plt.xlabel('# steps')
        plt.ylabel('W_shad')
        name = 'Integrator comparison'
        if constrained: name = name + ' (constrained)'
        else: name = name + ' (unconstrained)'
        plt.title(name)

    for constrained in [True, False]:
        print('Constraints: {}'.format(constrained))
        do_comparison(constrained)
        plt.savefig('shadow_work_comparison_{}.jpg'.format(constrained), bbox_inches='tight')