AtomNet: A Graph Neural Network Framework for Crystal Property Prediction Integrating Physical Priors

AtomNet is a crystal material property prediction framework based on physics-informed graph neural networks, developed by the JieCoa team and published in the npj Computational Materials journal. By introducing techniques like atomic electronegativity and polynomial feature engineering, it addresses the bottleneck of high computational cost in traditional DFT calculations, achieves high-precision predictions on multiple material datasets, and combines physical interpretability with engineering practicality. It has been open-sourced for researchers to use.

Crystal material property prediction is a core challenge in materials science. Traditional Density Functional Theory (DFT) calculations have high accuracy but are costly, making large-scale applications difficult. In recent years, Graph Neural Networks (GNNs) have shown great potential in this field, but effectively integrating physical prior knowledge remains a key challenge.

Framework Overview

AtomNet is an open-source framework whose core design idea is to maintain physical consistency in message passing and enhance the model's expressive power.

Physics-Informed Embedding

Introduce atomic electronegativity (Sanderson Electronegativity) as edge features, use periodic table properties to construct initial node features, and apply Radial Basis Function (RBF) kernels to process distance features.

Message Passing Architecture

Supports edge feature update control (e.g., the limitedUpdateEdge parameter improves band gap prediction performance), polynomial feature expansion, and neighbor sampling strategies.

Interpretability

Integrates the Captum library to visualize the contribution of atomic features to prediction results.

Datasets and Pre-trained Models

Supports datasets such as Jarvis DFT 3D 2021 and Materials Project 2018 (MEGNet), provides pre-trained models for five standard tasks, and data preprocessing code is in loader.py.

Experimental Configuration

Adjustable parameters include electronegativity_type (multiple RBF variants), envelope_type (weight function), usePolynomial (recommended to set to 3), etc. In the band gap prediction task, limitedUpdateEdge=3+4 layers of AtomNet Layer achieved the optimal MAE.

Training Tracking

Deeply integrates Weights & Biases (wandb) for experiment tracking, and offline training is also supported.

AtomNet provides researchers with an out-of-the-box tool:

• Accelerate material screening: Inference speed is several orders of magnitude faster than DFT
• Multi-task support: Covers properties like band gap, formation energy, and elastic modulus
• Interpretability: Helps understand the microscopic origin of material properties
• Easy to extend: Modular design facilitates integration of new physical constraints and feature methods

AtomNet is a successful practice of physics-informed machine learning in materials science, integrating chemical intuition with deep learning, and balancing accuracy and interpretability. It provides an academically rigorous and practical open-source tool for researchers in computational materials science and energy material development. The project code and pre-trained models have been open-sourced on GitHub and Zenodo, supporting Python 3.11+ and PyTorch 2.1+ environments.