xnn.common.deploy.lammps.LAMMPSWrapper#

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

Bases: Module

TorchScript-friendly bridge between LAMMPS tensors and the model.

Operates purely on tensors (no AtomicGraph dataclass) so the module is torch.jit.script-able, which is a hard requirement for use inside a LAMMPS pair style. Models expose a scriptable node_energy(...) core for exactly this purpose.

This wrapper defines the tensor ABI (application binary interface) that the LAMMPS pair style exchanges with the model: the argument order, dtypes and shapes of forward() are the contract between the C++ pair style and the serialized model, so they must not drift from what the pair style provides and consumes.

Parameters:

  • model (torch.nn.Module) – The trained model. Must expose a scriptable node_energy(atomic_numbers, edge_index, edge_vec) method returning per-atom (node) energies. A DispersionCorrection (D3 / D4) wrapper is unwrapped and its dispersion term added (the supplied neighbor list must then reach the wrapper’s cutoff and cover the whole system; the core sees only the edges within its own radius). LES needs per-atom features and is served by forward_lammps() instead.

  • cutoff (float) – Neighbor-list cutoff radius, stored (as a Python float) for serialization alongside the scripted module.

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

Variables:

  • model (torch.nn.Module) – The wrapped (core) model.

  • disp (torch.nn.Module) – The D4 head, or a null head.

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

  • core_cutoff (float) – The core model’s radius; edges beyond it are filtered before the core.

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

Compute energy and forces from LAMMPS-provided neighbor data.

Edge vectors are reconstructed as pos[dst] - pos[src] plus the periodic image shift cell_shifts @ cell. The model’s per-atom (node) energies are summed to the total energy, and forces are the negative gradient of that total energy with respect to the (grad- enabled) positions.

Parameters:

  • pos (Tensor) – Atomic positions, shape (n_atoms, 3).

  • edge_index (Tensor) – Neighbor (edge) list, shape (2, n_edges); row 0 is the source index and row 1 is the destination index of each edge.

  • cell_shifts (Tensor) – Integer periodic-image shift vectors per edge, shape (n_edges, 3); combined with cell to place neighbors in their correct periodic images.

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

  • cell (Tensor) – Simulation cell (lattice) matrix, shape (3, 3).

Returns:

A dictionary with keys "total_energy" (shape (1,)), "node_energy" (per-atom energies, shape (n_atoms,)) and "forces" (shape (n_atoms, 3)). Forces are zero when the autograd gradient is unavailable.

Return type:

dict of str to Tensor

On this page
Show Source