xnn.common.deploy.torchscript.TorchScriptPotential#

class xnn.common.deploy.torchscript.TorchScriptPotential(model, cutoff, total_charge=0.0)[source]#

Bases: Module

Tensor-only, scriptable potential: positions in, energy/forces/stress out.

Wraps a trained model (optionally LatentEwald- and/or D4Dispersion-wrapped) behind a fixed tensor ABI with no AtomicGraph and no Python-only constructs, so torch.jit.script produces a portable artifact. Forces come from -dE/dr and the stress from the symmetric-strain trick, matching ForceStressOutput.

Parameters:

  • model (torch.nn.Module) – The trained model. Its core must expose the scriptable tensor core node_features_energy(atomic_numbers, edge_index, edge_vec); LES and D4 wrappers are unwrapped automatically and carried over as heads.

  • cutoff (float) – Neighbor-list cutoff, baked in so forward() is self-contained. For a D4-wrapped model this is the wrapper’s (widened) cutoff; the core receives only the edges within its own radius.

  • total_charge (float, optional) – Net charge of the deployed system, used by the D4 EEQ charges; by default 0 (neutral).

Variables:

  • model (torch.nn.Module) – The short-range core.

  • lr (torch.nn.Module) – The long-range head (_LatentEwaldHead or _NoLongRange).

  • disp (torch.nn.Module) – The dispersion head (_DispersionHead or _NoDispersion).

  • cutoff (float) – The neighbor-list cutoff radius.

  • core_cutoff (float) – The core model’s own radius (<= cutoff).

  • has_long_range (bool) – Whether a long-range term is present.

  • has_dispersion (bool) – Whether a D4 dispersion term is present.

forward(pos, atomic_numbers, cell=None, pbc=None)[source]#

Whole-system entry point; builds its own neighbor list.

The general-purpose ABI, and the only correct one when the model carries a long-range term (see the module docstring).

Parameters:

  • pos (Tensor) – Cartesian positions, shape (N, 3).

  • atomic_numbers (Tensor) – Per-atom atomic numbers, shape (N,).

  • cell (Tensor or None, optional) – Lattice vectors as rows, shape (3, 3). None (the default) or all-zero means a molecular system.

  • pbc (Tensor or None, optional) – Boolean periodicity flags, shape (3,). None (the default) means non-periodic.

Returns:

energy / total_energy (1,), energy_sr (1,), energy_lr (1,), energy_disp (1,), node_energy (N,), forces (N, 3), stress (3, 3), virial (3, 3), latent_charges (N, n_channels) (empty without LES) and eeq_charges (N, 1) (empty without D4).

Return type:

dict of str to Tensor

forward_lammps(pos, edge_index, cell_shifts, atomic_numbers, cell)[source]#

Pair-style entry point; the caller supplies the neighbor list.

Argument order matches LAMMPSWrapper, so an existing pair style needs no change. Periodicity is inferred from cell being nonzero.

Parameters:

  • pos (Tensor) – Cartesian positions, shape (N, 3).

  • edge_index (Tensor) – Edge list [src, dst], shape (2, E).

  • cell_shifts (Tensor) – Integer image shift per edge, shape (E, 3).

  • atomic_numbers (Tensor) – Per-atom atomic numbers, shape (N,).

  • cell (Tensor) – Lattice vectors as rows, shape (3, 3).

Returns:

The same keys as forward().

Return type:

dict of str to Tensor

Notes

For a model with a long-range or dispersion term the supplied neighbor list must cover the whole system on a single rank (and reach the exported cutoff); a subdomain-local list does not reproduce the trained energy.