dgllife.utils.k_nearest_neighbors

dgllife.utils.k_nearest_neighbors(coordinates, neighbor_cutoff, max_num_neighbors=None, p_distance=2, self_loops=False)[source]

Find k nearest neighbors for each atom

We do not guarantee that the edges are sorted according to the distance between atoms.

Parameters
  • 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.

Returns

  • srcs (list of int) – Source nodes.

  • dsts (list of int) – Destination nodes, corresponding to srcs.

  • distances (list of float) – Distances between the end nodes, corresponding to srcs and dsts.

Examples

>>> 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]