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 aimmd
import aimmd.distributed as aimmdd

WARNING:asyncmd.slurm.process: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 (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),
         ]

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

4000

Run the TPS simulation for some more steps#

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

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()

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