Moment Graph Neural Network (MGNN) for molecular and materials potential learning.

MGNN is a rotation-invariant message-passing neural network that uses moment-based representations of 3D molecular graphs. This repository contains task-specific training and evaluation code for the systems reported in the MGNN paper, including QM9, MD17, 3BPA, 25-element high-entropy alloys, Li3PO4, and molecular spectra.

model/
  mgnn.py                 Shared MGNN implementation.
mgnn_QM9/
  main.py                 QM9 training and evaluation entry point.
  qm9.py                  QM9 dataset loader.
mgnn_MD17/
  main.py                 MD17 training entry point.
  test.py                 MD17 checkpoint evaluation.
  md17.py                 MD17 dataset loader.
mgnn_3BPA/
  main.py                 3BPA training and evaluation entry point.
  dataset_3bpa.py         3BPA dataset loader.
mgnn_25HEA/
  main.py                 25-element high-entropy alloy training entry point.
  HEA25.py                HEA25 dataset loader.
  mgnn.py                 HEA25 MGNN variant with PBC support.
mgnn_Li3PO4/
  create_dataset.py       Li3PO4 dataset conversion.
  main.py                 Li3PO4 train/test workflow.
mgnn_Spectra/
  train.py, test.py       Spectra model training and evaluation.
  ethanol_md*.py          Molecular dynamics/spectra utilities.

Create a Python environment with CUDA-enabled PyTorch if GPU training is needed. The exact versions depend on the target platform, but the code expects:

• torch
• torch_geometric
• torch_scatter
• torch_sparse
• torchmetrics
• numpy
• ase
• matscipy
• schnetpack
• pytorch_lightning
• tensorboard
• fairchem for the 25HEA periodic graph workflow

Install PyTorch Geometric packages using wheels that match your PyTorch and CUDA versions.

Run commands from the corresponding task directory so local imports and data paths resolve correctly.

cd mgnn_QM9
python main.py --debug --num_epochs 1

cd ../mgnn_MD17
python main.py --target aspirin --debug --num_epochs 1

cd ../mgnn_3BPA
python main.py --num_epochs 1

cd ../mgnn_25HEA
python main.py --debug --num_epochs 1

cd ../mgnn_Li3PO4
python create_dataset.py
python main.py --mode train

cd ../mgnn_Spectra
python train.py --datadir /path/to/data --workdir data

Use --ddp with torchrun for distributed PyTorch workflows where supported by the task script.

Some datasets are included in task folders, while others are downloaded or expected at user-provided paths. Check each task's dataset loader before running long jobs:

• mgnn_3BPA/dataset_3BPA/ contains 3BPA raw and processed files.
• mgnn_Li3PO4/data/ contains Li3PO4 data files.
• mgnn_25HEA/data/ is used by the HEA25 loader.
• QM9 and MD17 loaders can download or process data under the configured root.

Generated checkpoints, logs, processed caches, and TensorBoard runs should not be committed unless they are intentional small examples.

The model is consistently referred to as MGNN in this repository. Older names such as MTP and HMN refer to the same model family and have been unified to avoid ambiguity.

Before committing changes, run:

python -m compileall model mgnn_QM9 mgnn_MD17 mgnn_3BPA mgnn_25HEA mgnn_Li3PO4 mgnn_Spectra

For behavioral changes, also run the shortest relevant task command, such as --debug --num_epochs 1, and any task-specific test.py script.

If you use this code, please cite:

Chang, J., Zhu, S. MGNN: Moment Graph Neural Network for Universal Molecular Potentials. npj Computational Materials 11, 55 (2025). https://doi.org/10.1038/s41524-025-01541-5

@article{chang2025mgnn,
  title = {MGNN: Moment Graph Neural Network for Universal Molecular Potentials},
  author = {Chang, Jian and Zhu, Shuze},
  journal = {npj Computational Materials},
  volume = {11},
  pages = {55},
  year = {2025},
  doi = {10.1038/s41524-025-01541-5},
  url = {https://doi.org/10.1038/s41524-025-01541-5}
}