xnn#

Machine-Learning Interatomic Potentials in PyTorch#

xnn is a library of machine-learning interatomic potentials (MLIPs) for molecules and materials.

The x in xnn stands for the architecture family: graph neural networks (gnn), dense neural networks (dnn), convolutional neural networks on voxel grids (cnn), classical force field neural networks (ffnn), semiempirical quantum mechanics (sqmnn) and hybrid models (hybrid). Every model is implemented using the same PyTorch nn.Module interface. Thus, one data object, one trainer and one deployment path serve all model families.

Why xnn?#

xnn offers several benefits out-of-the-box:

  • Faithful implementations. Each model reproduces its reference code or paper to float32 and float64 round-off errors, verified block by block (Model Fidelity to Upstream Codes).

  • One pipeline. Shared data object, config schema, trainer and deployment to ASE, LAMMPS and the MolSSI Driver Interface.

  • Pre-trained and benchmark data in one line. from_pretrained() loads the MACE and AIMNet2 foundation models and your own checkpoints; load_dataset() downloads and converts standard datasets into ML-ready formats.

  • Fine-tuning, benchmarking and multi-GPU training from a config file and the xnn command.

  • Physics add-ons for any model. Latent Ewald long-range electrostatics and DFT-D3 / DFT-D4 dispersion wrap every potential.

  • Fast paths. Fused GPU kernels for the expensive blocks, switched on per run without changing the model.

Getting Started

Installing xnn and a first training run

To the Getting Started Guide

How-To Guides

Recipes for accomplishing common tasks

To the How-To Guides

Examples

Executed example notebooks, rendered with their outputs

To the Examples

User Guide

Reference information for using xnn

To the User Guide

Developer Guide

Extending xnn with new models and featurizers

To the Developer Guide

Background Information

The design of xnn and the models it implements

To the Background Information

API Reference

Documentation of the xnn Python API

To the API Reference

Equivariant GNNs with e3nn

A hands-on course for developing equivariant GNN interatomic potentials

To the Equivariant GNN Course

Models at a glance#

Family

Models

Graph neural networks (gnn)

MACE, NequIP, Allegro, CACE, AIMNet2, SchNet, DimeNet, DimeNet++, GemNet, PaiNN

Dense neural networks (dnn)

ANI (ANI-1, ANI-1x, ANI-1ccx, ANI-2x), PhysNet, HDNNP (generations 1 to 4)

Convolutional neural networks (cnn)

SE(3) steerable CNN, 3D CNN, spherical CNN

Classical force fields (ffnn)

ReaxFF / ReaxFF-nn, OPLS-AA / L-OPLS, DREIDING / X6

Semiempirical quantum chemistry (sqmnn)

PM3, PM6, PM6-ORG, PM7

Hybrid neural networks (hybrid)

BAMBOO, SpookyNet

Add-ons for any model (common)

LES long-range electrostatics, DFT-D3 and DFT-D4 dispersion

Each model matches its reference code or paper to round-off (Model Fidelity to Upstream Codes); the MACE and AIMNet2 foundation models load through the hub. Every model trains, evaluates and runs under ASE deployment interface. SchNet, DimeNet, GemNet, PaiNN, SpookyNet, NequIP, MACE, Allegro and AIMNet2 also export to TorchScript for LAMMPS, and any checkpoint can also serve as an MDI engine (Deployment).

xnn is developed by Mohammad Mostafanejad at The Molecular Sciences Software Institute (MolSSI). Visit the GitHub repository for the latest updates.