A Materials Foundation Model via Hybrid Invariant-Equivariant Architectures

[preprint]

The official implementation for the HIENet model from our preprint "A Materials Foundation Model via Hybrid Invariant-Equivariant Architectures"

Installation

Requirements

• Python >= 3.8
• PyTorch >= 1.12.0

We have tested the model with the following configuration:

• Python 3.9
• PyTorch 2.1.2

After installing PyTorch, you can install HIENet from source:

git clone https://github.com/divelab/AIRS.git
cd AIRS/OpenMat/HIENet
pip install .

Usage

ASE Calculator

HIENet provides an interface with the Atomic Simulation Environment (ASE) through its calculator implementation. You can use the pretrained model as follows:

from hienet.hienet_calculator import HIENetCalculator
from ase.io import read

# Load a structure from a file
atoms = read('your_structure.cif')

# Initialize the calculator with default settings
# This will use the pretrained model on the available device (GPU if available)
calc = HIENetCalculator()

# Attach the calculator to the atoms object
atoms.calc = calc

# Calculate energy, forces, and stress
energy = atoms.get_potential_energy()
forces = atoms.get_forces()
stress = atoms.get_stress()

Using a specific checkpoint

You can also load a specific model checkpoint:

# Load from a specific checkpoint file
calc = HIENetCalculator(model='path/to/your/checkpoint.pth')

# Specify device
calc = HIENetCalculator(device='cuda:0')  # or 'cpu' for CPU only

We provide a pretrained HIENet checkpoint under the checkpoints/ folder.

Citation

Please cite our preprint if you use our model or benchmark results. Thank you!

@misc{yan2025materialsfoundationmodelhybrid,
      title={A Materials Foundation Model via Hybrid Invariant-Equivariant Architectures}, 
      author={Keqiang Yan and Montgomery Bohde and Andrii Kryvenko and Ziyu Xiang and Kaiji Zhao and Siya Zhu and Saagar Kolachina and Doğuhan Sarıtürk and Jianwen Xie and Raymundo Arroyave and Xiaoning Qian and Xiaofeng Qian and Shuiwang Ji},
      year={2025},
      eprint={2503.05771},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2503.05771}, 
}

Acknowledgement

This project builds upon SevenNet as starter code, which is licensed under the GNU General Public License v3.0. In addition to adding our own files, we made significant changes to the following files to implement HIENet architecture within the SevenNet's training pipeline:

• hienet/model_build.py
• hienet/nn/linear.py
• hienet/nn/edge_embedding.py
• hienet/train/loss.py
• hienet/train/optim.py

We have also adjusted function and file names and made other slight changes in other files for consistency

K.Y., M.B., A.K. and S.J. acknowledge partial support from National Science Foundation (NSF) under grant IIS-2243850, ARPA-H under grant 1AY1AX000053, and National Institutes of Health under grant U01AG070112. K.Z., S.K. and X.F.Q. acknowledge partial support from NSF under awards CMMI-2226908 and DMR-2103842 and the donors of ACS Petroleum Research Fund under Grant 65502-ND10. R.A. and D.S. acknowledge support from ARO through Grant No. W911NF-22-2-0117. R.A., X.N.Q. and S.Z. acknowledge NSF through Grant No. 2119103 (DMREF). Z. X. and X.N.Q. acknowledge support from NSF through grants SHF-2215573 and IIS-2212419. Some computations were carried out at the Texas A&M High-Performance Research Computing (HPRC) facility. We gratefully acknowledge the support of Lambda Inc. for providing computing resources.

Contact

If you have any question, please contact keqiangyan@tamu.edu.