xnns 0.1.0__tar.gz

xnns-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 MolSSI AI
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

xnns-0.1.0/PKG-INFO

@@ -0,0 +1,422 @@
+Metadata-Version: 2.4
+Name: xnns
+Version: 0.1.0
+Summary: Machine-learning interatomic potentials in PyTorch for molecular and periodic systems
+Author-email: Mohammad Mostafanejad <smostafanejad@molssi.org>
+Maintainer-email: MolSSI AI <smostafanejad@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,SchNet,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: 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.2.4; extra == "ani"
+Requires-Dist: h5py>=3.0; extra == "ani"
+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: 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
+
+Machine-learning interatomic potentials for **molecules and materials** behind a
+single coherent PyTorch interface.
+
+## Installation
+
+From PyPI, where the distribution is called `xnns` (the import name and the
+command line are `xnn`):
+
+```
+pip install xnns                 # core (torch, numpy, pyyaml)
+pip install "xnns[gnn,ase]"      # + e3nn for NequIP/MACE/Allegro, + ASE calculator
+```
+
+From a clone, for development:
+
+```
+pip install -e .             # core (torch, numpy, pyyaml)
+pip install -e ".[ase]"      # + ASE calculator
+pip install -e ".[gnn]"      # + e3nn for NequIP/MACE/Allegro
+pip install -e ".[hydra]"    # + Hydra/OmegaConf config
+pip install -e ".[examples]" # + ASE, e3nn, mace-torch, nequip, jupyter (runs the notebooks)
+pip install -e ".[all]"
+```
+
+The Allegro reference implementation used by one fidelity notebook is not on
+PyPI; install it separately with
+`pip install "git+https://github.com/mir-group/allegro@v0.3.0"`.
+
+## Quick start
+
+```bash
+# builds trains a model on toy data and predicts on a test set
+python examples/quickstart.py
+
+# smoke tests for all models and families
+pytest tests/
+```
+
+```python
+from xnn.common.config import Config
+from xnn.common.data import AtomicDataset
+from xnn.common.train import Trainer
+
+cfg = Config()
+# E.g. schnet | hdnnp | ani | physnet | nequip | mace | allegro | cace | reaxff
+cfg.model.name = "nequip"
+# Unshared model specific hyperparameters go here
+cfg.model.extra = {"species": [1, 6, 8], "l_max": 2}
+# batch_size = 1 disables batch training
+cfg.data.batch_size = 16             
+# auto | cpu | cuda | cuda:0
+cfg.device = "auto"
+Trainer(cfg, AtomicDataset(structures, cfg.model.cutoff)).fit()
+```
+
+`structures` is a list of plain dictionaries (`pos`, `atomic_numbers`, and
+optionally, `cell`/`pbc` and `energy`/`forces`/`stress` targets). Data in any
+ASE-native format loads directly and targets can also be included. So, no
+pre-wrapping is required:
+
+```python
+# xnn supports ase-io native formats such as .extxyz, .cif, VASP, ...
+train_set = AtomicDataset.from_file("trajectory.extxyz", cutoff=4.0)
+```
+
+## Design principles
+
+Four ideas hold `xnn` together:
+
+1. **Unified data object.** Every model consumes an `AtomicGraph` and returns
+   `{"node_energy", "energy"}`. Molecular vs. periodic is invisible to models:
+   periodicity lives only in `edge_vectors()` (`r_ij = pos[dst] - pos[src] +
+   cell_shift @ cell`), keeping forces and stress differentiable.
+2. **Featurizers are first-class.** A `Featurizer` (subclass of `nn.Module`)
+   turns a graph into invariant descriptors (symmetry functions, AEV) or
+   equivariant edge attributes (spherical harmonics, Cartesian monomials).
+   Descriptor models (HDNNP, ANI) and GNNs (NequIP/MACE/Allegro/CACE) are thin
+   compositions over featurizers, so the featurization is reusable and
+   inspectable on its own.
+3. **Forces/stress in one place.** `ForceStressOutput` wraps any model and
+   differentiates energy w.r.t. positions (forces) and a symmetric strain
+   (stress). Models never implement them. The same wrapper idea powers
+   `LatentEwald` (Latent Ewald Summation, Cheng 2025): every model exposes
+   invariant `"node_features"`, so long-range electrostatics/dispersion can be
+   added to *any* short-range model with `extra: {long_range: {...}}` — a
+   faithful port of the CACE-LR reference implementation (see
+   `tests/test_les.py` and `examples/gnn/les/`).
+4. **Extensibility via registry + one config, three frontends.**
+   `@register_model("name")` + a `from_config` classmethod makes a model usable
+   from any of YAML / argparse / Hydra, which all funnel into one
+   `Config` dataclass.
+
+
+```python
+# Data (structures dict -> AtomicGraph)
+from xnn.common.data import AtomicDataset, build_neighbor_list
+ds = AtomicDataset(structures, cutoff=5.0)
+graph = ds[0]
+
+# Featurizers (AtomicGraph -> model inputs)
+from xnn.dnn.featurizers import AEV, RadialSymmetryFunctions
+from xnn.gnn.featurizers import SphericalHarmonicEdgeEmbedding
+# (N, D) invariant per-atom AEV
+descriptor = AEV(species=[1, 6, 8])(graph)
+# Equivariant edge attributes
+edges = SphericalHarmonicEdgeEmbedding(l_max=2)(graph)
+
+# Models (model inputs -> energy)
+from xnn.common.models import build_model, ForceStressOutput, available_models
+# Any registered model + autograd forces/stress
+model = ForceStressOutput(build_model(cfg.model))
+```
+
+## Config frontends (interchangeable)
+
+```python
+from xnn.common.config import from_yaml, from_argparse, from_hydra
+cfg = from_yaml("configs/train.yaml")
+cfg = from_argparse(["--config", "configs/train.yaml", "--set", "model.cutoff=6.0"])
+```
+
+CLI: `xnn train --config configs/train.yaml --set optim.epochs=50`
+(also `xnn benchmark --config configs/benchmark.yaml` and
+`xnn export --config ... --ckpt ... --to lammps|torchscript`).
+
+Model keys copied verbatim from an upstream code's yaml also work: a per-model
+key-translation registry (`xnn.common.config.translate`) rewrites the foreign
+spellings (MACE-CLI `r_max`/`num_radial_basis`/`atomic_numbers`/`E0s`, NequIP
+`num_layers`, ...) to the xnn canonical names at config-load time; the xnn
+spelling wins if both are given. Extend it for another code with
+`register_key_translation("name", {...})`.
+
+## Deployment
+
+```python
+from xnn.common.deploy import XNNCalculator, export_to_lammps
+atoms.calc = XNNCalculator(model, cutoff=5.0)          # ASE
+export_to_lammps(model, cutoff=5.0, path="deployed.pt") # TorchScript for LAMMPS
+```
+
+Pair the exported `.pt` with the matching C++ pair style (pair_nequip /
+pair_mace / pair_allegro pattern). The `LAMMPSWrapper` in
+`common/deploy/lammps.py` defines the tensor ABI. A model is exportable when it
+provides the scriptable `node_energy(atomic_numbers, edge_index, edge_vec)`
+core -- SchNet, NequIP, MACE and Allegro all do (the scripted models reproduce
+the eager ones to ~1e-15, verified in `tests/test_schnet.py` /
+`tests/test_mace.py` / `tests/test_nequip.py` / `tests/test_allegro.py`). For NequIP this required a scriptable, bit-exact stand-in
+for e3nn's `Gate` (`xnn.gnn.models.nequip._Gate`), which the e3nn 0.4.4
+original cannot do on torch 2.x.
+
+## Benchmarking
+
+Score a set of **pre-trained** models on one dataset with
+`xnn.common.benchmark` and tabulate their errors. A single config lists the
+`models` (each an architecture plus the `checkpoint` to load) and the
+`metrics` mapping, which ties each target quantity (`energy` / `forces` /
+`stress`) to the error metrics reported for it (`mae` / `mse` / `rmse`, or
+custom callables). Results are tabulated per model and
+written to CSV / JSON / Markdown (or a user-registered format). Energy can be
+scored per atom or, with `atomic_energies` (a `{Z: E0}` map or `average` to fit
+from data), as the physically meaningful atomization (interaction) energy.
+Benchmarking does not train — produce the checkpoints first with `xnn train`.
+
+```bash
+xnn benchmark --config configs/benchmark.yaml
+```
+
+```python
+from xnn.common.benchmark import from_yaml, run_benchmark
+rows = run_benchmark(from_yaml("configs/benchmark.yaml"))
+```
+
+The benchmark builds models with the same `Config` and model registry as a
+single run; new metrics and output formats plug in via `@register_metric` and
+`@register_writer`, mirroring `@register_model`.
+
+## Available models
+
+| Model | Family | Featurizers | State |
+|---|---|---|---|
+| SchNet | cnn | Gaussian RBF | Under development |
+| PhysNet | dnn | exp-Gaussian rbf + attention masks | Complete: Training, Evaluation, Deployment (ASE only) |
+| HDNNP | dnn | radial symmetry functions (G2) | Under development |
+| ANI | dnn | AEV (radial + angular symmetry functions) | Complete: Training, Evaluation, Deployment (ASE only) |
+| NequIP | gnn | spherical-harmonic edges | Complete: Training, Evaluation, Deployment (TorchScript, LAMMPS, ASE) |
+| MACE | gnn | spherical-harmonic edges | Complete: Training, Evaluation, Deployment (TorchScript, LAMMPS, ASE) |
+| CACE | gnn | Cartesian monomial edges | Complete: Training, Evaluation, Deployment (ASE only) |
+| Allegro | gnn | spherical-harmonic edges | Complete: Training, Evaluation, Deployment (TorchScript, LAMMPS, ASE) |
+| BAMBOO | hybrid | exp-normal rbf + multi-head edge attention | Complete: Training, Evaluation, Deployment (ASE only) |
+| ReaxFF / ReaxFF-nn | ffnn | bond orders + EEM charges (the force field is the model) | Complete: Training, Evaluation, Deployment (ASE only) |
+| OPLS / OPLS-AA / L-OPLS | ffnn | fixed valence topology (the force field is the model) | Complete: Training, Evaluation, Deployment (ASE only) |
+
+
+<details> <!-- Start Package layout -->
+<summary><h2 style="display:inline-block">Package layout</h2></summary>
+
+The package is organized **by model family** (`gnn`, `cnn`, `dnn`, `ffnn`, `hybrid`),
+with shared resources factored into the `common` modules and reusable
+transformer building blocks in `transformer`. An object (e.g., function, module
+etc.) lives with the model family that uses it, or with `common` if more than
+one family needs it. Of course, layers are designed as stand-alone entities and
+can be imported on their own.
+
+```
+src/xnn/
+├── __main__.py                 `python -m xnn` entry point
+├── common/                     shared across all model families
+│   ├── data/                   - common data abstractions
+│   │   └── …                     + AtomicGraph (the one data object), PBC neighbor list, AtomicDataset, ASE I/O
+│   ├── featurizers/            - common featurizers
+│   │   └── …                     + Featurizer base + shared basis functions (GaussianRBF, CosineCutoff)
+│   ├── config/                 - one dataclass schema; loaders for yaml / argparse / hydra
+│   │   └── …                     + schema, loaders, translate, coerce
+│   ├── models/                 - InteratomicPotential interface, registry, ForceStressOutput, ops (scatter_sum, shifted_softplus)
+│   │   └── …                     + base, registry, outputs, les, ops
+│   ├── train/                  - Trainer (batch + device aware), weighted energy/force/stress loss
+│   │   └── …                     + trainer, losses
+│   ├── benchmark/              - score pre-trained models on a dataset (metrics, atomization energy, report writers)
+│   │   └── …                     + config, runner, metrics, energy, report
+│   ├── deploy/                 - ASE Calculator, and LAMMPS/TorchScript export
+│   │   └── …                     + ase_calculator, lammps
+│   └── cli/                    - the `xnn` command-line interface
+│       └── main.py
+├── gnn/                        graph potentials
+│   ├── featurizers/            - GNN featurizers
+│   │   └── …                     + spherical, cartesian, radial, cutoff
+│   └── models/                 - base (GNNPotential, EquivariantGNN), blocks, nequip, mace, allegro, cace
+│       └── …                     + base, blocks, nequip, mace, allegro, cace
+├── cnn/                        continuous-filter conv net
+│   └── models/                 - schnet
+├── dnn/                        descriptor + per-element networks, and PhysNet
+│   ├── featurizers/            - DNN featurizers
+│   │   └── …                     + symmetry functions, AEV
+│   └── models/                 base (DescriptorPotential), hdnnp, ani, physnet and ported Grimme's D3
+│       └── …
+├── ffnn/                       learnable classical force fields
+│   ├── common/                 - frc (SEAMM .frc force-field files: reader, resolver, writer, registry),
+│   │                             typing (SMARTS atom typing), elements
+│   ├── data/                   - shipped parameter files: oplsaa.frc, lopls.frc, oplsaa_1996.frc, reaxff/*.frc
+│   └── models/                 - reaxff, ffield, opls, oplslib, topology
+│       └── …                     + ReaxFF / ReaxFF-nn reactive force field, OPLS / L-OPLS fixed-topology
+│                                   force field, the .frc <-> model parameter bridges
+├── transformer/                shared graph-transformer building blocks
+│   ├── attention.py            - EdgeMultiheadAttention (multi-head QKV attention on edges)
+│   └── featurizers/            - ExpNormalSmearing radial basis
+└── hybrid/                     GNN + transformer potentials with a physics energy split
+    └── models/                 - bamboo (BAMBOO graph equivariant transformer), dispersion (D3(CSO))
+        └── …
+```
+
+</details> <!-- End Package layout -->
+
+<details> <!-- Start Examples -->
+<summary><h2 style="display:inline-block">Examples</h2></summary>
+
+Runnable, pre-executed notebooks live under `examples/` (`pip install -e
+".[examples]"`). They are grouped by model family, where the leading letter of
+`<xnn>` names the architecture type (`gnn` graph, `cnn` convolutional, `dnn`
+deep/descriptor, `ffnn` force-field, `hybrid` mixed):
+
+- **Per-model training and MD** (`examples/<xnn>/<model>/`):
+  + `<model>_argon_train_test.ipynb` and `<model>_argon_density_md.ipynb` for
+    MACE, NequIP, Allegro, CACE and PhysNet.
+  + `schnet_rmd17_train.ipynb` and `schnet_ethanol_md.ipynb` for SchNet.
+  + `ani_rmd17_train.ipynb` plus the `ani1*_dataset.ipynb` / `ani2x_dataset.ipynb`
+    dataset walk-throughs for ANI.
+  + `les_molecular_dimers.ipynb` for the LES long-range wrapper (charged and
+    polar dimers) and `recreate_mace_architecture.ipynb` for a block-by-block
+    MACE rebuild.
+  + `bamboo_dimer_electrostatics.ipynb` and `bamboo_charge_analysis.ipynb` for
+    the hybrid BAMBOO model.
+- **Force fields** (`examples/ffnn/`):
+  + `reaxff/reaxff_rmd17_train_test.ipynb` trains the ReaxFF-nn reactive force
+    field on rMD17 and `reaxff/reaxff_md_bond_orders.ipynb` analyses its bond
+    orders, EEM charges and bond dissociation in MD. Both benchmark against the
+    published classical Chenoweth 2008 C/H/O field, loaded from the shipped
+    SEAMM `.frc` file (`ReaxFF("CHO_cho_2008")`).
+  + `opls/opls_conformational_energetics.ipynb` reproduces Table 1 of the 1996
+    OPLS-AA paper with a relaxed dihedral driver, typing every molecule from
+    coordinates with the SMARTS templates of `oplsaa.frc`.
+  + `opls/opls_lopls_torsion_refit.ipynb` re-derives the L-OPLS hydrocarbon
+    torsion refit of Siu et al. (2012) by gradient descent
+    (`trainable=("dihedral_v",)`) and writes the trained field back out as a
+    `.frc` file.
+- **Data and deployment:**
+  + `data/load_dataset_tutorial.ipynb` covers the one-line dataset hub used by
+    every training notebook (`load_dataset("argon_md")`, `load_dataset("rmd17",
+    ...)`, `load_dataset("ani1", ...)`, `load_dataset("lode_dimers",
+    subset="bio_scan")`).
+  + `deploy/mdi_argon_md.ipynb` and `deploy/mdi_argon_lammps.ipynb` drive a
+    trained model from an external MD code through the MDI engine.
+  + `examples/quickstart.py` is a minimal train/predict script on toy data.
+- **Fidelity checks** (`examples/fidelity_checks/<model>_verification.ipynb`):
+  + Block-by-block numerical comparisons against the upstream codes for MACE,
+    NequIP, Allegro, CACE, SchNet, PhysNet, ANI, BAMBOO and LES.
+  + OPLS is verified against OpenMM (an independent MD engine, optional
+    dependency) in `opls_verification.ipynb` and `tests/test_opls.py`.
+  + **ReaxFF is the exception:** the authors' reference implementation of
+    ReaxFF-nn is AGPL-licensed, so no verification notebook or test depending
+    on it (and no code derived from it) is distributed with this MIT-licensed
+    code base. The implementation follows the published equations, was checked
+    against that implementation during development without redistributing
+    anything from it, and ships with self-contained equation-by-equation tests
+    instead (`tests/test_reaxff.py`; see the fidelity notes in the
+    documentation).
+
+</details> <!-- End Examples -->
+
+<details> <!-- Start Extension points -->
+<summary><h2 style="display:inline-block">Extension points</h2></summary>
+
+- **Implement a new model:** 
+  + Pick your model family package (`gnn` / `cnn` / `dnn`/ `ffnn`,  or add one)
+    and add a module under `<family>/models/`.
+  + Subclass `xnn.common.models.InteratomicPotential` (or the family base,
+  e.g., `gnn.models.base.EquivariantGNN`) and implement its `forward(data)`
+  method.
+  + register your model implementation using `@register_model` decorator.
+  + Add a YAML config file for your model `configs/model/<name>.yaml`. 
+  + Import the family package so the model registers.
+  + For TorchScript/LAMMPS export, it is important to expose a scriptable
+  `node_energy(atomic_numbers, edge_index, edge_vec)` core (SchNet shows the
+  pattern; e3nn models need e3nn's JIT support for this).
+- **Add a new featurizer:**
+  + Subclass `xnn.common.featurizers.Featurizer` and implement `output_dim` and
+  `forward(data)`. Put the resulting featurizer module in the
+  `common/featurizers/` if shared by more than one model family or under the
+  using family's `featurizers/` if it is only used by that one model family.
+
+</details> <!-- End Extension points -->