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')
>>>
../_images/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]]
>>>