Table of Contents
- Introduction
- Installation
- Examples
- Pre-trained models
- Quick start using colab
- JARVIS-ALIGNN webapp
- Peformances on a few datasets
- Useful notes
- References
- How to contribute
- Correspondence
- Funding support
ALIGNN (Introduction)
The Atomistic Line Graph Neural Network (https://www.nature.com/articles/s41524-021-00650-1) introduces a new graph convolution layer that explicitly models both two and three body interactions in atomistic systems.
This is achieved by composing two edge-gated graph convolution layers, the first applied to the atomistic line graph L(g) (representing triplet interactions) and the second applied to the atomistic bond graph g (representing pair interactions).
The atomistic graph g consists of a node for each atom i (with atom/node representations hi), and one edge for each atom pair within a cutoff radius (with bond/pair representations eij).
The atomistic line graph L(g) represents relationships between atom triplets: it has nodes corresponding to bonds (sharing representations eij with those in g) and edges corresponding to bond angles (with angle/triplet representations tijk).
The line graph convolution updates the triplet representations and the pair representations; the direct graph convolution further updates the pair representations and the atom representations.
Installation
First create a conda environment: Install miniconda environment from https://conda.io/miniconda.html Based on your system requirements, you'll get a file something like 'Miniconda3-latest-XYZ'.
Now,
bash Miniconda3-latest-Linux-x86_64.sh (for linux)
bash Miniconda3-latest-MacOSX-x86_64.sh (for Mac)
Download 32/64 bit python 3.8 miniconda exe and install (for windows) Now, let's make a conda environment, say "version", choose other name as you like::
conda create --name version python=3.8
source activate version
Method 1 (using setup.py):
Now, let's install the package:
git clone https://github.com/usnistgov/alignn.git
cd alignn
python setup.py develop
For using GPUs/CUDA, install dgl-cu101 or dgl-cu111 based on the CUDA version available on your system, e.g.
pip install dgl-cu111
Method 2 (using pypi):
As an alternate method, ALIGNN can also be installed using pip
command as follows:
pip install alignn dgl-cu111
Examples
Dataset
Users can keep their structure files in POSCAR
, .cif
, .xyz
or .pdb
files in a directory. In the examples below we will use POSCAR format files. In the same directory, there should be an id_prop.csv
file.
In this directory, id_prop.csv
, the filenames, and correponding target values are kept in comma separated values (csv) format
.
Here is an example of training OptB88vdw bandgaps of 50 materials from JARVIS-DFT database. The example is created using the generate_sample_data_reg.py script. Users can modify the script for more than 50 data, or make their own dataset in this format. For list of available datasets see Databases.
The dataset in split in 80:10:10 as training-validation-test set (controlled by train_ratio, val_ratio, test_ratio
) . To change the split proportion and other parameters, change the config_example.json
file. If, users want to train on certain sets and val/test on another dataset, set n_train
, n_val
, n_test
manually in the config_example.json
and also set keep_data_order
as True there so that random shuffle is disabled.
A brief help guide (-h
) can be obtained as follows.
!train_folder.py -h
Regression example
Now, the model is trained as follows. Please increase the batch_size
parameter to something like 32 or 64 in config_example.json
for general trainings.
!train_folder.py --root_dir "alignn/examples/sample_data" --config "alignn/examples/sample_data/config_example.json" --output_dir=temp
Classification example
While the above example is for regression, the follwoing example shows a classification task for metal/non-metal based on the above bandgap values. We transform the dataset into 1 or 0 based on a threshold of 0.01 eV (controlled by the parameter, classification_threshold
) and train a similar classification model. Currently, the script allows binary classification tasks only.
!train_folder.py --root_dir "alignn/examples/sample_data" --classification_threshold 0.01 --config "alignn/examples/sample_data/config_example.json" --output_dir=temp
Multi-output model example
While the above example regression was for single-output values, we can train multi-output regression models as well. An example is given below for training formation energy per atom, bandgap and total energy per atom simulataneously. The script to generate the example data is provided in the script folder of the sample_data_multi_prop. Another example of training electron and phonon density of states is provided also.
!train_folder.py --root_dir "alignn/examples/sample_data_multi_prop" --config "alignn/examples/sample_data/config_example.json" --output_dir=temp
Automated model training
Users can try training using multiple example scripts to run multiple dataset (such as JARVIS-DFT, Materials project, QM9_JCTC etc.). Look into the alignn/scripts/train_*.py folder. This is done primarily to make the trainings more automated rather than making folder/ csv files etc. These scripts automatically download datasets from Databases in jarvis-tools and train several models. Make sure you specify your specific queuing system details in the scripts.
Using pre-trained models
All the trained models are distributed on figshare and this pretrained.py script can be applied to use them. These models can be used to directly make predictions.
A brief help section (-h
) is shown using:
!pretrained.py -h
An example of prediction formation energy per atom using JARVIS-DFT dataset trained model is shown below:
!pretrained.py --model_name jv_formation_energy_peratom_alignn --file_format poscar --file_path alignn/examples/sample_data/POSCAR-JVASP-10.vasp
Quick start using GoogleColab notebook example
The following notebook provides an example of 1) installing ALIGNN model, 2) training the example data and 3) using the pretrained models. For this example, you don't need to install alignn package on your local computer/cluster, it requires a gmail account to login. Learn more about Google colab here.
Web-app
A basic web-app is for direct-prediction available at JARVIS-ALIGNN app. Given atomistic structure in POSCAR format it predict formation energy, total energy per atom and bandgap using data trained on JARVIS-DFT dataset.
Performances
1) On QM9 dataset
2) On Materials project 2018 dataset
3) On JARVIS-DFT 2021 dataset (classification)
4) On JARVIS-DFT 2021 dataset (regression)
5) On hMOF dataset
6) On qMOF dataset
MAE on electronic bandgap 0.20 eV
7) On OMDB dataset
coming soon!
8) On HOPV dataset
coming soon!
9) On QETB dataset
coming soon!
10) On OpenCatalyst dataset
coming soon!
Useful notes (based on some of the queries we received)
- If you are using GPUs, make sure you have a compatible dgl-cuda version installed, for example: dgl-cu101 or dgl-cu111, so e.g.
pip install dgl-cu111
. - The undirected graph and its line graph is constructured in
jarvis-tools
package using jarvis.core.graphs - While comnventional '.cif' and '.pdb' files can be read using jarvis-tools, for complex files you might have to install
cif2cell
andpytraj
respectively i.e.pip install cif2cell==2.0.0a3
andconda install -c ambermd pytraj
. - Make sure you use
batch_size
as 32 or 64 for large datasets, and not 2 as given in the example config file, else it will take much longer to train, and performnce might drop a lot. - Note that
train_folder.py
andpretrained.py
in alignn folder are actually python executable scripts. So, even if you don't provide absolute path of these scripts, they should work.
References
Please see detailed publications list here.
How to contribute
For detailed instructions, please see Contribution instructions
Correspondence
Please report bugs as Github issues (https://github.com/usnistgov/alignn/issues) or email to [email protected].
Funding support
NIST-MGI (https://www.nist.gov/mgi).
Code of conduct
Please see Code of conduct <https://github.com/usnistgov/jarvis/blob/master/CODE_OF_CONDUCT.md>
__