Coverage for /builds/kinetik161/ase/ase/neighborlist.py: 97.35%

1import itertools 

2 

3import numpy as np 

4import scipy.sparse.csgraph as csgraph 

5from scipy import sparse as sp 

6from scipy.spatial import cKDTree 

7 

8from ase.cell import Cell 

9from ase.data import atomic_numbers, covalent_radii 

10from ase.geometry import (complete_cell, find_mic, minkowski_reduce, 

11 wrap_positions) 

12 

13 

14def natural_cutoffs(atoms, mult=1, **kwargs): 

15 """Generate a radial cutoff for every atom based on covalent radii. 

16 

17 The covalent radii are a reasonable cutoff estimation for bonds in 

18 many applications such as neighborlists, so function generates an 

19 atoms length list of radii based on this idea. 

20 

21 * atoms: An atoms object 

22 * mult: A multiplier for all cutoffs, useful for coarse grained adjustment 

23 * kwargs: Symbol of the atom and its corresponding cutoff, 

24 used to override the covalent radii 

25 """ 

26 return [kwargs.get(atom.symbol, covalent_radii[atom.number] * mult) 

27 for atom in atoms] 

28 

29 

30def build_neighbor_list(atoms, cutoffs=None, **kwargs): 

31 """Automatically build and update a NeighborList. 

32 

33 Parameters: 

34 

35 atoms : :class:`~ase.Atoms` object 

36 Atoms to build Neighborlist for. 

37 cutoffs: list of floats 

38 Radii for each atom. If not given it will be produced by calling 

39 :func:`ase.neighborlist.natural_cutoffs` 

40 kwargs: arbitrary number of options 

41 Will be passed to the constructor of 

42 :class:`~ase.neighborlist.NeighborList` 

43 

44 Returns: 

45 

46 return: :class:`~ase.neighborlist.NeighborList` 

47 A :class:`~ase.neighborlist.NeighborList` instance (updated). 

48 """ 

49 if cutoffs is None: 

50 cutoffs = natural_cutoffs(atoms) 

51 

52 nl = NeighborList(cutoffs, **kwargs) 

53 nl.update(atoms) 

54 

55 return nl 

56 

57 

58def get_distance_matrix(graph, limit=3): 

59 """Get Distance Matrix from a Graph. 

60 

61 Parameters: 

62 

63 graph: array, matrix or sparse matrix, 2 dimensions (N, N) 

64 Graph representation of the connectivity. 

65 See `scipy doc <https://docs.scipy.org/doc/scipy/reference/generated\ 

66/scipy.sparse.csgraph.dijkstra.html#scipy.sparse.csgraph.dijkstra>`_ 

67 for reference. 

68 limit: integer 

69 Maximum number of steps to analyze. For most molecular information, 

70 three should be enough. 

71 

72 Returns: 

73 

74 return: scipy.sparse.csr_matrix, shape (N, N) 

75 A scipy.sparse.csr_matrix. All elements that are not connected within 

76 *limit* steps are set to zero. 

77 

78 This is a potential memory bottleneck, as csgraph.dijkstra produces a 

79 dense output matrix. Here we replace all np.inf values with 0 and 

80 transform back to csr_matrix. 

81 Why not dok_matrix like the connectivity-matrix? Because row-picking 

82 is most likely and this is super fast with csr. 

83 """ 

84 mat = csgraph.dijkstra(graph, directed=False, limit=limit) 

85 mat[mat == np.inf] = 0 

86 return sp.csr_matrix(mat, dtype=np.int8) 

87 

88 

89def get_distance_indices(distanceMatrix, distance): 

90 """Get indices for each node that are distance or less away. 

91 

92 Parameters: 

93 

94 distanceMatrix: any one of scipy.sparse matrices (NxN) 

95 Matrix containing distance information of atoms. Get it e.g. with 

96 :func:`~ase.neighborlist.get_distance_matrix` 

97 distance: integer 

98 Number of steps to allow. 

99 

100 Returns: 

101 

102 return: list of length N 

103 List of length N. return[i] has all indices connected to item i. 

104 

105 The distance matrix only contains shortest paths, so when looking for 

106 distances longer than one, we need to add the lower values for cases 

107 where atoms are connected via a shorter path too. 

108 """ 

109 shape = distanceMatrix.get_shape() 

110 indices = [] 

111 # iterate over rows 

112 for i in range(shape[0]): 

113 row = distanceMatrix.getrow(i)[0] 

114 # find all non-zero 

115 found = sp.find(row) 

116 # screen for smaller or equal distance 

117 equal = np.where(found[-1] <= distance)[0] 

118 # found[1] contains the indexes 

119 indices.append([found[1][x] for x in equal]) 

120 return indices 

121 

122 

123def mic(dr, cell, pbc=True): 

124 """ 

125 Apply minimum image convention to an array of distance vectors. 

126 

127 Parameters: 

128 

129 dr : array_like 

130 Array of distance vectors. 

131 cell : array_like 

132 Simulation cell. 

133 pbc : array_like, optional 

134 Periodic boundary conditions in x-, y- and z-direction. Default is to 

135 assume periodic boundaries in all directions. 

136 

137 Returns: 

138 

139 dr : array 

140 Array of distance vectors, wrapped according to the minimum image 

141 convention. 

142 """ 

