1# Copyright (C) 2010, Jesper Friis 

2# (see accompanying license files for details). 

3 

4"""Utility tools for atoms/geometry manipulations. 

5 - convenient creation of slabs and interfaces of 

6different orientations. 

7 - detection of duplicate atoms / atoms within cutoff radius 

8""" 

9 

10import itertools 

11import numpy as np 

12from ase.geometry import complete_cell 

13from ase.geometry.minkowski_reduction import minkowski_reduce 

14from ase.utils import pbc2pbc 

15from ase.cell import Cell 

16 

17 

18def translate_pretty(fractional, pbc): 

19 """Translates atoms such that fractional positions are minimized.""" 

20 

21 for i in range(3): 

22 if not pbc[i]: 

23 continue 

24 

25 indices = np.argsort(fractional[:, i]) 

26 sp = fractional[indices, i] 

27 

28 widths = (np.roll(sp, 1) - sp) % 1.0 

29 fractional[:, i] -= sp[np.argmin(widths)] 

30 fractional[:, i] %= 1.0 

31 return fractional 

32 

33 

34def wrap_positions(positions, cell, pbc=True, center=(0.5, 0.5, 0.5), 

35 pretty_translation=False, eps=1e-7): 

36 """Wrap positions to unit cell. 

37 

38 Returns positions changed by a multiple of the unit cell vectors to 

39 fit inside the space spanned by these vectors. See also the 

40 :meth:`ase.Atoms.wrap` method. 

41 

42 Parameters: 

43 

44 positions: float ndarray of shape (n, 3) 

45 Positions of the atoms 

46 cell: float ndarray of shape (3, 3) 

47 Unit cell vectors. 

48 pbc: one or 3 bool 

49 For each axis in the unit cell decides whether the positions 

50 will be moved along this axis. 

51 center: three float 

52 The positons in fractional coordinates that the new positions 

53 will be nearest possible to. 

54 pretty_translation: bool 

55 Translates atoms such that fractional coordinates are minimized. 

56 eps: float 

57 Small number to prevent slightly negative coordinates from being 

58 wrapped. 

59 

60 Example: 

61 

62 >>> from ase.geometry import wrap_positions 

63 >>> wrap_positions([[-0.1, 1.01, -0.5]], 

64 ... [[1, 0, 0], [0, 1, 0], [0, 0, 4]], 

65 ... pbc=[1, 1, 0]) 

66 array([[ 0.9 , 0.01, -0.5 ]]) 

67 """ 

68 

69 if not hasattr(center, '__len__'): 

70 center = (center,) * 3 

71 

72 pbc = pbc2pbc(pbc) 

73 shift = np.asarray(center) - 0.5 - eps 

74 

75 # Don't change coordinates when pbc is False 

76 shift[np.logical_not(pbc)] = 0.0 

77 

78 assert np.asarray(cell)[np.asarray(pbc)].any(axis=1).all(), (cell, pbc) 

79 

80 cell = complete_cell(cell) 

81 fractional = np.linalg.solve(cell.T, 

82 np.asarray(positions).T).T - shift 

83 

84 if pretty_translation: 

85 fractional = translate_pretty(fractional, pbc) 

86 shift = np.asarray(center) - 0.5 

87 shift[np.logical_not(pbc)] = 0.0 

88 fractional += shift 

89 else: 

90 for i, periodic in enumerate(pbc): 

91 if periodic: 

92 fractional[:, i] %= 1.0 

93 fractional[:, i] += shift[i] 

94 

95 return np.dot(fractional, cell) 

96 

97 

98def get_layers(atoms, miller, tolerance=0.001): 

