Coverage for /builds/ase/ase/ase/neighborlist.py: 96.03%
378 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-08-02 00:12 +0000
1# fmt: off
3import itertools
5import numpy as np
6import scipy.sparse.csgraph as csgraph
7from scipy import sparse as sp
8from scipy.spatial import cKDTree
10from ase.cell import Cell
11from ase.data import atomic_numbers, covalent_radii
12from ase.geometry import (
13 complete_cell,
14 find_mic,
15 minkowski_reduce,
16 wrap_positions,
17)
18from ase.utils import deprecated
21def natural_cutoffs(atoms, mult=1, **kwargs):
22 """Generate a radial cutoff for every atom based on covalent radii.
24 The covalent radii are a reasonable cutoff estimation for bonds in
25 many applications such as neighborlists, so function generates an
26 atoms length list of radii based on this idea.
28 * atoms: An atoms object
29 * mult: A multiplier for all cutoffs, useful for coarse grained adjustment
30 * kwargs: Symbol of the atom and its corresponding cutoff,
31 used to override the covalent radii
32 """
33 return [kwargs.get(atom.symbol, covalent_radii[atom.number] * mult)
34 for atom in atoms]
37def build_neighbor_list(atoms, cutoffs=None, **kwargs):
38 """Automatically build and update a NeighborList.
40 Parameters:
42 atoms : :class:`~ase.Atoms` object
43 Atoms to build Neighborlist for.
44 cutoffs: list of floats
45 Radii for each atom. If not given it will be produced by calling
46 :func:`ase.neighborlist.natural_cutoffs`
47 kwargs: arbitrary number of options
48 Will be passed to the constructor of
49 :class:`~ase.neighborlist.NeighborList`
51 Returns:
53 return: :class:`~ase.neighborlist.NeighborList`
54 A :class:`~ase.neighborlist.NeighborList` instance (updated).
55 """
56 if cutoffs is None:
57 cutoffs = natural_cutoffs(atoms)
59 nl = NeighborList(cutoffs, **kwargs)
60 nl.update(atoms)
62 return nl
65def get_distance_matrix(graph, limit=3):
66 """Get Distance Matrix from a Graph.
68 Parameters:
70 graph: array, matrix or sparse matrix, 2 dimensions (N, N)
71 Graph representation of the connectivity.
72 See `scipy doc <https://docs.scipy.org/doc/scipy/reference/generated\
73/scipy.sparse.csgraph.dijkstra.html#scipy.sparse.csgraph.dijkstra>`_
74 for reference.
75 limit: integer
76 Maximum number of steps to analyze. For most molecular information,
77 three should be enough.
79 Returns:
81 return: scipy.sparse.csr_matrix, shape (N, N)
82 A scipy.sparse.csr_matrix. All elements that are not connected within
83 *limit* steps are set to zero.
85 This is a potential memory bottleneck, as csgraph.dijkstra produces a
86 dense output matrix. Here we replace all np.inf values with 0 and
87 transform back to csr_matrix.
88 Why not dok_matrix like the connectivity-matrix? Because row-picking
89 is most likely and this is super fast with csr.
90 """
91 mat = csgraph.dijkstra(graph, directed=False, limit=limit)
92 mat[mat == np.inf] = 0
93 return sp.csr_matrix(mat, dtype=np.int8)
96def get_distance_indices(distanceMatrix, distance):
97 """Get indices for each node that are distance or less away.
99 Parameters:
101 distanceMatrix: any one of scipy.sparse matrices (NxN)
102 Matrix containing distance information of atoms. Get it e.g. with
103 :func:`~ase.neighborlist.get_distance_matrix`
104 distance: integer
105 Number of steps to allow.
107 Returns:
109 return: list of length N
110 List of length N. return[i] has all indices connected to item i.
112 The distance matrix only contains shortest paths, so when looking for
113 distances longer than one, we need to add the lower values for cases
114 where atoms are connected via a shorter path too.
115 """
116 shape = distanceMatrix.get_shape()
117 indices = []
118 # iterate over rows
119 for i in range(shape[0]):
120 row = distanceMatrix.getrow(i)[0]
121 # find all non-zero
122 found = sp.find(row)
123 # screen for smaller or equal distance
124 equal = np.where(found[-1] <= distance)[0]
125 # found[1] contains the indexes
126 indices.append([found[1][x] for x in equal])
127 return indices
130def mic(dr, cell, pbc=True):
131 """
132 Apply minimum image convention to an array of distance vectors.
134 Parameters:
136 dr : array_like
137 Array of distance vectors.
138 cell : array_like
139 Simulation cell.
140 pbc : array_like, optional
141 Periodic boundary conditions in x-, y- and z-direction. Default is to
142 assume periodic boundaries in all directions.
144 Returns:
146 dr : array
147 Array of distance vectors, wrapped according to the minimum image
148 convention.
149 """
150 dr, _ = find_mic(dr, cell, pbc)
151 return dr
154def primitive_neighbor_list(quantities, pbc, cell, positions, cutoff,
155 numbers=None, self_interaction=False,
156 use_scaled_positions=False, max_nbins=1e6):
157 """Compute a neighbor list for an atomic configuration.
159 Atoms outside periodic boundaries are mapped into the box. Atoms
160 outside nonperiodic boundaries are included in the neighbor list
161 but complexity of neighbor list search for those can become n^2.
163 The neighbor list is sorted by first atom index 'i', but not by second
164 atom index 'j'.
166 Parameters:
168 quantities: str
169 Quantities to compute by the neighbor list algorithm. Each character
170 in this string defines a quantity. They are returned in a tuple of
171 the same order. Possible quantities are
173 * 'i' : first atom index
174 * 'j' : second atom index
175 * 'd' : absolute distance
176 * 'D' : distance vector
177 * 'S' : shift vector (number of cell boundaries crossed by the bond
178 between atom i and j). With the shift vector S, the
179 distances D between atoms can be computed from:
180 D = positions[j]-positions[i]+S.dot(cell)
181 pbc: array_like
182 3-tuple indicating giving periodic boundaries in the three Cartesian
183 directions.
184 cell: 3x3 matrix
185 Unit cell vectors.
186 positions: list of xyz-positions
187 Atomic positions. Anything that can be converted to an ndarray of
188 shape (n, 3) will do: [(x1,y1,z1), (x2,y2,z2), ...]. If
189 use_scaled_positions is set to true, this must be scaled positions.
190 cutoff: float or dict
191 Cutoff for neighbor search. It can be:
193 * A single float: This is a global cutoff for all elements.
194 * A dictionary: This specifies cutoff values for element
195 pairs. Specification accepts element numbers of symbols.
196 Example: {(1, 6): 1.1, (1, 1): 1.0, ('C', 'C'): 1.85}
197 * A list/array with a per atom value: This specifies the radius of
198 an atomic sphere for each atoms. If spheres overlap, atoms are
199 within each others neighborhood. See
200 :func:`~ase.neighborlist.natural_cutoffs`
201 for an example on how to get such a list.
202 self_interaction: bool
203 Return the atom itself as its own neighbor if set to true.
204 Default: False
205 use_scaled_positions: bool
206 If set to true, positions are expected to be scaled positions.
207 max_nbins: int
208 Maximum number of bins used in neighbor search. This is used to limit
209 the maximum amount of memory required by the neighbor list.
211 Returns:
213 i, j, ... : array
214 Tuple with arrays for each quantity specified above. Indices in `i`
215 are returned in ascending order 0..len(a)-1, but the order of (i,j)
216 pairs is not guaranteed.
218 """
220 # Naming conventions: Suffixes indicate the dimension of an array. The
221 # following convention is used here:
222 # c: Cartesian index, can have values 0, 1, 2
223 # i: Global atom index, can have values 0..len(a)-1
224 # xyz: Bin index, three values identifying x-, y- and z-component of a
225 # spatial bin that is used to make neighbor search O(n)
226 # b: Linearized version of the 'xyz' bin index
227 # a: Bin-local atom index, i.e. index identifying an atom *within* a
228 # bin
229 # p: Pair index, can have value 0 or 1
230 # n: (Linear) neighbor index
232 # Return empty neighbor list if no atoms are passed here
233 if len(positions) == 0:
234 empty_types = dict(i=(int, (0, )),
235 j=(int, (0, )),
236 D=(float, (0, 3)),
237 d=(float, (0, )),
238 S=(int, (0, 3)))
239 retvals = []
240 for i in quantities:
241 dtype, shape = empty_types[i]
242 retvals += [np.array([], dtype=dtype).reshape(shape)]
243 if len(retvals) == 1:
244 return retvals[0]
245 else:
246 return tuple(retvals)
248 # Compute reciprocal lattice vectors.
249 b1_c, b2_c, b3_c = np.linalg.pinv(cell).T
251 # Compute distances of cell faces.
252 l1 = np.linalg.norm(b1_c)
253 l2 = np.linalg.norm(b2_c)
254 l3 = np.linalg.norm(b3_c)
255 face_dist_c = np.array([1 / l1 if l1 > 0 else 1,
256 1 / l2 if l2 > 0 else 1,
257 1 / l3 if l3 > 0 else 1])
259 if isinstance(cutoff, dict):
260 max_cutoff = max(cutoff.values())
261 else:
262 if np.isscalar(cutoff):
263 max_cutoff = cutoff
264 else:
265 cutoff = np.asarray(cutoff)
266 max_cutoff = 2 * np.max(cutoff)
268 # We use a minimum bin size of 3 A
269 bin_size = max(max_cutoff, 3)
270 # Compute number of bins such that a sphere of radius cutoff fits into
271 # eight neighboring bins.
272 nbins_c = np.maximum((face_dist_c / bin_size).astype(int), [1, 1, 1])
273 nbins = np.prod(nbins_c)
274 # Make sure we limit the amount of memory used by the explicit bins.
275 while nbins > max_nbins:
276 nbins_c = np.maximum(nbins_c // 2, [1, 1, 1])
277 nbins = np.prod(nbins_c)
279 # Compute over how many bins we need to loop in the neighbor list search.
280 neigh_search_x, neigh_search_y, neigh_search_z = \
281 np.ceil(bin_size * nbins_c / face_dist_c).astype(int)
283 # If we only have a single bin and the system is not periodic, then we
284 # do not need to search neighboring bins
285 neigh_search_x = 0 if nbins_c[0] == 1 and not pbc[0] else neigh_search_x
286 neigh_search_y = 0 if nbins_c[1] == 1 and not pbc[1] else neigh_search_y
287 neigh_search_z = 0 if nbins_c[2] == 1 and not pbc[2] else neigh_search_z
289 # Sort atoms into bins.
290 if use_scaled_positions:
291 scaled_positions_ic = positions
292 positions = np.dot(scaled_positions_ic, cell)
293 else:
294 scaled_positions_ic = np.linalg.solve(complete_cell(cell).T,
295 positions.T).T
296 bin_index_ic = np.floor(scaled_positions_ic * nbins_c).astype(int)
297 cell_shift_ic = np.zeros_like(bin_index_ic)
299 for c in range(3):
300 if pbc[c]:
301 # (Note: np.divmod does not exist in older numpies)
302 cell_shift_ic[:, c], bin_index_ic[:, c] = \
303 divmod(bin_index_ic[:, c], nbins_c[c])
304 else:
305 bin_index_ic[:, c] = np.clip(bin_index_ic[:, c], 0, nbins_c[c] - 1)
307 # Convert Cartesian bin index to unique scalar bin index.
308 bin_index_i = (bin_index_ic[:, 0] +
309 nbins_c[0] * (bin_index_ic[:, 1] +
310 nbins_c[1] * bin_index_ic[:, 2]))
312 # atom_i contains atom index in new sort order.
313 atom_i = np.argsort(bin_index_i)
314 bin_index_i = bin_index_i[atom_i]
316 # Find max number of atoms per bin
317 max_natoms_per_bin = np.bincount(bin_index_i).max()
319 # Sort atoms into bins: atoms_in_bin_ba contains for each bin (identified
320 # by its scalar bin index) a list of atoms inside that bin. This list is
321 # homogeneous, i.e. has the same size *max_natoms_per_bin* for all bins.
322 # The list is padded with -1 values.
323 atoms_in_bin_ba = -np.ones([nbins, max_natoms_per_bin], dtype=int)
324 for i in range(max_natoms_per_bin):
325 # Create a mask array that identifies the first atom of each bin.
326 mask = np.append([True], bin_index_i[:-1] != bin_index_i[1:])
327 # Assign all first atoms.
328 atoms_in_bin_ba[bin_index_i[mask], i] = atom_i[mask]
330 # Remove atoms that we just sorted into atoms_in_bin_ba. The next
331 # "first" atom will be the second and so on.
332 mask = np.logical_not(mask)
333 atom_i = atom_i[mask]
334 bin_index_i = bin_index_i[mask]
336 # Make sure that all atoms have been sorted into bins.
337 assert len(atom_i) == 0
338 assert len(bin_index_i) == 0
340 # Now we construct neighbor pairs by pairing up all atoms within a bin or
341 # between bin and neighboring bin. atom_pairs_pn is a helper buffer that
342 # contains all potential pairs of atoms between two bins, i.e. it is a list
343 # of length max_natoms_per_bin**2.
344 atom_pairs_pn = np.indices((max_natoms_per_bin, max_natoms_per_bin),
345 dtype=int)
346 atom_pairs_pn = atom_pairs_pn.reshape(2, -1)
348 # Initialized empty neighbor list buffers.
349 first_at_neightuple_nn = []
350 secnd_at_neightuple_nn = []
351 cell_shift_vector_x_n = []
352 cell_shift_vector_y_n = []
353 cell_shift_vector_z_n = []
355 # This is the main neighbor list search. We loop over neighboring bins and
356 # then construct all possible pairs of atoms between two bins, assuming
357 # that each bin contains exactly max_natoms_per_bin atoms. We then throw
358 # out pairs involving pad atoms with atom index -1 below.
359 binz_xyz, biny_xyz, binx_xyz = np.meshgrid(np.arange(nbins_c[2]),
360 np.arange(nbins_c[1]),
361 np.arange(nbins_c[0]),
362 indexing='ij')
363 # The memory layout of binx_xyz, biny_xyz, binz_xyz is such that computing
364 # the respective bin index leads to a linearly increasing consecutive list.
365 # The following assert statement succeeds:
366 # b_b = (binx_xyz + nbins_c[0] * (biny_xyz + nbins_c[1] *
367 # binz_xyz)).ravel()
368 # assert (b_b == np.arange(np.prod(nbins_c))).all()
370 # First atoms in pair.
371 _first_at_neightuple_n = atoms_in_bin_ba[:, atom_pairs_pn[0]]
372 for dz in range(-neigh_search_z, neigh_search_z + 1):
373 for dy in range(-neigh_search_y, neigh_search_y + 1):
374 for dx in range(-neigh_search_x, neigh_search_x + 1):
375 # Bin index of neighboring bin and shift vector.
376 shiftx_xyz, neighbinx_xyz = divmod(binx_xyz + dx, nbins_c[0])
377 shifty_xyz, neighbiny_xyz = divmod(biny_xyz + dy, nbins_c[1])
378 shiftz_xyz, neighbinz_xyz = divmod(binz_xyz + dz, nbins_c[2])
379 neighbin_b = (neighbinx_xyz + nbins_c[0] *
380 (neighbiny_xyz + nbins_c[1] * neighbinz_xyz)
381 ).ravel()
383 # Second atom in pair.
384 _secnd_at_neightuple_n = \
385 atoms_in_bin_ba[neighbin_b][:, atom_pairs_pn[1]]
387 # Shift vectors.
388 _cell_shift_vector_x_n = \
389 np.resize(shiftx_xyz.reshape(-1, 1),
390 (max_natoms_per_bin**2, shiftx_xyz.size)).T
391 _cell_shift_vector_y_n = \
392 np.resize(shifty_xyz.reshape(-1, 1),
393 (max_natoms_per_bin**2, shifty_xyz.size)).T
394 _cell_shift_vector_z_n = \
395 np.resize(shiftz_xyz.reshape(-1, 1),
396 (max_natoms_per_bin**2, shiftz_xyz.size)).T
398 # We have created too many pairs because we assumed each bin
399 # has exactly max_natoms_per_bin atoms. Remove all surperfluous
400 # pairs. Those are pairs that involve an atom with index -1.
401 mask = np.logical_and(_first_at_neightuple_n != -1,
402 _secnd_at_neightuple_n != -1)
403 if mask.sum() > 0:
404 first_at_neightuple_nn += [_first_at_neightuple_n[mask]]
405 secnd_at_neightuple_nn += [_secnd_at_neightuple_n[mask]]
406 cell_shift_vector_x_n += [_cell_shift_vector_x_n[mask]]
407 cell_shift_vector_y_n += [_cell_shift_vector_y_n[mask]]
408 cell_shift_vector_z_n += [_cell_shift_vector_z_n[mask]]
410 # Flatten overall neighbor list.
411 first_at_neightuple_n = np.concatenate(first_at_neightuple_nn)
412 secnd_at_neightuple_n = np.concatenate(secnd_at_neightuple_nn)
413 cell_shift_vector_n = np.transpose([np.concatenate(cell_shift_vector_x_n),
414 np.concatenate(cell_shift_vector_y_n),
415 np.concatenate(cell_shift_vector_z_n)])
417 # Add global cell shift to shift vectors
418 cell_shift_vector_n += cell_shift_ic[first_at_neightuple_n] - \
419 cell_shift_ic[secnd_at_neightuple_n]
421 # Remove all self-pairs that do not cross the cell boundary.
422 if not self_interaction:
423 m = np.logical_not(np.logical_and(
424 first_at_neightuple_n == secnd_at_neightuple_n,
425 (cell_shift_vector_n == 0).all(axis=1)))
426 first_at_neightuple_n = first_at_neightuple_n[m]
427 secnd_at_neightuple_n = secnd_at_neightuple_n[m]
428 cell_shift_vector_n = cell_shift_vector_n[m]
430 # For nonperiodic directions, remove any bonds that cross the domain
431 # boundary.
432 for c in range(3):
433 if not pbc[c]:
434 m = cell_shift_vector_n[:, c] == 0
435 first_at_neightuple_n = first_at_neightuple_n[m]
436 secnd_at_neightuple_n = secnd_at_neightuple_n[m]
437 cell_shift_vector_n = cell_shift_vector_n[m]
439 # Sort neighbor list.
440 i = np.argsort(first_at_neightuple_n)
441 first_at_neightuple_n = first_at_neightuple_n[i]
442 secnd_at_neightuple_n = secnd_at_neightuple_n[i]
443 cell_shift_vector_n = cell_shift_vector_n[i]
445 # Compute distance vectors.
446 distance_vector_nc = positions[secnd_at_neightuple_n] - \
447 positions[first_at_neightuple_n] + \
448 cell_shift_vector_n.dot(cell)
449 abs_distance_vector_n = \
450 np.sqrt(np.sum(distance_vector_nc * distance_vector_nc, axis=1))
452 # We have still created too many pairs. Only keep those with distance
453 # smaller than max_cutoff.
454 mask = abs_distance_vector_n < max_cutoff
455 first_at_neightuple_n = first_at_neightuple_n[mask]
456 secnd_at_neightuple_n = secnd_at_neightuple_n[mask]
457 cell_shift_vector_n = cell_shift_vector_n[mask]
458 distance_vector_nc = distance_vector_nc[mask]
459 abs_distance_vector_n = abs_distance_vector_n[mask]
461 if isinstance(cutoff, dict) and numbers is not None:
462 # If cutoff is a dictionary, then the cutoff radii are specified per
463 # element pair. We now have a list up to maximum cutoff.
464 per_pair_cutoff_n = np.zeros_like(abs_distance_vector_n)
465 for (atomic_number1, atomic_number2), c in cutoff.items():
466 try:
467 atomic_number1 = atomic_numbers[atomic_number1]
468 except KeyError:
469 pass
470 try:
471 atomic_number2 = atomic_numbers[atomic_number2]
472 except KeyError:
473 pass
474 if atomic_number1 == atomic_number2:
475 mask = np.logical_and(
476 numbers[first_at_neightuple_n] == atomic_number1,
477 numbers[secnd_at_neightuple_n] == atomic_number2)
478 else:
479 mask = np.logical_or(
480 np.logical_and(
481 numbers[first_at_neightuple_n] == atomic_number1,
482 numbers[secnd_at_neightuple_n] == atomic_number2),
483 np.logical_and(
484 numbers[first_at_neightuple_n] == atomic_number2,
485 numbers[secnd_at_neightuple_n] == atomic_number1))
486 per_pair_cutoff_n[mask] = c
487 mask = abs_distance_vector_n < per_pair_cutoff_n
488 first_at_neightuple_n = first_at_neightuple_n[mask]
489 secnd_at_neightuple_n = secnd_at_neightuple_n[mask]
490 cell_shift_vector_n = cell_shift_vector_n[mask]
491 distance_vector_nc = distance_vector_nc[mask]
492 abs_distance_vector_n = abs_distance_vector_n[mask]
493 elif not np.isscalar(cutoff):
494 # If cutoff is neither a dictionary nor a scalar, then we assume it is
495 # a list or numpy array that contains atomic radii. Atoms are neighbors
496 # if their radii overlap.
497 mask = abs_distance_vector_n < \
498 cutoff[first_at_neightuple_n] + cutoff[secnd_at_neightuple_n]
499 first_at_neightuple_n = first_at_neightuple_n[mask]
500 secnd_at_neightuple_n = secnd_at_neightuple_n[mask]
501 cell_shift_vector_n = cell_shift_vector_n[mask]
502 distance_vector_nc = distance_vector_nc[mask]
503 abs_distance_vector_n = abs_distance_vector_n[mask]
505 # Assemble return tuple.
506 retvals = []
507 for q in quantities:
508 if q == 'i':
509 retvals += [first_at_neightuple_n]
510 elif q == 'j':
511 retvals += [secnd_at_neightuple_n]
512 elif q == 'D':
513 retvals += [distance_vector_nc]
514 elif q == 'd':
515 retvals += [abs_distance_vector_n]
516 elif q == 'S':
517 retvals += [cell_shift_vector_n]
518 else:
519 raise ValueError('Unsupported quantity specified.')
520 if len(retvals) == 1:
521 return retvals[0]
522 else:
523 return tuple(retvals)
526def neighbor_list(quantities, a, cutoff, self_interaction=False,
527 max_nbins=1e6):
528 """Compute a neighbor list for an atomic configuration.
530 Atoms outside periodic boundaries are mapped into the box. Atoms
531 outside nonperiodic boundaries are included in the neighbor list
532 but complexity of neighbor list search for those can become n^2.
534 The neighbor list is sorted by first atom index 'i', but not by second
535 atom index 'j'.
537 Parameters
538 ----------
539 quantities: str
540 Quantities to compute by the neighbor list algorithm. Each character
541 in this string defines a quantity. They are returned in a tuple of
542 the same order. Possible quantities are:
544 * 'i' : first atom index
545 * 'j' : second atom index
546 * 'd' : absolute distance
547 * 'D' : distance vector
548 * 'S' : shift vector (number of cell boundaries crossed by the bond
549 between atom i and j). With the shift vector S, the
550 distances D between atoms can be computed from:
551 D = a.positions[j]-a.positions[i]+S.dot(a.cell)
552 a: :class:`ase.Atoms`
553 Atomic configuration.
554 cutoff: float or dict
555 Cutoff for neighbor search. It can be:
557 * A single float: This is a global cutoff for all elements.
558 * A dictionary: This specifies cutoff values for element
559 pairs. Specification accepts element numbers of symbols.
560 Example: {(1, 6): 1.1, (1, 1): 1.0, ('C', 'C'): 1.85}
561 * A list/array with a per atom value: This specifies the radius of
562 an atomic sphere for each atoms. If spheres overlap, atoms are
563 within each others neighborhood. See
564 :func:`~ase.neighborlist.natural_cutoffs`
565 for an example on how to get such a list.
567 self_interaction: bool
568 Return the atom itself as its own neighbor if set to true.
569 Default: False
570 max_nbins: int
571 Maximum number of bins used in neighbor search. This is used to limit
572 the maximum amount of memory required by the neighbor list.
574 Returns
575 -------
576 i, j, ...: array
577 Tuple with arrays for each quantity specified above. Indices in `i`
578 are returned in ascending order 0..len(a), but the order of (i,j)
579 pairs is not guaranteed.
581 Examples
582 --------
584 >>> import numpy as np
585 >>> from ase.build import bulk, molecule
587 1. Coordination counting
589 >>> atoms = molecule('isobutane')
590 >>> i = neighbor_list('i', atoms, 1.85)
591 >>> coord = np.bincount(i, minlength=len(atoms))
593 2. Coordination counting with different cutoffs for each pair of species
595 >>> cutoff = {('H', 'H'): 1.1, ('C', 'H'): 1.3, ('C', 'C'): 1.85}
596 >>> i = neighbor_list('i', atoms, cutoff)
597 >>> coord = np.bincount(i, minlength=len(atoms))
599 3. Pair distribution function
601 >>> atoms = bulk('Cu', cubic=True) * 3
602 >>> atoms.rattle(0.5, rng=np.random.default_rng(42))
603 >>> cutoff = 5.0
604 >>> d = neighbor_list('d', atoms, cutoff)
605 >>> hist, bin_edges = np.histogram(d, bins=100, range=(0.0, cutoff))
606 >>> hist = hist / len(atoms) # per atom
607 >>> rho_mean = len(atoms) / atoms.cell.volume
608 >>> dv = 4.0 * np.pi * (bin_edges[1:] ** 3 - bin_edges[:-1] ** 3) / 3.0
609 >>> rho = hist / dv
610 >>> pdf = rho / rho_mean
612 4. Forces of a pair potential
614 >>> natoms = len(atoms)
615 >>> i, j, d, D = neighbor_list('ijdD', atoms, 5.0)
616 >>> # Lennard-Jones potential
617 >>> eps = 1.0
618 >>> sgm = 1.0
619 >>> epairs = 4.0 * eps * ((sgm / d) ** 12 - (sgm / d) ** 6)
620 >>> energy = 0.5 * epairs.sum() # correct double-counting
621 >>> dd = -4.0 * eps * (12 * (sgm / d) ** 13 - 6 * (sgm / d) ** 7) / sgm
622 >>> dd = (dd * (D.T / d)).T
623 >>> fx = -1.0 * np.bincount(i, weights=dd[:, 0], minlength=natoms)
624 >>> fy = -1.0 * np.bincount(i, weights=dd[:, 1], minlength=natoms)
625 >>> fz = -1.0 * np.bincount(i, weights=dd[:, 2], minlength=natoms)
627 5. Force-constant matrix of a pair potential
629 >>> i, j, d, D = neighbor_list('ijdD', atoms, 5.0)
630 >>> epairs = 1.0 / d # Coulomb potential
631 >>> forces = (D.T / d**3).T # (npairs, 3)
632 >>> # second derivative
633 >>> d2 = 3.0 * D[:, :, None] * D[:, None, :] / d[:, None, None] ** 5
634 >>> for k in range(3):
635 ... d2[:, k, k] -= 1.0 / d**3
636 >>> # force-constant matrix
637 >>> fc = np.zeros((natoms, 3, natoms, 3))
638 >>> for ia in range(natoms):
639 ... for ja in range(natoms):
640 ... fc[ia, :, ja, :] -= d2[(i == ia) & (j == ja), :, :].sum(axis=0)
641 >>> for ia in range(natoms):
642 ... fc[ia, :, ia, :] -= fc[ia].sum(axis=1)
644 """
645 return primitive_neighbor_list(quantities, a.pbc,
646 a.get_cell(complete=True),
647 a.positions, cutoff, numbers=a.numbers,
648 self_interaction=self_interaction,
649 max_nbins=max_nbins)
652def first_neighbors(natoms, first_atom):
653 """
654 Compute an index array pointing to the ranges within the neighbor list that
655 contain the neighbors for a certain atom.
657 Parameters:
659 natoms : int
660 Total number of atom.
661 first_atom : array_like
662 Array containing the first atom 'i' of the neighbor tuple returned
663 by the neighbor list.
665 Returns:
667 seed : array
668 Array containing pointers to the start and end location of the
669 neighbors of a certain atom. Neighbors of atom k have indices from s[k]
670 to s[k+1]-1.
671 """
672 if len(first_atom) == 0:
673 return np.zeros(natoms + 1, dtype=int)
674 # Create a seed array (which is returned by this function) populated with
675 # -1.
676 seed = -np.ones(natoms + 1, dtype=int)
678 first_atom = np.asarray(first_atom)
680 # Mask array contains all position where the number in the (sorted) array
681 # with first atoms (in the neighbor pair) changes.
682 mask = first_atom[:-1] != first_atom[1:]
684 # Seed array needs to start at 0
685 seed[first_atom[0]] = 0
686 # Seed array needs to stop at the length of the neighbor list
687 seed[-1] = len(first_atom)
688 # Populate all intermediate seed with the index of where the mask array is
689 # true, i.e. the index where the first_atom array changes.
690 seed[first_atom[1:][mask]] = (np.arange(len(mask)) + 1)[mask]
692 # Now fill all remaining -1 value with the value in the seed array right
693 # behind them. (There are no neighbor so seed[i] and seed[i+1] must point)
694 # to the same index.
695 mask = seed == -1
696 while mask.any():
697 seed[mask] = seed[np.arange(natoms + 1)[mask] + 1]
698 mask = seed == -1
699 return seed
702def get_connectivity_matrix(nl, sparse=True):
703 """Return connectivity matrix for a given NeighborList (dtype=numpy.int8).
705 A matrix of shape (nAtoms, nAtoms) will be returned.
706 Connected atoms i and j will have matrix[i,j] == 1, unconnected
707 matrix[i,j] == 0. If bothways=True the matrix will be symmetric,
708 otherwise not!
710 If *sparse* is True, a scipy csr matrix is returned.
711 If *sparse* is False, a numpy matrix is returned.
713 Note that the old and new neighborlists might give different results
714 for periodic systems if bothways=False.
716 Example:
718 Determine which molecule in a system atom 1 belongs to.
720 >>> from scipy import sparse
722 >>> from ase.build import molecule
723 >>> from ase.neighborlist import get_connectivity_matrix
724 >>> from ase.neighborlist import natural_cutoffs
725 >>> from ase.neighborlist import NeighborList
727 >>> mol = molecule('CH3CH2OH')
728 >>> cutoffs = natural_cutoffs(mol)
729 >>> neighbor_list = NeighborList(
730 ... cutoffs, self_interaction=False, bothways=True)
731 >>> neighbor_list.update(mol)
732 True
734 >>> matrix = neighbor_list.get_connectivity_matrix()
735 >>> # or: matrix = get_connectivity_matrix(neighbor_list.nl)
736 >>> n_components, component_list = sparse.csgraph.connected_components(
737 ... matrix)
738 >>> idx = 1
739 >>> molIdx = component_list[idx]
740 >>> print("There are {} molecules in the system".format(n_components))
741 There are 1 molecules in the system
742 >>> print("Atom {} is part of molecule {}".format(idx, molIdx))
743 Atom 1 is part of molecule 0
744 >>> molIdxs = [i for i in range(len(component_list))
745 ... if component_list[i] == molIdx]
746 >>> print("Atoms are part of molecule {}: {}".format(molIdx, molIdxs))
747 Atoms are part of molecule 0: [0, 1, 2, 3, 4, 5, 6, 7, 8]
748 """
750 nAtoms = len(nl.cutoffs)
752 if nl.nupdates <= 0:
753 raise RuntimeError(
754 'Must call update(atoms) on your neighborlist first!')
756 if sparse:
757 matrix = sp.dok_matrix((nAtoms, nAtoms), dtype=np.int8)
758 else:
759 matrix = np.zeros((nAtoms, nAtoms), dtype=np.int8)
761 for i in range(nAtoms):
762 for idx in nl.get_neighbors(i)[0]:
763 matrix[i, idx] = 1
765 return matrix
768class NewPrimitiveNeighborList:
769 """Neighbor list object. Wrapper around neighbor_list and first_neighbors.
771 cutoffs: list of float
772 List of cutoff radii - one for each atom. If the spheres (defined by
773 their cutoff radii) of two atoms overlap, they will be counted as
774 neighbors.
775 skin: float
776 If no atom has moved more than the skin-distance since the
777 last call to the
778 :meth:`~ase.neighborlist.NewPrimitiveNeighborList.update()`
779 method, then the neighbor list can be reused. This will save
780 some expensive rebuilds of the list, but extra neighbors outside
781 the cutoff will be returned.
782 sorted: bool
783 Sort neighbor list.
784 self_interaction: bool
785 Should an atom return itself as a neighbor?
786 bothways: bool
787 Return all neighbors. Default is to return only "half" of
788 the neighbors.
790 Example:
792 >>> from ase.build import bulk
793 >>> from ase.neighborlist import NewPrimitiveNeighborList
795 >>> nl = NewPrimitiveNeighborList([2.3, 1.7])
796 >>> atoms = bulk('Cu', 'fcc', a=3.6)
797 >>> nl.update(atoms.pbc, atoms.get_cell(), atoms.positions)
798 True
799 >>> indices, offsets = nl.get_neighbors(0)
800 """
802 def __init__(self, cutoffs, skin=0.3, sorted=False, self_interaction=True,
803 bothways=False, use_scaled_positions=False):
804 self.cutoffs = np.asarray(cutoffs) + skin
805 self.skin = skin
806 self.sorted = sorted
807 self.self_interaction = self_interaction
808 self.bothways = bothways
809 self.nupdates = 0
810 self.use_scaled_positions = use_scaled_positions
812 def update(self, pbc, cell, positions, numbers=None):
813 """Make sure the list is up to date."""
815 if self.nupdates == 0:
816 self.build(pbc, cell, positions, numbers=numbers)
817 return True
819 if ((self.pbc != pbc).any() or (self.cell != cell).any() or
820 ((self.positions - positions)**2).sum(1).max() > self.skin**2):
821 self.build(pbc, cell, positions, numbers=numbers)
822 return True
824 return False
826 def build(self, pbc, cell, positions, numbers=None):
827 """Build the list.
828 """
829 self.pbc = np.array(pbc, copy=True)
830 self.cell = np.array(cell, copy=True)
831 self.positions = np.array(positions, copy=True)
833 pair_first, pair_second, offset_vec = \
834 primitive_neighbor_list(
835 'ijS', pbc, cell, positions, self.cutoffs, numbers=numbers,
836 self_interaction=self.self_interaction,
837 use_scaled_positions=self.use_scaled_positions)
839 if len(positions) > 0 and not self.bothways:
840 offset_x, offset_y, offset_z = offset_vec.T
842 mask = offset_z > 0
843 mask &= offset_y == 0
844 mask |= offset_y > 0
845 mask &= offset_x == 0
846 mask |= offset_x > 0
847 mask |= (pair_first <= pair_second) & (offset_vec == 0).all(axis=1)
849 pair_first = pair_first[mask]
850 pair_second = pair_second[mask]
851 offset_vec = offset_vec[mask]
853 if len(positions) > 0 and self.sorted:
854 mask = np.argsort(pair_first * len(pair_first) +
855 pair_second)
856 pair_first = pair_first[mask]
857 pair_second = pair_second[mask]
858 offset_vec = offset_vec[mask]
860 self.pair_first = pair_first
861 self.pair_second = pair_second
862 self.offset_vec = offset_vec
864 # Compute the index array point to the first neighbor
865 self.first_neigh = first_neighbors(len(positions), pair_first)
867 self.nupdates += 1
869 def get_neighbors(self, a):
870 """Return neighbors of atom number a.
872 A list of indices and offsets to neighboring atoms is
873 returned. The positions of the neighbor atoms can be
874 calculated like this:
876 >>> from ase.build import bulk
877 >>> from ase.neighborlist import NewPrimitiveNeighborList
879 >>> nl = NewPrimitiveNeighborList([2.3, 1.7])
880 >>> atoms = bulk('Cu', 'fcc', a=3.6)
881 >>> nl.update(atoms.pbc, atoms.get_cell(), atoms.positions)
882 True
883 >>> indices, offsets = nl.get_neighbors(0)
884 >>> for i, offset in zip(indices, offsets):
885 ... print(atoms.positions[i] + offset @ atoms.get_cell())
886 ... # doctest: +SKIP
888 Notice that if get_neighbors(a) gives atom b as a neighbor,
889 then get_neighbors(b) will not return a as a neighbor - unless
890 bothways=True was used."""
892 return (self.pair_second[self.first_neigh[a]:self.first_neigh[a + 1]],
893 self.offset_vec[self.first_neigh[a]:self.first_neigh[a + 1]])
896class PrimitiveNeighborList:
897 """Neighbor list that works without Atoms objects.
899 This is less fancy, but can be used to avoid conversions between
900 scaled and non-scaled coordinates which may affect cell offsets
901 through rounding errors.
903 Attributes
904 ----------
905 nupdates : int
906 Number of updated times.
907 """
909 def __init__(self, cutoffs, skin=0.3, sorted=False, self_interaction=True,
910 bothways=False, use_scaled_positions=False):
911 self.cutoffs = np.asarray(cutoffs) + skin
912 self.skin = skin
913 self.sorted = sorted
914 self.self_interaction = self_interaction
915 self.bothways = bothways
916 self.nupdates = 0
917 self.use_scaled_positions = use_scaled_positions
919 def update(self, pbc, cell, coordinates):
920 """Make sure the list is up to date.
922 Returns
923 -------
924 bool
925 True if the neighbor list is updated.
926 """
928 if self.nupdates == 0:
929 self.build(pbc, cell, coordinates)
930 return True
932 if ((self.pbc != pbc).any() or (self.cell != cell).any() or (
933 (self.coordinates
934 - coordinates)**2).sum(1).max() > self.skin**2):
935 self.build(pbc, cell, coordinates)
936 return True
938 return False
940 def build(self, pbc, cell, coordinates):
941 """Build the list.
943 Coordinates are taken to be scaled or not according
944 to self.use_scaled_positions.
945 """
946 self.pbc = pbc = np.array(pbc, copy=True)
947 self.cell = cell = Cell(cell)
948 self.coordinates = coordinates = np.array(coordinates, copy=True)
950 if len(self.cutoffs) != len(coordinates):
951 raise ValueError('Wrong number of cutoff radii: {} != {}'
952 .format(len(self.cutoffs), len(coordinates)))
954 rcmax = self.cutoffs.max() if len(self.cutoffs) > 0 else 0.0
956 if self.use_scaled_positions:
957 positions0 = cell.cartesian_positions(coordinates)
958 else:
959 positions0 = coordinates
961 rcell, op = minkowski_reduce(cell, pbc)
962 positions = wrap_positions(positions0, rcell, pbc=pbc, eps=0)
964 natoms = len(positions)
965 self.nupdates += 1
966 if natoms == 0:
967 self.neighbors = []
968 self.displacements = []
969 return
971 tree = cKDTree(positions, copy_data=True)
972 offsets = cell.scaled_positions(positions - positions0)
973 offsets = offsets.round().astype(int)
975 N = _calc_expansion(rcell, pbc, rcmax)
977 neighbor_indices_a = [[] for _ in range(natoms)]
978 displacements_a = [[] for _ in range(natoms)]
980 for n1, n2, n3 in itertools.product(range(N[0] + 1),
981 range(-N[1], N[1] + 1),
982 range(-N[2], N[2] + 1)):
983 if n1 == 0 and (n2 < 0 or n2 == 0 and n3 < 0):
984 continue
986 displacement = (n1, n2, n3) @ rcell
987 shift0 = (n1, n2, n3) @ op
988 indices_all = tree.query_ball_point(
989 positions - displacement,
990 r=self.cutoffs + rcmax,
991 )
993 for a in range(natoms):
994 indices = indices_all[a]
996 if not indices:
997 continue
999 indices = np.array(indices)
1000 delta = positions[indices] + displacement - positions[a]
1001 distances = np.sqrt(np.add.reduce(delta**2, axis=1))
1002 cutoffs = self.cutoffs[indices] + self.cutoffs[a]
1003 i = indices[distances < cutoffs]
1004 if n1 == 0 and n2 == 0 and n3 == 0:
1005 if self.self_interaction:
1006 i = i[i >= a]
1007 else:
1008 i = i[i > a]
1010 neighbor_indices_a[a].append(i)
1012 disp = shift0 + offsets[i] - offsets[a]
1013 displacements_a[a].append(disp)
1015 self.neighbors = [np.concatenate(i) for i in neighbor_indices_a]
1016 self.displacements = [np.concatenate(d) for d in displacements_a]
1018 if self.bothways:
1019 neighbors2 = [[] for a in range(natoms)]
1020 displacements2 = [[] for a in range(natoms)]
1021 for a in range(natoms):
1022 for b, disp in zip(self.neighbors[a], self.displacements[a]):
1023 # avoid double counting of self interaction
1024 if a == b and (disp == 0).all():
1025 continue
1026 neighbors2[b].append(a)
1027 displacements2[b].append(-disp)
1028 for a in range(natoms):
1029 nbs = np.concatenate((self.neighbors[a], neighbors2[a]))
1030 disp = np.array(list(self.displacements[a]) + displacements2[a])
1031 # Force correct type and shape for case of no neighbors:
1032 self.neighbors[a] = nbs.astype(int)
1033 self.displacements[a] = disp.astype(int).reshape((-1, 3))
1035 if self.sorted:
1036 _sort_neighbors(self.neighbors, self.displacements)
1038 def get_neighbors(self, a):
1039 """Return neighbors of atom number a.
1041 A list of indices and offsets to neighboring atoms is
1042 returned. The positions of the neighbor atoms can be
1043 calculated like this:
1045 >>> from ase.build import bulk
1046 >>> from ase.neighborlist import NewPrimitiveNeighborList
1048 >>> nl = NewPrimitiveNeighborList([2.3, 1.7])
1049 >>> atoms = bulk('Cu', 'fcc', a=3.6)
1050 >>> nl.update(atoms.pbc, atoms.get_cell(), atoms.positions)
1051 True
1052 >>> indices, offsets = nl.get_neighbors(0)
1053 >>> for i, offset in zip(indices, offsets):
1054 ... print(atoms.positions[i] + offset @ atoms.get_cell())
1055 ... # doctest: +SKIP
1057 Notice that if get_neighbors(a) gives atom b as a neighbor,
1058 then get_neighbors(b) will not return a as a neighbor - unless
1059 bothways=True was used."""
1061 return self.neighbors[a], self.displacements[a]
1064def _calc_expansion(rcell, pbc, rcmax):
1065 r"""Calculate expansion to contain a sphere of radius `2.0 * rcmax`.
1067 This function determines the minimum supercell (parallelepiped) that
1068 contains a sphere of radius `2.0 * rcmax`. For this, `a_1` is projected
1069 onto the unit vector perpendicular to `a_2 \times a_3` (i.e. the unit
1070 vector along the direction `b_1`) to know how many `a_1`'s the supercell
1071 takes to contain the sphere.
1072 """
1073 ircell = np.linalg.pinv(rcell)
1074 vs = np.sqrt(np.add.reduce(ircell**2, axis=0))
1075 ns = np.where(pbc, np.ceil(2.0 * rcmax * vs), 0.0)
1076 return ns.astype(int)
1079def _sort_neighbors(neighbors, offsets):
1080 """Sort neighbors first by indices and then offsets."""
1081 natoms = len(neighbors)
1082 for a in range(natoms):
1083 keys = (
1084 offsets[a][:, 2],
1085 offsets[a][:, 1],
1086 offsets[a][:, 0],
1087 neighbors[a]
1088 )
1089 mask = np.lexsort(keys)
1090 neighbors[a] = neighbors[a][mask]
1091 offsets[a] = offsets[a][mask]
1094class NeighborList:
1095 """Neighbor list object.
1097 cutoffs: list of float
1098 List of cutoff radii - one for each atom. If the spheres
1099 (defined by their cutoff radii) of two atoms overlap, they
1100 will be counted as neighbors. See
1101 :func:`~ase.neighborlist.natural_cutoffs` for an example on
1102 how to get such a list.
1104 skin: float
1105 If no atom has moved more than the skin-distance since the
1106 last call to the
1107 :meth:`~ase.neighborlist.NeighborList.update()` method, then
1108 the neighbor list can be reused. This will save some
1109 expensive rebuilds of the list, but extra neighbors outside
1110 the cutoff will be returned.
1111 self_interaction: bool
1112 Should an atom return itself as a neighbor?
1113 bothways: bool
1114 Return all neighbors. Default is to return only "half" of
1115 the neighbors.
1116 primitive: class
1117 Define which implementation to use. Older and quadratically-scaling
1118 :class:`~ase.neighborlist.PrimitiveNeighborList` or newer and
1119 linearly-scaling :class:`~ase.neighborlist.NewPrimitiveNeighborList`.
1121 Example:
1123 >>> from ase.build import molecule
1124 >>> from ase.neighborlist import NeighborList
1126 >>> atoms = molecule("CO")
1127 >>> nl = NeighborList([0.76, 0.66])
1128 >>> nl.update(atoms)
1129 True
1130 >>> indices, offsets = nl.get_neighbors(0)
1132 """
1134 def __init__(self, cutoffs, skin=0.3, sorted=False, self_interaction=True,
1135 bothways=False, primitive=PrimitiveNeighborList):
1136 self.nl = primitive(cutoffs, skin, sorted,
1137 self_interaction=self_interaction,
1138 bothways=bothways)
1140 def update(self, atoms):
1141 """
1142 See :meth:`ase.neighborlist.PrimitiveNeighborList.update` or
1143 :meth:`ase.neighborlist.PrimitiveNeighborList.update`.
1144 """
1145 return self.nl.update(atoms.pbc, atoms.get_cell(complete=True),
1146 atoms.positions)
1148 def get_neighbors(self, a):
1149 """
1150 See :meth:`ase.neighborlist.PrimitiveNeighborList.get_neighbors` or
1151 :meth:`ase.neighborlist.PrimitiveNeighborList.get_neighbors`.
1152 """
1153 if self.nl.nupdates <= 0:
1154 raise RuntimeError('Must call update(atoms) on your neighborlist '
1155 'first!')
1157 return self.nl.get_neighbors(a)
1159 def get_connectivity_matrix(self, sparse=True):
1160 """
1161 See :func:`~ase.neighborlist.get_connectivity_matrix`.
1162 """
1163 return get_connectivity_matrix(self.nl, sparse)
1165 @property
1166 def nupdates(self):
1167 """Get number of updates."""
1168 return self.nl.nupdates
1170 @property
1171 @deprecated(
1172 'Use, e.g., `sum(_.size for _ in nl.neighbors)` '
1173 'for `bothways=False` and `self_interaction=False`.'
1174 )
1175 def nneighbors(self):
1176 """Get number of neighbors.
1178 .. deprecated:: 3.24.0
1179 """
1180 nneighbors = sum(indices.size for indices in self.nl.neighbors)
1181 if self.nl.self_interaction:
1182 nneighbors -= len(self.nl.neighbors)
1183 return nneighbors // 2 if self.nl.bothways else nneighbors
1185 @property
1186 @deprecated(
1187 'Use, e.g., `sum(_.any(1).sum() for _ in nl.displacements)` '
1188 'for `bothways=False` and `self_interaction=False`.'
1189 )
1190 def npbcneighbors(self):
1191 """Get number of pbc neighbors.
1193 .. deprecated:: 3.24.0
1194 """
1195 nneighbors = sum(
1196 offsets.any(axis=1).sum() for offsets in self.nl.displacements
1197 ) # sum up all neighbors that have non-zero supercell offsets
1198 return nneighbors // 2 if self.nl.bothways else nneighbors