Predicting Density of States via Multi-Modal Transformer

The offical source code for [Predicting Density of States via Multi-Modal Transformer](https://arxiv.org/abs/2303.07000) paper, accepted at ICLR 2023 Workshop on Machine Learning for Materials (ML4Materials)

Overview

The density of states (DOS) is a spectral property of materials, which provides fundamental insights on various characteristics of materials. In this paper, we propose to predict the density of states (DOS) by reflecting the nature of DOS: DOS determines the general distribution of states as a function of energy. Specifically, we integrate the heterogeneous information obtained from the crystal structure and the energies via multi-modal transformer, thereby modeling the complex relation- ships between the atoms in the crystal structure, and various energy levels. Exten- sive experiments on two types of DOS, i.e., phonon DOS and electron DOS, with various real-world scenarios demonstrate the superiority of DOSTransformer.

Phonon DOS Prediction

Dataset

You can dowload phonon dataset in this repository

Run model

Run main_phDOS.py for phonon DOS Prediction after downloading phonon dataset into data/processed

Electron DOS Prediction

Dataset

We build Electron DOS dataset consists of the materials and its electron DOS information which are collected from Materials Proejct
We converted raw files to pkl and made electronic DOS dataset by mat2graph.py

Run model

Run main_eDOS.py for electron DOS Prediction after building electron dataset.

Models

embedder eDOS

DOSTransformer.py: Our proposed model / graphnetwork.py: GraphNetwork using Energy Embedding
graphnetwork2.py: GraphNetwork not using Energy Embedding / mlp.py: Mlp using Energy Embedding
mlp2.py: Mlp not using Energy Embedding

embedder phDOS

DOSTransformer_phonon.py: Our proposed model / graphnetwork_phonon.py: GraphNetwork using Energy Embedding
graphnetwork2_phonon.py: GraphNetwork not using Energy Embedding / mlp_phonon.py: Mlp using Energy Embedding
mlp2_phonon.py: Mlp not using Energy Embedding

Hyperparameters

--layers: Number of GNN layers in DOSTransformer model

--transformer: Number pf Transformer layer in DOSTransformer

--embedder: Selecting embedder

--hidden: Size of hidden dim

--epochs: Number of epochs for training the model

--lr: Learning rate for training the model

--dataset: Selecting dataset for eDOS prediction (Random split, Crystal OOD, Element OOD, default dataset is Random split)

--es: Early Stopping Criteria

--eval: Evaluation Step