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