Massively parallel Transition Path Sampling

Massively parallel Transition Path Sampling#

Notebook 1: Setup and run TPS simulation#

This is the first of a series of example notebooks on massively parallel transition path sampling. We will use multiple samplers (each generating its own Markov chain) that are all steered by one central reaction coordinate model. This has the benefit that the central model learns from all chains/samplers, which is more time efficient since we can run all samplers at the same time (e.g. on an HPC cluster). And since the chains will diverge from each other the central model will also see multiple reaction mechanisms at the same time and learn all of them, i.e. the training samples are more diverse. Since all chains have equal weight you can still easily calculate ensemble averages over the transition path ensemble by weighting each Monte carlo state in each chain with equal weight. Note, that to start the sampling we do need an initial transition for each sampler/Markov chain (as opposed to TPS from equilibrium points). If you use different initial transitions for each sampler this can actually be a benefit as it enables you to estimate if the ensemble is converged by comparing the chains to each other (which you will see in the analysis notebook).

In this notebook we will use capped alanine dipetide as our example molecule and look at the transition between \(C7_{eq}\) and \(\alpha_R\) states. We will use the locally running GmxEngine and PyTrajectoryFunctionWrapper classes (such that you can run it on your workstation), but you can easily perform a massively parallel TPS on a HPC cluster running SLURM by using the SlurmGmxEngine and SlurmTrajectoryFunctionWrapper classes instead. However, in that case you will probably want to use a larger (and more interessting) system than capped alanine dipeptide :)

This notebook should be run on a multi-core workstation preferably with a GPU, otherwise you will have a very long coffee break and a very hot laptop.

Required knowledge/recommended reading: This notebooks assumes some familarity with the asyncmd (namely the [gromacs] engine and TrajectoryFunctionWrapper classes). Please see the example notebooks in asyncmd for an introduction.

Imports and set working directory#

%matplotlib inline

import os
import asyncio
# asyncmd for the engine and Trajectory class
import asyncmd
import asyncmd.gromacs as asyncgmx
from asyncmd import Trajectory
# aimmd for the TPS
import aimmd
import aimmd.distributed as aimmdd
# and pytorch for the reaction coordinate model
import torch

Could not initialize SLURM cluster handling. If you are sure SLURM (sinfo/sacct/etc) is available try calling `asyncmd.config.set_all_slurm_settings()` or `asyncmd.config.set_slurm_setting()` with the appropriate arguments.
WARNING:aimmd:dCGPy not found. SymReg will not be available.
WARNING:aimmd:Tensorflow/Keras not available

# setup working directory

scratch_dir = "."

workdir = os.path.join(scratch_dir, "TransitionPathSampling_ala")

if not os.path.isdir(workdir):
    os.mkdir(workdir)

Setup the TPS simulation#

To setup the TPS we need the following:

  • Decide how many Markov chain samplers we want to run in parallel

  • Create a storage file to save our models, trainset and other simulation results

  • Define the metastable states to know when to stop the integration (and what defines a transition)

  • Define the underlying dynamics (through defining a gromacs engine with its options)

  • Load our initial transitions, here we will use the three transitions that are part of the repository. Generating initial transitions can be facilitated in a multitude of ways (high temperature/pressure/etc, pulling, steered MD, …) and can/should take into account previous knowledge of the system. For an example on how to generate initial transitions if you have access (or can generate) configurations close to a putative transition state see the CommittorSimulation notebook.

  • Define the reaction coordinate model and the space it is learning in by choosing the descriptor_transform (which transforms from configurations to descriptor space)

  • Define the sampling scheme, i.e. how we generate new trials (here we will use two way shooting moves with random velocities).

  • Create a trainset into which we will add the simulation results (shooting outcomes).

  • Define the Tasks to run after a specified number of trials. These are used to e.g. train the reaction coordinate model or save the trainset, model and brain at specified intervals. They are similar to the openpathsampling concept of hooks.

  • Set the initial transitions for each Markov chain sampler

Number of Markov chains#

n_samplers = 4  # results in 2*n_samplers gmx engines

Create storage file#

storage = aimmd.Storage(os.path.join(workdir, "storage.h5"))

State definition#

alpha_R and C7_eq are python functions (using MDAnalysis) to claculate for each frame of a trajctory if it belongs to the respective state or not (returning an array of True and False values). We wrapp them using asyncmds TrajectoryFunctionWrapper to make them awaitable and run in parallel where possible (e.g. when the code releases the GIL or when you use the SlurmTrajectoryFunctionWrapper to submit the calculation via SLURM).

If you want to learn more about the TrajectoryFunctionWrappers have a look at the example notebooks and documentation ofa asyncmd.

