Coverage for /builds/ase/ase/ase/neighborlist.py : 97.36%
Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1import numpy as np
2import itertools
3from scipy import sparse as sp
4from scipy.spatial import cKDTree
5import scipy.sparse.csgraph as csgraph
7from ase.data import atomic_numbers, covalent_radii
8from ase.geometry import complete_cell, find_mic, wrap_positions
9from ase.geometry import minkowski_reduce
10from ase.cell import Cell
13def natural_cutoffs(atoms, mult=1, **kwargs):
14 """Generate a radial cutoff for every atom based on covalent radii.
16 The covalent radii are a reasonable cutoff estimation for bonds in
17 many applications such as neighborlists, so function generates an
18 atoms length list of radii based on this idea.
20 * atoms: An atoms object
21 * mult: A multiplier for all cutoffs, useful for coarse grained adjustment
22 * kwargs: Symbol of the atom and its corresponding cutoff,
23 used to override the covalent radii
24 """
25 return [kwargs.get(atom.symbol, covalent_radii[atom.number] * mult)
26 for atom in atoms]
29def build_neighbor_list(atoms, cutoffs=None, **kwargs):
30 """Automatically build and update a NeighborList.
32 Parameters:
34 atoms : :class:`~ase.Atoms` object
35 Atoms to build Neighborlist for.
36 cutoffs: list of floats
37 Radii for each atom. If not given it will be produced by calling
38 :func:`ase.neighborlist.natural_cutoffs`
39 kwargs: arbitrary number of options
40 Will be passed to the constructor of
41 :class:`~ase.neighborlist.NeighborList`
43 Returns:
45 return: :class:`~ase.neighborlist.NeighborList`
46 A :class:`~ase.neighborlist.NeighborList` instance (updated).
47 """
48 if cutoffs is None:
49 cutoffs = natural_cutoffs(atoms)
51 nl = NeighborList(cutoffs, **kwargs)
52 nl.update(atoms)
54 return nl
57def get_distance_matrix(graph, limit=3):
58 """Get Distance Matrix from a Graph.
60 Parameters:
62 graph: array, matrix or sparse matrix, 2 dimensions (N, N)
63 Graph representation of the connectivity.
64 See `scipy doc <https://docs.scipy.org/doc/scipy/reference/generated\
65/scipy.sparse.csgraph.dijkstra.html#scipy.sparse.csgraph.dijkstra>`_
66 for reference.
67 limit: integer
68 Maximum number of steps to analyze. For most molecular information,
69 three should be enough.
71 Returns:
73 return: scipy.sparse.csr_matrix, shape (N, N)
74 A scipy.sparse.csr_matrix. All elements that are not connected within
75 *limit* steps are set to zero.
77 This is a potential memory bottleneck, as csgraph.dijkstra produces a
78 dense output matrix. Here we replace all np.inf values with 0 and
79 transform back to csr_matrix.
80 Why not dok_matrix like the connectivity-matrix? Because row-picking
81 is most likely and this is super fast with csr.
82 """
83 mat = csgraph.dijkstra(graph, directed=False, limit=limit)
84 mat[mat == np.inf] = 0
85 return sp.csr_matrix(mat, dtype=np.int8)
88def get_distance_indices(distanceMatrix, distance):
89 """Get indices for each node that are distance or less away.
91 Parameters:
93 distanceMatrix: any one of scipy.sparse matrices (NxN)
94 Matrix containing distance information of atoms. Get it e.g. with
95 :func:`~ase.neighborlist.get_distance_matrix`
96 distance: integer
97 Number of steps to allow.
99 Returns:
101 return: list of length N
102 List of length N. return[i] has all indices connected to item i.
104 The distance matrix only contains shortest paths, so when looking for
105 distances longer than one, we need to add the lower values for cases
106 where atoms are connected via a shorter path too.
107 """
108 shape = distanceMatrix.get_shape()
109 indices = []
110 # iterate over rows
111 for i in range(shape[0]):
112 row = distanceMatrix.getrow(i)[0]
113 # find all non-zero
114 found = sp.find(row)
115 # screen for smaller or equal distance
116 equal = np.where(found[-1] <= distance)[0]
117 # found[1] contains the indexes
118 indices.append([found[1][x] for x in equal])
119 return indices
122def mic(dr, cell, pbc=True):
123 """
124 Apply minimum image convention to an array of distance vectors.
126 Parameters:
128 dr : array_like
129 Array of distance vectors.
130 cell : array_like
131 Simulation cell.
132 pbc : array_like, optional
133 Periodic boundary conditions in x-, y- and z-direction. Default is to
134 assume periodic boundaries in all directions.
136 Returns:
138 dr : array
139 Array of distance vectors, wrapped according to the minimum image
140 convention.
141 """
142 dr, _ = find_mic(dr, cell, pbc)
143 return dr
146def primitive_neighbor_list(quantities, pbc, cell, positions, cutoff,
147 numbers=None, self_interaction=False,
148 use_scaled_positions=False, max_nbins=1e6):
149 """Compute a neighbor list for an atomic configuration.
151 Atoms outside periodic boundaries are mapped into the box. Atoms
152 outside nonperiodic boundaries are included in the neighbor list
153 but complexity of neighbor list search for those can become n^2.
155 The neighbor list is sorted by first atom index 'i', but not by second
156 atom index 'j'.
158 Parameters:
160 quantities: str
161 Quantities to compute by the neighbor list algorithm. Each character
162 in this string defines a quantity. They are returned in a tuple of
163 the same order. Possible quantities are
165 * 'i' : first atom index
166 * 'j' : second atom index
167 * 'd' : absolute distance
168 * 'D' : distance vector
169 * 'S' : shift vector (number of cell boundaries crossed by the bond
170 between atom i and j). With the shift vector S, the
171 distances D between atoms can be computed from:
172 D = positions[j]-positions[i]+S.dot(cell)
173 pbc: array_like
174 3-tuple indicating giving periodic boundaries in the three Cartesian
175 directions.
176 cell: 3x3 matrix
177 Unit cell vectors.
178 positions: list of xyz-positions
179 Atomic positions. Anything that can be converted to an ndarray of
180 shape (n, 3) will do: [(x1,y1,z1), (x2,y2,z2), ...]. If
181 use_scaled_positions is set to true, this must be scaled positions.
182 cutoff: float or dict
183 Cutoff for neighbor search. It can be:
185 * A single float: This is a global cutoff for all elements.
186 * A dictionary: This specifies cutoff values for element
187 pairs. Specification accepts element numbers of symbols.
188 Example: {(1, 6): 1.1, (1, 1): 1.0, ('C', 'C'): 1.85}
189 * A list/array with a per atom value: This specifies the radius of
190 an atomic sphere for each atoms. If spheres overlap, atoms are
191 within each others neighborhood. See
192 :func:`~ase.neighborlist.natural_cutoffs`
193 for an example on how to get such a list.
194 self_interaction: bool
195 Return the atom itself as its own neighbor if set to true.
196 Default: False
197 use_scaled_positions: bool
198 If set to true, positions are expected to be scaled positions.
199 max_nbins: int
200 Maximum number of bins used in neighbor search. This is used to limit
201 the maximum amount of memory required by the neighbor list.
203 Returns:
205 i, j, ... : array
206 Tuple with arrays for each quantity specified above. Indices in `i`
207 are returned in ascending order 0..len(a)-1, but the order of (i,j)
208 pairs is not guaranteed.
210 """
212 # Naming conventions: Suffixes indicate the dimension of an array. The
213 # following convention is used here:
214 # c: Cartesian index, can have values 0, 1, 2
215 # i: Global atom index, can have values 0..len(a)-1
216 # xyz: Bin index, three values identifying x-, y- and z-component of a
217 # spatial bin that is used to make neighbor search O(n)
218 # b: Linearized version of the 'xyz' bin index
219 # a: Bin-local atom index, i.e. index identifying an atom *within* a
220 # bin
221 # p: Pair index, can have value 0 or 1
222 # n: (Linear) neighbor index
224 # Return empty neighbor list if no atoms are passed here
225 if len(positions) == 0:
226 empty_types = dict(i=(int, (0, )),
227 j=(int, (0, )),
228 D=(float, (0, 3)),
229 d=(float, (0, )),
230 S=(int, (0, 3)))
231 retvals = []
232 for i in quantities:
233 dtype, shape = empty_types[i]
234 retvals += [np.array([], dtype=dtype).reshape(shape)]
235 if len(retvals) == 1:
236 return retvals[0]
237 else:
238 return tuple(retvals)
240 # Compute reciprocal lattice vectors.
241 b1_c, b2_c, b3_c = np.linalg.pinv(cell).T
243 # Compute distances of cell faces.
244 l1 = np.linalg.norm(b1_c)
245 l2 = np.linalg.norm(b2_c)
246 l3 = np.linalg.norm(b3_c)
247 face_dist_c = np.array([1 / l1 if l1 > 0 else 1,
248 1 / l2 if l2 > 0 else 1,
249 1 / l3 if l3 > 0 else 1])
251 if isinstance(cutoff, dict):
252 max_cutoff = max(cutoff.values())
253 else:
254 if np.isscalar(cutoff):
255 max_cutoff = cutoff
256 else:
257 cutoff = np.asarray(cutoff)
258 max_cutoff = 2 * np.max(cutoff)
260 # We use a minimum bin size of 3 A
261 bin_size = max(max_cutoff, 3)
262 # Compute number of bins such that a sphere of radius cutoff fits into
263 # eight neighboring bins.
264 nbins_c = np.maximum((face_dist_c / bin_size).astype(int), [1, 1, 1])
265 nbins = np.prod(nbins_c)
266 # Make sure we limit the amount of memory used by the explicit bins.
267 while nbins > max_nbins:
268 nbins_c = np.maximum(nbins_c // 2, [1, 1, 1])
269 nbins = np.prod(nbins_c)
271 # Compute over how many bins we need to loop in the neighbor list search.
272 neigh_search_x, neigh_search_y, neigh_search_z = \
273 np.ceil(bin_size * nbins_c / face_dist_c).astype(int)
275 # If we only have a single bin and the system is not periodic, then we
276 # do not need to search neighboring bins
277 neigh_search_x = 0 if nbins_c[0] == 1 and not pbc[0] else neigh_search_x
278 neigh_search_y = 0 if nbins_c[1] == 1 and not pbc[1] else neigh_search_y
279 neigh_search_z = 0 if nbins_c[2] == 1 and not pbc[2] else neigh_search_z
281 # Sort atoms into bins.
282 if use_scaled_positions:
283 scaled_positions_ic = positions
284 positions = np.dot(scaled_positions_ic, cell)
285 else:
286 scaled_positions_ic = np.linalg.solve(complete_cell(cell).T,
287 positions.T).T
288 bin_index_ic = np.floor(scaled_positions_ic * nbins_c).astype(int)
289 cell_shift_ic = np.zeros_like(bin_index_ic)
291 for c in range(3):
292 if pbc[c]:
293 # (Note: np.divmod does not exist in older numpies)
294 cell_shift_ic[:, c], bin_index_ic[:, c] = \
295 divmod(bin_index_ic[:, c], nbins_c[c])
296 else:
297 bin_index_ic[:, c] = np.clip(bin_index_ic[:, c], 0, nbins_c[c] - 1)
299 # Convert Cartesian bin index to unique scalar bin index.
300 bin_index_i = (bin_index_ic[:, 0] +
301 nbins_c[0] * (bin_index_ic[:, 1] +
302 nbins_c[1] * bin_index_ic[:, 2]))
304 # atom_i contains atom index in new sort order.
305 atom_i = np.argsort(bin_index_i)
306 bin_index_i = bin_index_i[atom_i]
308 # Find max number of atoms per bin
309 max_natoms_per_bin = np.bincount(bin_index_i).max()
311 # Sort atoms into bins: atoms_in_bin_ba contains for each bin (identified
312 # by its scalar bin index) a list of atoms inside that bin. This list is
313 # homogeneous, i.e. has the same size *max_natoms_per_bin* for all bins.
314 # The list is padded with -1 values.
315 atoms_in_bin_ba = -np.ones([nbins, max_natoms_per_bin], dtype=int)
316 for i in range(max_natoms_per_bin):
317 # Create a mask array that identifies the first atom of each bin.
318 mask = np.append([True], bin_index_i[:-1] != bin_index_i[1:])
319 # Assign all first atoms.
320 atoms_in_bin_ba[bin_index_i[mask], i] = atom_i[mask]
322 # Remove atoms that we just sorted into atoms_in_bin_ba. The next
323 # "first" atom will be the second and so on.
324 mask = np.logical_not(mask)
325 atom_i = atom_i[mask]
326 bin_index_i = bin_index_i[mask]
328 # Make sure that all atoms have been sorted into bins.
329 assert len(atom_i) == 0
330 assert len(bin_index_i) == 0
332 # Now we construct neighbor pairs by pairing up all atoms within a bin or
333 # between bin and neighboring bin. atom_pairs_pn is a helper buffer that
334 # contains all potential pairs of atoms between two bins, i.e. it is a list
335 # of length max_natoms_per_bin**2.
336 atom_pairs_pn = np.indices((max_natoms_per_bin, max_natoms_per_bin),
337 dtype=int)
338 atom_pairs_pn = atom_pairs_pn.reshape(2, -1)
340 # Initialized empty neighbor list buffers.
341 first_at_neightuple_nn = []
342 secnd_at_neightuple_nn = []
343 cell_shift_vector_x_n = []
344 cell_shift_vector_y_n = []
345 cell_shift_vector_z_n = []
347 # This is the main neighbor list search. We loop over neighboring bins and
348 # then construct all possible pairs of atoms between two bins, assuming
349 # that each bin contains exactly max_natoms_per_bin atoms. We then throw
350 # out pairs involving pad atoms with atom index -1 below.
351 binz_xyz, biny_xyz, binx_xyz = np.meshgrid(np.arange(nbins_c[2]),
352 np.arange(nbins_c[1]),
353 np.arange(nbins_c[0]),
354 indexing='ij')
355 # The memory layout of binx_xyz, biny_xyz, binz_xyz is such that computing
356 # the respective bin index leads to a linearly increasing consecutive list.
357 # The following assert statement succeeds:
358 # b_b = (binx_xyz + nbins_c[0] * (biny_xyz + nbins_c[1] *
359 # binz_xyz)).ravel()
360 # assert (b_b == np.arange(np.prod(nbins_c))).all()
362 # First atoms in pair.
363 _first_at_neightuple_n = atoms_in_bin_ba[:, atom_pairs_pn[0]]
364 for dz in range(-neigh_search_z, neigh_search_z + 1):
365 for dy in range(-neigh_search_y, neigh_search_y + 1):
366 for dx in range(-neigh_search_x, neigh_search_x + 1):
367 # Bin index of neighboring bin and shift vector.
368 shiftx_xyz, neighbinx_xyz = divmod(binx_xyz + dx, nbins_c[0])
369 shifty_xyz, neighbiny_xyz = divmod(biny_xyz + dy, nbins_c[1])
370 shiftz_xyz, neighbinz_xyz = divmod(binz_xyz + dz, nbins_c[2])
371 neighbin_b = (neighbinx_xyz + nbins_c[0] *
372 (neighbiny_xyz + nbins_c[1] * neighbinz_xyz)
373 ).ravel()
375 # Second atom in pair.
376 _secnd_at_neightuple_n = \
377 atoms_in_bin_ba[neighbin_b][:, atom_pairs_pn[1]]
379 # Shift vectors.
380 _cell_shift_vector_x_n = \
381 np.resize(shiftx_xyz.reshape(-1, 1),
382 (max_natoms_per_bin**2, shiftx_xyz.size)).T
383 _cell_shift_vector_y_n = \
384 np.resize(shifty_xyz.reshape(-1, 1),
385 (max_natoms_per_bin**2, shifty_xyz.size)).T
386 _cell_shift_vector_z_n = \
387 np.resize(shiftz_xyz.reshape(-1, 1),
388 (max_natoms_per_bin**2, shiftz_xyz.size)).T
390 # We have created too many pairs because we assumed each bin
391 # has exactly max_natoms_per_bin atoms. Remove all surperfluous
392 # pairs. Those are pairs that involve an atom with index -1.
393 mask = np.logical_and(_first_at_neightuple_n != -1,
394 _secnd_at_neightuple_n != -1)
395 if mask.sum() > 0:
396 first_at_neightuple_nn += [_first_at_neightuple_n[mask]]
397 secnd_at_neightuple_nn += [_secnd_at_neightuple_n[mask]]
398 cell_shift_vector_x_n += [_cell_shift_vector_x_n[mask]]
399 cell_shift_vector_y_n += [_cell_shift_vector_y_n[mask]]
400 cell_shift_vector_z_n += [_cell_shift_vector_z_n[mask]]
402 # Flatten overall neighbor list.
403 first_at_neightuple_n = np.concatenate(first_at_neightuple_nn)
404 secnd_at_neightuple_n = np.concatenate(secnd_at_neightuple_nn)
405 cell_shift_vector_n = np.transpose([np.concatenate(cell_shift_vector_x_n),
406 np.concatenate(cell_shift_vector_y_n),
407 np.concatenate(cell_shift_vector_z_n)])
409 # Add global cell shift to shift vectors
410 cell_shift_vector_n += cell_shift_ic[first_at_neightuple_n] - \
411 cell_shift_ic[secnd_at_neightuple_n]
413 # Remove all self-pairs that do not cross the cell boundary.
414 if not self_interaction:
415 m = np.logical_not(np.logical_and(
416 first_at_neightuple_n == secnd_at_neightuple_n,
417 (cell_shift_vector_n == 0).all(axis=1)))
418 first_at_neightuple_n = first_at_neightuple_n[m]
419 secnd_at_neightuple_n = secnd_at_neightuple_n[m]
420 cell_shift_vector_n = cell_shift_vector_n[m]
422 # For nonperiodic directions, remove any bonds that cross the domain
423 # boundary.
424 for c in range(3):
425 if not pbc[c]:
426 m = cell_shift_vector_n[:, c] == 0
427 first_at_neightuple_n = first_at_neightuple_n[m]
428 secnd_at_neightuple_n = secnd_at_neightuple_n[m]
429 cell_shift_vector_n = cell_shift_vector_n[m]
431 # Sort neighbor list.
432 i = np.argsort(first_at_neightuple_n)
433 first_at_neightuple_n = first_at_neightuple_n[i]
434 secnd_at_neightuple_n = secnd_at_neightuple_n[i]
435 cell_shift_vector_n = cell_shift_vector_n[i]
437 # Compute distance vectors.
438 distance_vector_nc = positions[secnd_at_neightuple_n] - \
439 positions[first_at_neightuple_n] + \
440 cell_shift_vector_n.dot(cell)
441 abs_distance_vector_n = \
442 np.sqrt(np.sum(distance_vector_nc * distance_vector_nc, axis=1))
444 # We have still created too many pairs. Only keep those with distance
445 # smaller than max_cutoff.
446 mask = abs_distance_vector_n < max_cutoff
447 first_at_neightuple_n = first_at_neightuple_n[mask]
448 secnd_at_neightuple_n = secnd_at_neightuple_n[mask]
449 cell_shift_vector_n = cell_shift_vector_n[mask]
450 distance_vector_nc = distance_vector_nc[mask]
451 abs_distance_vector_n = abs_distance_vector_n[mask]
453 if isinstance(cutoff, dict) and numbers is not None:
454 # If cutoff is a dictionary, then the cutoff radii are specified per
455 # element pair. We now have a list up to maximum cutoff.
456 per_pair_cutoff_n = np.zeros_like(abs_distance_vector_n)
457 for (atomic_number1, atomic_number2), c in cutoff.items():
458 try:
459 atomic_number1 = atomic_numbers[atomic_number1]
460 except KeyError:
461 pass
462 try:
463 atomic_number2 = atomic_numbers[atomic_number2]
464 except KeyError:
465 pass
466 if atomic_number1 == atomic_number2:
467 mask = np.logical_and(
468 numbers[first_at_neightuple_n] == atomic_number1,
469 numbers[secnd_at_neightuple_n] == atomic_number2)
470 else:
471 mask = np.logical_or(
472 np.logical_and(
473 numbers[first_at_neightuple_n] == atomic_number1,
474 numbers[secnd_at_neightuple_n] == atomic_number2),
475 np.logical_and(
476 numbers[first_at_neightuple_n] == atomic_number2,
477 numbers[secnd_at_neightuple_n] == atomic_number1))
478 per_pair_cutoff_n[mask] = c
479 mask = abs_distance_vector_n < per_pair_cutoff_n
480 first_at_neightuple_n = first_at_neightuple_n[mask]
481 secnd_at_neightuple_n = secnd_at_neightuple_n[mask]
482 cell_shift_vector_n = cell_shift_vector_n[mask]
483 distance_vector_nc = distance_vector_nc[mask]
484 abs_distance_vector_n = abs_distance_vector_n[mask]
485 elif not np.isscalar(cutoff):
486 # If cutoff is neither a dictionary nor a scalar, then we assume it is
487 # a list or numpy array that contains atomic radii. Atoms are neighbors
488 # if their radii overlap.
489 mask = abs_distance_vector_n < \
490 cutoff[first_at_neightuple_n] + cutoff[secnd_at_neightuple_n]
491 first_at_neightuple_n = first_at_neightuple_n[mask]
492 secnd_at_neightuple_n = secnd_at_neightuple_n[mask]
493 cell_shift_vector_n = cell_shift_vector_n[mask]
494 distance_vector_nc = distance_vector_nc[mask]
495 abs_distance_vector_n = abs_distance_vector_n[mask]
497 # Assemble return tuple.
498 retvals = []
499 for q in quantities:
500 if q == 'i':
501 retvals += [first_at_neightuple_n]
502 elif q == 'j':
503 retvals += [secnd_at_neightuple_n]
504 elif q == 'D':
505 retvals += [distance_vector_nc]
506 elif q == 'd':
507 retvals += [abs_distance_vector_n]
508 elif q == 'S':
509 retvals += [cell_shift_vector_n]
510 else:
511 raise ValueError('Unsupported quantity specified.')
512 if len(retvals) == 1:
513 return retvals[0]
514 else:
515 return tuple(retvals)
518def neighbor_list(quantities, a, cutoff, self_interaction=False,
519 max_nbins=1e6):
520 """Compute a neighbor list for an atomic configuration.
522 Atoms outside periodic boundaries are mapped into the box. Atoms
523 outside nonperiodic boundaries are included in the neighbor list
524 but complexity of neighbor list search for those can become n^2.
526 The neighbor list is sorted by first atom index 'i', but not by second
527 atom index 'j'.
529 Parameters:
531 quantities: str
532 Quantities to compute by the neighbor list algorithm. Each character
533 in this string defines a quantity. They are returned in a tuple of
534 the same order. Possible quantities are:
536 * 'i' : first atom index
537 * 'j' : second atom index
538 * 'd' : absolute distance
539 * 'D' : distance vector
540 * 'S' : shift vector (number of cell boundaries crossed by the bond
541 between atom i and j). With the shift vector S, the
542 distances D between atoms can be computed from:
543 D = a.positions[j]-a.positions[i]+S.dot(a.cell)
544 a: :class:`ase.Atoms`
545 Atomic configuration.
546 cutoff: float or dict
547 Cutoff for neighbor search. It can be:
549 * A single float: This is a global cutoff for all elements.
550 * A dictionary: This specifies cutoff values for element
551 pairs. Specification accepts element numbers of symbols.
552 Example: {(1, 6): 1.1, (1, 1): 1.0, ('C', 'C'): 1.85}
553 * A list/array with a per atom value: This specifies the radius of
554 an atomic sphere for each atoms. If spheres overlap, atoms are
555 within each others neighborhood. See
556 :func:`~ase.neighborlist.natural_cutoffs`
557 for an example on how to get such a list.
559 self_interaction: bool
560 Return the atom itself as its own neighbor if set to true.
561 Default: False
562 max_nbins: int
563 Maximum number of bins used in neighbor search. This is used to limit
564 the maximum amount of memory required by the neighbor list.
566 Returns:
568 i, j, ...: array
569 Tuple with arrays for each quantity specified above. Indices in `i`
570 are returned in ascending order 0..len(a), but the order of (i,j)
571 pairs is not guaranteed.
573 Examples:
575 Examples assume Atoms object *a* and numpy imported as *np*.
577 1. Coordination counting::
579 i = neighbor_list('i', a, 1.85)
580 coord = np.bincount(i)
582 2. Coordination counting with different cutoffs for each pair of species::
584 i = neighbor_list('i', a,
585 {('H', 'H'): 1.1, ('C', 'H'): 1.3, ('C', 'C'): 1.85})
586 coord = np.bincount(i)
588 3. Pair distribution function::
590 d = neighbor_list('d', a, 10.00)
591 h, bin_edges = np.histogram(d, bins=100)
592 pdf = h/(4*np.pi/3*(
593 bin_edges[1:]**3 - bin_edges[:-1]**3)) * a.get_volume()/len(a)
595 4. Pair potential::
597 i, j, d, D = neighbor_list('ijdD', a, 5.0)
598 energy = (-C/d**6).sum()
599 forces = (6*C/d**5 * (D/d).T).T
600 forces_x = np.bincount(j, weights=forces[:, 0], minlength=len(a)) - \
601 np.bincount(i, weights=forces[:, 0], minlength=len(a))
602 forces_y = np.bincount(j, weights=forces[:, 1], minlength=len(a)) - \
603 np.bincount(i, weights=forces[:, 1], minlength=len(a))
604 forces_z = np.bincount(j, weights=forces[:, 2], minlength=len(a)) - \
605 np.bincount(i, weights=pair_forces[:, 2], minlength=len(a))
607 5. Dynamical matrix for a pair potential stored in a block sparse format::
609 from scipy.sparse import bsr_matrix
610 i, j, dr, abs_dr = neighbor_list('ijDd', atoms)
611 energy = (dr.T / abs_dr).T
612 dynmat = -(dde * (energy.reshape(-1, 3, 1)
613 * energy.reshape(-1, 1, 3)).T).T \
614 -(de / abs_dr * (np.eye(3, dtype=energy.dtype) - \
615 (energy.reshape(-1, 3, 1) * energy.reshape(-1, 1, 3))).T).T
616 dynmat_bsr = bsr_matrix((dynmat, j, first_i),
617 shape=(3*len(a), 3*len(a)))
619 dynmat_diag = np.empty((len(a), 3, 3))
620 for x in range(3):
621 for y in range(3):
622 dynmat_diag[:, x, y] = -np.bincount(i, weights=dynmat[:, x, y])
624 dynmat_bsr += bsr_matrix((dynmat_diag, np.arange(len(a)),
625 np.arange(len(a) + 1)),
626 shape=(3 * len(a), 3 * len(a)))
628 """
629 return primitive_neighbor_list(quantities, a.pbc,
630 a.get_cell(complete=True),
631 a.positions, cutoff, numbers=a.numbers,
632 self_interaction=self_interaction,
633 max_nbins=max_nbins)
636def first_neighbors(natoms, first_atom):
637 """
638 Compute an index array pointing to the ranges within the neighbor list that
639 contain the neighbors for a certain atom.
641 Parameters:
643 natoms : int
644 Total number of atom.
645 first_atom : array_like
646 Array containing the first atom 'i' of the neighbor tuple returned
647 by the neighbor list.
649 Returns:
651 seed : array
652 Array containing pointers to the start and end location of the
653 neighbors of a certain atom. Neighbors of atom k have indices from s[k]
654 to s[k+1]-1.
655 """
656 if len(first_atom) == 0:
657 return np.zeros(natoms + 1, dtype=int)
658 # Create a seed array (which is returned by this function) populated with
659 # -1.
660 seed = -np.ones(natoms + 1, dtype=int)
662 first_atom = np.asarray(first_atom)
664 # Mask array contains all position where the number in the (sorted) array
665 # with first atoms (in the neighbor pair) changes.
666 mask = first_atom[:-1] != first_atom[1:]
668 # Seed array needs to start at 0
669 seed[first_atom[0]] = 0
670 # Seed array needs to stop at the length of the neighbor list
671 seed[-1] = len(first_atom)
672 # Populate all intermediate seed with the index of where the mask array is
673 # true, i.e. the index where the first_atom array changes.
674 seed[first_atom[1:][mask]] = (np.arange(len(mask)) + 1)[mask]
676 # Now fill all remaining -1 value with the value in the seed array right
677 # behind them. (There are no neighbor so seed[i] and seed[i+1] must point)
678 # to the same index.
679 mask = seed == -1
680 while mask.any():
681 seed[mask] = seed[np.arange(natoms + 1)[mask] + 1]
682 mask = seed == -1
683 return seed
686def get_connectivity_matrix(nl, sparse=True):
687 """Return connectivity matrix for a given NeighborList (dtype=numpy.int8).
689 A matrix of shape (nAtoms, nAtoms) will be returned.
690 Connected atoms i and j will have matrix[i,j] == 1, unconnected
691 matrix[i,j] == 0. If bothways=True the matrix will be symmetric,
692 otherwise not!
694 If *sparse* is True, a scipy csr matrix is returned.
695 If *sparse* is False, a numpy matrix is returned.
697 Note that the old and new neighborlists might give different results
698 for periodic systems if bothways=False.
700 Example:
702 Determine which molecule in a system atom 1 belongs to.
704 >>> from ase import neighborlist
705 >>> from ase.build import molecule
706 >>> from scipy import sparse
707 >>> mol = molecule('CH3CH2OH')
708 >>> cutOff = neighborlist.natural_cutoffs(mol)
709 >>> neighborList = neighborlist.NeighborList(
710 ... cutOff, self_interaction=False, bothways=True)
711 >>> neighborList.update(mol)
712 >>> matrix = neighborList.get_connectivity_matrix()
713 >>> #or: matrix = neighborlist.get_connectivity_matrix(neighborList.nl)
714 >>> n_components, component_list = sparse.csgraph.connected_components(
715 ... matrix)
716 >>> idx = 1
717 >>> molIdx = component_list[idx]
718 >>> print("There are {} molecules in the system".format(n_components))
719 >>> print("Atom {} is part of molecule {}".format(idx, molIdx))
720 >>> molIdxs = [i for i in range(len(component_list))
721 ... if component_list[i] == molIdx]
722 >>> print("Atoms are part of molecule {}: {}".format(molIdx, molIdxs))
723 """
725 nAtoms = len(nl.cutoffs)
727 if nl.nupdates <= 0:
728 raise RuntimeError(
729 'Must call update(atoms) on your neighborlist first!')
731 if sparse:
732 matrix = sp.dok_matrix((nAtoms, nAtoms), dtype=np.int8)
733 else:
734 matrix = np.zeros((nAtoms, nAtoms), dtype=np.int8)
736 for i in range(nAtoms):
737 for idx in nl.get_neighbors(i)[0]:
738 matrix[i, idx] = 1
740 return matrix
743class NewPrimitiveNeighborList:
744 """Neighbor list object. Wrapper around neighbor_list and first_neighbors.
746 cutoffs: list of float
747 List of cutoff radii - one for each atom. If the spheres (defined by
748 their cutoff radii) of two atoms overlap, they will be counted as
749 neighbors.
750 skin: float
751 If no atom has moved more than the skin-distance since the
752 last call to the
753 :meth:`~ase.neighborlist.NewPrimitiveNeighborList.update()`
754 method, then the neighbor list can be reused. This will save
755 some expensive rebuilds of the list, but extra neighbors outside
756 the cutoff will be returned.
757 sorted: bool
758 Sort neighbor list.
759 self_interaction: bool
760 Should an atom return itself as a neighbor?
761 bothways: bool
762 Return all neighbors. Default is to return only "half" of
763 the neighbors.
765 Example::
767 nl = NeighborList([2.3, 1.7])
768 nl.update(atoms)
769 indices, offsets = nl.get_neighbors(0)
770 """
772 def __init__(self, cutoffs, skin=0.3, sorted=False, self_interaction=True,
773 bothways=False, use_scaled_positions=False):
774 self.cutoffs = np.asarray(cutoffs) + skin
775 self.skin = skin
776 self.sorted = sorted
777 self.self_interaction = self_interaction
778 self.bothways = bothways
779 self.nupdates = 0
780 self.use_scaled_positions = use_scaled_positions
781 self.nneighbors = 0
782 self.npbcneighbors = 0
784 def update(self, pbc, cell, positions, numbers=None):
785 """Make sure the list is up to date."""
787 if self.nupdates == 0:
788 self.build(pbc, cell, positions, numbers=numbers)
789 return True
791 if ((self.pbc != pbc).any() or (self.cell != cell).any() or
792 ((self.positions - positions)**2).sum(1).max() > self.skin**2):
793 self.build(pbc, cell, positions, numbers=numbers)
794 return True
796 return False
798 def build(self, pbc, cell, positions, numbers=None):
799 """Build the list.
800 """
801 self.pbc = np.array(pbc, copy=True)
802 self.cell = np.array(cell, copy=True)
803 self.positions = np.array(positions, copy=True)
805 pair_first, pair_second, offset_vec = \
806 primitive_neighbor_list(
807 'ijS', pbc, cell, positions, self.cutoffs, numbers=numbers,
808 self_interaction=self.self_interaction,
809 use_scaled_positions=self.use_scaled_positions)
811 if len(positions) > 0 and not self.bothways:
812 offset_x, offset_y, offset_z = offset_vec.T
814 mask = offset_z > 0
815 mask &= offset_y == 0
816 mask |= offset_y > 0
817 mask &= offset_x == 0
818 mask |= offset_x > 0
819 mask |= (pair_first <= pair_second) & (offset_vec == 0).all(axis=1)
821 pair_first = pair_first[mask]
822 pair_second = pair_second[mask]
823 offset_vec = offset_vec[mask]
825 if len(positions) > 0 and self.sorted:
826 mask = np.argsort(pair_first * len(pair_first) +
827 pair_second)
828 pair_first = pair_first[mask]
829 pair_second = pair_second[mask]
830 offset_vec = offset_vec[mask]
832 self.pair_first = pair_first
833 self.pair_second = pair_second
834 self.offset_vec = offset_vec
836 # Compute the index array point to the first neighbor
837 self.first_neigh = first_neighbors(len(positions), pair_first)
839 self.nupdates += 1
841 def get_neighbors(self, a):
842 """Return neighbors of atom number a.
844 A list of indices and offsets to neighboring atoms is
845 returned. The positions of the neighbor atoms can be
846 calculated like this:
848 >>> indices, offsets = nl.get_neighbors(42)
849 >>> for i, offset in zip(indices, offsets):
850 >>> print(atoms.positions[i] + dot(offset, atoms.get_cell()))
852 Notice that if get_neighbors(a) gives atom b as a neighbor,
853 then get_neighbors(b) will not return a as a neighbor - unless
854 bothways=True was used."""
856 return (self.pair_second[self.first_neigh[a]:self.first_neigh[a + 1]],
857 self.offset_vec[self.first_neigh[a]:self.first_neigh[a + 1]])
860class PrimitiveNeighborList:
861 """Neighbor list that works without Atoms objects.
863 This is less fancy, but can be used to avoid conversions between
864 scaled and non-scaled coordinates which may affect cell offsets
865 through rounding errors.
866 """
868 def __init__(self, cutoffs, skin=0.3, sorted=False, self_interaction=True,
869 bothways=False, use_scaled_positions=False):
870 self.cutoffs = np.asarray(cutoffs) + skin
871 self.skin = skin
872 self.sorted = sorted
873 self.self_interaction = self_interaction
874 self.bothways = bothways
875 self.nupdates = 0
876 self.use_scaled_positions = use_scaled_positions
877 self.nneighbors = 0
878 self.npbcneighbors = 0
880 def update(self, pbc, cell, coordinates):
881 """Make sure the list is up to date."""
883 if self.nupdates == 0:
884 self.build(pbc, cell, coordinates)
885 return True
887 if ((self.pbc != pbc).any() or (self.cell != cell).any() or (
888 (self.coordinates
889 - coordinates)**2).sum(1).max() > self.skin**2):
890 self.build(pbc, cell, coordinates)
891 return True
893 return False
895 def build(self, pbc, cell, coordinates):
896 """Build the list.
898 Coordinates are taken to be scaled or not according
899 to self.use_scaled_positions.
900 """
901 self.pbc = pbc = np.array(pbc, copy=True)
902 self.cell = cell = Cell(cell)
903 self.coordinates = coordinates = np.array(coordinates, copy=True)
905 if len(self.cutoffs) != len(coordinates):
906 raise ValueError('Wrong number of cutoff radii: {0} != {1}'
907 .format(len(self.cutoffs), len(coordinates)))
909 if len(self.cutoffs) > 0:
910 rcmax = self.cutoffs.max()
911 else:
912 rcmax = 0.0
914 if self.use_scaled_positions:
915 positions0 = cell.cartesian_positions(coordinates)
916 else:
917 positions0 = coordinates
919 rcell, op = minkowski_reduce(cell, pbc)
920 positions = wrap_positions(positions0, rcell, pbc=pbc, eps=0)
922 natoms = len(positions)
923 self.nneighbors = 0
924 self.npbcneighbors = 0
925 self.neighbors = [np.empty(0, int) for a in range(natoms)]
926 self.displacements = [np.empty((0, 3), int) for a in range(natoms)]
927 self.nupdates += 1
928 if natoms == 0:
929 return
931 N = []
932 ircell = np.linalg.pinv(rcell)
933 for i in range(3):
934 if self.pbc[i]:
935 v = ircell[:, i]
936 h = 1 / np.linalg.norm(v)
937 n = int(2 * rcmax / h) + 1
938 else:
939 n = 0
940 N.append(n)
942 tree = cKDTree(positions, copy_data=True)
943 offsets = cell.scaled_positions(positions - positions0)
944 offsets = offsets.round().astype(int)
946 for n1, n2, n3 in itertools.product(range(0, N[0] + 1),
947 range(-N[1], N[1] + 1),
948 range(-N[2], N[2] + 1)):
949 if n1 == 0 and (n2 < 0 or n2 == 0 and n3 < 0):
950 continue
952 displacement = (n1, n2, n3) @ rcell
953 for a in range(natoms):
955 indices = tree.query_ball_point(positions[a] - displacement,
956 r=self.cutoffs[a] + rcmax)
957 if not len(indices):
958 continue
960 indices = np.array(indices)
961 delta = positions[indices] + displacement - positions[a]
962 cutoffs = self.cutoffs[indices] + self.cutoffs[a]
963 i = indices[np.linalg.norm(delta, axis=1) < cutoffs]
964 if n1 == 0 and n2 == 0 and n3 == 0:
965 if self.self_interaction:
966 i = i[i >= a]
967 else:
968 i = i[i > a]
970 self.nneighbors += len(i)
971 self.neighbors[a] = np.concatenate((self.neighbors[a], i))
973 disp = (n1, n2, n3) @ op + offsets[i] - offsets[a]
974 self.npbcneighbors += disp.any(1).sum()
975 self.displacements[a] = np.concatenate((self.displacements[a],
976 disp))
978 if self.bothways:
979 neighbors2 = [[] for a in range(natoms)]
980 displacements2 = [[] for a in range(natoms)]
981 for a in range(natoms):
982 for b, disp in zip(self.neighbors[a], self.displacements[a]):
983 neighbors2[b].append(a)
984 displacements2[b].append(-disp)
985 for a in range(natoms):
986 nbs = np.concatenate((self.neighbors[a], neighbors2[a]))
987 disp = np.array(list(self.displacements[a]) + displacements2[a])
988 # Force correct type and shape for case of no neighbors:
989 self.neighbors[a] = nbs.astype(int)
990 self.displacements[a] = disp.astype(int).reshape((-1, 3))
992 if self.sorted:
993 for a, i in enumerate(self.neighbors):
994 mask = (i < a)
995 if mask.any():
996 j = i[mask]
997 offsets = self.displacements[a][mask]
998 for b, offset in zip(j, offsets):
999 self.neighbors[b] = np.concatenate(
1000 (self.neighbors[b], [a]))
1001 self.displacements[b] = np.concatenate(
1002 (self.displacements[b], [-offset]))
1003 mask = np.logical_not(mask)
1004 self.neighbors[a] = self.neighbors[a][mask]
1005 self.displacements[a] = self.displacements[a][mask]
1007 def get_neighbors(self, a):
1008 """Return neighbors of atom number a.
1010 A list of indices and offsets to neighboring atoms is
1011 returned. The positions of the neighbor atoms can be
1012 calculated like this::
1014 indices, offsets = nl.get_neighbors(42)
1015 for i, offset in zip(indices, offsets):
1016 print(atoms.positions[i] + offset @ atoms.get_cell())
1018 Notice that if get_neighbors(a) gives atom b as a neighbor,
1019 then get_neighbors(b) will not return a as a neighbor - unless
1020 bothways=True was used."""
1022 return self.neighbors[a], self.displacements[a]
1025class NeighborList:
1026 """Neighbor list object.
1028 cutoffs: list of float
1029 List of cutoff radii - one for each atom. If the spheres
1030 (defined by their cutoff radii) of two atoms overlap, they
1031 will be counted as neighbors. See
1032 :func:`~ase.neighborlist.natural_cutoffs` for an example on
1033 how to get such a list.
1035 skin: float
1036 If no atom has moved more than the skin-distance since the
1037 last call to the
1038 :meth:`~ase.neighborlist.NeighborList.update()` method, then
1039 the neighbor list can be reused. This will save some
1040 expensive rebuilds of the list, but extra neighbors outside
1041 the cutoff will be returned.
1042 self_interaction: bool
1043 Should an atom return itself as a neighbor?
1044 bothways: bool
1045 Return all neighbors. Default is to return only "half" of
1046 the neighbors.
1047 primitive: class
1048 Define which implementation to use. Older and quadratically-scaling
1049 :class:`~ase.neighborlist.PrimitiveNeighborList` or newer and
1050 linearly-scaling :class:`~ase.neighborlist.NewPrimitiveNeighborList`.
1052 Example::
1054 nl = NeighborList([2.3, 1.7])
1055 nl.update(atoms)
1056 indices, offsets = nl.get_neighbors(0)
1058 """
1060 def __init__(self, cutoffs, skin=0.3, sorted=False, self_interaction=True,
1061 bothways=False, primitive=PrimitiveNeighborList):
1062 self.nl = primitive(cutoffs, skin, sorted,
1063 self_interaction=self_interaction,
1064 bothways=bothways)
1066 def update(self, atoms):
1067 """
1068 See :meth:`ase.neighborlist.PrimitiveNeighborList.update` or
1069 :meth:`ase.neighborlist.PrimitiveNeighborList.update`.
1070 """
1071 return self.nl.update(atoms.pbc, atoms.get_cell(complete=True),
1072 atoms.positions)
1074 def get_neighbors(self, a):
1075 """
1076 See :meth:`ase.neighborlist.PrimitiveNeighborList.get_neighbors` or
1077 :meth:`ase.neighborlist.PrimitiveNeighborList.get_neighbors`.
1078 """
1079 if self.nl.nupdates <= 0:
1080 raise RuntimeError('Must call update(atoms) on your neighborlist '
1081 'first!')
1083 return self.nl.get_neighbors(a)
1085 def get_connectivity_matrix(self, sparse=True):
1086 """
1087 See :func:`~ase.neighborlist.get_connectivity_matrix`.
1088 """
1089 return get_connectivity_matrix(self.nl, sparse)
1091 @property
1092 def nupdates(self):
1093 """Get number of updates."""
1094 return self.nl.nupdates
1096 @property
1097 def nneighbors(self):
1098 """Get number of neighbors."""
1099 return self.nl.nneighbors
1101 @property
1102 def npbcneighbors(self):
1103 """Get number of pbc neighbors."""
1104 return self.nl.npbcneighbors