Metadata-Version: 2.4
Name: xnns
Version: 0.13.0
Summary: Machine-learning interatomic potentials in PyTorch for molecular and periodic systems
Author-email: Mohammad Mostafanejad <smostafanejad@molssi.org>
Maintainer-email: Mohammad Mostafanejad <sina.mostafanejad@gmail.com>, MolSSI AI <ai@molssi.org>
License-Expression: MIT
Project-URL: Homepage, https://github.com/molssi-ai/xnn
Project-URL: Documentation, https://molssi-ai.github.io/xnn/
Project-URL: Source, https://github.com/molssi-ai/xnn
Project-URL: Issues, https://github.com/molssi-ai/xnn/issues
Project-URL: Changelog, https://github.com/molssi-ai/xnn/releases
Keywords: interatomic potentials,machine learning,molecular dynamics,MACE,NequIP,Allegro,CACE,AIMNet2,SchNet,PaiNN,3D steerable CNN,PhysNet,ANI,ReaxFF,OPLS,pytorch,LAMMPS,ASE
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Chemistry
Classifier: Topic :: Scientific/Engineering :: Physics
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0
Requires-Dist: numpy
Requires-Dist: pyyaml
Requires-Dist: tqdm
Provides-Extra: ase
Requires-Dist: ase>=3.22; extra == "ase"
Provides-Extra: hub
Requires-Dist: h5py>=3.0; extra == "hub"
Provides-Extra: gnn
Requires-Dist: e3nn>=0.4.4; extra == "gnn"
Provides-Extra: mdi
Requires-Dist: pymdi>=1.4; extra == "mdi"
Provides-Extra: vesin
Requires-Dist: vesin-torch>=0.3; extra == "vesin"
Provides-Extra: hydra
Requires-Dist: hydra-core>=1.3; extra == "hydra"
Requires-Dist: omegaconf>=2.3; extra == "hydra"
Provides-Extra: tensorboard
Requires-Dist: tensorboard>=2.10; extra == "tensorboard"
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"
Provides-Extra: docs
Requires-Dist: sphinx>=7; extra == "docs"
Requires-Dist: pydata-sphinx-theme>=0.15; extra == "docs"
Requires-Dist: sphinx-design; extra == "docs"
Requires-Dist: sphinx-copybutton; extra == "docs"
Requires-Dist: sphinx-togglebutton; extra == "docs"
Requires-Dist: myst-nb>=1.1; extra == "docs"
Provides-Extra: examples
Requires-Dist: ase>=3.22; extra == "examples"
Requires-Dist: e3nn>=0.4.4; extra == "examples"
Requires-Dist: h5py>=3.0; extra == "examples"
Requires-Dist: mace-torch==0.3.16; extra == "examples"
Requires-Dist: nequip==0.6.2; extra == "examples"
Requires-Dist: matplotlib>=3.8; extra == "examples"
Requires-Dist: jupyter; extra == "examples"
Requires-Dist: ipykernel; extra == "examples"
Requires-Dist: nbconvert; extra == "examples"
Requires-Dist: nbformat; extra == "examples"
Requires-Dist: openmm>=8.0; extra == "examples"
Requires-Dist: rdkit>=2023.9; extra == "examples"
Provides-Extra: ani
Requires-Dist: torchani==2.9.0; extra == "ani"
Requires-Dist: h5py>=3.0; extra == "ani"
Provides-Extra: d4
Requires-Dist: dftd4>=3.7; extra == "d4"
Requires-Dist: tomli; extra == "d4"
Provides-Extra: d3
Requires-Dist: dftd3>=1.0; extra == "d3"
Provides-Extra: ffnn
Requires-Dist: rdkit>=2023.9; extra == "ffnn"
Provides-Extra: opls
Requires-Dist: openmm>=8.0; extra == "opls"
Requires-Dist: rdkit>=2023.9; extra == "opls"
Provides-Extra: all
Requires-Dist: ase>=3.22; extra == "all"
Requires-Dist: e3nn>=0.4.4; extra == "all"
Requires-Dist: vesin-torch>=0.3; extra == "all"
Requires-Dist: hydra-core>=1.3; extra == "all"
Requires-Dist: omegaconf>=2.3; extra == "all"
Requires-Dist: tensorboard>=2.10; extra == "all"
Requires-Dist: pytest; extra == "all"
Requires-Dist: mace-torch==0.3.16; extra == "all"
Requires-Dist: nequip==0.6.2; extra == "all"
Requires-Dist: matplotlib>=3.8; extra == "all"
Requires-Dist: jupyter; extra == "all"
Requires-Dist: ipykernel; extra == "all"
Requires-Dist: nbconvert; extra == "all"
Requires-Dist: nbformat; extra == "all"
Dynamic: license-file