# state functions
from state_funcs_mda import alpha_R, C7_eq

wrapped_alphaR = asyncmd.trajectory.PyTrajectoryFunctionWrapper(alpha_R)
wrapped_C7_eq = asyncmd.trajectory.PyTrajectoryFunctionWrapper(C7_eq)

Underlying dynamics#

We will use the localy running GmxEngine from asyncmd (again have a look at the examples and documentation there to learn more), you could however easily run the TPS on a HPC cluster by using the SlurmGmxEngine.

In aimmd.distributed the propagation options are specified using the MDEngineSpec dataclass for all TPS simulations and also for committor simulations.

# Define the engine(s) for the PathMovers
# (they will all be the same)
gro = "gmx_infiles/conf.gro"
top = "gmx_infiles/topol_amber99sbildn.top"
ndx = "gmx_infiles/index.ndx"
mdp = asyncgmx.MDP("gmx_infiles/md.mdp")

# use VV integrator
mdp["integrator"] = "md-vv"

gmx_engine_kwargs = {"mdconfig": mdp,
                     "gro_file": gro,
                     "top_file": top,
                     "ndx_file": ndx,
                     "mdrun_extra_args": "-nt 2",
                     "output_traj_type": "XTC",
                     }

engine_spec = aimmdd.MDEngineSpec(
                    engine_cls=asyncgmx.GmxEngine,
                    engine_kwargs=gmx_engine_kwargs,
                    max_steps=500 * 10**5,  # 500 steps * dt (2 fs) = 1 ps
                    # NOTE: choose this as short as possible!
                    #       since ala is super-small and commits fast we should make sure
                    #       that most trials reach a state in the first part
                    #       this in turn makes sure that we do not call gromacs multiple times per trial (saving setup time)
                    #       but still ensures that the resulting trajectories are not too long and large
                    #       it also reduces the time needed per step (we need at least walltime_per_part hours per step)
                    #walltime_per_part=0.000015625,  # 0.055125 s per part
                    walltime_per_part=0.00003125,  # 0.1125 s per part
                    #walltime_per_part=0.0000625,  # 0.225 s per part
                    #walltime_per_part=0.000125,  # 0.45 s per part
                    #walltime_per_part=0.001,  # 3.6 s per part
                    #walltime_per_part=0.004,  # 14.4 s per part
                    )

Load initial transitions#

tp_lb = Trajectory(structure_file="gmx_infiles/ala_300K_amber99sb-ildn.tpr", trajectory_files="gmx_infiles/TP_low_barrier_300K_amber99sbildn.trr")
# Note that `tp_lb` actually contains a few frames inside of the alpha_R state, because it was generated with a slightly less forgiving state definition
# It is therefore not used (commented out) as an initial transition below, you can however uncomment it to see that the TPS simulation can (gracefully) handle
# this situation by prininting an error message and retrying with another frame 
tp_short = Trajectory(structure_file="gmx_infiles/ala_300K_amber99sb-ildn.tpr", trajectory_files="gmx_infiles/TP_short_low_barrier_300K_amber99sbildn.trr")
tp_short2 = Trajectory(structure_file="gmx_infiles/ala_300K_amber99sb-ildn.tpr", trajectory_files="gmx_infiles/TP_short2_low_barrier_300K_amber99sbildn.trr")

Define reaction coordinate model and descriptor_transform#

# import descriptor_transform for the model
# descriptor_func_ic gives us an internal coordinate representation (i.e. bond lengths, angles and dihedrals)
# descriptor_func_psi_phi gives us the ψ and φ dihedral angles (we use it to project to a 2d space in which we can look at the TPE)
from state_funcs_mda import descriptor_func_ic, descriptor_func_psi_phi

# and as usual wrapp them to become awaitable
wrapped_transform = asyncmd.trajectory.PyTrajectoryFunctionWrapper(descriptor_func_ic, call_kwargs={"molecule_selection": "protein"})
wrapped_psi_phi = asyncmd.trajectory.PyTrajectoryFunctionWrapper(descriptor_func_psi_phi)

# get the descriptors for them to infer the number of inputs for our model
# we would only need one but since we can execute them in parallel anyway...
descriptors_for_tps = await asyncio.gather(*(wrapped_transform(tp) for tp in [tp_short, tp_short2, tp_lb]))

# model architecture definition
# we use a pyramidal ResNet as described in "Machine-guided path sampling to discover mechanisms of molecular self-organization" (Nat.Comput.Sci 2023)