99 """Returns two arrays describing which layer each atom belongs 

100 to and the distance between the layers and origo. 

101 

102 Parameters: 

103 

104 miller: 3 integers 

105 The Miller indices of the planes. Actually, any direction 

106 in reciprocal space works, so if a and b are two float 

107 vectors spanning an atomic plane, you can get all layers 

108 parallel to this with miller=np.cross(a,b). 

109 tolerance: float 

110 The maximum distance in Angstrom along the plane normal for 

111 counting two atoms as belonging to the same plane. 

112 

113 Returns: 

114 

115 tags: array of integres 

116 Array of layer indices for each atom. 

117 levels: array of floats 

118 Array of distances in Angstrom from each layer to origo. 

119 

120 Example: 

121 

122 >>> import numpy as np 

123 >>> from ase.spacegroup import crystal 

124 >>> atoms = crystal('Al', [(0,0,0)], spacegroup=225, cellpar=4.05) 

125 >>> np.round(atoms.positions, decimals=5) 

126 array([[ 0. , 0. , 0. ], 

127 [ 0. , 2.025, 2.025], 

128 [ 2.025, 0. , 2.025], 

129 [ 2.025, 2.025, 0. ]]) 

130 >>> get_layers(atoms, (0,0,1)) # doctest: +ELLIPSIS 

131 (array([0, 1, 1, 0]...), array([ 0. , 2.025])) 

132 """ 

133 miller = np.asarray(miller) 

134 

135 metric = np.dot(atoms.cell, atoms.cell.T) 

136 c = np.linalg.solve(metric.T, miller.T).T 

137 miller_norm = np.sqrt(np.dot(c, miller)) 

138 d = np.dot(atoms.get_scaled_positions(), miller) / miller_norm 

139 

140 keys = np.argsort(d) 

141 ikeys = np.argsort(keys) 

142 mask = np.concatenate(([True], np.diff(d[keys]) > tolerance)) 

143 tags = np.cumsum(mask)[ikeys] 

144 if tags.min() == 1: 

145 tags -= 1 

146 

147 levels = d[keys][mask] 

148 return tags, levels 

149 

150 

151def naive_find_mic(v, cell): 

152 """Finds the minimum-image representation of vector(s) v. 

153 Safe to use for (pbc.all() and (norm(v_mic) < 0.5 * min(cell.lengths()))). 

154 Can otherwise fail for non-orthorhombic cells. 

155 Described in: 

156 W. Smith, "The Minimum Image Convention in Non-Cubic MD Cells", 1989, 

157 http://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.57.1696.""" 

158 f = Cell(cell).scaled_positions(v) 

159 f -= np.floor(f + 0.5) 

160 vmin = f @ cell 

161 vlen = np.linalg.norm(vmin, axis=1) 

162 return vmin, vlen 

163 

164 

165def general_find_mic(v, cell, pbc=True): 

166 """Finds the minimum-image representation of vector(s) v. Using the 

167 Minkowski reduction the algorithm is relatively slow but safe for any cell. 

168 """ 

169 

170 cell = complete_cell(cell) 

171 rcell, _ = minkowski_reduce(cell, pbc=pbc) 

172 positions = wrap_positions(v, rcell, pbc=pbc, eps=0) 

173 

174 # In a Minkowski-reduced cell we only need to test nearest neighbors, 

175 # or "Voronoi-relevant" vectors. These are a subset of combinations of 

176 # [-1, 0, 1] of the reduced cell vectors. 

177 

178 # Define ranges [-1, 0, 1] for periodic directions and [0] for aperiodic 

179 # directions. 

180 ranges = [np.arange(-1 * p, p + 1) for p in pbc] 

181 

182 # Get Voronoi-relevant vectors. 

183 # Pre-pend (0, 0, 0) to resolve issue #772 

184 hkls = np.array([(0, 0, 0)] + list(itertools.product(*ranges))) 

185 vrvecs = hkls @ rcell 

186 

187 # Map positions into neighbouring cells. 

188 x = positions + vrvecs[:, None] 

189 

190 # Find minimum images 

191 lengths = np.linalg.norm(x, axis=2) 

192 indices = np.argmin(lengths, axis=0) 

193 vmin = x[indices, np.arange(len(positions)), :] 

194 vlen = lengths[indices, np.arange(len(positions))] 

195 return vmin, vlen 

196 

197 

198def find_mic(v, cell, pbc=True): 

199 """Finds the minimum-image representation of vector(s) v using either one 

200 of two find mic algorithms depending on the given cell, v and pbc.""" 

201 

202 cell = Cell(cell) 

203 pbc = cell.any(1) & pbc2pbc(pbc) 

204 dim = np.sum(pbc) 

205 v = np.asarray(v) 

206 single = v.ndim == 1 

207 v = np.atleast_2d(v) 

208 

209 if dim > 0: 

