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
import numpy as np
import matplotlib.pyplot as plt
import MDAnalysis as mda
# 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.nn.functional as F
import torch
/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
scratch_dir = "."
workdir = os.path.join(scratch_dir, "TransitionPathSampling_ala")
if not os.path.isdir(workdir):
os.mkdir(workdir)
Setup logging to a file (optional)#
The next few cells are just to show you how to configure pythons logging module to write to a logfile in the directory where we do the simulation. It is not necessary to run aimmd but it might be helpful to find out what went wrong if something does.
# setup logging
# executing this file sets the variable LOGCONFIG, which is a dictionary of logging presets
%run ../resources/logconf.py
# have a look at the default logging level (the level used for the root logger)
print(LOGCONFIG["loggers"][""])
# have a look at the logger for aimmd
print(LOGCONFIG["loggers"]["aimmd"])
# and have a look at the log-level for the filehandler
print(LOGCONFIG["handlers"]["stdf"])
# the last two should both be `INFO`
{'level': 'WARN', 'handlers': ['stdf', 'warnout']}
{'level': 'INFO'}
{'class': 'logging.FileHandler', 'level': 'INFO', 'mode': 'w', 'filename': 'simulation.log', 'formatter': 'standardFormatter'}
# OPTIONAL: more logging to file
level = "INFO"
LOGCONFIG["handlers"]["stdf"]["level"] = level
LOGCONFIG["loggers"]["aimmd"]["level"] = level
LOGCONFIG["loggers"]["asyncmd"] = {"level": level}
# 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
import logging.config
LOGCONFIG["handlers"]["stdf"]["filename"] = os.path.join(workdir, "simulation_pathsampling.log")
LOGCONFIG["handlers"]["initf"]["filename"] = os.path.join(workdir, "initlog_pathsampling.log")
logging.config.dictConfig(LOGCONFIG)
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 = 5 # 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.
# 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")
gmx_engine_kwargs = {"mdconfig": mdp,
"gro_file": gro,
"top_file": top,
"ndx_file": ndx,
"output_traj_type": "XTC",
#"mdrun_extra_args": "-nt 2",
# use this for gmx sans (thread) MPI
"mdrun_extra_args": "-ntomp 2",
}
gmx_engine_cls = asyncgmx.GmxEngine
Load initial transitions#
# 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_tp, descriptors_for_tp_short, descriptors_for_tp_short2 = await asyncio.gather(wrapped_transform(tp_lb),
wrapped_transform(tp_short),
wrapped_transform(tp_short2),
)
# 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_tp.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.
# wrapp 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,
cache_file=storage,
)
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.
# we will use the reaction coodrinate 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 parmeters for now
# One notable feature is the density adaption, which will correct for the denisty of points on transition projected into the space of the committor,
# it is enabled by default
spselector = aimmdd.spselectors.RCModelSPSelectorFromTraj()
# 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],
'engine_cls': gmx_engine_cls,
'engine_kwargs': gmx_engine_kwargs,
# 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 gromasc 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
'T': mdp["ref-t"][0],
"sp_selector": spselector, # use the spselctor we have defined above
"max_steps": 500 * 10**5, # 500 steps * dt (2 fs) = 1 ps
}
]
# 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],
# engine_cls=gmx_engine_cls,
# engine_kwargs=gmx_engine_kwargs,
# engine_config=mdp,
# walltime_per_part=0.00003125,
# T=mdp["ref-t"][0],
# sp_selector=spselector,
# max_steps=500 * 10**5,
# )
# ] 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 and the DensityCollectionTask will run last 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), the SaveTask (which saves the reaction coordinate model and training set periodically), and the DensityCollectionTask (which regularly updates the estimate of the density of shooting points projected to committor space) will be an integral part of most/all TPS simulations.
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 DensityCollectionTask takes care of updating the estimate of the density of shooting points
# projected into committor space
# It needs to know what the ensemble is we shoot from (e.g. "custom" for a self-defined set of shooting points
# or the default "p_x_TP" if we shoot from previous transitions)
aimmdd.pathsampling.DensityCollectionTask(model=model,
first_collection=100,
recreate_interval=250,
),
# 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.
# seed initial transitions
brain.seed_initial_paths(trajectories=[#tp_lb,
tp_short, tp_short2],
weights=[#1.,
1., 1.],
)
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 that spend more time around the transition state than 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 remidy 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 next line at the beginning of a TPS simulation:
[s.always_accept_next_TP = True for s in brain.samplers]
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.
Run the TPS simulation#
import time
n_steps = 10000
start = time.time()
await brain.run_for_n_steps(n_steps,
print_progress=500, # print a short info on simulation progress every 500 steps
)
end = time.time()
print("-" * 100)
print(f"Running for {n_steps} cummulative MCSteps took {end-start} s (= {(end-start)/60} min).")
Starting simulation at Tue Jul 23 18:57:23 2024.
Tue Jul 23 19:05:13 2024: 500 (of 10000) steps done. Produced 146 new accepts so far. Estimated end time is Tue Jul 23 21:34:03 2024.
Tue Jul 23 19:13:35 2024: 1000 (of 10000) steps done. Produced 307 new accepts so far. Estimated end time is Tue Jul 23 21:39:23 2024.
Tue Jul 23 19:22:30 2024: 1500 (of 10000) steps done. Produced 446 new accepts so far. Estimated end time is Tue Jul 23 21:44:48 2024.
Tue Jul 23 19:31:38 2024: 2000 (of 10000) steps done. Produced 595 new accepts so far. Estimated end time is Tue Jul 23 21:48:35 2024.
Tue Jul 23 19:41:18 2024: 2500 (of 10000) steps done. Produced 717 new accepts so far. Estimated end time is Tue Jul 23 21:53:03 2024.
Tue Jul 23 19:51:37 2024: 3000 (of 10000) steps done. Produced 843 new accepts so far. Estimated end time is Tue Jul 23 21:58:08 2024.
Tue Jul 23 20:01:46 2024: 3500 (of 10000) steps done. Produced 985 new accepts so far. Estimated end time is Tue Jul 23 22:01:20 2024.
Tue Jul 23 20:12:28 2024: 4000 (of 10000) steps done. Produced 1134 new accepts so far. Estimated end time is Tue Jul 23 22:05:04 2024.
Tue Jul 23 20:23:17 2024: 4500 (of 10000) steps done. Produced 1269 new accepts so far. Estimated end time is Tue Jul 23 22:08:15 2024.
Tue Jul 23 20:34:09 2024: 5000 (of 10000) steps done. Produced 1416 new accepts so far. Estimated end time is Tue Jul 23 22:10:55 2024.
Tue Jul 23 20:45:08 2024: 5500 (of 10000) steps done. Produced 1560 new accepts so far. Estimated end time is Tue Jul 23 22:13:17 2024.
Tue Jul 23 20:56:46 2024: 6000 (of 10000) steps done. Produced 1688 new accepts so far. Estimated end time is Tue Jul 23 22:16:21 2024.
Tue Jul 23 21:08:12 2024: 6500 (of 10000) steps done. Produced 1835 new accepts so far. Estimated end time is Tue Jul 23 22:18:38 2024.
Tue Jul 23 21:21:29 2024: 7000 (of 10000) steps done. Produced 1960 new accepts so far. Estimated end time is Tue Jul 23 22:23:15 2024.
Tue Jul 23 21:35:37 2024: 7500 (of 10000) steps done. Produced 2079 new accepts so far. Estimated end time is Tue Jul 23 22:28:21 2024.
Tue Jul 23 21:49:10 2024: 8000 (of 10000) steps done. Produced 2213 new accepts so far. Estimated end time is Tue Jul 23 22:32:06 2024.
Tue Jul 23 22:03:02 2024: 8500 (of 10000) steps done. Produced 2354 new accepts so far. Estimated end time is Tue Jul 23 22:35:47 2024.
Tue Jul 23 22:17:40 2024: 9000 (of 10000) steps done. Produced 2500 new accepts so far. Estimated end time is Tue Jul 23 22:39:55 2024.
Tue Jul 23 22:35:15 2024: 9500 (of 10000) steps done. Produced 2643 new accepts so far. Estimated end time is Tue Jul 23 22:46:43 2024.
Tue Jul 23 23:03:26 2024: 10000 (of 10000) steps done. Produced 2778 new accepts so far. Estimated end time is Tue Jul 23 23:03:26 2024.
----------------------------------------------------------------------------------------------------
Running for 10000 cummulative MCSteps took 14762.756368637085 s (= 246.04593947728475 min).
brain.total_steps
10000
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()