n_lay_pyramid = 5  # number of resunits
n_unit_top = 10  # number of units in the last layer before the log_predictor
dropout_base = 0.3  # dropot fraction in the first layer (will be reduced going to the top)
n_unit_base = cv_ndim = descriptors_for_tps[0].shape[1]  # input dimension
# the factor by which we reduce the number of units per layer (the width) and the dropout fraction
fact = (n_unit_top / n_unit_base)**(1./(n_lay_pyramid))

# create a list of modules to build our pytorch reaction coodrinate model from
modules = []

for i in range(1, n_lay_pyramid + 1):
    modules += [aimmd.pytorch.networks.FFNet(n_in=max(n_unit_top, int(n_unit_base * fact**(i-1))),
                                             n_hidden=[max(n_unit_top, int(n_unit_base * fact**i))],  # 1 hidden layer network
                                             activation=torch.nn.Identity(),
                                             dropout={"0": dropout_base * fact**i}
                                             )
                ]
    print(f"ResUnit {i} is {max(n_unit_top, int(n_unit_base * fact**(i)))} units wide.")
    print(f"Dropout before it is {dropout_base * fact**i}.")
    modules += [aimmd.pytorch.networks.ResNet(n_units=max(n_unit_top, int(n_unit_base * fact**i)),
                                              n_blocks=1)
                ]

torch_model = aimmd.pytorch.networks.ModuleStack(n_out=1,  # using a single output we will predict only p_B and use a binomial loss
                                                           # we could have also used n_out=n_states to use a multinomial loss and predict all states,
                                                           # but this is probably only worthwhile if n_states > 2 as it would increase the number of free parameters in the NN
                                                 modules=modules,  # modules is a list of initialized torch.nn.Modules from arcd.pytorch.networks
                                                 )

# move model to GPU if CUDA is available
if torch.cuda.is_available():
    torch_model = torch_model.to('cuda')

# choose and initialize an optimizer to train the model
optimizer = torch.optim.Adam(torch_model.parameters(), lr=1e-3)

ResUnit 1 is 66 units wide.
Dropout before it is 0.18674307214231128.
ResUnit 2 is 41 units wide.
Dropout before it is 0.1162432499771616.
ResUnit 3 is 25 units wide.
Dropout before it is 0.07235873872180604.
ResUnit 4 is 16 units wide.
Dropout before it is 0.045041643884176266.
ResUnit 5 is 10 units wide.
Dropout before it is 0.02803738317757008.

# wrap the pytorch neural network model in a RCModel class,
# these classes know how to decide if they should train in a self-consistent way
# and they also know how to transform from configurations to descriptors space (because they know about the descriptor_transform) 
# Here we take an ExpectedEfficiencyPytorchRCModel,
# this RCmodel scales the learning rate by the expected efficiency factor (1 - n_TP_true / n_TP_expected)**2
model = aimmd.pytorch.EEScalePytorchRCModelAsync(nnet=torch_model,
                                                 optimizer=optimizer,
                                                 states=[wrapped_C7_eq, wrapped_alphaR],
                                                 ee_params={'lr_0': 1e-3,
                                                            'lr_min': 5e-5,  # lr_min = lr_0 / 20 is a good choice empirically
                                                            'epochs_per_train': 3,
                                                            'window': 100,
                                                            'batch_size': 8192,
                                                           },
                                                 descriptor_transform=wrapped_transform,
                                                 )

Define the sampling scheme#

We will first define the shooting point selection, the shooting point selector (or equivalently the choosen selection scheme) determines the acceptance probability for each new trial because it decides on how we select the new shooting points from the last accepted transition.

We will then use the selector to setup our sampling scheme, which is very simple here as it only consists of one mover (the two way shooting mover). However you can use an arbitray number of movers (potentially defining a probability for each of them), in that case each mover will be picked with the given probability to generate the next trial.

Shooting point selection#

The shooting point selection (scheme) can be influenced by choosing a shooting point selector. It determines if we, e.g., select a random shooting point from the last trajectory(using a UniformSPSelector), select a shooting point from the last trajectory according to the learned committor (RCModelSelectorFromTraj), or even if we shoot from a set of configurations with known equilibrium weight (RCModelSelectorFromEQ).

In addition to selecting the shooting point, the shooting point selectors that select according to a committor model also take care of the density adaption. Density adaption flattens the density of points on transitions when projected into the space of the committor and thereby corrects for asymetric barriers (see also the documentation here).