# xnn

[![Tests](https://github.com/molssi-ai/xnn/actions/workflows/tests.yml/badge.svg)](https://github.com/molssi-ai/xnn/actions/workflows/tests.yml)
[![codecov](https://codecov.io/gh/molssi-ai/xnn/graph/badge.svg)](https://codecov.io/gh/molssi-ai/xnn)

Machine-learning interatomic potentials for molecules and materials behind one
PyTorch interface. 

**x** in **xnn** refers to the neural network model architectures: **g** stands
for graph, **d** for dense, **c** for convolutional, **ff** for classical
force-fields, **sqm** for semi-empirical quantum mechanics, **hybrid** for mixed models, all sharing one data object, one
trainer and one deployment path. This categorization is reflected in the package
structure.

**Documentation:** https://molssi-ai.github.io/xnn/

## Installation

```
pip install xnns                  # core (torch, numpy, pyyaml)
pip install "xnns[gnn,ase]"       # + e3nn for NequIP/MACE/Allegro, + ASE calculator
pip install -e ".[all]"           # from a clone: everything, for development
```

The distribution is `xnns`; the import name and the command line are `xnn`.

## Design principles

1. **One data object.** Every model takes an `AtomicGraph` and returns
   `{"node_energy", "energy"}`. Periodicity lives only in the edge vectors, so
   molecules and crystals look the same to a model.
2. **Featurizers are first-class.** Descriptors (symmetry functions, AEV) and
   equivariant edge features (spherical harmonics, Cartesian monomials) are
   standalone modules that models compose and you can inspect on their own.
3. **Forces and stress in one place.** `ForceStressOutput` differentiates any
   model's energy for forces and stress; models never implement them. Wrappers
   in the same style add long-range electrostatics (LES) and DFT-D3 / DFT-D4
   dispersion to any model.
4. **One config, three frontends.** `@register_model` exposes a model to the
   YAML, argparse and Hydra interfaces, all of which fill one `Config`
   dataclass.

## Load data

Structures are plain dictionaries (`pos`, `atomic_numbers`, optional `cell`,
`pbc` and `energy` / `forces` / `stress` targets), any ASE-readable file, or a
dataset from the hub:

```python
from xnn.common.data import AtomicDataset, load_dataset

# From a list of structures dict
data = AtomicDataset(structures, cutoff=5.0)

# From a file (e.g., .extxyz, .cif, VASP, ...)
data = AtomicDataset.from_file("trajectory.extxyz", cutoff=5.0)

# From the xnn data hub (e.g., RMD17 dataset)
data = load_dataset("rmd17", molecule="aspirin", split="train", cutoff=5.0)
```

## Train

```python
from xnn.common.config import Config
from xnn.common.train import Trainer

# Initialize the configuration
cfg = Config()

# Set the model and data configuration
cfg.model.name = "nequip"
cfg.model.extra = {"species": [1, 6, 8], "l_max": 2}
cfg.data.batch_size = 16

# Train the model
Trainer(cfg, data).fit()
```

or from the command line: `xnn train --config configs/train.yaml --set optim.epochs=50`.

## Pretrained models

```python
from xnn.common.models import from_pretrained, list_models

# Foundation models, xnn models, local cache
list_models()

# MACE-MP / MACE-OFF, converted once and cached
model = from_pretrained("mace-off23-small")

# Your own checkpoints
model = from_pretrained("runs/exp/best.pt")
```

## Fine-tune

Any pretrained model continues training through the same `Trainer`: naive
fine-tuning, LoRA, multi-head replay (with original or pseudolabelled replay
data), layer freezing and model-aware reference-energy reestimation, all from
the config:

```yaml
model:
  pretrained: mace-off23-small
  cutoff: 5.0
  lora: {rank: 16}                      # or heads: [pt_head, Default] with data.replay_path
  atomic_energies: estimated
```

## Benchmark

Score pre-trained checkpoints on one dataset (energy / force / stress MAE, MSE
or RMSE, per atom or as atomization energies), tabulated per model and written
to CSV, JSON or Markdown:

```bash
xnn benchmark --config configs/benchmark.yaml
```

## Deploy

```python
from xnn.common.deploy import XNNCalculator, export_to_lammps

# Export as ASE calculator
atoms.calc = XNNCalculator(model, cutoff=5.0)

# Export to LAMMPS (TorchScript)
export_to_lammps(model, cutoff=5.0, path="deployed.pt")
```

`xnn mdi --ckpt runs/exp/best.pt` serves a model to an external MD code through
the MolSSI Driver Interface.

## Models

| Family | Models |
|---|---|
| Generic add-ons (`common`) | LES long-range electrostatics (optionally charge-constrained or charge-equilibrated); DFT-D3 (zero, BJ, mzero, op damping) and DFT-D4 dispersion |
| Graph neural networks (`gnn`) | SchNet, DimeNet, DimeNet++, GemNet (-Q, -T, -dQ, -dT), PaiNN, NequIP, MACE, Allegro, CACE, AIMNet2 |
| Dense neural networks (`dnn`) | ANI (ANI-1, ANI-1x, ANI-1ccx, ANI-2x), PhysNet, HDNNP (generations 1 to 4) |
| Convolutional neural networks (`cnn`) | 3D steerable-/ conventional- / spherical-CNN |
| Force field neural networks (`ffnn`) | ReaxFF / ReaxFF-nn, OPLS-AA / L-OPLS, DREIDING / X6 |
| Hybrid neural networks (`hybrid`) | BAMBOO, SpookyNet |
| Semiempirical quantum mechanics neural networks (`sqmnn`) | PM3, PM6, PM6-ORG, PM7 (classics or trainable) |

Each model is checked against its reference code or paper
(`examples/fidelity_checks/`). SchNet, DimeNet, GemNet, PaiNN, SpookyNet, NequIP, MACE, Allegro and AIMNet2
also export to TorchScript / LAMMPS; every model deploys through ASE and the MDI
engine.

## Examples and tests

Executed notebooks for every model are under `examples/`
(`pip install -e ".[examples]"`); `examples/quickstart.py` is a minimal script.
Run the tests with `pytest tests/`.

## Contributing

Bug reports, fixes and new models are welcome. Read [CONTRIBUTING.md](https://github.com/molssi-ai/xnn/blob/main/CONTRIBUTING.md)
first: it explains how to report a bug with a reproducible example and what a pull request
needs to be merged.

## License

MIT

## How to cite

If you use xnn in your work, please cite it as:

```bibtex
@software{Mostafanejad:2026:xnn,
  author  = {Mostafanejad, Mohammad},
  title   = {xnn: Machine-learning interatomic potentials for molecules and materials},
  year    = {2026},
  version = {0.11.0},
  url     = {https://github.com/molssi-ai/xnn}
}
```

**Author:** Mohammad Mostafanejad, The Molecular Sciences Software Institute
(MolSSI)
