Massively parallel Transition Path Sampling

Massively parallel Transition Path Sampling#

Notebook 2: Continue TPS simulation#

This is the second of a series of example notebooks on massively parallel transition path sampling. Obviously you need to have run the TPS simulation in the first notebook to continue it here, so please do so if you have not yet already.

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
import numpy as np
import matplotlib.pyplot as plt
import mdtraj as mdt
import MDAnalysis as mda
import aimmd
import aimmd.distributed as aimmdd

/home/think/.conda/envs/aimmd_dev/lib/python3.10/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm
Could not initialize SLURM cluster handling. If you are sure SLURM (sinfo/sacct/etc) is available try calling `asyncmd.config.set_slurm_settings()` with the appropriate arguments.
Tensorflow/Keras not available

# setup working directory (this should be the same as in the previous notebook)
scratch_dir = "."

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

Continue the TPS simulation#

To continue the TPS simulation we need to:

  • open the storage file from last time

  • load the reaction coodrinate model and trainset from it

  • redefine the brain tasks using the storage, reaction coordinate model and trainset

  • and finaly load the brain using the reaction coodrinate model and tasks

The loaded brain will then be exactly in the state we left it at the end of the last notebook.

Open the storage (in append mode)#

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

Load the reaction coordinate model#

model = storage.rcmodels["model_to_continue_with"]

Load the trainset#

trainset = storage.load_trainset()

Redefine the tasks#

tasks = [aimmdd.pathsampling.TrainingTask(model=model, trainset=trainset),
         aimmdd.pathsampling.SaveTask(storage=storage, model=model, trainset=trainset),
         aimmdd.pathsampling.DensityCollectionTask(model=model,
                                                   first_collection=100,
                                                   recreate_interval=500,
                                                   ),
         ]

Load the brain#

brain = storage.load_brain(model=model, tasks=tasks)

# this will be exactly the number of steps we left of at
brain.total_steps

10009

Run the TPS simulation for some more steps#

# we can just run the simulation a bit longer
await brain.run_for_n_steps(9)

brain.total_steps

10018

Save the last model, trainset and brain to storage#

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

storage.close()
Contents