example_cKDTree.py

import scipy.spatial as spatial
import numpy as np

configuration = np.random.rand((1024, 3))
point_tree = spatial.cKDTree(configuration)
r_cut = 0.5
i_atom = 0
neighbors_indices = point_tree.query_ball_point(configuration[i_atom], r_cut) 

periodize.py

import numpy as np

def periodize_configuration(configuration, r_cut, dimensions):
  """Periodically replicate the atoms in a rectangular box that are distance <= r_cut of the faces.

  Parameters
    configuration: np.array of shape (n_atoms, 3)
        coordinates of the atoms to be periodized

    r_cut: float
        cutoff radius

    dimensions: np.array of shape (3,) (or list length 3)
        dimensions of the periodic rectangle

  Returns
    periodized_configuration: np.array of shape (n_atoms_periodized, 3)

    initial_atom_ids: np.array of shape (n_atoms_periodized, )
        ids of the periodized atoms in the initial configuration
  """
  periodized_configuration = []
  initial_atom_ids = []

  x_translation = np.array([[dimensions[0], 0, 0]], dtype=configuration.dtype)
  y_translation = np.array([[0, dimensions[1], 0]], dtype=configuration.dtype)
  z_translation = np.array([[0, 0, dimensions[2]]], dtype=configuration.dtype)

  mask_true = np.ones(configuration.shape[0], dtype=bool)

  for i_x, mask_x in [(-1., configuration[:, 0] > (dimensions[0] - r_cut)), (0., mask_true), (1., configuration[:, 0] < r_cut)]:
    for i_y, mask_y in [(-1., configuration[:, 1] > (dimensions[1] - r_cut)), (0., mask_true), (1., configuration[:, 1] < r_cut)]:
      for i_z, mask_z in [(-1., configuration[:, 2] > (dimensions[2] - r_cut)), (0., mask_true), (1., configuration[:, 2] < r_cut)]:
        mask = mask_x * mask_y * mask_z
        initial_atom_ids.append(np.nonzero(mask)[0])
        periodized_configuration.append(configuration[mask] + i_x*x_translation + i_y*y_translation + i_z*z_translation)

  periodized_configuration = np.concatenate(periodized_configuration, axis=0)
  initial_atom_ids = np.concatenate(initial_atom_ids, axis=0)

  return periodized_configuration, initial_atom_ids