Skip to content

Graph neural network prediction of electronic Hamiltonians in atomic orbital representation with many body messages

License

Notifications You must be signed in to change notification settings

maurergroup/MACE-H

Repository files navigation

MACE-H

This code is released as the MACE-H model in the paper Equivariant Electronic Hamiltonian Prediction with Many-Body Message Passing (arXiv:2508.15108).

This package is developed based upon DeepH-E3, and share the similar workflow and mostly the same data management pipeline. If you have used DeepH-E3 before, MACE-H will be very easy to implement.

The paper includes the 2D material datasets from previous DeepH-E3 work and our generated bulk Au dataset. For the 2D material datasets, which are generated by OpenMx package and can be found in the repositries released by DeepH-E3 group, i.e., Dataset 1, Dataset 2 and Dataset 3. For our bulk Au dataset, it is calculated by FHI-aims with proper basis sets transformation and cutoff radius truncation.

The model config files and other necessary file to reproduce the data in teh paper can be found in Zenodo.

Installation

To use MACE-H, the python environment is required to implement the model and the Julia environment is used to diagonalized the reciprocal space Hamiltonian matrices and calculate the electronic properties.

Python

The python environment can be enabled by either installation using provided python environment yaml file or loading precompiled container. The python environment yaml file is provide in Zenodo. To facilitate our users, we compiled two containers, one for Nvidia machine and the otehr for Rocm machine, which are provided in Zenodo

Julia

The julia environmet is used to diagonalize the matrices in the reciprocal space and study the electronic strucure in the postprocessing stage. The following packages are requred:

  • Distributed
  • LinearAlgebra

To ensure the Julia environment is loaded correctly, please specify the path of Julia environment in Band.py using "os.environ["JULIA_PROJECT"] = "/path/to/project"".

Usage

The simple demons of invoking the package can be found in the ./sh directory with example config fie in Zenodo. The overall usage is very simialr to DeepH-E3.

Preprocess

The generation of the DFT dataset depends on which DFT pakcages you are using. For DFT packages with pseodo-potential methods that only consider the valence orbials and comparatively smaller cutoff radius, such as OpenMX or Abacus, you can directly using the data pipeline of DeepH-E3. However, in terms of full-electronic packages that adopts larger cutoff-radius, especially when elements of higher periodic numbers are involved, we should use the provied basis set transformation to exclude the excessive core orbitals and sparcify the matrix according to the distance.

We can either preprocess the processed dataset to graph data using the deephe3-preprocess.py code and preprocess.ini config file in DeepH-E3 or directly specify the arguemets in the data section in train.ini, which are also the same as DeepH-E3. Note that when no config file where fed to the python implementation code, the default config file in ./maceh/default_configs will be automatically loaded. The argument not specified in the fed config file will employ which in the defalt one.

Model training

After specified corresponding config file, you can train the model using

${python_path} ./deephe3-train.py ./configs/train.ini | tee -a sh/log_train.txt

as the slurm bash file in the ./sh directory.

Model inference

Given the trained model, you can get the predicted Hamiltonian matrices (i.e., the "Hamiltonians_pred.h5" file) by

/path/to/deephe3-eval.py eval.ini

The relationship between the fed config file and the default one is same as the train.ini yaml file.

The model inference tool in ./inference_tools directory are used to calulate the band structure and density of states (DOS), and when the DFT reference data (i.e., the "hamiltonians.h5" file) is provided, to calculate the metrics of eigenvalue error and electronic entropy error. To call the code, please use

${python_path} $MACEH_DIR/inference_tools/Band.py -args

The more detailed specification can be found in the example bash file in Zenodo. Here, we also suggested that we use another json file called SI.json (check ./inference_tools/SI.json as an example) specifying the number of valence orbitals and k-path for band structure, since the default number of orbitals in this regards are not always correct without knowing the exact potential files of the elements.

Citation

@article{qian2025equivariant,
  title={Equivariant Electronic Hamiltonian Prediction with Many-Body Message Passing},
  author={Qian, Chen and Vitartas, Valdas and Kermode, James and Maurer, Reinhard J},
  journal={arXiv preprint arXiv:2508.15108},
  year={2025}
}

Since our packages are developed as extension of DeepH-E3, so we also suggest to cite the DeepH-E3 paper.

@article{gong2023general,
  title={General framework for E (3)-equivariant neural network representation of density functional theory Hamiltonian},
  author={Gong, Xiaoxun and Li, He and Zou, Nianlong and Xu, Runzhang and Duan, Wenhui and Xu, Yong},
  journal={Nature Communications},
  volume={14},
  number={1},
  pages={2848},
  year={2023},
  publisher={Nature Publishing Group UK London}
}

About

Graph neural network prediction of electronic Hamiltonians in atomic orbital representation with many body messages

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published