myung-group/BAM_MPtrj_v1

BAM_MPtrj

BAM (Bayesian Atoms Modeling) pretrained on the Materials Project Trajectory (MPtrj) dataset.

This model is a Bayesian E(3) Equivariant Machine Learning Potential based on the RACE (Restratification of Atoms with Combined Encoding) architecture. It provides uncertainty-aware energy and force predictions for atomistic simulations of inorganic materials.

Model Description

BAM_MPtrj_v1 is trained on the MPtrj dataset, which contains DFT-calculated energies and forces from the Materials Project. The model uses E(3)-equivariant message passing with iterative restratification to achieve ab initio-level accuracy while providing robust uncertainty quantification.

Key Features

• E(3) Equivariance: Maintains rotational and translational symmetry for physically consistent predictions
• RACE Architecture: Iterative Restratification of Atoms with Combined Encoding for improved message passing
• Joint Energy-Force NLL Loss: Novel loss function explicitly modeling uncertainty in both energies and interatomic forces
• Uncertainty Quantification: Comprehensive uncertainty estimation for active learning, calibration, and out-of-distribution detection
• Scalable: Designed for large-scale atomistic simulations with GPU acceleration

Intended Uses

• Energy and force prediction for inorganic materials
• Molecular dynamics simulations
• Uncertainty-aware atomistic simulations
• Active learning for efficient data acquisition
• Out-of-distribution detection in materials discovery

Training Details

Training Data

The model is trained on the Materials Project Trajectory (MPtrj) dataset, which includes DFT-calculated energies, forces, and stresses from relaxation trajectories across diverse inorganic materials.

Architecture

• Model type: RACE
• Equivariance: E(3) equivariant via e3nn
• Framework: PyTorch + PyTorch Geometric

How to Use

Installation

git clone https://github.com/myung-group/BAM-torch
cd BAM-torch
pip install "torch<=2.8"
python install_deps.py
pip install -e .

Inference

import json
import torch
from bam_torch.predicting.evaluator import Evaluator
from bam_torch.utils import find_input_json

input_json_path = find_input_json()
with open(input_json_path) as f:
    json_data = json.load(f)

evaluator = Evaluator(json_data)
evaluator.evaluate()

Citation

If you use this model in your research, please cite:

@article{bam2025,
  title={Bayesian E(3)-Equivariant Interatomic Potential with Iterative Restratification of Many-body Message Passing},
  author={Soohaeng Yoo Willow, Tae Hyeon Park, Gi Beom Sim, Sung Wook Moon, Seung Kyu Min, D. ChangMo Yang, Hyun Woo Kim, Juho Lee, Chang Woo Myung},
  journal={arXiv:2510.03046},
  year={2025}
}

More Information

• GitHub Repository: https://github.com/myung-group/BAM-torch
• Paper: arXiv:2510.03046

License

This model is released under the MIT License.