MAUL (Modified ANI with Uncertainity Limits)

Maul is a set of Bayesian neural networks built based upon the DNN ensemble ANI. This library contains the key tools to build and train a Maul neural potential. The neural potential offers a predicted energy together with an estimate of the model (epistemic) uncertaininty based upon Monte Carlo sampling over parameter distributions.

Library

The library consists of five core modules built on top of pytorch, torchani and Bayesian torch. The core modules are:

1. network_classes.py
2. network_data_loader.py
3. network_train.py
4. network_test.py
5. network_utilities.py

Firstly, network_classes.py defines a super class for ANI style networks, enabling efficient automated building of these networks and methods to track their status. The core class here is the ANI() class. From this module one can define a standard ANI potential and in one command convert the potential to Maul style BNN neural potential.

Next, we have newtork_data_loader.py. This module as it says is for loading data in a consistent way. The module is designed around the HDF5 files from the ANI1 and ANI-1x datasets. The module allows the data used to be mutated in a consistent way by passing a --mutate_index INT which becomes a seed to repeatably partition the trianing and validation sets such that over an ensemble each member has a different sub set.

Network_train.py and network_test.py have functionality for training and testing the neural potentials The modules can train and test DNN neural potentials and BNN neural potentials.

Finally, network_utilities.py is an important module conatining most of the plotting tools for analysis of performance.

Install

Build a new environment using conda or venv

python -m venv

or

conda create -n maul-bnnani python=3.9

conda activate maul-bnnani

Install the Maul bnnani library

python setup.py install

Args

optional arguments:

-h, --help show this help message and exit

Featurization options:

--radial_cutoff RADIAL_CUTOFF Radial cutoff distance (default: 5.2)

--theta_max_rad THETA_MAX_RAD Theta distance cutoff (default: 3.33794218)

--angular_cutoff ANGULAR_CUTOFF Angular cutoff in radians (default: 3.5)

--etar ETAR Eta in the radial function determines the guassians width (default: 16.0)

--etaa ETAA Eta in the angular function determines the guassians width for the radial portion (default: 8.0)

--zeta ZETA Zeta in the angular function determines the guassians width for the angular portion (default: 32.0)

--radial_steps RADIAL_STEPS Number of radial shifts moves the peak of the gaussian (default: 16)

--angular_radial_steps ANGULAR_RADIAL_STEPS Number of angular shifts moves the peak of the gaussian (default: 4)

--theta_steps THETA_STEPS Angle steps (default: 8)

--species_order [SPECIES_ORDER [SPECIES_ORDER ...]] Species order at atomic symbol (default: ['H', 'C', 'N', 'O'])

--self_energies [SELF_ENERGIES [SELF_ENERGIES ...]] Self energies in the order of species order (default: None)

Data options:

--data str Data directory where HDF5 file of training and testing data can be found (default: ../../data/ANI-1_release)

--no_reload Tell the code to reload data don't use previous loading data (default: False)

--checkpoint_filename CHECKPOINT_FILENAME initial weights path (default: ani_model_checkpoint.pt)

--learning_curve_csv_filename LEARNING_CURVE_CSV_FILENAME best weights path (default: learning_curve.csv)

--batch_size BATCH_SIZE Batch size for training (default: 1024)

--learning_rate LEARNING_RATE Learning rate (default: 1e-06)

--train_size TRAIN_SIZE Training dataset size (default: 0.8)

--validation_size VALIDATION_SIZE Validation dataset size (default: 0.1)

--random_seed RANDOM_SEED random seed tp initialize (default: 15234)

--forces Train set has forces and use them (default: False)

--force_scalar FORCE_SCALAR the weight to apply to the forces part of the loss when training (default: 1.0)

--max_epochs MAX_EPOCHS maximum number of epochs (default: 1000)

--species_indicies SPECIES_INDICIES [SPECIES_INDICIES ...] Whether to read species by atomic number (ani 1x data) or label mapped to network index (ani 1 data) (default: ['periodic_table'])

--early_stopping_learning_rate EARLY_STOPPING_LEARNING_RATE Early stopping if this learning rate is met (default: 1e-09)

BNN prior options:

--params_prior PARAMS_PRIOR 'ani' or Json file with prior parameters stored under appropriate keys for setting each individually (default: None)

--prior_mu PRIOR_MU Prior mean (default: 0.0)

--prior_sigma PRIOR_SIGMA Prior standard deviation (default: 1.0)

--posterior_mu_init POSTERIOR_MU_INIT Posterior mean initialization (default: 0.0)

--posterior_rho_init POSTERIOR_RHO_INIT Posterior denisty initialization (default: -3.0)

--reparameterization Use the reparameterization version of the bayesian layers (default: False)

--flipout Use the flipout version of the bayesian layers (default: False)

--moped Use the moped priors (default: False)

--moped_delta MOPED_DELTA MOPED delta value (default: 0.2)

--moped_enable Initialize mu/sigma from the dnn weights (default: False)

--moped_init_model MOPED_INIT_MODEL DNN model to intialize MOPED method filename and path /path/to/dnn_v1_moped_init.pt (default: dnn_v1_moped_init.pt)

--bayesian_depth BAYESIAN_DEPTH Use bayesian layers to a depth of n (None means all) (default: None)

--lml Whether to replace the MSE loss with LML loss for the BNN (default: False)

Training options:

--use_schedulers Use lr schedulers for adam and sgd (default: False)

--train_mc_runs TRAIN_MC_RUNS Number of Monte Carlo runs during training of the Bayesian potential (default: 2)

--test_mc_runs TEST_MC_RUNS Number of Monte Carlo runs during testing of the Bayesian potential (default: 2)

--kl_weight KL_WEIGHT KL factor (default: 0.0)

--load_pretrained_parameters_dnn LOAD_PRETRAINED_PARAMETERS_DNN Path and file name to load pre-trained DNN parameters from (default: None)

--load_pretrained_parameters_bnn LOAD_PRETRAINED_PARAMETERS_BNN Path and file name to load pre-trained BNN parameters from (default: None)

--pretrained_model_key PRETRAINED_MODEL_KEY if there is a key the model weights and biases are stored give it here (default: None)

--update_bnn_variance UPDATE_BNN_VARIANCE If you load a pretrained BNN but want to update the varience of the parameters set a fraction between 0.0 and 1.0. (default: None)

--name NAME Name for the model sor reference later on (default: my)

--mutate_index MUTATE_INDEX Permute dataset swapping examples in training and validation set using this index as a randomseed. None will not permute the data set beyond the normal loading (default: None)

--no_mu_opt Whether to omit training the mean weights of the BNN (default: False)

--no_sig_opt Whether to omit training the sigma weights of the BNN (default: False)

Running options:

--train_ani_dnn Train a DNN (default: False)

--train_ani_bnn Train a BNN (default: False)

--ensemble ENSEMBLE Number of concurrent models to train as an ensemble prediction (default: 1)

--set_rho_explicit If params prior is being used will set rho values using the params prior (default: False)

--set_prior_explicit If params prior is being used will set prior values using the params prior (default: False)

--use_cuaev Use the cuda ave codes note these are installed separately from ani library (default: False)

--summary_plot_only in the training only output summary plots for each epoch (default: True)

--parallel Run using data parallel pytorch stratergy (default: False)

Logging arguments: --loglevel LOGLEVEL log level (default: INFO)