Installation

Requirements

xnn requires Python 3.10 or later. The core package depends only on torch (>= 2.0), numpy and pyyaml; everything else is optional.

Installing from PyPI

The distribution on PyPI is called xnns; the import package and the command line are xnn:

pip install xnns                 # core
pip install "xnns[gnn,ase]"      # with e3nn (NequIP / MACE / Allegro) and ASE

Installing from source

Clone the repository and install with pip:

git clone https://github.com/molssi-ai/xnn.git
cd xnn
pip install -e .

Optional extras

The optional dependencies are organized as extras, so you install only what you need:

pip install -e .              # core (torch, numpy, pyyaml)
pip install -e ".[gnn]"      # + e3nn, for NequIP / MACE / Allegro
pip install -e ".[ase]"      # + ASE calculator support
pip install -e ".[hydra]"    # + Hydra / OmegaConf config frontend
pip install -e ".[dev]"      # + pytest, for running the test suite
pip install -e ".[examples]" # everything the example notebooks need
pip install -e ".[all]"      # all of the above

Extra	Adds	Needed for
gnn	e3nn >= 0.4.4	the E(3)-equivariant models: NequIP, MACE, Allegro
ase	ase >= 3.22	XNNCalculator and reading structure files with the xnn command line
hydra	hydra-core, omegaconf	the from_hydra() config frontend
dev	pytest	running pytest tests/
examples	ase, e3nn, mace-torch, nequip, matplotlib, jupyter	the validation notebooks in examples/gnn/, which benchmark xnn against the reference implementations. The Allegro reference (mir-group/allegro v0.3.0) is not on PyPI, so install it separately: pip install "git+https://github.com/mir-group/allegro@v0.3.0"

Note

The examples extra pins mace-torch==0.3.16 and nequip==0.6.2, which in turn pin e3nn==0.4.4. xnn itself runs fine on that pin (all tests pass), so the extras can coexist in one environment.

GPU installation with uv

pyproject.toml carries a uv configuration that reproduces the GPU environment the example notebooks were built in (torch 2.5.1+cu121 from the PyTorch cu121 wheel index, for CUDA 12.x drivers):

uv sync --extra all

Verifying the installation

Run the test suite (requires the dev extra, and gnn for the equivariant-model tests):

pytest tests/

or check quickly from Python:

import xnn
from xnn.common.models import available_models

print(xnn.__version__)
print(available_models())   # ['allegro', 'ani', 'bamboo', 'cace', 'hdnnp', 'mace', 'nequip', 'opls', 'physnet', 'reaxff', 'schnet']

The GNN models (NequIP, MACE, Allegro, etc.) only appear in the registry when e3nn is installed.