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