22 dimensional toy system setup with a keras/tensorflow RCmodel¶

%matplotlib inline

import os
import aimmd
import numpy as np
import matplotlib.pyplot as plt
import openpathsampling as paths
# convenience for the toy dynamics
import openpathsampling.engines.toy as toys
from functools import reduce

# load a helper class for visualization of the toy system
%run resources/toy_plot_helpers.py

# setup logging
# executing this file sets the variable LOGCONFIG, which is a dictionary of logging presets for openpathsampling
%run resources/logconf.py

# set working directory
# set to None to use current folder (where the notebook is located)
wdir = '/home/tb/hejung/DATA/aimmd_scratch/SimData_keras_toy_22dim'  # relative path
if wdir is not None:
    if not os.path.exists(wdir):
        # create working dir
        os.mkdir(wdir)
    # and change to it
    os.chdir(wdir)

# you can either modify single values or use it as is to get the same setup as in the OPS default logging config file
# you could e.g. do LOGCONF['handlers']['stdf']['filename'] = new_name to change the filename of the log
# the default is to create 'simulation.log' and 'initialization.log' in the current working directory
# having them in a dictionary makes it easy to e.g. change the logfilename as such:
LOGCONFIG['handlers']['stdf']['filename'] = 'toy_keras_simulation.log'
LOGCONFIG['handlers']['initf']['filename'] = 'toy_keras_init.log'

import logging.config
logging.config.dictConfig(LOGCONFIG)

# create a topology with 2 + n_harmonics cartesian dimensions
topology=toys.Topology(n_spatial=2 + n_harmonics,
                       # set the mass for every dimension to 1
                       masses=np.array([1.0 for _ in range(2 + n_harmonics)]),
                       # use the PES we defined above
                       pes=pes,
                       # use one particle, since they would not interact anyway
                       n_atoms=1
                       )

# standard openpathsampling toy integrator
integ = toys.LangevinBAOABIntegrator(dt=0.02, temperature=0.1, gamma=2.5)

# create the openpathsampling toy engine
options={'integ': integ,
         'n_frames_max': 5000,
         'n_steps_per_frame': 1
         }
toy_eng = toys.Engine(options=options,
                      topology=topology
                      )
toy_eng.initialized = True

# create a template snapshot, such that the toy engine knows about its dimensions, masses and PES
template = toys.Snapshot(coordinates=np.array([[-0.75, -0.5] + [0. for _ in range(n_harmonics)]]), 
                         velocities=np.array([[0.0, 0.0] + [0. for _ in range(n_harmonics)]]),
                         engine=toy_eng
                         )
toy_eng.current_snapshot = template

def circle(snapshot, center):
    """
    Return (radial) distance of snapshot from center.
    """
    import math
    return math.sqrt((snapshot.xyz[0][0]-center[0])**2 + (snapshot.xyz[0][1]-center[1])**2)

# Order parameters to define the states
# these just measure how far we are from the center of the states
opA = paths.CoordinateFunctionCV(name="opA", f=circle, center=[-0.75, -0.5])
opB = paths.CoordinateFunctionCV(name="opB", f=circle, center=[0.75, 0.5])

# State volumes in order parameter space
stateA = paths.CVDefinedVolume(opA, 0.0, 0.15).named('StateA')
stateB = paths.CVDefinedVolume(opB, 0.0, 0.15).named('StateB')

# collective variable to transform OPS snapshots to model descriptor space,
# i.e. the space in which the model learns
# for this simple toy system we can just feed the cartesian coordinates of the system
descriptor_transform = paths.FunctionCV('descriptor_transform', lambda s: s.coordinates[0], cv_wrap_numpy_array=True).with_diskcache()

initAB = paths.Trajectory([toys.Snapshot(coordinates=np.array([[-0.75 + i/700., -0.5 + i/1000] + [0. for _ in range(n_harmonics)]]), 
                                         velocities=np.array([[1.0, 0.0]  + [0. for _ in range(n_harmonics)]]),
                                         engine=toy_eng
                                        )
                          for i in range(1001)])

# to project N dimensional trajectories to first 2 (relevant) dimesnsions for ploting
def tra_to_2d(tra):
    """
    Extracts the two non harmonic degrees of freedom of a trajectory for ploting.
    """
    snap_list = []
    for s in tra:
        snap_list.append(toys.Snapshot(coordinates=np.array(s.coordinates[:,:2])))
    return paths.Trajectory(snap_list)

from tensorflow.keras import optimizers, regularizers

# create aimmd storage
aimmd_store = aimmd.Storage('aimmd_storage.h5')