210 naive_find_mic_is_safe = False 

211 if dim == 3: 

212 vmin, vlen = naive_find_mic(v, cell) 

213 # naive find mic is safe only for the following condition 

214 if (vlen < 0.5 * min(cell.lengths())).all(): 

215 naive_find_mic_is_safe = True # hence skip Minkowski reduction 

216 

217 if not naive_find_mic_is_safe: 

218 vmin, vlen = general_find_mic(v, cell, pbc=pbc) 

219 else: 

220 vmin = v.copy() 

221 vlen = np.linalg.norm(vmin, axis=1) 

222 

223 if single: 

224 return vmin[0], vlen[0] 

225 else: 

226 return vmin, vlen 

227 

228 

229def conditional_find_mic(vectors, cell, pbc): 

230 """Return list of vector arrays and corresponding list of vector lengths 

231 for a given list of vector arrays. The minimum image convention is applied 

232 if cell and pbc are set. Can be used like a simple version of get_distances. 

233 """ 

234 if (cell is None) != (pbc is None): 

235 raise ValueError("cell or pbc must be both set or both be None") 

236 if cell is not None: 

237 mics = [find_mic(v, cell, pbc) for v in vectors] 

238 vectors, vector_lengths = zip(*mics) 

239 else: 

240 vector_lengths = np.linalg.norm(vectors, axis=2) 

241 return [np.asarray(v) for v in vectors], vector_lengths 

242 

243 

244def get_angles(v0, v1, cell=None, pbc=None): 

245 """Get angles formed by two lists of vectors. 

246 

247 Calculate angle in degrees between vectors v0 and v1 

248 

249 Set a cell and pbc to enable minimum image 

250 convention, otherwise angles are taken as-is. 

251 """ 

252 (v0, v1), (nv0, nv1) = conditional_find_mic([v0, v1], cell, pbc) 

253 

254 if (nv0 <= 0).any() or (nv1 <= 0).any(): 

255 raise ZeroDivisionError('Undefined angle') 

256 v0n = v0 / nv0[:, np.newaxis] 

257 v1n = v1 / nv1[:, np.newaxis] 

258 # We just normalized the vectors, but in some cases we can get 

259 # bad things like 1+2e-16. These we clip away: 

260 angles = np.arccos(np.einsum('ij,ij->i', v0n, v1n).clip(-1.0, 1.0)) 

261 return np.degrees(angles) 

262 

263 

264def get_angles_derivatives(v0, v1, cell=None, pbc=None): 

265 """Get derivatives of angles formed by two lists of vectors (v0, v1) w.r.t. 

266 Cartesian coordinates in degrees. 

267 

268 Set a cell and pbc to enable minimum image 

269 convention, otherwise derivatives of angles are taken as-is. 

270 

271 There is a singularity in the derivatives for sin(angle) -> 0 for which 

272 a ZeroDivisionError is raised. 

273 

274 Derivative output format: [[dx_a0, dy_a0, dz_a0], [...], [..., dz_a2]. 

275 """ 

276 (v0, v1), (nv0, nv1) = conditional_find_mic([v0, v1], cell, pbc) 

277 

278 angles = np.radians(get_angles(v0, v1, cell=cell, pbc=pbc)) 

279 sin_angles = np.sin(angles) 

280 cos_angles = np.cos(angles) 

281 if (sin_angles == 0.).any(): # identify singularities 

282 raise ZeroDivisionError('Singularity for angle derivative') 

283 

284 product = nv0 * nv1 

285 deriv_d0 = (-(v1 / product[:, np.newaxis] # derivatives by atom 0 

286 - np.einsum('ij,i->ij', v0, cos_angles / nv0**2)) 

287 / sin_angles[:, np.newaxis]) 

288 deriv_d2 = (-(v0 / product[:, np.newaxis] # derivatives by atom 2 

289 - np.einsum('ij,i->ij', v1, cos_angles / nv1**2)) 

290 / sin_angles[:, np.newaxis]) 

291 deriv_d1 = -(deriv_d0 + deriv_d2) # derivatives by atom 1 

292 derivs = np.stack((deriv_d0, deriv_d1, deriv_d2), axis=1) 

293 return np.degrees(derivs) 

294 

295 

296def get_dihedrals(v0, v1, v2, cell=None, pbc=None): 