# we will use the reaction coordinate model to select SPs biased towards its current guess of the transition state
# from the last accepted transition
# the RCModelSelectorFromTraj has a couple of options to customize its behavior (e.g. to decide how sharply peaked the selection should be)
# but we will leave it at its default parameters for now
# One notable feature is the density adaption, which will correct for the density of points on transition projected into the space of the committor,
# it is enabled by default
spselector_cls = aimmdd.spselectors.RCModelSPSelectorFromTraj
# None is the default value for ``sp_selector_kwargs``, it is equivalent to an empty dictionary,
# i.e., passing None will initialize the spselector with its default arguments
spselector_kwargs = None

Pathmovers#

A Pathmover moves from one transition path to the next, i.e. it generates a new trial transition path which can then be accepted or rejected. All shooting path movers need a shooting point selector to function.

A TPS scheme can be defined via multiple PathMovers, each with a weight. Each move then begins by drawing a random mover (according to their weights), which then generates the new trial trajectory. This enables the mixing of different move types and can increase exploration of the path space.

# setup a list of movers
# since we want to create n_sampler identical samplers it is easiest to use the `Brain.samplers_from_moverlist()` function
# This function will create n_sampler identical PathChainSamplers where the movers for each sampler are
# specified by movers_cls (a list of mover classes) and movers_kwargs (a dict with keyword arguments used for initialization of the movers)
movers_cls = [aimmdd.pathmovers.TwoWayShootingPathMover]
movers_kwargs = [{'states': [wrapped_alphaR, wrapped_C7_eq],
                  'md_engine_spec': engine_spec,
                  'temperature': mdp["ref-t"][0],
                  # use the spselector (and possibly arguments) we have defined above
                  "sp_selector_cls": spselector_cls,
                  "sp_selector_kwargs": spselector_kwargs,
                  # Many additional settings for the shooting movers can be set via kwargs
                  # these are usually not needed, but can influence details like
                  # the names of the output trajectories and files, or if temporary
                  # trajectories are removed (the default is to remove them)
                  #"remove_temporary_trajectories": False,
                  }
                 ]

# Note that for full flexibility of the sampling scheme setup we could also use a list of lists with initialized movers
# then however we need the outermost list to be of length n_sampler as shown below
#movers = [[aimmdd.TwoWayShootingPathMover(states=[wrapped_C7_eq, wrapped_alphaR],
#                                          md_engine_spec=engine_spec,
#                                          temperature=mdp["ref-t"][0],
#                                          sp_selector_cls=spselector_cls,
#                                          sp_selector_kwargs=spselector_kwargs,
#                                         )
#           ] for i in range(n_samplers)
#         ]

Trainset#

trainset = aimmd.TrainSet(n_states=2)

Brain tasks#

Each task will be run after every finished Monte Carlo step with the step and the index of the sampler it came from as arguments. You can also define your own tasks to modify the behaviour of the TPS simulation easily (openpathsampling users should think of hooks). Note that each user defined Task should subclass aimmd.distributed.pathsampling.BrainTask. Note also that the tasks will be run in the order that they are in the list, i.e. for the list below the TrainingTask will be run first for every Monte Carlo step.

A number of required and (potentially) useful Tasks are already included with aimmd.distributed, e.g. the TrainingTask (which adds new shooting results to the training set and perdiodically asks the reaction coordinate model to train) and the SaveTask (which saves the reaction coordinate model and training set periodically).

A noteworthy Task class for longrunning TPS simulations is the StorageCheckpointTask which creates a copy of the aimmd.Storage used during the simulation at a given interval of finished Monte Carlo steps. The copy of the storage file can then be openend and accessed while the TPS simulation is still running, thereby enabeling preliminary analyses of a long-running TPS simulation.

tasks = [
    # the TrainingTask takes care of training the model (or better: reminding the model to decide if it wants to train)
    aimmdd.pathsampling.TrainingTask(model=model, trainset=trainset),
    # the SaveTask saves the model, trainset and brain to storage at specified interval (default=100 steps) during the simulation
    aimmdd.pathsampling.SaveTask(storage=storage, model=model, trainset=trainset),
    # the StorageCheckpointTask is commented out because this simulation should not run soo long
    #aimmdd.pathsampling.StorageCheckpointTask(storage=storage,  # the storage to checkpoint
    #                                          # increase the checkpoint interval to 100 MCSteps
    #                                          interval=100,
    #                                          # and leave the options that control the checkpoint
    #                                          # naming at their default values
    #                                          checkpoint_suffix=".ckpt",
    #                                          checkpoint_prev_suffix="_prev",
    #                                          )
         ]

# this is the 'easy' way to setup n_sampler identical samplers (using `samplers_from_moverlist` as described above)
brain = aimmdd.Brain.samplers_from_moverlist(model=model, workdir=workdir, storage=storage,
                                             n_sampler=n_samplers,
                                             movers_cls=movers_cls, movers_kwargs=movers_kwargs,
                                             samplers_use_same_stepcollection=False,
                                             tasks=tasks)
                                             # Note that we left mover_weights=None at its default, this results
                                             # in uniform weights for all movers


