xnn.common.deploy.torchscript#
Self-contained TorchScript export: a trained model as a standalone .pt.
The artifact produced here is driven purely by tensors and carries its own
neighbor list, so a consumer needs only libtorch / torch.jit.load –
no xnn import, no Python model code, no config file. That is what makes a
checkpoint usable from LAMMPS, i-PI, OpenMM, a C++ driver or any other MD
package.
Two entry points are exported on the same module:
forward(pos, atomic_numbers, cell, pbc)The whole-system ABI. The module builds its own neighbor list from the cutoff baked into it at export time. This is the general-purpose entry point and the only correct one for models carrying a long-range term.
forward_lammps(pos, edge_index, cell_shifts, atomic_numbers, cell)The pair-style ABI, matching
LAMMPSWrapper(and hence pair_nequip / pair_mace / pair_allegro): the caller supplies the neighbor list. Cheaper, because the MD engine already has a neighbor list, but see the long-range caveat below.
Calling conventions#
One structure per call (inputs are (N, 3), with no batch dimension).
Positions may be float32 or float64 whatever the weights’ dtype: the module
computes in its own dtype and answers in the caller’s. Wrapping the call in
torch.no_grad() is fine – forces come from autograd, so the module
re-enables grad internally and restores the caller’s mode afterwards – but
torch.inference_mode() raises, because tensors created under it can never
participate in autograd.
Add-on terms: LES long-range and D4 dispersion#
Two wrappers can sit on top of a core model, in either order or together:
LatentEwald (an Ewald energy over latent
charges) and DispersionCorrection (the DFT-D3 /
DFT-D4 corrections; D4’s EEQ charges couple every atom of the structure
through a dense linear system). Both are global: they do not decompose into a
local, per-domain neighbor list. A model exported with either must be driven
with the whole system on one rank (forward, or forward_lammps with a
full-system neighbor list) – an MPI-decomposed pair style that only ever sees
its own subdomain plus ghosts cannot reproduce the trained energy.
export_torchscript_potential() records long_range and dispersion
in the archive metadata so a consumer can check.
The dispersion wrapper also widens the neighbor list (its cutoffs default to the
upstream 60 / 40 / 30 / 25 bohr): the export bakes in the wrapper’s cutoff and
hands the core model only the edges within the core’s own radius, exactly as
the eager wrapper does. The total charge of the system is fixed at export
time (total_charge, default neutral), since neither tensor ABI carries it;
the LES head reads it for its net-charge constraint and charge solve, whose
fragment charges come from the ion table alone (no label at deployment).
Core ABIs#
A core model exposes one of two scriptable tensor methods. Most models
provide node_features_energy(atomic_numbers, edge_index, edge_vec)
(features and per-atom energies from the local graph alone). Models that
predict partial charges and add their own electrostatics
(AIMNet2) provide
node_features_energy_charges(atomic_numbers, edge_index, edge_vec, pos,
cell, pbc, total_charge, spin_multiplicity), which also receives the
positions, the cell and the charge state (the latter fixed at export, like
total_charge above, through spin_multiplicity) and returns the
charges as well; the artifact then carries them under "charges". Their
Coulomb sum is global (all pairs, or an Ewald sum), so the same whole-system
rule as for the add-on terms applies.
Functions
|
|
|
Script a trained model to a standalone |
|
Peel the LES / D4 wrappers off a model. |
Classes
|
Tensor-only, scriptable potential: positions in, energy/forces/stress out. |