From 2615798088a22f4af4aa839b42e6e2d723c66d9f Mon Sep 17 00:00:00 2001 From: Sebastian Kloth Date: Tue, 26 Dec 2023 11:07:30 +0100 Subject: [PATCH] Added dnn to next_neighbors --- src/mdevaluate/atoms.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/mdevaluate/atoms.py b/src/mdevaluate/atoms.py index 22582a4..77bebac 100644 --- a/src/mdevaluate/atoms.py +++ b/src/mdevaluate/atoms.py @@ -251,17 +251,20 @@ 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( - query_atoms, box, thickness=distance_upper_bound+0.1, index=True, **kwargs + query_atoms, box, thickness=distance_upper_bound + 0.1, index=True, **kwargs ) 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:] -