xnn.common.deploy.torchscript.build_neighbor_list_ts#

xnn.common.deploy.torchscript.build_neighbor_list_ts(pos, cutoff, cell, pbc)[source]#

torch.jit-able neighbor list, matching the reference implementation.

Semantically identical to build_neighbor_list() (same dst = i / src = j convention, same negated cell_shifts, same wrapping of out-of-cell positions with the removed image offsets folded back into the shifts), rewritten without itertools and without .tolist() so it compiles under TorchScript and can live inside the exported artifact.

A structure is treated as molecular when no periodic flag is set or the cell is all zeros, in which case image enumeration is skipped.

Parameters:

pos (Tensor) – Cartesian positions, shape (N, 3).
cutoff (float) – Neighbor cutoff radius.
cell (Tensor) – Lattice vectors as rows, shape (3, 3); all-zero for a molecule.
pbc (Tensor) – Boolean periodicity flags, shape (3,).

Returns:

edge_index (Tensor) – Edge list [src, dst], shape (2, E).
cell_shifts (Tensor) – Integer image shift per edge, shape (E, 3).

Return type:

Tuple[Tensor, Tensor]

Notes

Brute force over (S, N, N) pairs, like the reference implementation it mirrors; memory grows as O(S N^2). Fine for the molecular and modest periodic systems this export targets, but for large cells supply the neighbor list from the MD engine via forward_lammps instead.