MT-PotNet

Open-source code for "A Material Property Prediction Method Using Multi-Task Learning Strategy to Tackle Data Scarcity"

Environment Setup

• We'll use conda to install dependencies and set up the environment. We recommend using the Python 3.9 Miniconda installer.
• After installing conda, run

conda env create -f environment.yml

• Then run below code to activate the environment

conda activate mtpotnet

• We are using torch-geometric for the implementation of our graph neural networks. Installation of torch-geometric via conda seems problematic so we install it explicitly using pip as

pip install torch-scatter==2.1.0 torch-sparse==0.6.16 torch-cluster==1.6.0 torch-spline-conv==1.2.1 -f https://data.pyg.org/whl/torch-1.12.1%2Bcu116.html

and

pip install torch-geometric

• Be aware that we are using an old version of JARVIS toolkits jarvis-tools==2022.9.16. The newest JARVIS toolkits will contain new versions of datasets that include more data than the one we present in the paper.

Running Summation Algorithm

• To run the summation algorithm, please run below commands in order to install the algorithm package (remember to replace the TARGET_PATH with your own destination.)

cd functions
tar xzvf gsl-latest.tar.gz
cd gsl-2.7.1
./configure --prefix=TARGET_PATH
make
make install

• Then edit ~/.bashrc by adding

export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:TARGET_PATH/lib/

and

source ~/.bashrc

or create another terminal.

• Now we back to functions directory and run

python setup.py build_ext --inplace

• Then the algorithm is installed as Cython package. A simple way to test if it is successfully installed is to run below in the root directory.

python test_algorithm.py

Train and Evaluate Models

In this code base, the datasets are directly provided by JARVIS toolkits and there is no need to download the JARVIS or Materials Project dataset from the official site. To change between different datasets and among different properties, go to potnet.yaml and set the corresponding entries such as

dataset: dft_3d_2021
# JARVIS dataset: dft_3d_2021

To train our model, use the script

python main.py --config configs/potnet.yaml --output_dir xxx --checkpoint xxx

Here, output_dir denotes the output directory of checkpoints and processed data files, and checkpoint denotes the path of a checkpoint meaning restarting training from a certain checkpoint. One can also omit checkpoint in this script. Note that after training, the code will conduct the evaluation for the last epoch. It is recommended to do an evaluation below by specifying a checkpoint from the best 5 saved in the checkpoint directory.

To evaluate our model, use the script

python main.py --config configs/potnet.yaml --output_dir xxx --checkpoint xxx --testing

and here checkpoint denotes the path of a checkpoint and testing denotes enabling the evaluation phase.

Train on Custom Dataset

We are supporting custom datasets in the same format as datasets in JARVIS Leaderboard [4]. Once the corresponding dataset is prepared, use the script

python main.py --config configs/potnet.yaml --output_dir xxx --checkpoint xxx --data_root xxx

and change the target name as target in potnet.yaml such like

dataset:dft_3d_2021
target: target

And here data_root denotes the path of the custom dataset, where dataset information dataset_info.json and targets id_prop.csv are included. Then our code will read the dataset information from dataset_info.json and id_prop.csv in the dataset directory and then read the data from data_root. Note that a dataset in the JARVIS Leaderboard can be generated by jarvis_populate_data.py. To generate the same dataset format based on jarvis_populate_data.py to accommodate our code, it is recommended to

• Generate the crystal structures that can be read by JARVIS toolkits and their corresponding properties to predict
• Predefine the train-val-test split of your dataset and write the crystal ids and targets in order in id_prop.csv

Pretrained Models

We provide preprocessed files and pretrained models in this google drive. Right now we only provide the checkpoint for formation energy per atom of the JARVIS dataset. Please stay tuned for more pretrained models! To use these files, specify --output_dir such like

python main.py --config configs/potnet.yaml --output_dir formation_energy_peratom --checkpoint xxx

then the processed file will be read automatically.

Acknowledgement

The underlying training part is based on PotNet.