Source code for ase.calculators.neighborlist

from math import sqrt

import numpy as np


[docs]class NeighborList: """Neighbor list object. cutoffs: list of float List of cutoff radii - one for each atom. If the spheres (defined by their cutoff radii) of two atoms overlap, they will be counted as neighbors. skin: float If no atom has moved more than the skin-distance since the last call to the ``update()`` method, then the neighbor list can be reused. This will save some expensive rebuilds of the list, but extra neighbors outside the cutoff will be returned. self_interaction: bool Should an atom return itself as a neighbor? bothways: bool Return all neighbors. Default is to return only "half" of the neighbors. Example:: nl = NeighborList([2.3, 1.7]) nl.update(atoms) indices, offsets = nl.get_neighbors(0) """ def __init__(self, cutoffs, skin=0.3, sorted=False, self_interaction=True, bothways=False): self.cutoffs = np.asarray(cutoffs) + skin self.skin = skin self.sorted = sorted self.self_interaction = self_interaction self.bothways = bothways self.nupdates = 0
[docs] def update(self, atoms): """Make sure the list is up to date.""" if self.nupdates == 0: self.build(atoms) return True if ((self.pbc != atoms.get_pbc()).any() or (self.cell != atoms.get_cell()).any() or ((self.positions - atoms.get_positions())**2).sum(1).max() > self.skin**2): self.build(atoms) return True return False
[docs] def build(self, atoms): """Build the list.""" self.positions = atoms.get_positions() self.pbc = atoms.get_pbc() self.cell = atoms.get_cell() if len(self.cutoffs) > 0: rcmax = self.cutoffs.max() else: rcmax = 0.0 icell = np.linalg.inv(self.cell) scaled = np.dot(self.positions, icell) scaled0 = scaled.copy() N = [] for i in range(3): if self.pbc[i]: scaled0[:, i] %= 1.0 v = icell[:, i] h = 1 / sqrt(np.dot(v, v)) n = int(2 * rcmax / h) + 1 else: n = 0 N.append(n) offsets = (scaled0 - scaled).round().astype(int) positions0 = np.dot(scaled0, self.cell) natoms = len(atoms) indices = np.arange(natoms) self.nneighbors = 0 self.npbcneighbors = 0 self.neighbors = [np.empty(0, int) for a in range(natoms)] self.displacements = [np.empty((0, 3), int) for a in range(natoms)] for n1 in range(0, N[0] + 1): for n2 in range(-N[1], N[1] + 1): for n3 in range(-N[2], N[2] + 1): if n1 == 0 and (n2 < 0 or n2 == 0 and n3 < 0): continue displacement = np.dot((n1, n2, n3), self.cell) for a in range(natoms): d = positions0 + displacement - positions0[a] i = indices[(d**2).sum(1) < (self.cutoffs + self.cutoffs[a])**2] if n1 == 0 and n2 == 0 and n3 == 0: if self.self_interaction: i = i[i >= a] else: i = i[i > a] self.nneighbors += len(i) self.neighbors[a] = np.concatenate( (self.neighbors[a], i)) disp = np.empty((len(i), 3), int) disp[:] = (n1, n2, n3) disp += offsets[i] - offsets[a] self.npbcneighbors += disp.any(1).sum() self.displacements[a] = np.concatenate( (self.displacements[a], disp)) if self.bothways: neighbors2 = [[] for a in range(natoms)] displacements2 = [[] for a in range(natoms)] for a in range(natoms): for b, disp in zip(self.neighbors[a], self.displacements[a]): neighbors2[b].append(a) displacements2[b].append(-disp) for a in range(natoms): # Force neighbors to be integer array self.neighbors[a] = np.array(np.concatenate((self.neighbors[a], neighbors2[a])), int) self.displacements[a] = np.array(list(self.displacements[a]) + displacements2[a]) if self.sorted: for a, i in enumerate(self.neighbors): mask = (i < a) if mask.any(): j = i[mask] offsets = self.displacements[a][mask] for b, offset in zip(j, offsets): self.neighbors[b] = np.concatenate( (self.neighbors[b], [a])) self.displacements[b] = np.concatenate( (self.displacements[b], [-offset])) mask = np.logical_not(mask) self.neighbors[a] = self.neighbors[a][mask] self.displacements[a] = self.displacements[a][mask] self.nupdates += 1
[docs] def get_neighbors(self, a): """Return neighbors of atom number a. A list of indices and offsets to neighboring atoms is returned. The positions of the neighbor atoms can be calculated like this:: indices, offsets = nl.get_neighbors(42) for i, offset in zip(indices, offsets): print(atoms.positions[i] + dot(offset, atoms.get_cell())) Notice that if get_neighbors(a) gives atom b as a neighbor, then get_neighbors(b) will not return a as a neighbor - unless bothways=True was used.""" return self.neighbors[a], self.displacements[a]