143 dr, _ = find_mic(dr, cell, pbc) 

144 return dr 

145 

146 

147def primitive_neighbor_list(quantities, pbc, cell, positions, cutoff, 

148 numbers=None, self_interaction=False, 

149 use_scaled_positions=False, max_nbins=1e6): 

150 """Compute a neighbor list for an atomic configuration. 

151 

152 Atoms outside periodic boundaries are mapped into the box. Atoms 

153 outside nonperiodic boundaries are included in the neighbor list 

154 but complexity of neighbor list search for those can become n^2. 

155 

156 The neighbor list is sorted by first atom index 'i', but not by second 

157 atom index 'j'. 

158 

159 Parameters: 

160 

161 quantities: str 

162 Quantities to compute by the neighbor list algorithm. Each character 

163 in this string defines a quantity. They are returned in a tuple of 

164 the same order. Possible quantities are 

165 

166 * 'i' : first atom index 

167 * 'j' : second atom index 

168 * 'd' : absolute distance 

169 * 'D' : distance vector 

170 * 'S' : shift vector (number of cell boundaries crossed by the bond 

171 between atom i and j). With the shift vector S, the 

172 distances D between atoms can be computed from: 

173 D = positions[j]-positions[i]+S.dot(cell) 

174 pbc: array_like 

175 3-tuple indicating giving periodic boundaries in the three Cartesian 

176 directions. 

177 cell: 3x3 matrix 

178 Unit cell vectors. 

179 positions: list of xyz-positions 

180 Atomic positions. Anything that can be converted to an ndarray of 

181 shape (n, 3) will do: [(x1,y1,z1), (x2,y2,z2), ...]. If 

182 use_scaled_positions is set to true, this must be scaled positions. 

183 cutoff: float or dict 

184 Cutoff for neighbor search. It can be: 

185 

186 * A single float: This is a global cutoff for all elements. 

187 * A dictionary: This specifies cutoff values for element 

188 pairs. Specification accepts element numbers of symbols. 

189 Example: {(1, 6): 1.1, (1, 1): 1.0, ('C', 'C'): 1.85} 

190 * A list/array with a per atom value: This specifies the radius of 

191 an atomic sphere for each atoms. If spheres overlap, atoms are 

192 within each others neighborhood. See 

193 :func:`~ase.neighborlist.natural_cutoffs` 

194 for an example on how to get such a list. 

195 self_interaction: bool 

196 Return the atom itself as its own neighbor if set to true. 

197 Default: False 

198 use_scaled_positions: bool 

199 If set to true, positions are expected to be scaled positions. 

200 max_nbins: int 

201 Maximum number of bins used in neighbor search. This is used to limit 

202 the maximum amount of memory required by the neighbor list. 

203 

204 Returns: 

205 

206 i, j, ... : array 

207 Tuple with arrays for each quantity specified above. Indices in `i` 

208 are returned in ascending order 0..len(a)-1, but the order of (i,j) 

209 pairs is not guaranteed. 

210 

211 """ 

212 

213 # Naming conventions: Suffixes indicate the dimension of an array. The 

214 # following convention is used here: 

215 # c: Cartesian index, can have values 0, 1, 2 

216 # i: Global atom index, can have values 0..len(a)-1 

217 # xyz: Bin index, three values identifying x-, y- and z-component of a 

218 # spatial bin that is used to make neighbor search O(n) 

219 # b: Linearized version of the 'xyz' bin index 

220 # a: Bin-local atom index, i.e. index identifying an atom *within* a 

221 # bin 

222 # p: Pair index, can have value 0 or 1 

223 # n: (Linear) neighbor index 

224 

225 # Return empty neighbor list if no atoms are passed here 

226 if len(positions) == 0: 

227 empty_types = dict(i=(int, (0, )), 

228 j=(int, (0, )), 

229 D=(float, (0, 3)), 

230 d=(float, (0, )), 

231 S=(int, (0, 3))) 

232 retvals = [] 

233 for i in quantities: 

234 dtype, shape = empty_types[i] 

235 retvals += [np.array([], dtype=dtype).reshape(shape)] 

236 if len(retvals) == 1: 

237 return retvals[0] 

238 else: 

239 return tuple(retvals) 

240 

241 # Compute reciprocal lattice vectors. 

242 b1_c, b2_c, b3_c = np.linalg.pinv(cell).T 

243 

244 # Compute distances of cell faces. 

245 l1 = np.linalg.norm(b1_c) 

246 l2 = np.linalg.norm(b2_c) 

247 l3 = np.linalg.norm(b3_c) 

248 face_dist_c = np.array([1 / l1 if l1 > 0 else 1, 

249 1 / l2 if l2 > 0 else 1, 

250 1 / l3 if l3 > 0 else 1]) 

251 

252 if isinstance(cutoff, dict): 

253 max_cutoff = max(cutoff.values()) 

254 else: 

255 if np.isscalar(cutoff): 

256 max_cutoff = cutoff 

257 else: 

258 cutoff = np.asarray(cutoff) 

259 max_cutoff = 2 * np.max(cutoff) 

260 

261 # We use a minimum bin size of 3 A 

262 bin_size = max(max_cutoff, 3) 

