barebones_integrator_benchmark.py

def get_total_energy(context):
    state = context.getState(getEnergy=True)
    ke, pe = state.getKineticEnergy(), state.getPotentialEnergy()
    return ke + pe

def measure_shadow_work(simulation, n_steps):
    '''Given a `simulation` that uses an integrator that accumulates heat exchange with bath,
    apply the integrator for n_steps and return the change in energy - the heat.
    '''
    get_energy = lambda : get_total_energy(simulation.context)
    get_heat = lambda : simulation.integrator.getGlobalVariableByName('heat')

    E_0 = get_energy()
    Q_0 = get_heat()

    simulation.step(n_steps)

    E_1 = get_energy()
    Q_1 = get_heat()

    delta_E = E_1 - E_0
    delta_Q = Q_1 - Q_0

    W_shad = delta_E.value_in_unit(unit.kilojoule_per_mole) - delta_Q

    return W_shad

def randomization_midpoint_operator(simulation, temperature):
    simulation.context.setVelocitiesToTemperature(temperature)

def reversal_midpoint_operator(simulation):
    simulation.context.setVelocities(- simulation.context.getState(getVelocities=True).getVelocities(asNumpy=True))

if __name__ == '__main__':
    from tqdm import tqdm
    from simtk import unit
    import simtk.openmm as mm
    from simtk.openmm import app
    import numpy as np

    temperature = 298.0 * unit.kelvin

    burn_in_length = 1000 # in steps
    n_samples = 100 # number of samples to collect
    ghmc_thinning = 1000 # number of GHMC steps between samples
    protocol_length = 1000 # length of protocol
    midpoint_operator = reversal_midpoint_operator
    #midpoint_operator = lambda simulation: randomization_midpoint_operator(simulation, temperature)

    #platform = 'OpenCL'
    platform = 'Reference'
    #platform = 'CUDA'

    from openmmtools.integrators import GHMCIntegrator
    ghmc = GHMCIntegrator(temperature)

    # unbiased simulation
    from openmmtools.testsystems import WaterBox
    testsystem = WaterBox(constrained=True)

    from openmmtools.testsystems import AlanineDipeptideVacuum, AlanineDipeptideImplicit

    testsystem = AlanineDipeptideVacuum(constraints=app.AllBonds)
    #testsystem = AlanineDipeptideVacuum(constraints=None)
    #testsystem = AlanineDipeptideImplicit(constraints=None)

    (system, positions) = testsystem.system, testsystem.positions

    if platform == 'Reference':
        platform = mm.Platform.getPlatformByName('Reference')
    elif platform == 'OpenCL':
        platform = mm.Platform.getPlatformByName('OpenCL')
        platform.setPropertyDefaultValue('OpenCLPrecision', 'mixed')
    elif platform == 'CUDA':
        platform = mm.Platform.getPlatformByName('CUDA')
        platform.setPropertyDefaultValue('CUDAPrecision', 'mixed')

    unbiased_simulation = app.Simulation(testsystem.topology, system, ghmc, platform)
    unbiased_simulation.context.setPositions(positions)
    unbiased_simulation.context.setVelocitiesToTemperature(temperature)

    print('"Burning in" unbiased GHMC sampler')
    unbiased_simulation.step(burn_in_length)
    print('Done!')

    print('Collecting statistics!')

    def benchmark(simulation, unbiased_simulation, n_samples, M):
        W_shads_F = []
        W_shads_R = []

        for i in tqdm(range(n_samples)):
            unbiased_simulation.step(ghmc_thinning)

            simulation.context.setPositions(unbiased_simulation.context.getState(getPositions=True).getPositions())
            simulation.context.setVelocities(unbiased_simulation.context.getState(getVelocities=True).getVelocities())

            W_shads_F.append(measure_shadow_work(simulation, protocol_length))
            midpoint_operator(simulation)
            W_shads_R.append(measure_shadow_work(simulation, protocol_length))

        DeltaF_neq = 0.5 * (np.mean(W_shads_F) - np.mean(W_shads_R))
        N_eff = n_samples # for now, assuming independent
        sq_uncertainty = (np.std(W_shads_F)**2 + np.std(W_shads_R)**2 - np.cov(W_shads_F, W_shads_R)[0,1]) / (4 * N_eff)

        print(DeltaF_neq, sq_uncertainty)
        return DeltaF_neq, sq_uncertainty

    timestep = 2.5 * unit.femtoseconds

    from barebones_lsi import LangevinSplittingIntegrator

    for scheme in ['RVOVR', 'OVRVO', 'VRORV', 'VRRRORRRV']:

        print(scheme)
        lsi = LangevinSplittingIntegrator(scheme, timestep=timestep, temperature=temperature)
        simulation = app.Simulation(testsystem.topology, system, lsi, platform)
        simulation.context.setPositions(positions)
        simulation.context.setVelocitiesToTemperature(temperature)

        benchmark(simulation, unbiased_simulation, n_samples, M=protocol_length)