k_nearest_neighbors(coordinates, neighbor_cutoff, max_num_neighbors=None, p_distance=2, self_loops=False)¶
Find k nearest neighbors for each atom
We do not guarantee that the edges are sorted according to the distance between atoms.
coordinates (numpy.ndarray of shape (N, D)) – The coordinates of atoms in the molecule. N for the number of atoms and D for the dimensions of the coordinates.
neighbor_cutoff (float) – If the distance between a pair of nodes is larger than neighbor_cutoff, they will not be considered as neighboring nodes.
max_num_neighbors (int or None.) – If not None, then this specifies the maximum number of neighbors allowed for each atom. Default to None.
p_distance (int) – We compute the distance between neighbors using Minkowski (\(l_p\)) distance. When
p_distance = 1, Minkowski distance is equivalent to Manhattan distance. When
p_distance = 2, Minkowski distance is equivalent to the standard Euclidean distance. Default to 2.
self_loops (bool) – Whether to allow a node to be its own neighbor. Default to False.
srcs (list of int) – Source nodes.
dsts (list of int) – Destination nodes, corresponding to
distances (list of float) – Distances between the end nodes, corresponding to
>>> from dgllife.utils import get_mol_3d_coordinates, k_nearest_neighbors >>> from rdkit import Chem >>> from rdkit.Chem import AllChem
>>> mol = Chem.MolFromSmiles('CC1(C(N2C(S1)C(C2=O)NC(=O)CC3=CC=CC=C3)C(=O)O)C') >>> AllChem.EmbedMolecule(mol) >>> AllChem.MMFFOptimizeMolecule(mol) >>> coords = get_mol_3d_coordinates(mol) >>> srcs, dsts, dists = k_nearest_neighbors(coords, neighbor_cutoff=1.25) >>> print(srcs) [8, 7, 11, 10, 20, 19] >>> print(dsts) [7, 8, 10, 11, 19, 20] >>> print(dists) [1.2084666104583117, 1.2084666104583117, 1.226457824344217, 1.226457824344217, 1.2230522248065987, 1.2230522248065987]