amep.order.k_nearest_neighbors#
- amep.order.k_nearest_neighbors(coords: ndarray, box_boundary: ndarray, other_coords: ndarray | None = None, k: int = 1, pbc: bool = True, rmax: float = inf, enforce_nd: int | None = None) ndarray#
Calculates the \(k\) nearest neighbors and returns the number of nearest neighbors (can be different from \(k\) if a distance cutoff is given), the \(k\) nearest-neighbor distances, and the \(k\) nearest-neighbor indices.
- Parameters:
coords (np.ndarray of shape (N,3)) – Array of coordinates/points to search for neighbors.
box_boundary (np.ndarray of shape (3,2)) – Boundary of the simulation box in the form of np.array([[xmin, xmax], [ymin, ymax], [zmin, zmax]]).
other_coords (np.ndarray of shape (M,3) or None, optional) – Array of points to search for neighbors of. The default is None (use coords).
rmax (float, optional) – Distance cutoff. Only neighbors with distance smaller than the cutoff are returned. The default is np.inf.
pbc (boolean, optional) – If True, periodic boundary conditions are used. The default is True.
k (int, optional) – Number of nearest neighbors to determine. The default is 1.
enforce_nd (int, None, optional) – enforces the number of dimensions. 2 for 2d, 3 for 3d. If None is supplied, a best guess is used by checking if all particles have the same dimension in the last coordinate. See utils.dimension()
- Returns:
nnn (np.ndarray of shape (N,)) – Number of nearest neighbors.
distances (np.ndarray of shape (N,)) – Array of lists containing the nearest-neighbor distances.
neighbors (np.ndarray of shape (N,)) – Array of lists containing the nearest-neighbor indices.
Examples
Create a test setup of a box and some particles at certain positions:
>>> import amep >>> import numpy as np >>> coords = np.array([[0,0,0],[1,0,0],[0,1,0],[-1,0,0],[0,-2,0]]) >>> box_boundary = np.array([[-4,4],[-4,4],[-0.5,0.5]]) >>> fig, axs = amep.plot.new(figsize=(3,3)) >>> amep.plot.box(axs, box_boundary, ls='--') >>> axs.plot(coords[:,0], coords[:,1], "x", c="k") >>> axs.set_xlabel(r'$x$') >>> axs.set_ylabel(r'$y$') >>> fig.savefig('./figures/order/order-k_nearest_neighbors.png') >>>
Calculate k nearest neighbors without periodic boundary conditions:
>>> nnn, distances, neighbors = amep.order.k_nearest_neighbors( ... coords, box_boundary, k=5, pbc=False ... ) >>> print(nnn) # number of nearest neighbors [4 4 4 4 4] >>> print(distances) # nearest-neighbor distances [[1.0 1.0 1.0 2.0] [1.0 1.4142135623730951 2.0 2.23606797749979] [1.0 1.4142135623730951 1.4142135623730951 3.0] [1.0 1.4142135623730951 2.0 2.23606797749979] [2.0 2.23606797749979 2.23606797749979 3.0]] >>> print(neighbors) # nearest neighbor ids [[2 1 3 4] [0 2 3 4] [0 1 3 4] [0 2 1 4] [0 1 3 2]] >>>
Calculate k nearest neighbors with periodic boundary conditions:
>>> nnn, distances, neighbors = amep.order.k_nearest_neighbors( ... coords, box_boundary, k=5, pbc=True ... ) >>> print(nnn) # number of nearest neighbors [5 5 5 5 5] >>> print(distances) # nearest-neighbor distances [[1.0 1.0 1.0 2.0 6.0] [1.0 1.4142135623730951 2.0 2.23606797749979 6.0] [1.0 1.4142135623730951 1.4142135623730951 3.0 5.0] [1.0 1.4142135623730951 2.0 2.23606797749979 6.0] [2.0 2.23606797749979 2.23606797749979 3.0 5.0]] >>> print(neighbors) # nearest neighbor ids [[3 1 2 4 4] [0 2 3 4 3] [0 3 1 4 4] [0 2 1 4 1] [0 3 1 2 2]] >>>
Calculate k nearest neighbors of all coords in other_coords:
>>> nnn, distances, neighbors = amep.order.k_nearest_neighbors( ... coords[1:], box_boundary, k=5, pbc=False, other_coords=coords[0:1] ... ) >>> print(nnn) # number of nearest neighbors [1 1 1 1] >>> print(distances) # nearest-neighbor distances [[1.0] [1.0] [1.0] [2.0]] >>> print(neighbors) # nearest neighbor ids [[0] [0] [0] [0]] >>>