297 """Get dihedral angles formed by three lists of vectors. 

298 

299 Calculate dihedral angle (in degrees) between the vectors a0->a1, 

300 a1->a2 and a2->a3, written as v0, v1 and v2. 

301 

302 Set a cell and pbc to enable minimum image 

303 convention, otherwise angles are taken as-is. 

304 """ 

305 (v0, v1, v2), (_, nv1, _) = conditional_find_mic([v0, v1, v2], cell, pbc) 

306 

307 v1n = v1 / nv1[:, np.newaxis] 

308 # v, w: projection of v0, v2 onto plane perpendicular to v1 

309 v = -v0 - np.einsum('ij,ij,ik->ik', -v0, v1n, v1n) 

310 w = v2 - np.einsum('ij,ij,ik->ik', v2, v1n, v1n) 

311 

312 # formula returns 0 for undefined dihedrals; prefer ZeroDivisionError 

313 undefined_v = np.all(v == 0.0, axis=1) 

314 undefined_w = np.all(w == 0.0, axis=1) 

315 if np.any([undefined_v, undefined_w]): 

316 raise ZeroDivisionError('Undefined dihedral') 

317 

318 x = np.einsum('ij,ij->i', v, w) 

319 y = np.einsum('ij,ij->i', np.cross(v1n, v, axis=1), w) 

320 dihedrals = np.arctan2(y, x) # dihedral angle in [-pi, pi] 

321 dihedrals[dihedrals < 0.] += 2 * np.pi # dihedral angle in [0, 2*pi] 

322 return np.degrees(dihedrals) 

323 

324 

325def get_dihedrals_derivatives(v0, v1, v2, cell=None, pbc=None): 

326 """Get derivatives of dihedrals formed by three lists of vectors 

327 (v0, v1, v2) w.r.t Cartesian coordinates in degrees. 

328 

329 Set a cell and pbc to enable minimum image 

330 convention, otherwise dihedrals are taken as-is. 

331 

332 Derivative output format: [[dx_a0, dy_a0, dz_a0], ..., [..., dz_a3]]. 

333 """ 

334 (v0, v1, v2), (nv0, nv1, nv2) = conditional_find_mic([v0, v1, v2], cell, 

335 pbc) 

336 

337 v0 /= nv0[:, np.newaxis] 

338 v1 /= nv1[:, np.newaxis] 

339 v2 /= nv2[:, np.newaxis] 

340 normal_v01 = np.cross(v0, v1, axis=1) 

341 normal_v12 = np.cross(v1, v2, axis=1) 

342 cos_psi01 = np.einsum('ij,ij->i', v0, v1) # == np.sum(v0 * v1, axis=1) 

343 sin_psi01 = np.sin(np.arccos(cos_psi01)) 

344 cos_psi12 = np.einsum('ij,ij->i', v1, v2) 

345 sin_psi12 = np.sin(np.arccos(cos_psi12)) 

346 if (sin_psi01 == 0.).any() or (sin_psi12 == 0.).any(): 

347 raise ZeroDivisionError('Undefined derivative for undefined dihedral') 

348 

349 deriv_d0 = -normal_v01 / (nv0 * sin_psi01**2)[:, np.newaxis] # by atom 0 

350 deriv_d3 = normal_v12 / (nv2 * sin_psi12**2)[:, np.newaxis] # by atom 3 

351 deriv_d1 = (((nv1 + nv0 * cos_psi01) / nv1)[:, np.newaxis] * -deriv_d0 

352 + (cos_psi12 * nv2 / nv1)[:, np.newaxis] * deriv_d3) # by a1 

353 deriv_d2 = (-((nv1 + nv2 * cos_psi12) / nv1)[:, np.newaxis] * deriv_d3 

354 - (cos_psi01 * nv0 / nv1)[:, np.newaxis] * -deriv_d0) # by a2 

355 derivs = np.stack((deriv_d0, deriv_d1, deriv_d2, deriv_d3), axis=1) 

356 return np.degrees(derivs) 

357 

358 

359def get_distances(p1, p2=None, cell=None, pbc=None): 

