alanine_diverse_starting_structure_experiment.py

n_threads = 32

from simtk.openmm import app
import simtk.openmm as mm
import openmmtools
from simtk.unit import *
from openmmtools.integrators import kB
import mdtraj as md
from msmbuilder.example_datasets import AlanineDipeptide
import os
import numpy as np

# quick benchmark indicated that, for this very small system in implicit solvent,
# the reference implementation is about fast as the GPU implementation.
# this is convenient, since I can then parallelize 1 simulation per CPU core.
# (I have access to way more CPU cores than GPUs.)

n_runs = 4
n_clones = n_threads / n_runs

temperature = 300 * kelvin
timestep = 2.0 * femtoseconds


target_clone_length = 10 * nanosecond
n_steps = int(target_clone_length / timestep)
reporter_interval = 500

# get diverse starting configurations, by global minimization
ala_trajs = AlanineDipeptide().get().trajectories
top = ala_trajs[0].top


forcefield = app.ForceField('amber99sbildn.xml', 'amber99_obc.xml')
system = forcefield.createSystem(top.to_openmm(), nonbondedMethod=app.CutoffNonPeriodic,
                                 nonbondedCutoff=1.0 * nanometers, constraints=app.HBonds)
integrator = openmmtools.integrators.GradientDescentMinimizationIntegrator(initial_step_size=0.02 * angstroms)

platform = mm.Platform.getPlatformByName('Reference')

simulation = app.Simulation(top.to_openmm(), system, integrator, platform=platform)
simulation.context.setPositions(ala_trajs[0].xyz[0])

def get_starting_configurations(n_trial = 100):
    def get_positions():
        pos = simulation.context.getState(getPositions=True).getPositions()
        unit = nanometer
        X = np.array(pos.value_in_unit(unit))
        return X

    def perturb_positions(sigma=0.1):
        X = get_positions()
        noise = np.random.randn(*X.shape) * sigma
        simulation.context.setPositions(X + noise)

    configurations = []
    Es = []
    for i in range(n_trial):
        perturb_positions(sigma=1.0)
        simulation.step(2000)
        configurations.append(get_positions())
        Es.append(simulation.context.getState(getEnergy=True).getPotentialEnergy())

    traj = md.Trajectory(configurations, ala_trajs[0].top)
    runs = traj[np.argsort(Es)][:n_runs]

    clones = []
    for run in runs:
        for i in range(n_clones):
            clones.append(run.xyz[0])

    return clones

def simulate_clone(i):
    # create system
    forcefield = app.ForceField('amber99sbildn.xml', 'amber99_obc.xml')
    system = forcefield.createSystem(top.to_openmm(), nonbondedMethod=app.CutoffNonPeriodic,
                                     nonbondedCutoff=1.0 * nanometers, constraints=app.HBonds)
    integrator = openmmtools.integrators.GHMCIntegrator(temperature=temperature, timestep=timestep)

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

    # set starting configuration and
    starting_conf = starting_confs[i]
    simulation.context.setPositions(starting_conf)
    simulation.context.setVelocitiesToTemperature(temperature)

    # add a reporter
    name = 'alanine_trajectory_{0}'.format(i)
    simulation.reporters.append(app.DCDReporter(name + '.dcd', reporter_interval))

    # to-do for the VVVR comparison: also add a shadow work reporter here, but for now not important

    # simulate
    simulation.step(n_steps)

    # get acceptance rate
    acceptance_rate = integrator.getGlobalVariableByName('naccept') / integrator.getGlobalVariableByName('ntrials')

    # delete the reporter
    del (simulation.reporters[:])

    # process the trajectory into an HDF5 file
    traj = md.load(name + '.dcd', top=top)
    traj.save_hdf5(name + '.h5')

    # delete the temporary dcd file
    os.remove(name + '.dcd')

    return acceptance_rate


from multiprocessing import Pool
from time import time

starting_confs = get_starting_configurations(100)

t = time()
pool = Pool(n_threads)
acceptance_rates = pool.map(simulate_clone, range(len(starting_confs)))

print('Wall clock time: {0:.3f} s'.format(time() - t))
print('Simulation time: {0:.3f} ns'.format((n_clones * n_steps * timestep).value_in_unit(nanosecond)))
print('Average acceptance rate: {0:.3f}%'.format(100 * np.mean(acceptance_rates)))