Loading and saving RCModels

Loading and saving RCModels#

This notebook will show you:

  • how to load any aimmd.RCModel by hand

  • how to save any aimmd.RCModel by hand

Note: This notebook depends on files created in either one of the the notebooks 1_Toy_pytorch_simple_setup.ipynb and 1_Toy_tensorflow_simple_setup.ipynb. Please do one of them first if you have not already.

%matplotlib inline

import os
import aimmd
import numpy as np
import matplotlib.pyplot as plt
import openpathsampling as paths

# set this to False if you want to use the files from the pytorch notebook instead
use_tf_notebook_files = True

# change to the working directory of choice
# (should be the same as in the first notebook)
if use_tf_notebook_files:
    wdir = '/home/tb/hejung/DATA/aimmd_scratch/SimData_keras_toy_22dim/'
else:
    wdir = '/home/tb/hejung/DATA/aimmd_scratch/SimData_pytorch_toy_22dim/'

#wdir = None
if wdir is not None:
    os.chdir(wdir)

Loading#

You can load any RCModel from the aimmd.Storage. This is useful for example to do analysis of the intermediate models that have been saved.

aimmd_storage = aimmd.Storage('aimmd_storage.h5', "a")  # open storage in append mode

The aimmd.Storage exposes a dictionary like interface for all models at storage.rcmodels.

print(aimmd_storage.rcmodels.keys())  # will be a KeysView
print()
print("Available models in the storage are:")
for k in aimmd_storage.rcmodels:
    print(k)

KeysView(<aimmd.base.storage.RCModelRack object at 0x7f1d3eb53b50>)

Available models in the storage are:
RCModel_after_step_1000
RCModel_after_step_500
most_recent

There should always be a most_recent model if you used the aimmd.Storage with a TPS simulation. This key will always point to the model (state) at the point in time when the simulation has finished its last TPS step. This should make it easy to always get the most recent model for restarting/continuing a simulation.

# load the most recent model
model = aimmd_storage.rcmodels["most_recent"]

Note that the model will have forgotten about the descriptor_transform and the states (if they are openpathsampling objects)#

All openpathsampling objects are saved in the openpathsampling.Storage (because this storage knows how to save them and always will). Currently this affects the states (if they are openpathsampling.Volume objects) and the descriptor_transform (if it is a openpathsampling.CollectiveVariable or sublcass thereof).

We can however easily readd them to the model in the current python session by calling model.complete_from_ops_storage(store) where store is the openpathsampling.Storage containing the descriptor_transform collective variable and and the state volumes.

Note: If states or descriptor_transform had been pure python objects/functions, they would have been be pickled and saved with the model to the aimmd.Storage.

print(model.states)  # these will be strings indicating the names of the ops volumes
print(model.descriptor_transform)  # same for the descriptor transform

['StateA', 'StateB']
descriptor_transform

# open ops storage for reading
if use_tf_notebook_files:
    storage = paths.Storage("22dim_toy_keras.nc", "r")  # read only
else:
    storage = paths.Storage('22dim_toy_pytorch.nc', 'r')

# complete the model
model.complete_from_ops_storage(storage)
# now states and descriptor transform should be ops objects again
print(model.states)
print(model.descriptor_transform)

[<openpathsampling.volume.CVDefinedVolume object at 0x7f17a3e19190>, <openpathsampling.volume.CVDefinedVolume object at 0x7f17a3e19730>]
<openpathsampling.collectivevariable.FunctionCV object at 0x7f1c7675f370>

# check that it worked
print(model)  # should print an initialized aimmd model
print(model(storage.snapshots[-1]))  # predicts the commitment probability for the last shooting snapshot in the ops storage

<aimmd.keras.rcmodel.EEScaleKerasRCModel object at 0x7f1c76f03490>
[[0.14685754]]

Saving#

Of course you can also save any aimmd.RCModel you want at any time. Simply (re)set the key with the desired name to any aimmd.RCModel.

As noted above, in case the model has an attached OPS CollectiveVariable as descriptor_transform or Volume objects as states, it will save their names instead, such that you can later only completely reconstruct the model if passing an openpathsampling.Storage that contains the/any CV with that name.

aimmd_storage.rcmodels["save_test"] = model

# list all keys to see that it is there
list(aimmd_storage.rcmodels.keys())

['RCModel_after_step_1000',
 'RCModel_after_step_500',
 'most_recent',
 'save_test']

Arbitrary python functions for descriptor_transform and states#

# now lets test saving arbitrary python functions as descriptor_transform
def test_func(x, y):
    return x + y

# this will break our model as it changes the call signature of descriptor_transform, only for demonstration purposes
model.descriptor_transform = test_func

aimmd_storage.rcmodels["save_test2"] = model

# load it to see if it worked
loaded_model = aimmd_storage.rcmodels["save_test2"]

# will return 2
loaded_model.descriptor_transform(1, 1)

2

aimmd_storage.close()  # close the aimmd.Storage, important to make sure everything is flushed from cache to disk

storage.close()  # close the ops storage (should not be necessary if open in read-only, but better be save than sorry...)
Contents