360 """Return distance matrix of every position in p1 with every position in p2 

361 

362 If p2 is not set, it is assumed that distances between all positions in p1 

363 are desired. p2 will be set to p1 in this case. 

364 

365 Use set cell and pbc to use the minimum image convention. 

366 """ 

367 p1 = np.atleast_2d(p1) 

368 if p2 is None: 

369 np1 = len(p1) 

370 ind1, ind2 = np.triu_indices(np1, k=1) 

371 D = p1[ind2] - p1[ind1] 

372 else: 

373 p2 = np.atleast_2d(p2) 

374 D = (p2[np.newaxis, :, :] - p1[:, np.newaxis, :]).reshape((-1, 3)) 

375 

376 (D, ), (D_len, ) = conditional_find_mic([D], cell=cell, pbc=pbc) 

377 

378 if p2 is None: 

379 Dout = np.zeros((np1, np1, 3)) 

380 Dout[(ind1, ind2)] = D 

381 Dout -= np.transpose(Dout, axes=(1, 0, 2)) 

382 

383 Dout_len = np.zeros((np1, np1)) 

384 Dout_len[(ind1, ind2)] = D_len 

385 Dout_len += Dout_len.T 

386 return Dout, Dout_len 

387 

388 # Expand back to matrix indexing 

389 D.shape = (-1, len(p2), 3) 

390 D_len.shape = (-1, len(p2)) 

391 

392 return D, D_len 

393 

394 

395def get_distances_derivatives(v0, cell=None, pbc=None): 

396 """Get derivatives of distances for all vectors in v0 w.r.t. Cartesian 

397 coordinates in Angstrom. 

398 

399 Set cell and pbc to use the minimum image convention. 

400 

401 There is a singularity for distances -> 0 for which a ZeroDivisionError is 

402 raised. 

403 Derivative output format: [[dx_a0, dy_a0, dz_a0], [dx_a1, dy_a1, dz_a1]]. 

404 """ 

405 (v0, ), (dists, ) = conditional_find_mic([v0], cell, pbc) 

406 

407 if (dists <= 0.).any(): # identify singularities 

408 raise ZeroDivisionError('Singularity for distance derivative') 

409 

410 derivs_d0 = np.einsum('i,ij->ij', -1. / dists, v0) # derivatives by atom 0 

411 derivs_d1 = -derivs_d0 # derivatives by atom 1 

412 derivs = np.stack((derivs_d0, derivs_d1), axis=1) 

413 return derivs 

414 

415 

416def get_duplicate_atoms(atoms, cutoff=0.1, delete=False): 

417 """Get list of duplicate atoms and delete them if requested. 

418 

419 Identify all atoms which lie within the cutoff radius of each other. 

420 Delete one set of them if delete == True. 

421 """ 

422 from scipy.spatial.distance import pdist 

423 dists = pdist(atoms.get_positions(), 'sqeuclidean') 

424 dup = np.nonzero(dists < cutoff**2) 

425 rem = np.array(_row_col_from_pdist(len(atoms), dup[0])) 

426 if delete: 

427 if rem.size != 0: 

428 del atoms[rem[:, 0]] 

429 else: 

430 return rem 

431 

432 

433def _row_col_from_pdist(dim, i): 

434 """Calculate the i,j index in the square matrix for an index in a 

435 condensed (triangular) matrix. 

436 """ 

437 i = np.array(i) 

438 b = 1 - 2 * dim 

439 x = (np.floor((-b - np.sqrt(b**2 - 8 * i)) / 2)).astype(int) 

440 y = (i + x * (b + x + 2) / 2 + 1).astype(int) 

441 if i.shape: 

442 return list(zip(x, y)) 

443 else: 

444 return [(x, y)] 

445 

446 

447def permute_axes(atoms, permutation): 

448 """Permute axes of unit cell and atom positions. Considers only cell and 

449 atomic positions. Other vector quantities such as momenta are not 

450 modified.""" 

451 assert (np.sort(permutation) == np.arange(3)).all() 

452 

453 permuted = atoms.copy() 

454 scaled = permuted.get_scaled_positions() 

455 permuted.set_cell(permuted.cell.permute_axes(permutation), 

456 scale_atoms=False) 

457 permuted.set_scaled_positions(scaled[:, permutation]) 

458 permuted.set_pbc(permuted.pbc[permutation]) 

459 return permuted