Added dnn to next_neighbors

This commit is contained in:
Sebastian Kloth 2023-12-26 11:07:30 +01:00
parent 97e29e6de3
commit 2615798088

View File

@ -251,7 +251,9 @@ def next_neighbors(
atoms = atoms % np.diag(box)
tree = KDTree(atoms, boxsize=np.diag(box))
distances, indices = tree.query(
query_atoms, number_of_neighbors, distance_upper_bound=distance_upper_bound
query_atoms,
number_of_neighbors + dnn,
distance_upper_bound=distance_upper_bound,
)
else:
atoms_pbc, atoms_pbc_index = pbc_points(
@ -259,9 +261,10 @@ def next_neighbors(
)
tree = KDTree(atoms_pbc)
distances, indices = tree.query(
query_atoms, number_of_neighbors, distance_upper_bound=distance_upper_bound
query_atoms,
number_of_neighbors + dnn,
distance_upper_bound=distance_upper_bound,
)
indices = atoms_pbc_index[indices]
return distances[:, dnn:], indices[:, dnn:]