# and this would be the full __init__ call to the brain (given you defined `movers` as above commented out) 
# it gives you full flexibility of setting up every PathChainSamplers individually
#brain = aimmdd.Brain(model=model, workdir=workdir, storage=storage, movers=movers, mover_weights=[[1.], [1.], [1.]], tasks=tasks)

Seed initial transitions for each Markov chain#

Here we will use a convinience function of the brain that takes a list of transitions (and optionaly a list of weights) and seeds all Markov chains from the given transitions (and weights) randomly. If no weights are given a uniform distribution is used.

# set initial paths for each sampler directly
for traj, sampler in zip((tp_short, tp_short, tp_short2, tp_short2),
                         brain.samplers):
    sampler.set_step_zero_from_trajectory(trajectory=traj)

# or seed the initial paths at random
#brain.seed_initial_paths(trajectories=[#tp_lb,
#                                       tp_short, tp_short2],
#                        )

# have a look at what the initial path for each sampler is
for s in brain.samplers:
    print(s.current_step.path)

Trajectory(trajectory_files=gmx_infiles/TP_short_low_barrier_300K_amber99sbildn.trr, structure_file=gmx_infiles/ala_300K_amber99sb-ildn.tpr)
Trajectory(trajectory_files=gmx_infiles/TP_short_low_barrier_300K_amber99sbildn.trr, structure_file=gmx_infiles/ala_300K_amber99sb-ildn.tpr)
Trajectory(trajectory_files=gmx_infiles/TP_short2_low_barrier_300K_amber99sbildn.trr, structure_file=gmx_infiles/ala_300K_amber99sb-ildn.tpr)
Trajectory(trajectory_files=gmx_infiles/TP_short2_low_barrier_300K_amber99sbildn.trr, structure_file=gmx_infiles/ala_300K_amber99sb-ildn.tpr)

NOTE: Depending on how your initial transitions were generated they can be very different from a transition that is typical for the equilibrium transition path ensemble you are sampling. This is e.g. the case if your initial transition was generated using pulling or other non-equilibrium work protocols and can result in very low acceptance rates at the beginning of the TPS simulation.#

To understand why especially initial transition that are shorter than the typical transition (or transitions on which a larger fraction of the path is spent around the transition state than on typical transitions) result in low acceptance rates please have a look at the acceptance criterion detailed e.g. in the aimmd publication (doi:10.1038/s43588-023-00428-z).

To remedy this situation the PathChainSamplers have a boolean attribute (always_accept_next_TP), which can be used to force the acceptance of the next generated transition independent of the acceptance criterion. It can (and needs) to be set for every sampler individually, because every sampler can have a different initial transition (and sampling scheme). To ensure that the first produced transition in every sampler is accepted you could e.g. use something similiar to the following before starting the TPS simulation:

for s in brain.samplers:
    s.always_accept_next_TP = True

Note however that this essentially means that your Markov chain starts after the first (forced) accept, i.e. you should start any analyses earliest after/with the forced accept.

#for s in brain.samplers:
#    s.always_accept_next_TP = True

Run the TPS simulation#

n_steps = 4000
await brain.run_for_n_steps(n_steps)

/dpool/coding/projects/aimmd/aimmd/pytorch/rcmodel.py:343: UserWarning: Converting a tensor with requires_grad=True to a scalar may lead to unexpected behavior.
Consider using tensor.detach() first. (Triggered internally at /pytorch/torch/csrc/autograd/generated/python_variable_methods.cpp:836.)
  total_loss += float(loss)

# you can also run for a specified number of accepts,
# independent on how many steps that might take
#await brain.run_for_n_accepts(200)

brain.total_steps

4000

Save the last model, trainset and brain to storage#

This enables us to do the analysis in a different notebook and/or continue the TPS simulation from the last step easily.

Note that the SaveTask defined above also saves the reaction coordinate model, the trainset and the brain, but only at the specified interval, so it is always good to save the last state. For simulations where single steps are already quite costly it might be worth sacrificing some disk space at setting the interval=1 to save after every finished step.

# save the last model
storage.rcmodels["model_to_continue_with"] = model
storage.save_trainset(trainset)  # the trainset
storage.save_brain(brain)  # and the brain

# close the storage to make sure all writes are synced
storage.close()

WARNING:asyncmd.config:Deregistered global writeable h5py cache. No TrajectoryFunction values will be cached until a new h5py cache has been registered.