torch_nl

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.

How to

instal with pip

pip install torch-nl

use the neighborlist

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
)

Benchmarks

Periodic structure