# we take an ExpectedEfficiencyKerasRCModel,
# this RCmodel scales the learning rate by the expected efficiency factor (1 - n_TP_true / n_TP_expected)**2
model = aimmd.keras.EEScaleKerasRCModel(nnet=resnet, # any keras model predicting a float/a vector of floats,
                                                     # each output will be interpreted as the log probability to reach a certain state
                                                     # if the model only predicts one float it is assumed to be p_B
                                        states=[stateA, stateB],
                                        cache_file=aimmd_store,
                                        descriptor_transform=descriptor_transform,  # the function transforming snapshots to descriptors,
                                                                                    # i.e. the space in which the model learns)
                                        ee_params={'lr_0': 0.005,  # lr = lr_0 * (1 - n_TP_true / n_TP_expected)**2
                                                   'lr_min': 0.0001,  # do not train if lr < lr_min
                                                   'epochs_per_train': 2,  # train for 2 epochs every time we train
                                                   'interval': 5,  # attempt to train after every 5rd MCStep
                                                   'window': 50,  # average expected efficiency factor over 50 MCSteps
                                                   },
                                        )

trainset = aimmd.TrainSet(n_states=2)

trainhook = aimmd.ops.TrainingHook(model, trainset)
storehook = aimmd.ops.AimmdStorageHook(aimmd_store, model, trainset)
densityhook = aimmd.ops.DensityCollectionHook(model)

selector = aimmd.ops.RCModelSelector(model=model,  # always takes a RCModel
                                     # we can greatly speed up rejecting/accepting trial TPS by passing the list of states
                                     # this enables testing if a trial TP is a TP and calculating the potentialy costly
                                     # transformation from Cartesian to descriptor space only if neccessary (i.e. if it is a TP and we need to calculate the MC acceptance)
                                     # if we were lazy and knew that the transformation was fast, we could also explicitly pass None
                                     states=[stateA, stateB],
                                     # new shooting points are selected with p_sel(SP) ~ p_lorentz(model.z_sel(SP))
                                     # we could also choose 'gaussian'
                                     distribution='lorentzian',
                                     # softness of the selection distribtion,
                                     # lower values result in a sharper concentration around the predicted transition state,
                                     # higher values result in a more uniform selection
                                     scale=1.0,
                                    )

# velocity randomizer setup
beta = integ.beta
modifier = paths.RandomVelocities(beta=beta, engine=toy_eng)
# shooting strategy: TwoWayShooting
tw_strategy = paths.strategies.TwoWayShootingStrategy(modifier=modifier, selector=selector, engine=toy_eng, group='TwoWayShooting')
# transition network, tells us which transitions to sample
tps = paths.TPSNetwork.from_states_all_to_all([stateA, stateB])
# move scheme, tells us how to do the sampling
move_scheme = paths.MoveScheme(network=tps)
move_scheme.append(tw_strategy)
# This should always be the last added strategy when constructing custom move_schemes in openpathsampling,
# it sorts all previously appended strategies into a sane hirachy when building the move tree
move_scheme.append(paths.strategies.OrganizeByMoveGroupStrategy())
move_scheme.build_move_decision_tree()

# create a storage file
storage = paths.Storage('22dim_toy_keras.nc', 'w', template=template)

# create a PathSampling
sampler = paths.PathSampling(storage=storage,
                             sample_set=initial_conditions,
                             move_scheme=move_scheme
                             )

sampler.attach_hook(trainhook)
sampler.attach_hook(storehook)
sampler.attach_hook(densityhook)

# get the number of accepts from OPS storage
accepts = []
for step in storage.steps:
    if step.change.canonical.accepted:
        accepts.append(1.)
    else:
        accepts.append(0.)

# we can also have a look at the models prediction in x/y
x = np.linspace(-1.2, 1.2, 240)
y = np.linspace(-1., 1., 200)

# we take the hypersurface where all oscillators are 0
oscis = [0. for _ in range(n_harmonics)]
# do one prediction for one big vector and unwrap it afterwards
# this is much faster than looping over the single coordinate points as we can make use of the GPU/CPU vector operations...!
coord = np.array([[xv, yv]  + oscis for yv in y for xv in x])
p = model(coord, use_transform=False)
p = p.reshape((len(y), len(x)))
pred = np.stack([1. - p, p], axis=-1)

# perform a HIPR analysis of the model over the trainset
hipr = aimmd.analysis.HIPRanalysis(model,
                                   trainset,
                                   # setting n_redraw=5 will do 5 passses over the randomized training set for every descriptor dimension and average the results 
                                   n_redraw=25,
                                   )

# the HIPR analysis always returns a tuple: the average loss for every coordinate and the standard deviation over the redraws 
hipr_plus_losses, hipr_plus_stds = hipr.do_hipr_plus(n_redraw=50)  # hipr plus keeps the distribution of the input descriptors invariant and only permutates them
# the n_redraw argument given to do_hipr_plus()/do_hipr() takes precedence over the one given at init time in case you realize you need a less noisy estimate later

# sync and close the OPS storage
# not strictly neccessary but a good habit to make sure everything you expect is in there
storage.sync_all()
storage.close()

aimmd_store.close()