Provide a pytorch implementation of a naive (compute_neighborlist_n2
) and a linked cell (compute_neighborlist
) neighbor list that are compatible with TorchScript.
Their correctness is tested against ASE's implementation.
Note that contrary to ASE, the atoms positions are assumed to be wrapped inside the unit cell.
pip install torch-nl
from torch_nl import compute_neighborlist, ase2data
from ase.build import bulk, molecule
frames = [bulk("Si", "diamond", a=6, cubic=True), molecule("CH3CH2NH2")]
pos, cell, pbc, batch, n_atoms = ase2data(frames)
mapping, batch_mapping, shifts_idx = compute_neighborlist(
cutoff, pos, cell, pbc, batch, self_interaction
)