felixmusil / torch_nl Goto Github PK

from ase.neighborlist import primitive_neighbor_list
from torch_nl import compute_neighborlist
import torch

# cell of 10 x 10 x 10
cell = torch.eye(3)*10.0
pbc = torch.tensor([True, True, True],dtype=bool)

cutoff = 4.5 # a bit smaller than half box length

# make 10 points that are spread out over many periodic images
x = torch.randn(10,3)*100.0


# ASE
i, j, S = primitive_neighbor_list(
    quantities="ijS",
    pbc=pbc.numpy(),
    cell=cell.numpy(),
    positions=x.numpy(),
    cutoff=cutoff,
    self_interaction=True,  
    use_scaled_positions=False,  
)


# torch nl
tnl_ij, _, tnl_S = compute_neighborlist(cutoff=cutoff, pos=x, cell=cell, pbc=pbc, batch=torch.zeros(10,dtype=torch.int), self_interaction=True)
tnl_i = tnl_ij[0]
tnl_j = tnl_ij[1]

#They should give the same 
i=list(i)
j=list(j)

tnl_i = [ int(a) for a in tnl_i]
tnl_j = [ int(a) for a in tnl_j]

print("ASE:",[i,j])
print("torch-nl:",[tnl_i, tnl_j])

assert( i == tnl_i)

ASE: [[0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 6, 6, 7, 7, 7, 7, 7, 7, 7, 8, 8, 8, 9, 9, 9, 9, 9], [6, 7, 2, 8, 1, 3, 4, 0, 7, 1, 3, 0, 3, 7, 9, 5, 0, 2, 4, 1, 3, 0, 2, 4, 3, 4, 0, 2, 9, 5, 4, 9, 7, 5, 2, 0, 6, 1, 9, 7, 5, 2, 8, 0, 0, 8, 7, 9, 7, 5, 2, 4]]
torch-nl: [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]]
Traceback (most recent call last):
    assert( i == tnl_i)
AssertionError