Prediction of high fidelity band structures using graph neural networks with node-level and graph-level features

Installation

The model uses the same environment as ALIGNN; otherwise, one can use conda env create -f environment.yml to install the environment. The code can be installed using pip install .

Training script

python scripts/train.py \
        --root_dir="datasets/mnc_datasets" \
        --sub_dir="example" \
        --dataset_name="mnxc_processed_020725.pkl" \
        --config_name="config_example.json"

root_dir: a directory of datasets.

sub_dir: a director of the train, val, and test files in csv, and config files. The csv files contain indices for train, val, and test, queried from the dataset.

dataset_name: a dataset name, stored in root_dir.

config_name: a config file to config the model, stored in sub_dir.

mnxc_processed_020725.pkl is provided in this repo. Installing git lfs, then using git lfs pull to download this file.

Prediction script

python scripts/predict.py \
        --root_dir="datasets/mnc_datasets" \
        --sub_dir="for_prediction" \
        --dataset_name="processed_file_for_prediction.pkl" \
        --prefix="prefix" \
        --config_name="config_example.json"

prefix: a prefix of an index file for querying data from the dataset. The file should be named as, for example, prefix_ids.csv.

Configuration

In a config file,

"output_dir": an output directory for storing output files.

"model": model config

• "atom_input_features": set 0 for not using default atomic features; otherwise, the number should be according the atom_features set above (92 for CGCNN).
• "extra_atom_features": atomic features, defined as a dictionary whose keys are feature names and values are feature dimensions. Set null if this feature is not used.
• "global_features": graph-level features, defined similar to "extra_atom_features".
• "embed_global_features": embedding graph-level features using MLP.