Coverage for /builds/debichem-team/python-ase/ase/neighborlist.py: 95.77%
378 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-03-06 04:00 +0000
1import itertools
3import numpy as np
4import scipy.sparse.csgraph as csgraph
5from scipy import sparse as sp
6from scipy.spatial import cKDTree
8from ase.cell import Cell
9from ase.data import atomic_numbers, covalent_radii
10from ase.geometry import (
11 complete_cell,
12 find_mic,
13 minkowski_reduce,
14 wrap_positions,
15)
16from ase.utils import deprecated
19def natural_cutoffs(atoms, mult=1, **kwargs):
20 """Generate a radial cutoff for every atom based on covalent radii.
22 The covalent radii are a reasonable cutoff estimation for bonds in
23 many applications such as neighborlists, so function generates an
24 atoms length list of radii based on this idea.
26 * atoms: An atoms object
27 * mult: A multiplier for all cutoffs, useful for coarse grained adjustment
28 * kwargs: Symbol of the atom and its corresponding cutoff,
29 used to override the covalent radii
30 """
31 return [kwargs.get(atom.symbol, covalent_radii[atom.number] * mult)
32 for atom in atoms]
35def build_neighbor_list(atoms, cutoffs=None, **kwargs):
36 """Automatically build and update a NeighborList.
38 Parameters:
40 atoms : :class:`~ase.Atoms` object
41 Atoms to build Neighborlist for.
42 cutoffs: list of floats
43 Radii for each atom. If not given it will be produced by calling
44 :func:`ase.neighborlist.natural_cutoffs`
45 kwargs: arbitrary number of options
46 Will be passed to the constructor of
47 :class:`~ase.neighborlist.NeighborList`
49 Returns:
51 return: :class:`~ase.neighborlist.NeighborList`
52 A :class:`~ase.neighborlist.NeighborList` instance (updated).
53 """
54 if cutoffs is None:
55 cutoffs = natural_cutoffs(atoms)
57 nl = NeighborList(cutoffs, **kwargs)
58 nl.update(atoms)
60 return nl
63def get_distance_matrix(graph, limit=3):
64 """Get Distance Matrix from a Graph.
66 Parameters:
68 graph: array, matrix or sparse matrix, 2 dimensions (N, N)
69 Graph representation of the connectivity.
70 See `scipy doc <https://docs.scipy.org/doc/scipy/reference/generated\
71/scipy.sparse.csgraph.dijkstra.html#scipy.sparse.csgraph.dijkstra>`_
72 for reference.
73 limit: integer
74 Maximum number of steps to analyze. For most molecular information,
75 three should be enough.
77 Returns:
79 return: scipy.sparse.csr_matrix, shape (N, N)
80 A scipy.sparse.csr_matrix. All elements that are not connected within
81 *limit* steps are set to zero.
83 This is a potential memory bottleneck, as csgraph.dijkstra produces a
84 dense output matrix. Here we replace all np.inf values with 0 and
85 transform back to csr_matrix.
86 Why not dok_matrix like the connectivity-matrix? Because row-picking
87 is most likely and this is super fast with csr.
88 """
89 mat = csgraph.dijkstra(graph, directed=False, limit=limit)
90 mat[mat == np.inf] = 0
91 return sp.csr_matrix(mat, dtype=np.int8)
94def get_distance_indices(distanceMatrix, distance):
95 """Get indices for each node that are distance or less away.
97 Parameters:
99 distanceMatrix: any one of scipy.sparse matrices (NxN)
100 Matrix containing distance information of atoms. Get it e.g. with
101 :func:`~ase.neighborlist.get_distance_matrix`
102 distance: integer
103 Number of steps to allow.
105 Returns:
107 return: list of length N
108 List of length N. return[i] has all indices connected to item i.
110 The distance matrix only contains shortest paths, so when looking for
111 distances longer than one, we need to add the lower values for cases
112 where atoms are connected via a shorter path too.
113 """
114 shape = distanceMatrix.get_shape()
115 indices = []
116 # iterate over rows
117 for i in range(shape[0]):
118 row = distanceMatrix.getrow(i)[0]
119 # find all non-zero
120 found = sp.find(row)
121 # screen for smaller or equal distance
122 equal = np.where(found[-1] <= distance)[0]
123 # found[1] contains the indexes
124 indices.append([found[1][x] for x in equal])
125 return indices
128def mic(dr, cell, pbc=True):
129 """
130 Apply minimum image convention to an array of distance vectors.
132 Parameters:
134 dr : array_like
135 Array of distance vectors.
136 cell : array_like
137 Simulation cell.
138 pbc : array_like, optional
139 Periodic boundary conditions in x-, y- and z-direction. Default is to
140 assume periodic boundaries in all directions.
142 Returns:
144 dr : array
145 Array of distance vectors, wrapped according to the minimum image
146 convention.
147 """
148 dr, _ = find_mic(dr, cell, pbc)
149 return dr
152def primitive_neighbor_list(quantities, pbc, cell, positions, cutoff,
153 numbers=None, self_interaction=False,
154 use_scaled_positions=False, max_nbins=1e6):
155 """Compute a neighbor list for an atomic configuration.
157 Atoms outside periodic boundaries are mapped into the box. Atoms
158 outside nonperiodic boundaries are included in the neighbor list
159 but complexity of neighbor list search for those can become n^2.
161 The neighbor list is sorted by first atom index 'i', but not by second
162 atom index 'j'.
164 Parameters:
166 quantities: str
167 Quantities to compute by the neighbor list algorithm. Each character
168 in this string defines a quantity. They are returned in a tuple of
169 the same order. Possible quantities are
171 * 'i' : first atom index
172 * 'j' : second atom index
173 * 'd' : absolute distance
174 * 'D' : distance vector
175 * 'S' : shift vector (number of cell boundaries crossed by the bond
176 between atom i and j). With the shift vector S, the
177 distances D between atoms can be computed from:
178 D = positions[j]-positions[i]+S.dot(cell)
179 pbc: array_like
180 3-tuple indicating giving periodic boundaries in the three Cartesian
181 directions.
182 cell: 3x3 matrix
183 Unit cell vectors.
184 positions: list of xyz-positions
185 Atomic positions. Anything that can be converted to an ndarray of
186 shape (n, 3) will do: [(x1,y1,z1), (x2,y2,z2), ...]. If
187 use_scaled_positions is set to true, this must be scaled positions.
188 cutoff: float or dict
189 Cutoff for neighbor search. It can be:
191 * A single float: This is a global cutoff for all elements.
192 * A dictionary: This specifies cutoff values for element
193 pairs. Specification accepts element numbers of symbols.
194 Example: {(1, 6): 1.1, (1, 1): 1.0, ('C', 'C'): 1.85}
195 * A list/array with a per atom value: This specifies the radius of
196 an atomic sphere for each atoms. If spheres overlap, atoms are
197 within each others neighborhood. See
198 :func:`~ase.neighborlist.natural_cutoffs`
199 for an example on how to get such a list.
200 self_interaction: bool
201 Return the atom itself as its own neighbor if set to true.
202 Default: False
203 use_scaled_positions: bool
204 If set to true, positions are expected to be scaled positions.
205 max_nbins: int
206 Maximum number of bins used in neighbor search. This is used to limit
207 the maximum amount of memory required by the neighbor list.
209 Returns:
211 i, j, ... : array
212 Tuple with arrays for each quantity specified above. Indices in `i`
213 are returned in ascending order 0..len(a)-1, but the order of (i,j)
214 pairs is not guaranteed.
216 """
218 # Naming conventions: Suffixes indicate the dimension of an array. The
219 # following convention is used here:
220 # c: Cartesian index, can have values 0, 1, 2
221 # i: Global atom index, can have values 0..len(a)-1
222 # xyz: Bin index, three values identifying x-, y- and z-component of a
223 # spatial bin that is used to make neighbor search O(n)
224 # b: Linearized version of the 'xyz' bin index
225 # a: Bin-local atom index, i.e. index identifying an atom *within* a
226 # bin
227 # p: Pair index, can have value 0 or 1
228 # n: (Linear) neighbor index
230 # Return empty neighbor list if no atoms are passed here
231 if len(positions) == 0:
232 empty_types = dict(i=(int, (0, )),
233 j=(int, (0, )),
234 D=(float, (0, 3)),
235 d=(float, (0, )),
236 S=(int, (0, 3)))
237 retvals = []
238 for i in quantities:
239 dtype, shape = empty_types[i]
240 retvals += [np.array([], dtype=dtype).reshape(shape)]
241 if len(retvals) == 1:
242 return retvals[0]
243 else:
244 return tuple(retvals)
246 # Compute reciprocal lattice vectors.
247 b1_c, b2_c, b3_c = np.linalg.pinv(cell).T
249 # Compute distances of cell faces.
250 l1 = np.linalg.norm(b1_c)
251 l2 = np.linalg.norm(b2_c)
252 l3 = np.linalg.norm(b3_c)
253 face_dist_c = np.array([1 / l1 if l1 > 0 else 1,
254 1 / l2 if l2 > 0 else 1,
255 1 / l3 if l3 > 0 else 1])
257 if isinstance(cutoff, dict):
258 max_cutoff = max(cutoff.values())
259 else:
260 if np.isscalar(cutoff):
261 max_cutoff = cutoff
262 else:
263 cutoff = np.asarray(cutoff)
264 max_cutoff = 2 * np.max(cutoff)
266 # We use a minimum bin size of 3 A
267 bin_size = max(max_cutoff, 3)
268 # Compute number of bins such that a sphere of radius cutoff fits into
269 # eight neighboring bins.
270 nbins_c = np.maximum((face_dist_c / bin_size).astype(int), [1, 1, 1])
271 nbins = np.prod(nbins_c)
272 # Make sure we limit the amount of memory used by the explicit bins.
273 while nbins > max_nbins:
274 nbins_c = np.maximum(nbins_c // 2, [1, 1, 1])
275 nbins = np.prod(nbins_c)
277 # Compute over how many bins we need to loop in the neighbor list search.
278 neigh_search_x, neigh_search_y, neigh_search_z = \
279 np.ceil(bin_size * nbins_c / face_dist_c).astype(int)
281 # If we only have a single bin and the system is not periodic, then we
282 # do not need to search neighboring bins
283 neigh_search_x = 0 if nbins_c[0] == 1 and not pbc[0] else neigh_search_x
284 neigh_search_y = 0 if nbins_c[1] == 1 and not pbc[1] else neigh_search_y
285 neigh_search_z = 0 if nbins_c[2] == 1 and not pbc[2] else neigh_search_z
287 # Sort atoms into bins.
288 if use_scaled_positions:
289 scaled_positions_ic = positions
290 positions = np.dot(scaled_positions_ic, cell)
291 else:
292 scaled_positions_ic = np.linalg.solve(complete_cell(cell).T,
293 positions.T).T
294 bin_index_ic = np.floor(scaled_positions_ic * nbins_c).astype(int)
295 cell_shift_ic = np.zeros_like(bin_index_ic)
297 for c in range(3):
298 if pbc[c]:
299 # (Note: np.divmod does not exist in older numpies)
300 cell_shift_ic[:, c], bin_index_ic[:, c] = \
301 divmod(bin_index_ic[:, c], nbins_c[c])
302 else:
303 bin_index_ic[:, c] = np.clip(bin_index_ic[:, c], 0, nbins_c[c] - 1)
305 # Convert Cartesian bin index to unique scalar bin index.
306 bin_index_i = (bin_index_ic[:, 0] +
307 nbins_c[0] * (bin_index_ic[:, 1] +
308 nbins_c[1] * bin_index_ic[:, 2]))
310 # atom_i contains atom index in new sort order.
311 atom_i = np.argsort(bin_index_i)
312 bin_index_i = bin_index_i[atom_i]
314 # Find max number of atoms per bin
315 max_natoms_per_bin = np.bincount(bin_index_i).max()
317 # Sort atoms into bins: atoms_in_bin_ba contains for each bin (identified
318 # by its scalar bin index) a list of atoms inside that bin. This list is
319 # homogeneous, i.e. has the same size *max_natoms_per_bin* for all bins.
320 # The list is padded with -1 values.
321 atoms_in_bin_ba = -np.ones([nbins, max_natoms_per_bin], dtype=int)
322 for i in range(max_natoms_per_bin):
323 # Create a mask array that identifies the first atom of each bin.
324 mask = np.append([True], bin_index_i[:-1] != bin_index_i[1:])
325 # Assign all first atoms.
326 atoms_in_bin_ba[bin_index_i[mask], i] = atom_i[mask]
328 # Remove atoms that we just sorted into atoms_in_bin_ba. The next
329 # "first" atom will be the second and so on.
330 mask = np.logical_not(mask)
331 atom_i = atom_i[mask]
332 bin_index_i = bin_index_i[mask]
334 # Make sure that all atoms have been sorted into bins.
335 assert len(atom_i) == 0
336 assert len(bin_index_i) == 0
338 # Now we construct neighbor pairs by pairing up all atoms within a bin or
339 # between bin and neighboring bin. atom_pairs_pn is a helper buffer that
340 # contains all potential pairs of atoms between two bins, i.e. it is a list
341 # of length max_natoms_per_bin**2.
342 atom_pairs_pn = np.indices((max_natoms_per_bin, max_natoms_per_bin),
343 dtype=int)
344 atom_pairs_pn = atom_pairs_pn.reshape(2, -1)
346 # Initialized empty neighbor list buffers.
347 first_at_neightuple_nn = []
348 secnd_at_neightuple_nn = []
349 cell_shift_vector_x_n = []
350 cell_shift_vector_y_n = []
351 cell_shift_vector_z_n = []
353 # This is the main neighbor list search. We loop over neighboring bins and
354 # then construct all possible pairs of atoms between two bins, assuming
355 # that each bin contains exactly max_natoms_per_bin atoms. We then throw
356 # out pairs involving pad atoms with atom index -1 below.
357 binz_xyz, biny_xyz, binx_xyz = np.meshgrid(np.arange(nbins_c[2]),
358 np.arange(nbins_c[1]),
359 np.arange(nbins_c[0]),
360 indexing='ij')
361 # The memory layout of binx_xyz, biny_xyz, binz_xyz is such that computing
362 # the respective bin index leads to a linearly increasing consecutive list.
363 # The following assert statement succeeds:
364 # b_b = (binx_xyz + nbins_c[0] * (biny_xyz + nbins_c[1] *
365 # binz_xyz)).ravel()
366 # assert (b_b == np.arange(np.prod(nbins_c))).all()
368 # First atoms in pair.
369 _first_at_neightuple_n = atoms_in_bin_ba[:, atom_pairs_pn[0]]
370 for dz in range(-neigh_search_z, neigh_search_z + 1):
371 for dy in range(-neigh_search_y, neigh_search_y + 1):
372 for dx in range(-neigh_search_x, neigh_search_x + 1):
373 # Bin index of neighboring bin and shift vector.
374 shiftx_xyz, neighbinx_xyz = divmod(binx_xyz + dx, nbins_c[0])
375 shifty_xyz, neighbiny_xyz = divmod(biny_xyz + dy, nbins_c[1])
376 shiftz_xyz, neighbinz_xyz = divmod(binz_xyz + dz, nbins_c[2])
377 neighbin_b = (neighbinx_xyz + nbins_c[0] *
378 (neighbiny_xyz + nbins_c[1] * neighbinz_xyz)
379 ).ravel()
381 # Second atom in pair.
382 _secnd_at_neightuple_n = \
383 atoms_in_bin_ba[neighbin_b][:, atom_pairs_pn[1]]
385 # Shift vectors.
386 _cell_shift_vector_x_n = \
387 np.resize(shiftx_xyz.reshape(-1, 1),
388 (max_natoms_per_bin**2, shiftx_xyz.size)).T
389 _cell_shift_vector_y_n = \
390 np.resize(shifty_xyz.reshape(-1, 1),
391 (max_natoms_per_bin**2, shifty_xyz.size)).T
392 _cell_shift_vector_z_n = \
393 np.resize(shiftz_xyz.reshape(-1, 1),
394 (max_natoms_per_bin**2, shiftz_xyz.size)).T
396 # We have created too many pairs because we assumed each bin
397 # has exactly max_natoms_per_bin atoms. Remove all surperfluous
398 # pairs. Those are pairs that involve an atom with index -1.
399 mask = np.logical_and(_first_at_neightuple_n != -1,
400 _secnd_at_neightuple_n != -1)
401 if mask.sum() > 0:
402 first_at_neightuple_nn += [_first_at_neightuple_n[mask]]
403 secnd_at_neightuple_nn += [_secnd_at_neightuple_n[mask]]
404 cell_shift_vector_x_n += [_cell_shift_vector_x_n[mask]]
405 cell_shift_vector_y_n += [_cell_shift_vector_y_n[mask]]
406 cell_shift_vector_z_n += [_cell_shift_vector_z_n[mask]]
408 # Flatten overall neighbor list.
409 first_at_neightuple_n = np.concatenate(first_at_neightuple_nn)
410 secnd_at_neightuple_n = np.concatenate(secnd_at_neightuple_nn)
411 cell_shift_vector_n = np.transpose([np.concatenate(cell_shift_vector_x_n),
412 np.concatenate(cell_shift_vector_y_n),
413 np.concatenate(cell_shift_vector_z_n)])
415 # Add global cell shift to shift vectors
416 cell_shift_vector_n += cell_shift_ic[first_at_neightuple_n] - \
417 cell_shift_ic[secnd_at_neightuple_n]
419 # Remove all self-pairs that do not cross the cell boundary.
420 if not self_interaction:
421 m = np.logical_not(np.logical_and(
422 first_at_neightuple_n == secnd_at_neightuple_n,
423 (cell_shift_vector_n == 0).all(axis=1)))
424 first_at_neightuple_n = first_at_neightuple_n[m]
425 secnd_at_neightuple_n = secnd_at_neightuple_n[m]
426 cell_shift_vector_n = cell_shift_vector_n[m]
428 # For nonperiodic directions, remove any bonds that cross the domain
429 # boundary.
430 for c in range(3):
431 if not pbc[c]:
432 m = cell_shift_vector_n[:, c] == 0
433 first_at_neightuple_n = first_at_neightuple_n[m]
434 secnd_at_neightuple_n = secnd_at_neightuple_n[m]
435 cell_shift_vector_n = cell_shift_vector_n[m]
437 # Sort neighbor list.
438 i = np.argsort(first_at_neightuple_n)
439 first_at_neightuple_n = first_at_neightuple_n[i]
440 secnd_at_neightuple_n = secnd_at_neightuple_n[i]
441 cell_shift_vector_n = cell_shift_vector_n[i]
443 # Compute distance vectors.
444 distance_vector_nc = positions[secnd_at_neightuple_n] - \
445 positions[first_at_neightuple_n] + \
446 cell_shift_vector_n.dot(cell)
447 abs_distance_vector_n = \
448 np.sqrt(np.sum(distance_vector_nc * distance_vector_nc, axis=1))
450 # We have still created too many pairs. Only keep those with distance
451 # smaller than max_cutoff.
452 mask = abs_distance_vector_n < max_cutoff
453 first_at_neightuple_n = first_at_neightuple_n[mask]
454 secnd_at_neightuple_n = secnd_at_neightuple_n[mask]
455 cell_shift_vector_n = cell_shift_vector_n[mask]
456 distance_vector_nc = distance_vector_nc[mask]
457 abs_distance_vector_n = abs_distance_vector_n[mask]
459 if isinstance(cutoff, dict) and numbers is not None:
460 # If cutoff is a dictionary, then the cutoff radii are specified per
461 # element pair. We now have a list up to maximum cutoff.
462 per_pair_cutoff_n = np.zeros_like(abs_distance_vector_n)
463 for (atomic_number1, atomic_number2), c in cutoff.items():
464 try:
465 atomic_number1 = atomic_numbers[atomic_number1]
466 except KeyError:
467 pass
468 try:
469 atomic_number2 = atomic_numbers[atomic_number2]
470 except KeyError:
471 pass
472 if atomic_number1 == atomic_number2:
473 mask = np.logical_and(
474 numbers[first_at_neightuple_n] == atomic_number1,
475 numbers[secnd_at_neightuple_n] == atomic_number2)
476 else:
477 mask = np.logical_or(
478 np.logical_and(
479 numbers[first_at_neightuple_n] == atomic_number1,
480 numbers[secnd_at_neightuple_n] == atomic_number2),
481 np.logical_and(
482 numbers[first_at_neightuple_n] == atomic_number2,
483 numbers[secnd_at_neightuple_n] == atomic_number1))
484 per_pair_cutoff_n[mask] = c
485 mask = abs_distance_vector_n < per_pair_cutoff_n
486 first_at_neightuple_n = first_at_neightuple_n[mask]
487 secnd_at_neightuple_n = secnd_at_neightuple_n[mask]
488 cell_shift_vector_n = cell_shift_vector_n[mask]
489 distance_vector_nc = distance_vector_nc[mask]
490 abs_distance_vector_n = abs_distance_vector_n[mask]
491 elif not np.isscalar(cutoff):
492 # If cutoff is neither a dictionary nor a scalar, then we assume it is
493 # a list or numpy array that contains atomic radii. Atoms are neighbors
494 # if their radii overlap.
495 mask = abs_distance_vector_n < \
496 cutoff[first_at_neightuple_n] + cutoff[secnd_at_neightuple_n]
497 first_at_neightuple_n = first_at_neightuple_n[mask]
498 secnd_at_neightuple_n = secnd_at_neightuple_n[mask]
499 cell_shift_vector_n = cell_shift_vector_n[mask]
500 distance_vector_nc = distance_vector_nc[mask]
501 abs_distance_vector_n = abs_distance_vector_n[mask]
503 # Assemble return tuple.
504 retvals = []
505 for q in quantities:
506 if q == 'i':
507 retvals += [first_at_neightuple_n]
508 elif q == 'j':
509 retvals += [secnd_at_neightuple_n]
510 elif q == 'D':
511 retvals += [distance_vector_nc]
512 elif q == 'd':
513 retvals += [abs_distance_vector_n]
514 elif q == 'S':
515 retvals += [cell_shift_vector_n]
516 else:
517 raise ValueError('Unsupported quantity specified.')
518 if len(retvals) == 1:
519 return retvals[0]
520 else:
521 return tuple(retvals)
524def neighbor_list(quantities, a, cutoff, self_interaction=False,
525 max_nbins=1e6):
526 """Compute a neighbor list for an atomic configuration.
528 Atoms outside periodic boundaries are mapped into the box. Atoms
529 outside nonperiodic boundaries are included in the neighbor list
530 but complexity of neighbor list search for those can become n^2.
532 The neighbor list is sorted by first atom index 'i', but not by second
533 atom index 'j'.
535 Parameters:
537 quantities: str
538 Quantities to compute by the neighbor list algorithm. Each character
539 in this string defines a quantity. They are returned in a tuple of
540 the same order. Possible quantities are:
542 * 'i' : first atom index
543 * 'j' : second atom index
544 * 'd' : absolute distance
545 * 'D' : distance vector
546 * 'S' : shift vector (number of cell boundaries crossed by the bond
547 between atom i and j). With the shift vector S, the
548 distances D between atoms can be computed from:
549 D = a.positions[j]-a.positions[i]+S.dot(a.cell)
550 a: :class:`ase.Atoms`
551 Atomic configuration.
552 cutoff: float or dict
553 Cutoff for neighbor search. It can be:
555 * A single float: This is a global cutoff for all elements.
556 * A dictionary: This specifies cutoff values for element
557 pairs. Specification accepts element numbers of symbols.
558 Example: {(1, 6): 1.1, (1, 1): 1.0, ('C', 'C'): 1.85}
559 * A list/array with a per atom value: This specifies the radius of
560 an atomic sphere for each atoms. If spheres overlap, atoms are
561 within each others neighborhood. See
562 :func:`~ase.neighborlist.natural_cutoffs`
563 for an example on how to get such a list.
565 self_interaction: bool
566 Return the atom itself as its own neighbor if set to true.
567 Default: False
568 max_nbins: int
569 Maximum number of bins used in neighbor search. This is used to limit
570 the maximum amount of memory required by the neighbor list.
572 Returns:
574 i, j, ...: array
575 Tuple with arrays for each quantity specified above. Indices in `i`
576 are returned in ascending order 0..len(a), but the order of (i,j)
577 pairs is not guaranteed.
579 Examples:
581 Examples assume Atoms object *a* and numpy imported as *np*.
583 1. Coordination counting::
585 i = neighbor_list('i', a, 1.85)
586 coord = np.bincount(i)
588 2. Coordination counting with different cutoffs for each pair of species::
590 i = neighbor_list('i', a,
591 {('H', 'H'): 1.1, ('C', 'H'): 1.3, ('C', 'C'): 1.85})
592 coord = np.bincount(i)
594 3. Pair distribution function::
596 d = neighbor_list('d', a, 10.00)
597 h, bin_edges = np.histogram(d, bins=100)
598 pdf = h/(4*np.pi/3*(
599 bin_edges[1:]**3 - bin_edges[:-1]**3)) * a.get_volume()/len(a)
601 4. Pair potential::
603 i, j, d, D = neighbor_list('ijdD', a, 5.0)
604 energy = (-C/d**6).sum()
605 forces = (6*C/d**5 * (D/d).T).T
606 forces_x = np.bincount(j, weights=forces[:, 0], minlength=len(a)) - \
607 np.bincount(i, weights=forces[:, 0], minlength=len(a))
608 forces_y = np.bincount(j, weights=forces[:, 1], minlength=len(a)) - \
609 np.bincount(i, weights=forces[:, 1], minlength=len(a))
610 forces_z = np.bincount(j, weights=forces[:, 2], minlength=len(a)) - \
611 np.bincount(i, weights=pair_forces[:, 2], minlength=len(a))
613 5. Dynamical matrix for a pair potential stored in a block sparse format::
615 from scipy.sparse import bsr_matrix
616 i, j, dr, abs_dr = neighbor_list('ijDd', atoms)
617 energy = (dr.T / abs_dr).T
618 dynmat = -(dde * (energy.reshape(-1, 3, 1)
619 * energy.reshape(-1, 1, 3)).T).T \
620 -(de / abs_dr * (np.eye(3, dtype=energy.dtype) - \
621 (energy.reshape(-1, 3, 1) * energy.reshape(-1, 1, 3))).T).T
622 dynmat_bsr = bsr_matrix((dynmat, j, first_i),
623 shape=(3*len(a), 3*len(a)))
625 dynmat_diag = np.empty((len(a), 3, 3))
626 for x in range(3):
627 for y in range(3):
628 dynmat_diag[:, x, y] = -np.bincount(i, weights=dynmat[:, x, y])
630 dynmat_bsr += bsr_matrix((dynmat_diag, np.arange(len(a)),
631 np.arange(len(a) + 1)),
632 shape=(3 * len(a), 3 * len(a)))
634 """
635 return primitive_neighbor_list(quantities, a.pbc,
636 a.get_cell(complete=True),
637 a.positions, cutoff, numbers=a.numbers,
638 self_interaction=self_interaction,
639 max_nbins=max_nbins)
642def first_neighbors(natoms, first_atom):
643 """
644 Compute an index array pointing to the ranges within the neighbor list that
645 contain the neighbors for a certain atom.
647 Parameters:
649 natoms : int
650 Total number of atom.
651 first_atom : array_like
652 Array containing the first atom 'i' of the neighbor tuple returned
653 by the neighbor list.
655 Returns:
657 seed : array
658 Array containing pointers to the start and end location of the
659 neighbors of a certain atom. Neighbors of atom k have indices from s[k]
660 to s[k+1]-1.
661 """
662 if len(first_atom) == 0:
663 return np.zeros(natoms + 1, dtype=int)
664 # Create a seed array (which is returned by this function) populated with
665 # -1.
666 seed = -np.ones(natoms + 1, dtype=int)
668 first_atom = np.asarray(first_atom)
670 # Mask array contains all position where the number in the (sorted) array
671 # with first atoms (in the neighbor pair) changes.
672 mask = first_atom[:-1] != first_atom[1:]
674 # Seed array needs to start at 0
675 seed[first_atom[0]] = 0
676 # Seed array needs to stop at the length of the neighbor list
677 seed[-1] = len(first_atom)
678 # Populate all intermediate seed with the index of where the mask array is
679 # true, i.e. the index where the first_atom array changes.
680 seed[first_atom[1:][mask]] = (np.arange(len(mask)) + 1)[mask]
682 # Now fill all remaining -1 value with the value in the seed array right
683 # behind them. (There are no neighbor so seed[i] and seed[i+1] must point)
684 # to the same index.
685 mask = seed == -1
686 while mask.any():
687 seed[mask] = seed[np.arange(natoms + 1)[mask] + 1]
688 mask = seed == -1
689 return seed
692def get_connectivity_matrix(nl, sparse=True):
693 """Return connectivity matrix for a given NeighborList (dtype=numpy.int8).
695 A matrix of shape (nAtoms, nAtoms) will be returned.
696 Connected atoms i and j will have matrix[i,j] == 1, unconnected
697 matrix[i,j] == 0. If bothways=True the matrix will be symmetric,
698 otherwise not!
700 If *sparse* is True, a scipy csr matrix is returned.
701 If *sparse* is False, a numpy matrix is returned.
703 Note that the old and new neighborlists might give different results
704 for periodic systems if bothways=False.
706 Example:
708 Determine which molecule in a system atom 1 belongs to.
710 >>> from scipy import sparse
712 >>> from ase.build import molecule
713 >>> from ase.neighborlist import get_connectivity_matrix
714 >>> from ase.neighborlist import natural_cutoffs
715 >>> from ase.neighborlist import NeighborList
717 >>> mol = molecule('CH3CH2OH')
718 >>> cutoffs = natural_cutoffs(mol)
719 >>> neighbor_list = NeighborList(
720 ... cutoffs, self_interaction=False, bothways=True)
721 >>> neighbor_list.update(mol)
722 True
724 >>> matrix = neighbor_list.get_connectivity_matrix()
725 >>> # or: matrix = get_connectivity_matrix(neighbor_list.nl)
726 >>> n_components, component_list = sparse.csgraph.connected_components(
727 ... matrix)
728 >>> idx = 1
729 >>> molIdx = component_list[idx]
730 >>> print("There are {} molecules in the system".format(n_components))
731 There are 1 molecules in the system
732 >>> print("Atom {} is part of molecule {}".format(idx, molIdx))
733 Atom 1 is part of molecule 0
734 >>> molIdxs = [i for i in range(len(component_list))
735 ... if component_list[i] == molIdx]
736 >>> print("Atoms are part of molecule {}: {}".format(molIdx, molIdxs))
737 Atoms are part of molecule 0: [0, 1, 2, 3, 4, 5, 6, 7, 8]
738 """
740 nAtoms = len(nl.cutoffs)
742 if nl.nupdates <= 0:
743 raise RuntimeError(
744 'Must call update(atoms) on your neighborlist first!')
746 if sparse:
747 matrix = sp.dok_matrix((nAtoms, nAtoms), dtype=np.int8)
748 else:
749 matrix = np.zeros((nAtoms, nAtoms), dtype=np.int8)
751 for i in range(nAtoms):
752 for idx in nl.get_neighbors(i)[0]:
753 matrix[i, idx] = 1
755 return matrix
758class NewPrimitiveNeighborList:
759 """Neighbor list object. Wrapper around neighbor_list and first_neighbors.
761 cutoffs: list of float
762 List of cutoff radii - one for each atom. If the spheres (defined by
763 their cutoff radii) of two atoms overlap, they will be counted as
764 neighbors.
765 skin: float
766 If no atom has moved more than the skin-distance since the
767 last call to the
768 :meth:`~ase.neighborlist.NewPrimitiveNeighborList.update()`
769 method, then the neighbor list can be reused. This will save
770 some expensive rebuilds of the list, but extra neighbors outside
771 the cutoff will be returned.
772 sorted: bool
773 Sort neighbor list.
774 self_interaction: bool
775 Should an atom return itself as a neighbor?
776 bothways: bool
777 Return all neighbors. Default is to return only "half" of
778 the neighbors.
780 Example:
782 >>> from ase.build import bulk
783 >>> from ase.neighborlist import NewPrimitiveNeighborList
785 >>> nl = NewPrimitiveNeighborList([2.3, 1.7])
786 >>> atoms = bulk('Cu', 'fcc', a=3.6)
787 >>> nl.update(atoms.pbc, atoms.get_cell(), atoms.positions)
788 True
789 >>> indices, offsets = nl.get_neighbors(0)
790 """
792 def __init__(self, cutoffs, skin=0.3, sorted=False, self_interaction=True,
793 bothways=False, use_scaled_positions=False):
794 self.cutoffs = np.asarray(cutoffs) + skin
795 self.skin = skin
796 self.sorted = sorted
797 self.self_interaction = self_interaction
798 self.bothways = bothways
799 self.nupdates = 0
800 self.use_scaled_positions = use_scaled_positions
802 def update(self, pbc, cell, positions, numbers=None):
803 """Make sure the list is up to date."""
805 if self.nupdates == 0:
806 self.build(pbc, cell, positions, numbers=numbers)
807 return True
809 if ((self.pbc != pbc).any() or (self.cell != cell).any() or
810 ((self.positions - positions)**2).sum(1).max() > self.skin**2):
811 self.build(pbc, cell, positions, numbers=numbers)
812 return True
814 return False
816 def build(self, pbc, cell, positions, numbers=None):
817 """Build the list.
818 """
819 self.pbc = np.array(pbc, copy=True)
820 self.cell = np.array(cell, copy=True)
821 self.positions = np.array(positions, copy=True)
823 pair_first, pair_second, offset_vec = \
824 primitive_neighbor_list(
825 'ijS', pbc, cell, positions, self.cutoffs, numbers=numbers,
826 self_interaction=self.self_interaction,
827 use_scaled_positions=self.use_scaled_positions)
829 if len(positions) > 0 and not self.bothways:
830 offset_x, offset_y, offset_z = offset_vec.T
832 mask = offset_z > 0
833 mask &= offset_y == 0
834 mask |= offset_y > 0
835 mask &= offset_x == 0
836 mask |= offset_x > 0
837 mask |= (pair_first <= pair_second) & (offset_vec == 0).all(axis=1)
839 pair_first = pair_first[mask]
840 pair_second = pair_second[mask]
841 offset_vec = offset_vec[mask]
843 if len(positions) > 0 and self.sorted:
844 mask = np.argsort(pair_first * len(pair_first) +
845 pair_second)
846 pair_first = pair_first[mask]
847 pair_second = pair_second[mask]
848 offset_vec = offset_vec[mask]
850 self.pair_first = pair_first
851 self.pair_second = pair_second
852 self.offset_vec = offset_vec
854 # Compute the index array point to the first neighbor
855 self.first_neigh = first_neighbors(len(positions), pair_first)
857 self.nupdates += 1
859 def get_neighbors(self, a):
860 """Return neighbors of atom number a.
862 A list of indices and offsets to neighboring atoms is
863 returned. The positions of the neighbor atoms can be
864 calculated like this:
866 >>> from ase.build import bulk
867 >>> from ase.neighborlist import NewPrimitiveNeighborList
869 >>> nl = NewPrimitiveNeighborList([2.3, 1.7])
870 >>> atoms = bulk('Cu', 'fcc', a=3.6)
871 >>> nl.update(atoms.pbc, atoms.get_cell(), atoms.positions)
872 True
873 >>> indices, offsets = nl.get_neighbors(0)
874 >>> for i, offset in zip(indices, offsets):
875 ... print(
876 ... atoms.positions[i] + offset @ atoms.get_cell()
877 ... ) # doctest: +ELLIPSIS
878 [3.6 ... 0. ]
880 Notice that if get_neighbors(a) gives atom b as a neighbor,
881 then get_neighbors(b) will not return a as a neighbor - unless
882 bothways=True was used."""
884 return (self.pair_second[self.first_neigh[a]:self.first_neigh[a + 1]],
885 self.offset_vec[self.first_neigh[a]:self.first_neigh[a + 1]])
888class PrimitiveNeighborList:
889 """Neighbor list that works without Atoms objects.
891 This is less fancy, but can be used to avoid conversions between
892 scaled and non-scaled coordinates which may affect cell offsets
893 through rounding errors.
895 Attributes
896 ----------
897 nupdates : int
898 Number of updated times.
899 """
901 def __init__(self, cutoffs, skin=0.3, sorted=False, self_interaction=True,
902 bothways=False, use_scaled_positions=False):
903 self.cutoffs = np.asarray(cutoffs) + skin
904 self.skin = skin
905 self.sorted = sorted
906 self.self_interaction = self_interaction
907 self.bothways = bothways
908 self.nupdates = 0
909 self.use_scaled_positions = use_scaled_positions
911 def update(self, pbc, cell, coordinates):
912 """Make sure the list is up to date.
914 Returns
915 -------
916 bool
917 True if the neighbor list is updated.
918 """
920 if self.nupdates == 0:
921 self.build(pbc, cell, coordinates)
922 return True
924 if ((self.pbc != pbc).any() or (self.cell != cell).any() or (
925 (self.coordinates
926 - coordinates)**2).sum(1).max() > self.skin**2):
927 self.build(pbc, cell, coordinates)
928 return True
930 return False
932 def build(self, pbc, cell, coordinates):
933 """Build the list.
935 Coordinates are taken to be scaled or not according
936 to self.use_scaled_positions.
937 """
938 self.pbc = pbc = np.array(pbc, copy=True)
939 self.cell = cell = Cell(cell)
940 self.coordinates = coordinates = np.array(coordinates, copy=True)
942 if len(self.cutoffs) != len(coordinates):
943 raise ValueError('Wrong number of cutoff radii: {} != {}'
944 .format(len(self.cutoffs), len(coordinates)))
946 rcmax = self.cutoffs.max() if len(self.cutoffs) > 0 else 0.0
948 if self.use_scaled_positions:
949 positions0 = cell.cartesian_positions(coordinates)
950 else:
951 positions0 = coordinates
953 rcell, op = minkowski_reduce(cell, pbc)
954 positions = wrap_positions(positions0, rcell, pbc=pbc, eps=0)
956 natoms = len(positions)
957 self.nupdates += 1
958 if natoms == 0:
959 self.neighbors = []
960 self.displacements = []
961 return
963 tree = cKDTree(positions, copy_data=True)
964 offsets = cell.scaled_positions(positions - positions0)
965 offsets = offsets.round().astype(int)
967 N = _calc_expansion(rcell, pbc, rcmax)
969 neighbor_indices_a = [[] for _ in range(natoms)]
970 displacements_a = [[] for _ in range(natoms)]
972 for n1, n2, n3 in itertools.product(range(N[0] + 1),
973 range(-N[1], N[1] + 1),
974 range(-N[2], N[2] + 1)):
975 if n1 == 0 and (n2 < 0 or n2 == 0 and n3 < 0):
976 continue
978 displacement = (n1, n2, n3) @ rcell
979 shift0 = (n1, n2, n3) @ op
980 indices_all = tree.query_ball_point(
981 positions - displacement,
982 r=self.cutoffs + rcmax,
983 )
985 for a in range(natoms):
986 indices = indices_all[a]
988 if not indices:
989 continue
991 indices = np.array(indices)
992 delta = positions[indices] + displacement - positions[a]
993 distances = np.sqrt(np.add.reduce(delta**2, axis=1))
994 cutoffs = self.cutoffs[indices] + self.cutoffs[a]
995 i = indices[distances < cutoffs]
996 if n1 == 0 and n2 == 0 and n3 == 0:
997 if self.self_interaction:
998 i = i[i >= a]
999 else:
1000 i = i[i > a]
1002 neighbor_indices_a[a].append(i)
1004 disp = shift0 + offsets[i] - offsets[a]
1005 displacements_a[a].append(disp)
1007 self.neighbors = [np.concatenate(i) for i in neighbor_indices_a]
1008 self.displacements = [np.concatenate(d) for d in displacements_a]
1010 if self.bothways:
1011 neighbors2 = [[] for a in range(natoms)]
1012 displacements2 = [[] for a in range(natoms)]
1013 for a in range(natoms):
1014 for b, disp in zip(self.neighbors[a], self.displacements[a]):
1015 # avoid double counting of self interaction
1016 if a == b and (disp == 0).all():
1017 continue
1018 neighbors2[b].append(a)
1019 displacements2[b].append(-disp)
1020 for a in range(natoms):
1021 nbs = np.concatenate((self.neighbors[a], neighbors2[a]))
1022 disp = np.array(list(self.displacements[a]) + displacements2[a])
1023 # Force correct type and shape for case of no neighbors:
1024 self.neighbors[a] = nbs.astype(int)
1025 self.displacements[a] = disp.astype(int).reshape((-1, 3))
1027 if self.sorted:
1028 _sort_neighbors(self.neighbors, self.displacements)
1030 def get_neighbors(self, a):
1031 """Return neighbors of atom number a.
1033 A list of indices and offsets to neighboring atoms is
1034 returned. The positions of the neighbor atoms can be
1035 calculated like this:
1037 >>> from ase.build import bulk
1038 >>> from ase.neighborlist import NewPrimitiveNeighborList
1040 >>> nl = NewPrimitiveNeighborList([2.3, 1.7])
1041 >>> atoms = bulk('Cu', 'fcc', a=3.6)
1042 >>> nl.update(atoms.pbc, atoms.get_cell(), atoms.positions)
1043 True
1044 >>> indices, offsets = nl.get_neighbors(0)
1045 >>> for i, offset in zip(indices, offsets):
1046 ... print(
1047 ... atoms.positions[i] + offset @ atoms.get_cell()
1048 ... ) # doctest: +ELLIPSIS
1049 [3.6 ... 0. ]
1051 Notice that if get_neighbors(a) gives atom b as a neighbor,
1052 then get_neighbors(b) will not return a as a neighbor - unless
1053 bothways=True was used."""
1055 return self.neighbors[a], self.displacements[a]
1058def _calc_expansion(rcell, pbc, rcmax):
1059 r"""Calculate expansion to contain a sphere of radius `2.0 * rcmax`.
1061 This function determines the minimum supercell (parallelepiped) that
1062 contains a sphere of radius `2.0 * rcmax`. For this, `a_1` is projected
1063 onto the unit vector perpendicular to `a_2 \times a_3` (i.e. the unit
1064 vector along the direction `b_1`) to know how many `a_1`'s the supercell
1065 takes to contain the sphere.
1066 """
1067 ircell = np.linalg.pinv(rcell)
1068 vs = np.sqrt(np.add.reduce(ircell**2, axis=0))
1069 ns = np.where(pbc, np.ceil(2.0 * rcmax * vs), 0.0)
1070 return ns.astype(int)
1073def _sort_neighbors(neighbors, offsets):
1074 """Sort neighbors first by indices and then offsets."""
1075 natoms = len(neighbors)
1076 for a in range(natoms):
1077 keys = (
1078 offsets[a][:, 2],
1079 offsets[a][:, 1],
1080 offsets[a][:, 0],
1081 neighbors[a]
1082 )
1083 mask = np.lexsort(keys)
1084 neighbors[a] = neighbors[a][mask]
1085 offsets[a] = offsets[a][mask]
1088class NeighborList:
1089 """Neighbor list object.
1091 cutoffs: list of float
1092 List of cutoff radii - one for each atom. If the spheres
1093 (defined by their cutoff radii) of two atoms overlap, they
1094 will be counted as neighbors. See
1095 :func:`~ase.neighborlist.natural_cutoffs` for an example on
1096 how to get such a list.
1098 skin: float
1099 If no atom has moved more than the skin-distance since the
1100 last call to the
1101 :meth:`~ase.neighborlist.NeighborList.update()` method, then
1102 the neighbor list can be reused. This will save some
1103 expensive rebuilds of the list, but extra neighbors outside
1104 the cutoff will be returned.
1105 self_interaction: bool
1106 Should an atom return itself as a neighbor?
1107 bothways: bool
1108 Return all neighbors. Default is to return only "half" of
1109 the neighbors.
1110 primitive: class
1111 Define which implementation to use. Older and quadratically-scaling
1112 :class:`~ase.neighborlist.PrimitiveNeighborList` or newer and
1113 linearly-scaling :class:`~ase.neighborlist.NewPrimitiveNeighborList`.
1115 Example:
1117 >>> from ase.build import molecule
1118 >>> from ase.neighborlist import NeighborList
1120 >>> atoms = molecule("CO")
1121 >>> nl = NeighborList([0.76, 0.66])
1122 >>> nl.update(atoms)
1123 True
1124 >>> indices, offsets = nl.get_neighbors(0)
1126 """
1128 def __init__(self, cutoffs, skin=0.3, sorted=False, self_interaction=True,
1129 bothways=False, primitive=PrimitiveNeighborList):
1130 self.nl = primitive(cutoffs, skin, sorted,
1131 self_interaction=self_interaction,
1132 bothways=bothways)
1134 def update(self, atoms):
1135 """
1136 See :meth:`ase.neighborlist.PrimitiveNeighborList.update` or
1137 :meth:`ase.neighborlist.PrimitiveNeighborList.update`.
1138 """
1139 return self.nl.update(atoms.pbc, atoms.get_cell(complete=True),
1140 atoms.positions)
1142 def get_neighbors(self, a):
1143 """
1144 See :meth:`ase.neighborlist.PrimitiveNeighborList.get_neighbors` or
1145 :meth:`ase.neighborlist.PrimitiveNeighborList.get_neighbors`.
1146 """
1147 if self.nl.nupdates <= 0:
1148 raise RuntimeError('Must call update(atoms) on your neighborlist '
1149 'first!')
1151 return self.nl.get_neighbors(a)
1153 def get_connectivity_matrix(self, sparse=True):
1154 """
1155 See :func:`~ase.neighborlist.get_connectivity_matrix`.
1156 """
1157 return get_connectivity_matrix(self.nl, sparse)
1159 @property
1160 def nupdates(self):
1161 """Get number of updates."""
1162 return self.nl.nupdates
1164 @property
1165 @deprecated(
1166 'Use, e.g., `sum(_.size for _ in nl.neighbors)` '
1167 'for `bothways=False` and `self_interaction=False`.'
1168 )
1169 def nneighbors(self):
1170 """Get number of neighbors.
1172 .. deprecated:: 3.24.0
1173 """
1174 nneighbors = sum(indices.size for indices in self.nl.neighbors)
1175 if self.nl.self_interaction:
1176 nneighbors -= len(self.nl.neighbors)
1177 return nneighbors // 2 if self.nl.bothways else nneighbors
1179 @property
1180 @deprecated(
1181 'Use, e.g., `sum(_.any(1).sum() for _ in nl.displacements)` '
1182 'for `bothways=False` and `self_interaction=False`.'
1183 )
1184 def npbcneighbors(self):
1185 """Get number of pbc neighbors.
1187 .. deprecated:: 3.24.0
1188 """
1189 nneighbors = sum(
1190 offsets.any(axis=1).sum() for offsets in self.nl.displacements
1191 ) # sum up all neighbors that have non-zero supercell offsets
1192 return nneighbors // 2 if self.nl.bothways else nneighbors