263 # Compute number of bins such that a sphere of radius cutoff fits into 

264 # eight neighboring bins. 

265 nbins_c = np.maximum((face_dist_c / bin_size).astype(int), [1, 1, 1]) 

266 nbins = np.prod(nbins_c) 

267 # Make sure we limit the amount of memory used by the explicit bins. 

268 while nbins > max_nbins: 

269 nbins_c = np.maximum(nbins_c // 2, [1, 1, 1]) 

270 nbins = np.prod(nbins_c) 

271 

272 # Compute over how many bins we need to loop in the neighbor list search. 

273 neigh_search_x, neigh_search_y, neigh_search_z = \ 

274 np.ceil(bin_size * nbins_c / face_dist_c).astype(int) 

275 

276 # If we only have a single bin and the system is not periodic, then we 

277 # do not need to search neighboring bins 

278 neigh_search_x = 0 if nbins_c[0] == 1 and not pbc[0] else neigh_search_x 

279 neigh_search_y = 0 if nbins_c[1] == 1 and not pbc[1] else neigh_search_y 

280 neigh_search_z = 0 if nbins_c[2] == 1 and not pbc[2] else neigh_search_z 

281 

282 # Sort atoms into bins. 

283 if use_scaled_positions: 

284 scaled_positions_ic = positions 

285 positions = np.dot(scaled_positions_ic, cell) 

286 else: 

287 scaled_positions_ic = np.linalg.solve(complete_cell(cell).T, 

288 positions.T).T 

289 bin_index_ic = np.floor(scaled_positions_ic * nbins_c).astype(int) 

290 cell_shift_ic = np.zeros_like(bin_index_ic) 

291 

292 for c in range(3): 

293 if pbc[c]: 

294 # (Note: np.divmod does not exist in older numpies) 

295 cell_shift_ic[:, c], bin_index_ic[:, c] = \ 

296 divmod(bin_index_ic[:, c], nbins_c[c]) 

297 else: 

298 bin_index_ic[:, c] = np.clip(bin_index_ic[:, c], 0, nbins_c[c] - 1) 

299 

300 # Convert Cartesian bin index to unique scalar bin index. 

301 bin_index_i = (bin_index_ic[:, 0] + 

302 nbins_c[0] * (bin_index_ic[:, 1] + 

303 nbins_c[1] * bin_index_ic[:, 2])) 

304 

305 # atom_i contains atom index in new sort order. 

306 atom_i = np.argsort(bin_index_i) 

307 bin_index_i = bin_index_i[atom_i] 

308 

309 # Find max number of atoms per bin 

310 max_natoms_per_bin = np.bincount(bin_index_i).max() 

311 

312 # Sort atoms into bins: atoms_in_bin_ba contains for each bin (identified 

313 # by its scalar bin index) a list of atoms inside that bin. This list is 

314 # homogeneous, i.e. has the same size *max_natoms_per_bin* for all bins. 

315 # The list is padded with -1 values. 

316 atoms_in_bin_ba = -np.ones([nbins, max_natoms_per_bin], dtype=int) 

317 for i in range(max_natoms_per_bin): 

318 # Create a mask array that identifies the first atom of each bin. 

319 mask = np.append([True], bin_index_i[:-1] != bin_index_i[1:]) 

320 # Assign all first atoms. 

321 atoms_in_bin_ba[bin_index_i[mask], i] = atom_i[mask] 

322 

323 # Remove atoms that we just sorted into atoms_in_bin_ba. The next 

324 # "first" atom will be the second and so on. 

325 mask = np.logical_not(mask) 

326 atom_i = atom_i[mask] 

327 bin_index_i = bin_index_i[mask] 

328 

329 # Make sure that all atoms have been sorted into bins. 

330 assert len(atom_i) == 0 

331 assert len(bin_index_i) == 0 

332 

333 # Now we construct neighbor pairs by pairing up all atoms within a bin or 

334 # between bin and neighboring bin. atom_pairs_pn is a helper buffer that 

335 # contains all potential pairs of atoms between two bins, i.e. it is a list 

336 # of length max_natoms_per_bin**2. 

337 atom_pairs_pn = np.indices((max_natoms_per_bin, max_natoms_per_bin), 

338 dtype=int) 

339 atom_pairs_pn = atom_pairs_pn.reshape(2, -1) 

340 

341 # Initialized empty neighbor list buffers. 

342 first_at_neightuple_nn = [] 

343 secnd_at_neightuple_nn = [] 

344 cell_shift_vector_x_n = [] 

345 cell_shift_vector_y_n = [] 

346 cell_shift_vector_z_n = [] 

347 

348 # This is the main neighbor list search. We loop over neighboring bins and 

349 # then construct all possible pairs of atoms between two bins, assuming 

350 # that each bin contains exactly max_natoms_per_bin atoms. We then throw 

351 # out pairs involving pad atoms with atom index -1 below. 

352 binz_xyz, biny_xyz, binx_xyz = np.meshgrid(np.arange(nbins_c[2]), 

353 np.arange(nbins_c[1]), 

354 np.arange(nbins_c[0]), 

355 indexing='ij') 

356 # The memory layout of binx_xyz, biny_xyz, binz_xyz is such that computing 

357 # the respective bin index leads to a linearly increasing consecutive list. 

358 # The following assert statement succeeds: 

359 # b_b = (binx_xyz + nbins_c[0] * (biny_xyz + nbins_c[1] * 

360 # binz_xyz)).ravel() 

361 # assert (b_b == np.arange(np.prod(nbins_c))).all() 

362 

363 # First atoms in pair. 

364 _first_at_neightuple_n = atoms_in_bin_ba[:, atom_pairs_pn[0]] 

365 for dz in range(-neigh_search_z, neigh_search_z + 1): 

366 for dy in range(-neigh_search_y, neigh_search_y + 1): 

367 for dx in range(-neigh_search_x, neigh_search_x + 1): 

368 # Bin index of neighboring bin and shift vector. 

369 shiftx_xyz, neighbinx_xyz = divmod(binx_xyz + dx, nbins_c[0]) 

370 shifty_xyz, neighbiny_xyz = divmod(biny_xyz + dy, nbins_c[1]) 

371 shiftz_xyz, neighbinz_xyz = divmod(binz_xyz + dz, nbins_c[2]) 

372 neighbin_b = (neighbinx_xyz + nbins_c[0] * 

373 (neighbiny_xyz + nbins_c[1] * neighbinz_xyz) 

374 ).ravel() 

375 

376 # Second atom in pair. 

377 _secnd_at_neightuple_n = \ 

378 atoms_in_bin_ba[neighbin_b][:, atom_pairs_pn[1]] 

379 

380 # Shift vectors. 

381 _cell_shift_vector_x_n = \ 

382 np.resize(shiftx_xyz.reshape(-1, 1), 

383 (max_natoms_per_bin**2, shiftx_xyz.size)).T 

384 _cell_shift_vector_y_n = \ 

385 np.resize(shifty_xyz.reshape(-1, 1), 

386 (max_natoms_per_bin**2, shifty_xyz.size)).T 

387 _cell_shift_vector_z_n = \ 

388 np.resize(shiftz_xyz.reshape(-1, 1), 

389 (max_natoms_per_bin**2, shiftz_xyz.size)).T 

390 

391 # We have created too many pairs because we assumed each bin 

392 # has exactly max_natoms_per_bin atoms. Remove all surperfluous 

393 # pairs. Those are pairs that involve an atom with index -1. 

394 mask = np.logical_and(_first_at_neightuple_n != -1, 

395 _secnd_at_neightuple_n != -1) 

396 if mask.sum() > 0: 

397 first_at_neightuple_nn += [_first_at_neightuple_n[mask]] 

398 secnd_at_neightuple_nn += [_secnd_at_neightuple_n[mask]] 

399 cell_shift_vector_x_n += [_cell_shift_vector_x_n[mask]] 

400 cell_shift_vector_y_n += [_cell_shift_vector_y_n[mask]] 

401 cell_shift_vector_z_n += [_cell_shift_vector_z_n[mask]] 

402 

403 # Flatten overall neighbor list. 

404 first_at_neightuple_n = np.concatenate(first_at_neightuple_nn) 

405 secnd_at_neightuple_n = np.concatenate(secnd_at_neightuple_nn) 

406 cell_shift_vector_n = np.transpose([np.concatenate(cell_shift_vector_x_n), 

407 np.concatenate(cell_shift_vector_y_n), 

408 np.concatenate(cell_shift_vector_z_n)]) 

409 

410 # Add global cell shift to shift vectors 

411 cell_shift_vector_n += cell_shift_ic[first_at_neightuple_n] - \ 

412 cell_shift_ic[secnd_at_neightuple_n] 

413 

414 # Remove all self-pairs that do not cross the cell boundary. 

415 if not self_interaction: 

416 m = np.logical_not(np.logical_and( 

417 first_at_neightuple_n == secnd_at_neightuple_n, 

418 (cell_shift_vector_n == 0).all(axis=1))) 

419 first_at_neightuple_n = first_at_neightuple_n[m] 

420 secnd_at_neightuple_n = secnd_at_neightuple_n[m] 

421 cell_shift_vector_n = cell_shift_vector_n[m] 

422 

423 # For nonperiodic directions, remove any bonds that cross the domain 

424 # boundary. 

425 for c in range(3): 

426 if not pbc[c]: 

427 m = cell_shift_vector_n[:, c] == 0 

428 first_at_neightuple_n = first_at_neightuple_n[m] 

429 secnd_at_neightuple_n = secnd_at_neightuple_n[m] 

430 cell_shift_vector_n = cell_shift_vector_n[m] 

431 

432 # Sort neighbor list. 

433 i = np.argsort(first_at_neightuple_n) 

434 first_at_neightuple_n = first_at_neightuple_n[i] 

435 secnd_at_neightuple_n = secnd_at_neightuple_n[i] 

436 cell_shift_vector_n = cell_shift_vector_n[i] 

437 

438 # Compute distance vectors. 

439 distance_vector_nc = positions[secnd_at_neightuple_n] - \ 

440 positions[first_at_neightuple_n] + \ 

441 cell_shift_vector_n.dot(cell) 

442 abs_distance_vector_n = \ 

443 np.sqrt(np.sum(distance_vector_nc * distance_vector_nc, axis=1)) 

444 

445 # We have still created too many pairs. Only keep those with distance 

446 # smaller than max_cutoff. 

447 mask = abs_distance_vector_n < max_cutoff 

448 first_at_neightuple_n = first_at_neightuple_n[mask] 

449 secnd_at_neightuple_n = secnd_at_neightuple_n[mask] 

450 cell_shift_vector_n = cell_shift_vector_n[mask] 

451 distance_vector_nc = distance_vector_nc[mask] 

452 abs_distance_vector_n = abs_distance_vector_n[mask] 

453 

454 if isinstance(cutoff, dict) and numbers is not None: 

455 # If cutoff is a dictionary, then the cutoff radii are specified per 

456 # element pair. We now have a list up to maximum cutoff. 

457 per_pair_cutoff_n = np.zeros_like(abs_distance_vector_n) 

458 for (atomic_number1, atomic_number2), c in cutoff.items(): 

459 try: 

460 atomic_number1 = atomic_numbers[atomic_number1] 

461 except KeyError: 

462 pass 

463 try: 

464 atomic_number2 = atomic_numbers[atomic_number2] 

465 except KeyError: 

466 pass 

467 if atomic_number1 == atomic_number2: 

468 mask = np.logical_and( 

469 numbers[first_at_neightuple_n] == atomic_number1, 

470 numbers[secnd_at_neightuple_n] == atomic_number2) 

471 else: 

472 mask = np.logical_or( 

473 np.logical_and( 

474 numbers[first_at_neightuple_n] == atomic_number1, 

475 numbers[secnd_at_neightuple_n] == atomic_number2), 

476 np.logical_and( 

477 numbers[first_at_neightuple_n] == atomic_number2, 

478 numbers[secnd_at_neightuple_n] == atomic_number1)) 

479 per_pair_cutoff_n[mask] = c 

480 mask = abs_distance_vector_n < per_pair_cutoff_n 

481 first_at_neightuple_n = first_at_neightuple_n[mask] 

482 secnd_at_neightuple_n = secnd_at_neightuple_n[mask] 

483 cell_shift_vector_n = cell_shift_vector_n[mask] 

484 distance_vector_nc = distance_vector_nc[mask] 

485 abs_distance_vector_n = abs_distance_vector_n[mask] 

486 elif not np.isscalar(cutoff): 

487 # If cutoff is neither a dictionary nor a scalar, then we assume it is 

488 # a list or numpy array that contains atomic radii. Atoms are neighbors 

489 # if their radii overlap. 

490 mask = abs_distance_vector_n < \ 

491 cutoff[first_at_neightuple_n] + cutoff[secnd_at_neightuple_n] 

492 first_at_neightuple_n = first_at_neightuple_n[mask] 

493 secnd_at_neightuple_n = secnd_at_neightuple_n[mask] 

494 cell_shift_vector_n = cell_shift_vector_n[mask] 

495 distance_vector_nc = distance_vector_nc[mask] 

496 abs_distance_vector_n = abs_distance_vector_n[mask] 

497 

498 # Assemble return tuple. 

499 retvals = [] 

500 for q in quantities: 

501 if q == 'i': 

502 retvals += [first_at_neightuple_n] 

503 elif q == 'j': 

504 retvals += [secnd_at_neightuple_n] 

505 elif q == 'D': 

506 retvals += [distance_vector_nc] 

507 elif q == 'd': 

508 retvals += [abs_distance_vector_n] 

509 elif q == 'S': 

510 retvals += [cell_shift_vector_n] 

511 else: 

512 raise ValueError('Unsupported quantity specified.') 

513 if len(retvals) == 1: 

514 return retvals[0] 

515 else: 

516 return tuple(retvals) 

517 

518 

519def neighbor_list(quantities, a, cutoff, self_interaction=False, 

520 max_nbins=1e6): 

521 """Compute a neighbor list for an atomic configuration. 

522 

523 Atoms outside periodic boundaries are mapped into the box. Atoms 

524 outside nonperiodic boundaries are included in the neighbor list 

525 but complexity of neighbor list search for those can become n^2. 

526 

527 The neighbor list is sorted by first atom index 'i', but not by second 

528 atom index 'j'. 

529 

530 Parameters: 

531 

532 quantities: str 

533 Quantities to compute by the neighbor list algorithm. Each character 

534 in this string defines a quantity. They are returned in a tuple of 

535 the same order. Possible quantities are: 

536 

537 * 'i' : first atom index 

538 * 'j' : second atom index 

539 * 'd' : absolute distance 

540 * 'D' : distance vector 

541 * 'S' : shift vector (number of cell boundaries crossed by the bond 

542 between atom i and j). With the shift vector S, the 

543 distances D between atoms can be computed from: 

544 D = a.positions[j]-a.positions[i]+S.dot(a.cell) 

545 a: :class:`ase.Atoms` 

546 Atomic configuration. 

547 cutoff: float or dict 

548 Cutoff for neighbor search. It can be: 

549 

550 * A single float: This is a global cutoff for all elements. 

551 * A dictionary: This specifies cutoff values for element 

552 pairs. Specification accepts element numbers of symbols. 

553 Example: {(1, 6): 1.1, (1, 1): 1.0, ('C', 'C'): 1.85} 

554 * A list/array with a per atom value: This specifies the radius of 

555 an atomic sphere for each atoms. If spheres overlap, atoms are 

556 within each others neighborhood. See 

557 :func:`~ase.neighborlist.natural_cutoffs` 

558 for an example on how to get such a list. 

559 

560 self_interaction: bool 

561 Return the atom itself as its own neighbor if set to true. 

562 Default: False 

563 max_nbins: int 

564 Maximum number of bins used in neighbor search. This is used to limit 

565 the maximum amount of memory required by the neighbor list. 

566 

567 Returns: 

568 

569 i, j, ...: array 

570 Tuple with arrays for each quantity specified above. Indices in `i` 

571 are returned in ascending order 0..len(a), but the order of (i,j) 

572 pairs is not guaranteed. 

573 

574 Examples: 

575 

576 Examples assume Atoms object *a* and numpy imported as *np*. 

577 

578 1. Coordination counting:: 

579 

580 i = neighbor_list('i', a, 1.85) 

581 coord = np.bincount(i) 

582 

583 2. Coordination counting with different cutoffs for each pair of species:: 

584 

585 i = neighbor_list('i', a, 

586 {('H', 'H'): 1.1, ('C', 'H'): 1.3, ('C', 'C'): 1.85}) 

587 coord = np.bincount(i) 

588 

589 3. Pair distribution function:: 

590 

591 d = neighbor_list('d', a, 10.00) 

592 h, bin_edges = np.histogram(d, bins=100) 

593 pdf = h/(4*np.pi/3*( 

594 bin_edges[1:]**3 - bin_edges[:-1]**3)) * a.get_volume()/len(a) 

595 

596 4. Pair potential:: 

597 

598 i, j, d, D = neighbor_list('ijdD', a, 5.0) 

599 energy = (-C/d**6).sum() 

600 forces = (6*C/d**5 * (D/d).T).T 

601 forces_x = np.bincount(j, weights=forces[:, 0], minlength=len(a)) - \ 

602 np.bincount(i, weights=forces[:, 0], minlength=len(a)) 

603 forces_y = np.bincount(j, weights=forces[:, 1], minlength=len(a)) - \ 

604 np.bincount(i, weights=forces[:, 1], minlength=len(a)) 

605 forces_z = np.bincount(j, weights=forces[:, 2], minlength=len(a)) - \ 

606 np.bincount(i, weights=pair_forces[:, 2], minlength=len(a)) 

607 

608 5. Dynamical matrix for a pair potential stored in a block sparse format:: 

609 

610 from scipy.sparse import bsr_matrix 

611 i, j, dr, abs_dr = neighbor_list('ijDd', atoms) 

612 energy = (dr.T / abs_dr).T 

613 dynmat = -(dde * (energy.reshape(-1, 3, 1) 

614 * energy.reshape(-1, 1, 3)).T).T \ 

615 -(de / abs_dr * (np.eye(3, dtype=energy.dtype) - \ 

616 (energy.reshape(-1, 3, 1) * energy.reshape(-1, 1, 3))).T).T 

617 dynmat_bsr = bsr_matrix((dynmat, j, first_i), 

618 shape=(3*len(a), 3*len(a))) 

619 

620 dynmat_diag = np.empty((len(a), 3, 3)) 

621 for x in range(3): 

622 for y in range(3): 

623 dynmat_diag[:, x, y] = -np.bincount(i, weights=dynmat[:, x, y]) 

624 

625 dynmat_bsr += bsr_matrix((dynmat_diag, np.arange(len(a)), 

626 np.arange(len(a) + 1)), 

627 shape=(3 * len(a), 3 * len(a))) 

628 

629 """ 

630 return primitive_neighbor_list(quantities, a.pbc, 

631 a.get_cell(complete=True), 

632 a.positions, cutoff, numbers=a.numbers, 

633 self_interaction=self_interaction, 

634 max_nbins=max_nbins) 

635 

636 

637def first_neighbors(natoms, first_atom): 

638 """ 

639 Compute an index array pointing to the ranges within the neighbor list that 

640 contain the neighbors for a certain atom. 

641 

642 Parameters: 

643 

644 natoms : int 

645 Total number of atom. 

646 first_atom : array_like 

647 Array containing the first atom 'i' of the neighbor tuple returned 

648 by the neighbor list. 

649 

650 Returns: 

651 

652 seed : array 

653 Array containing pointers to the start and end location of the 

654 neighbors of a certain atom. Neighbors of atom k have indices from s[k] 

655 to s[k+1]-1. 

656 """ 

657 if len(first_atom) == 0: 

658 return np.zeros(natoms + 1, dtype=int) 

659 # Create a seed array (which is returned by this function) populated with 

660 # -1. 

661 seed = -np.ones(natoms + 1, dtype=int) 

662 

663 first_atom = np.asarray(first_atom) 

664 

665 # Mask array contains all position where the number in the (sorted) array 

666 # with first atoms (in the neighbor pair) changes. 

667 mask = first_atom[:-1] != first_atom[1:] 

668 

669 # Seed array needs to start at 0 

670 seed[first_atom[0]] = 0 

671 # Seed array needs to stop at the length of the neighbor list 

672 seed[-1] = len(first_atom) 

673 # Populate all intermediate seed with the index of where the mask array is 

674 # true, i.e. the index where the first_atom array changes. 

675 seed[first_atom[1:][mask]] = (np.arange(len(mask)) + 1)[mask] 

676 

677 # Now fill all remaining -1 value with the value in the seed array right 

678 # behind them. (There are no neighbor so seed[i] and seed[i+1] must point) 

679 # to the same index. 

680 mask = seed == -1 

681 while mask.any(): 

682 seed[mask] = seed[np.arange(natoms + 1)[mask] + 1] 

683 mask = seed == -1 

684 return seed 

685 

686 

687def get_connectivity_matrix(nl, sparse=True): 

688 """Return connectivity matrix for a given NeighborList (dtype=numpy.int8). 

689 

690 A matrix of shape (nAtoms, nAtoms) will be returned. 

691 Connected atoms i and j will have matrix[i,j] == 1, unconnected 

692 matrix[i,j] == 0. If bothways=True the matrix will be symmetric, 

693 otherwise not! 

694 

695 If *sparse* is True, a scipy csr matrix is returned. 

696 If *sparse* is False, a numpy matrix is returned. 

697 

698 Note that the old and new neighborlists might give different results 

699 for periodic systems if bothways=False. 

700 

701 Example: 

702 

703 Determine which molecule in a system atom 1 belongs to. 

704 

705 >>> from scipy import sparse 

706 

707 >>> from ase.build import molecule 

708 >>> from ase.neighborlist import get_connectivity_matrix 

709 >>> from ase.neighborlist import natural_cutoffs 

710 >>> from ase.neighborlist import NeighborList 

711 

712 >>> mol = molecule('CH3CH2OH') 

713 >>> cutoffs = natural_cutoffs(mol) 

714 >>> neighbor_list = NeighborList( 

715 ... cutoffs, self_interaction=False, bothways=True) 

716 >>> neighbor_list.update(mol) 

717 True 

718 

719 >>> matrix = neighbor_list.get_connectivity_matrix() 

720 >>> # or: matrix = get_connectivity_matrix(neighbor_list.nl) 

721 >>> n_components, component_list = sparse.csgraph.connected_components( 

722 ... matrix) 

723 >>> idx = 1 

724 >>> molIdx = component_list[idx] 

725 >>> print("There are {} molecules in the system".format(n_components)) 

726 There are 1 molecules in the system 

727 >>> print("Atom {} is part of molecule {}".format(idx, molIdx)) 

728 Atom 1 is part of molecule 0 

729 >>> molIdxs = [i for i in range(len(component_list)) 

730 ... if component_list[i] == molIdx] 

731 >>> print("Atoms are part of molecule {}: {}".format(molIdx, molIdxs)) 

732 Atoms are part of molecule 0: [0, 1, 2, 3, 4, 5, 6, 7, 8] 

733 """ 

734 

735 nAtoms = len(nl.cutoffs) 

736 

737 if nl.nupdates <= 0: 

738 raise RuntimeError( 

739 'Must call update(atoms) on your neighborlist first!') 

740 

741 if sparse: 

742 matrix = sp.dok_matrix((nAtoms, nAtoms), dtype=np.int8) 

743 else: 

744 matrix = np.zeros((nAtoms, nAtoms), dtype=np.int8) 

745 

746 for i in range(nAtoms): 

747 for idx in nl.get_neighbors(i)[0]: 

748 matrix[i, idx] = 1 

749 

750 return matrix 

751 

752 

753class NewPrimitiveNeighborList: 

754 """Neighbor list object. Wrapper around neighbor_list and first_neighbors. 

755 

756 cutoffs: list of float 

757 List of cutoff radii - one for each atom. If the spheres (defined by 

758 their cutoff radii) of two atoms overlap, they will be counted as 

759 neighbors. 

760 skin: float 

761 If no atom has moved more than the skin-distance since the 

762 last call to the 

763 :meth:`~ase.neighborlist.NewPrimitiveNeighborList.update()` 

764 method, then the neighbor list can be reused. This will save 

765 some expensive rebuilds of the list, but extra neighbors outside 

766 the cutoff will be returned. 

767 sorted: bool 

768 Sort neighbor list. 

769 self_interaction: bool 

770 Should an atom return itself as a neighbor? 

771 bothways: bool 

772 Return all neighbors. Default is to return only "half" of 

773 the neighbors. 

774 

775 Example: 

776 

777 >>> from ase.build import bulk 

778 >>> from ase.neighborlist import NewPrimitiveNeighborList 

779 

780 >>> nl = NewPrimitiveNeighborList([2.3, 1.7]) 

781 >>> atoms = bulk('Cu', 'fcc', a=3.6) 

782 >>> nl.update(atoms.pbc, atoms.get_cell(), atoms.positions) 

783 True 

784 >>> indices, offsets = nl.get_neighbors(0) 

785 """ 

786 

787 def __init__(self, cutoffs, skin=0.3, sorted=False, self_interaction=True, 

788 bothways=False, use_scaled_positions=False): 

789 self.cutoffs = np.asarray(cutoffs) + skin 

790 self.skin = skin 

791 self.sorted = sorted 

792 self.self_interaction = self_interaction 

793 self.bothways = bothways 

794 self.nupdates = 0 

795 self.use_scaled_positions = use_scaled_positions 

796 self.nneighbors = 0 

797 self.npbcneighbors = 0 

798 

799 def update(self, pbc, cell, positions, numbers=None): 

800 """Make sure the list is up to date.""" 

801 

802 if self.nupdates == 0: 

803 self.build(pbc, cell, positions, numbers=numbers) 

804 return True 

805 

806 if ((self.pbc != pbc).any() or (self.cell != cell).any() or 

807 ((self.positions - positions)**2).sum(1).max() > self.skin**2): 

808 self.build(pbc, cell, positions, numbers=numbers) 

809 return True 

810 

811 return False 

812 

813 def build(self, pbc, cell, positions, numbers=None): 

814 """Build the list. 

815 """ 

816 self.pbc = np.array(pbc, copy=True) 

817 self.cell = np.array(cell, copy=True) 

818 self.positions = np.array(positions, copy=True) 

819 

820 pair_first, pair_second, offset_vec = \ 

821 primitive_neighbor_list( 

822 'ijS', pbc, cell, positions, self.cutoffs, numbers=numbers, 

823 self_interaction=self.self_interaction, 

824 use_scaled_positions=self.use_scaled_positions) 

825 

826 if len(positions) > 0 and not self.bothways: 

827 offset_x, offset_y, offset_z = offset_vec.T 

828 

829 mask = offset_z > 0 

830 mask &= offset_y == 0 

831 mask |= offset_y > 0 

832 mask &= offset_x == 0 

833 mask |= offset_x > 0 

834 mask |= (pair_first <= pair_second) & (offset_vec == 0).all(axis=1) 

835 

836 pair_first = pair_first[mask] 

837 pair_second = pair_second[mask] 

838 offset_vec = offset_vec[mask] 

839 

840 if len(positions) > 0 and self.sorted: 

841 mask = np.argsort(pair_first * len(pair_first) + 

842 pair_second) 

843 pair_first = pair_first[mask] 

844 pair_second = pair_second[mask] 

845 offset_vec = offset_vec[mask] 

846 

847 self.pair_first = pair_first 

848 self.pair_second = pair_second 

849 self.offset_vec = offset_vec 

850 

851 # Compute the index array point to the first neighbor 

852 self.first_neigh = first_neighbors(len(positions), pair_first) 

853 

854 self.nupdates += 1 

855 

856 def get_neighbors(self, a): 

857 """Return neighbors of atom number a. 

858 

859 A list of indices and offsets to neighboring atoms is 

860 returned. The positions of the neighbor atoms can be 

861 calculated like this: 

862 

863 >>> from ase.build import bulk 

864 >>> from ase.neighborlist import NewPrimitiveNeighborList 

865 

866 >>> nl = NewPrimitiveNeighborList([2.3, 1.7]) 

867 >>> atoms = bulk('Cu', 'fcc', a=3.6) 

868 >>> nl.update(atoms.pbc, atoms.get_cell(), atoms.positions) 

869 True 

870 >>> indices, offsets = nl.get_neighbors(0) 

871 >>> for i, offset in zip(indices, offsets): 

872 ... print( 

873 ... atoms.positions[i] + offset @ atoms.get_cell() 

874 ... ) # doctest: +ELLIPSIS 

875 [3.6 ... 0. ] 

876 

877 Notice that if get_neighbors(a) gives atom b as a neighbor, 

878 then get_neighbors(b) will not return a as a neighbor - unless 

879 bothways=True was used.""" 

880 

881 return (self.pair_second[self.first_neigh[a]:self.first_neigh[a + 1]], 

882 self.offset_vec[self.first_neigh[a]:self.first_neigh[a + 1]]) 

883 

884 

885class PrimitiveNeighborList: 

886 """Neighbor list that works without Atoms objects. 

887 

888 This is less fancy, but can be used to avoid conversions between 

889 scaled and non-scaled coordinates which may affect cell offsets 

890 through rounding errors. 

891 """ 

892 

893 def __init__(self, cutoffs, skin=0.3, sorted=False, self_interaction=True, 

894 bothways=False, use_scaled_positions=False): 

895 self.cutoffs = np.asarray(cutoffs) + skin 

896 self.skin = skin 

897 self.sorted = sorted 

898 self.self_interaction = self_interaction 

899 self.bothways = bothways 

900 self.nupdates = 0 

901 self.use_scaled_positions = use_scaled_positions 

902 self.nneighbors = 0 

903 self.npbcneighbors = 0 

904 

905 def update(self, pbc, cell, coordinates): 

906 """Make sure the list is up to date.""" 

907 

908 if self.nupdates == 0: 

909 self.build(pbc, cell, coordinates) 

910 return True 

911 

912 if ((self.pbc != pbc).any() or (self.cell != cell).any() or ( 

913 (self.coordinates 

914 - coordinates)**2).sum(1).max() > self.skin**2): 

915 self.build(pbc, cell, coordinates) 

916 return True 

917 

918 return False 

919 

920 def build(self, pbc, cell, coordinates): 

921 """Build the list. 

922 

923 Coordinates are taken to be scaled or not according 

924 to self.use_scaled_positions.