TorchScript-able neighbor lists implementations (linear and quadratic scaling) for molecular modeling
Project description
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
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
torch_nl-0.3.tar.gz
(11.6 kB
view hashes)
Built Distribution
torch_nl-0.3-py3-none-any.whl
(12.9 kB
view hashes)