xnn.common.deploy.mdi_engine#

MDI engine serving any trained xnn model to an external driver.

The MolSSI Driver Interface (MDI) lets simulation codes (LAMMPS, QCEngine, SEAMM, …) drive an external “engine” through a small command protocol. MDIEngine implements the engine side for xnn: the driver sends the system (>NATOMS, >ELEMENTS, >CELL, >COORDS, optionally >TOTCHARGE) and requests results (<ENERGY, <FORCES, <STRESS), and the engine evaluates the wrapped model each time the geometry changes.

The engine is model agnostic: it speaks to the model exclusively through the AtomicGraph contract shared by every model in the library (the same contract XNNCalculator uses), so any registered family (MACE, NequIP, Allegro, CACE, SchNet, ANI, PhysNet, HDNNP, BAMBOO, …) works unchanged. Graphs are built with the library’s own structure_to_graph(); forces and stress come from the ForceStressOutput wrapper.

Typical use, from a trainer checkpoint:

from xnn.common.deploy import MDIEngine
engine = MDIEngine.from_checkpoint("runs/exp/best.pt", device="cuda")
engine.run("-role ENGINE -name xnn -method TCP -port 8021 -hostname localhost")

or from the command line (see main()):

xnn mdi --ckpt runs/exp/best.pt -mdi "-role ENGINE -name xnn -method TCP ..."

A checkpoint trained without dispersion can be served with a D3 / D4 correction added on top (--dispersion d4 or a YAML mapping such as "{name: d4, cutoff_pair: 12.0, switch_width_pair: 2.0}"), and the system’s net charge (used by D4’s EEQ charges and by charge-aware models) is set with --total-charge or by the driver through >TOTCHARGE.

Requires the pymdi package (pip install pymdi); MPI communication additionally requires mpi4py.

Units: MDI communicates in atomic units (Bohr / Hartree) while xnn models follow the library’s ASE-style convention of angstrom / eV (the units of the training data). The engine converts at the boundary in both directions.

Functions

main([argv])

Command-line entry point serving a checkpoint as an MDI engine.

Classes

MDIEngine(model, cutoff[, device, total_charge])

MDI engine exposing a trained xnn model to an external driver.
Show Source