<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Obtaining an equilibrated initial trajectory
This is the first file to run in the alanine dipeptide TPS example. This teaches you how to
Step1:  Setting up high-temperature engine
Now we set things up for the OpenMM simulation. We will need a openmm.System object and an openmm.Integrator object.
To learn more about OpenMM, read the OpenMM documentation. The code we use here is based on output from the convenient web-based OpenMM builder.
Step2:  The storage file will need a template snapshot. In addition, the OPS OpenMM-based Engine has a few properties and options that are set by these dictionaries.
Step3:  Defining states
First we define the CVs using the md.compute_dihedrals function. Then we define our states using PeriodicCVDefinedVolume (since our CVs are periodic.)
Step4:  Getting a first trajectory
We'll use the VisitAllStatesEnsemble to create a high-temperature trajectory that visits all of our states. Here we only have 2 states, but this approach generalizes to multiple states.
Step5:  Plotting the trajectory
Step6:  Setting up the normal temperature engine
We'll create another engine that uses a 300K integrator, and equilibrate to a 300K path from the 500K path.
Step7:  Using TPS to equilibrate
This is, again, a simple path sampling setup. We use the same TPSNetwork we'll use later, and only shooting moves. One the initial conditions are correctly set up, we run one step at a time until the initial trajectory is decorrelated.
This setup of a path sampler always consists of defining a network and a move_scheme. See toy model notebooks for further discussion.
Step8:  The move scheme's initial_conditions_from_trajectories method will take whatever input you give it and try to generate valid initial conditions. This includes looking for subtrajectories that satisfy the ensembles you're sampling, as well as reversing the trajectory. If it succeeds, you'll see the message "No missing ensembles. No extra ensembles." It's also a good practice to use scheme.assert_initial_conditions to ensure that your initial conditions are valid. This will halt a script if the initial conditions are not valid.
Step9:  We can create a StepVisualizer2D, which plays the simulation as it is running. It requires CVs for the $x$ direction and the $y$ direction, as well as bounds of the plot. You can set the background attribute to an existing matplotlib Figure, and the simulation will plot on top of that. See the toy model MSTIS example for an example of that.
This isn't necessary, and isn't even a useful for long simulations (which typically won't be in an interactive notebook), but it can be fun!
<ASSISTANT_TASK:>
Python Code:
from __future__ import print_function
%matplotlib inline
import matplotlib.pyplot as plt
import openpathsampling as paths
import openpathsampling.engines.openmm as omm
from simtk.openmm import app
import simtk.openmm as mm
import simtk.unit as unit
from openmmtools.integrators import VVVRIntegrator
import mdtraj as md
import numpy as np
import os
initial_pdb = os.path.join("..", "resources", "AD_initial_frame.pdb")
Explanation: Obtaining an equilibrated initial trajectory
This is the first file to run in the alanine dipeptide TPS example. This teaches you how to:
set up an engine using OpenMM
set up states using MDTraj-based collective variables
obtain a initial trajectory using high temperature MD
equilibrate by using shooting moves until the first decorrelated trajectory
We assume at this point that you are familiar with the basic concepts of OPS. If you find this file confusing, we recommend working through the toy model examples.
End of explanation
# this cell is all OpenMM specific
forcefield = app.ForceField('amber96.xml', 'tip3p.xml')
pdb = app.PDBFile(initial_pdb)
system = forcefield.createSystem(
    pdb.topology, 
    nonbondedMethod=app.PME, 
    nonbondedCutoff=1.0*unit.nanometers,
    constraints=app.HBonds, 
    rigidWater=True,
    ewaldErrorTolerance=0.0005
)
hi_T_integrator = VVVRIntegrator(
    500*unit.kelvin, 
    1.0/unit.picoseconds, 
    2.0*unit.femtoseconds)
hi_T_integrator.setConstraintTolerance(0.00001)
Explanation: Setting up high-temperature engine
Now we set things up for the OpenMM simulation. We will need a openmm.System object and an openmm.Integrator object.
To learn more about OpenMM, read the OpenMM documentation. The code we use here is based on output from the convenient web-based OpenMM builder.
End of explanation
template = omm.snapshot_from_pdb(initial_pdb)
openmm_properties = {'OpenCLPrecision': 'mixed'}
engine_options = {
    'n_steps_per_frame': 10,
    'n_frames_max': 20000
}
hi_T_engine = omm.Engine(
    template.topology, 
    system, 
    hi_T_integrator, 
    openmm_properties=openmm_properties,
    options=engine_options
)
hi_T_engine.name = '500K'
hi_T_engine.current_snapshot = template
hi_T_engine.minimize()
Explanation: The storage file will need a template snapshot. In addition, the OPS OpenMM-based Engine has a few properties and options that are set by these dictionaries.
End of explanation
# define the CVs
psi = paths.MDTrajFunctionCV(
    "psi", md.compute_dihedrals,
    topology=template.topology, indices=[[6,8,14,16]]
).enable_diskcache()
phi = paths.MDTrajFunctionCV(
    "phi", md.compute_dihedrals,
    topology=template.topology, indices=[[4,6,8,14]]
).enable_diskcache()
# define the states
deg = 180.0/np.pi
C_7eq = (
    paths.PeriodicCVDefinedVolume(phi, lambda_min=-180/deg, lambda_max=0/deg, 
                                  period_min=-np.pi, period_max=np.pi) &
    paths.PeriodicCVDefinedVolume(psi, lambda_min=100/deg, lambda_max=200/deg,
                                  period_min=-np.pi, period_max=np.pi)
).named("C_7eq")
# similarly, without bothering with the labels:
alpha_R = (
    paths.PeriodicCVDefinedVolume(phi, -180/deg, 0/deg, -np.pi, np.pi) &
    paths.PeriodicCVDefinedVolume(psi, -100/deg, 0/deg, -np.pi, np.pi)
).named("alpha_R")
Explanation: Defining states
First we define the CVs using the md.compute_dihedrals function. Then we define our states using PeriodicCVDefinedVolume (since our CVs are periodic.)
End of explanation
visit_all = paths.VisitAllStatesEnsemble([C_7eq, alpha_R])
trajectory = hi_T_engine.generate(hi_T_engine.current_snapshot, [visit_all.can_append])
Explanation: Getting a first trajectory
We'll use the VisitAllStatesEnsemble to create a high-temperature trajectory that visits all of our states. Here we only have 2 states, but this approach generalizes to multiple states.
End of explanation
# create a network so we can use its ensemble to obtain an initial trajectory
# use all-to-all because initial traj can be A->B or B->A; will be reversed
tmp_network = paths.TPSNetwork.from_states_all_to_all([C_7eq, alpha_R])
# take the subtrajectory matching the ensemble (only one ensemble, only one subtraj)
subtrajectories = []
for ens in tmp_network.analysis_ensembles:
    subtrajectories += ens.split(trajectory)
print(subtrajectories)
plt.plot(phi(trajectory), psi(trajectory), 'k.')
plt.plot(phi(subtrajectories[0]), psi(subtrajectories[0]), 'r')
Explanation: Plotting the trajectory
End of explanation
integrator = VVVRIntegrator(
    300*unit.kelvin, 
    1.0/unit.picoseconds, 
    2.0*unit.femtoseconds
)
integrator.setConstraintTolerance(0.00001)
engine = omm.Engine(
    template.topology, 
    system, 
    integrator, 
    openmm_properties=openmm_properties,
    options=engine_options
).named('300K')
Explanation: Setting up the normal temperature engine
We'll create another engine that uses a 300K integrator, and equilibrate to a 300K path from the 500K path.
End of explanation
network = paths.TPSNetwork(initial_states=C_7eq,
                           final_states=alpha_R).named('tps_network')
scheme = paths.OneWayShootingMoveScheme(network, 
                                        selector=paths.UniformSelector(),
                                        engine=engine).named('equil_scheme')
Explanation: Using TPS to equilibrate
This is, again, a simple path sampling setup. We use the same TPSNetwork we'll use later, and only shooting moves. One the initial conditions are correctly set up, we run one step at a time until the initial trajectory is decorrelated.
This setup of a path sampler always consists of defining a network and a move_scheme. See toy model notebooks for further discussion.
End of explanation
# make subtrajectories into initial conditions (trajectories become a sampleset)
initial_conditions = scheme.initial_conditions_from_trajectories(subtrajectories)
# check that initial conditions are valid and complete (raise AssertionError otherwise)
scheme.assert_initial_conditions(initial_conditions)
Explanation: The move scheme's initial_conditions_from_trajectories method will take whatever input you give it and try to generate valid initial conditions. This includes looking for subtrajectories that satisfy the ensembles you're sampling, as well as reversing the trajectory. If it succeeds, you'll see the message "No missing ensembles. No extra ensembles." It's also a good practice to use scheme.assert_initial_conditions to ensure that your initial conditions are valid. This will halt a script if the initial conditions are not valid.
End of explanation
visualizer = paths.StepVisualizer2D(network, phi, psi, xlim=[-3.14, 3.14], 
                                    ylim=[-3.14, 3.14])
sampler = paths.PathSampling(storage=paths.Storage("ad_tps_equil.nc", "w", template),
                             move_scheme=scheme,
                             sample_set=initial_conditions)
sampler.live_visualizer = visualizer
# initially, these trajectories are correlated (actually, identical)
# once decorrelated, we have a (somewhat) reasonable 300K trajectory
initial_conditions[0].trajectory.is_correlated(sampler.sample_set[0].trajectory)
sampler.run_until_decorrelated()
# now we have decorrelated: no frames are in common
initial_conditions[0].trajectory.is_correlated(sampler.sample_set[0].trajectory)
# run an extra 10 to decorrelate a little futher
sampler.run(10)
sampler.storage.close()
Explanation: We can create a StepVisualizer2D, which plays the simulation as it is running. It requires CVs for the $x$ direction and the $y$ direction, as well as bounds of the plot. You can set the background attribute to an existing matplotlib Figure, and the simulation will plot on top of that. See the toy model MSTIS example for an example of that.
This isn't necessary, and isn't even a useful for long simulations (which typically won't be in an interactive notebook), but it can be fun!
End of explanation
<END_TASK>