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, ops, les (Latent Ewald), dispersion (shared D3/D4 machinery), d4 (DFT-D4), d3 (DFT-D3)
│   ├── 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
│       └── …
├── 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 and DREIDING 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))
        └── …

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.
  • mace_foundation_molecules.ipynb and mace_foundation_materials.ipynb for the pretrained MACE foundation models (MACE.from_foundation()): MACE-OFF23 on torsions / hydrogen bonds / Trainer fine-tuning, and equations of state across the MACE-MP generations. Every published checkpoint's conversion is verified in examples/fidelity_checks/mace_foundation_verification.ipynb.
  • 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.
  • common/d4/d4_paper_examples.ipynb and common/d4/d4_benchmark.ipynb for the DFT-D4 dispersion add-on (paper reproductions, benchmark against dftd4, deployment of a D4-corrected MLIP); common/d3/ likewise for DFT-D3 (against simple-dftd3).
• 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.
  • dreiding/dreiding_conformational_energetics.ipynb reproduces the rotational barriers of Table XI and the conformational energies of Table XII of the 1990 DREIDING paper (14 molecules, mean difference ~0.01 kcal/mol from the paper's own calculated values) from relaxed scans.
  • dreiding/dreiding_refit_aromatics.ipynb refits DREIDING's generators on benzene against rMD17 PBE forces and tests whether they transfer to naphthalene and toluene, neither of which is trained on.
• 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, LES, DFT-D4 and DFT-D3.
  • OPLS is verified against OpenMM (an independent MD engine, optional dependency) in opls_verification.ipynb and tests/test_opls.py.
  • DREIDING is verified term by term against LAMMPS's DREIDING styles (an independent MD engine, optional dependency) in dreiding_verification.ipynb, and against the paper's own tables in tests/test_dreiding.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).

• 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.