Coverage for /builds/debichem-team/python-ase/ase/geometry/geometry.py: 97.56%
205 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-03-06 04:00 +0000
1# Copyright (C) 2010, Jesper Friis
2# (see accompanying license files for details).
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"""
10import itertools
12import numpy as np
14from ase.cell import Cell
15from ase.geometry import complete_cell
16from ase.geometry.minkowski_reduction import minkowski_reduce
17from ase.utils import pbc2pbc
20def translate_pretty(fractional, pbc):
21 """Translates atoms such that fractional positions are minimized."""
23 for i in range(3):
24 if not pbc[i]:
25 continue
27 indices = np.argsort(fractional[:, i])
28 sp = fractional[indices, i]
30 widths = (np.roll(sp, 1) - sp) % 1.0
31 fractional[:, i] -= sp[np.argmin(widths)]
32 fractional[:, i] %= 1.0
33 return fractional
36def wrap_positions(positions, cell, pbc=True, center=(0.5, 0.5, 0.5),
37 pretty_translation=False, eps=1e-7):
38 """Wrap positions to unit cell.
40 Returns positions changed by a multiple of the unit cell vectors to
41 fit inside the space spanned by these vectors. See also the
42 :meth:`ase.Atoms.wrap` method.
44 Parameters:
46 positions: float ndarray of shape (n, 3)
47 Positions of the atoms
48 cell: float ndarray of shape (3, 3)
49 Unit cell vectors.
50 pbc: one or 3 bool
51 For each axis in the unit cell decides whether the positions
52 will be moved along this axis.
53 center: three float
54 The positons in fractional coordinates that the new positions
55 will be nearest possible to.
56 pretty_translation: bool
57 Translates atoms such that fractional coordinates are minimized.
58 eps: float
59 Small number to prevent slightly negative coordinates from being
60 wrapped.
62 Example:
64 >>> from ase.geometry import wrap_positions
65 >>> wrap_positions([[-0.1, 1.01, -0.5]],
66 ... [[1, 0, 0], [0, 1, 0], [0, 0, 4]],
67 ... pbc=[1, 1, 0])
68 array([[ 0.9 , 0.01, -0.5 ]])
69 """
71 if not hasattr(center, '__len__'):
72 center = (center,) * 3
74 pbc = pbc2pbc(pbc)
75 shift = np.asarray(center) - 0.5 - eps
77 # Don't change coordinates when pbc is False
78 shift[np.logical_not(pbc)] = 0.0
80 assert np.asarray(cell)[np.asarray(pbc)].any(axis=1).all(), (cell, pbc)
82 cell = complete_cell(cell)
83 fractional = np.linalg.solve(cell.T,
84 np.asarray(positions).T).T - shift
86 if pretty_translation:
87 fractional = translate_pretty(fractional, pbc)
88 shift = np.asarray(center) - 0.5
89 shift[np.logical_not(pbc)] = 0.0
90 fractional += shift
91 else:
92 for i, periodic in enumerate(pbc):
93 if periodic:
94 fractional[:, i] %= 1.0
95 fractional[:, i] += shift[i]
97 return np.dot(fractional, cell)
100def get_layers(atoms, miller, tolerance=0.001):
101 """Returns two arrays describing which layer each atom belongs
102 to and the distance between the layers and origo.
104 Parameters:
106 miller: 3 integers
107 The Miller indices of the planes. Actually, any direction
108 in reciprocal space works, so if a and b are two float
109 vectors spanning an atomic plane, you can get all layers
110 parallel to this with miller=np.cross(a,b).
111 tolerance: float
112 The maximum distance in Angstrom along the plane normal for
113 counting two atoms as belonging to the same plane.
115 Returns:
117 tags: array of integres
118 Array of layer indices for each atom.
119 levels: array of floats
120 Array of distances in Angstrom from each layer to origo.
122 Example:
124 >>> import numpy as np
126 >>> from ase.spacegroup import crystal
127 >>> from ase.geometry.geometry import get_layers
129 >>> atoms = crystal('Al', [(0,0,0)], spacegroup=225, cellpar=4.05)
130 >>> np.round(atoms.positions, decimals=5) # doctest: +NORMALIZE_WHITESPACE
131 array([[ 0. , 0. , 0. ],
132 [ 0. , 2.025, 2.025],
133 [ 2.025, 0. , 2.025],
134 [ 2.025, 2.025, 0. ]])
135 >>> get_layers(atoms, (0,0,1)) # doctest: +ELLIPSIS
136 (array([0, 1, 1, 0]...), array([ 0. , 2.025]))
137 """
138 miller = np.asarray(miller)
140 metric = np.dot(atoms.cell, atoms.cell.T)
141 c = np.linalg.solve(metric.T, miller.T).T
142 miller_norm = np.sqrt(np.dot(c, miller))
143 d = np.dot(atoms.get_scaled_positions(), miller) / miller_norm
145 keys = np.argsort(d)
146 ikeys = np.argsort(keys)
147 mask = np.concatenate(([True], np.diff(d[keys]) > tolerance))
148 tags = np.cumsum(mask)[ikeys]
149 if tags.min() == 1:
150 tags -= 1
152 levels = d[keys][mask]
153 return tags, levels
156def naive_find_mic(v, cell):
157 """Finds the minimum-image representation of vector(s) v.
158 Safe to use for (pbc.all() and (norm(v_mic) < 0.5 * min(cell.lengths()))).
159 Can otherwise fail for non-orthorhombic cells.
160 Described in:
161 W. Smith, "The Minimum Image Convention in Non-Cubic MD Cells", 1989,
162 http://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.57.1696."""
163 f = Cell(cell).scaled_positions(v)
164 f -= np.floor(f + 0.5)
165 vmin = f @ cell
166 vlen = np.linalg.norm(vmin, axis=1)
167 return vmin, vlen
170def general_find_mic(v, cell, pbc=True):
171 """Finds the minimum-image representation of vector(s) v. Using the
172 Minkowski reduction the algorithm is relatively slow but safe for any cell.
173 """
175 cell = complete_cell(cell)
176 rcell, _ = minkowski_reduce(cell, pbc=pbc)
177 positions = wrap_positions(v, rcell, pbc=pbc, eps=0)
179 # In a Minkowski-reduced cell we only need to test nearest neighbors,
180 # or "Voronoi-relevant" vectors. These are a subset of combinations of
181 # [-1, 0, 1] of the reduced cell vectors.
183 # Define ranges [-1, 0, 1] for periodic directions and [0] for aperiodic
184 # directions.
185 ranges = [np.arange(-1 * p, p + 1) for p in pbc]
187 # Get Voronoi-relevant vectors.
188 # Pre-pend (0, 0, 0) to resolve issue #772
189 hkls = np.array([(0, 0, 0)] + list(itertools.product(*ranges)))
190 vrvecs = hkls @ rcell
192 # Map positions into neighbouring cells.
193 x = positions + vrvecs[:, None]
195 # Find minimum images
196 lengths = np.linalg.norm(x, axis=2)
197 indices = np.argmin(lengths, axis=0)
198 vmin = x[indices, np.arange(len(positions)), :]
199 vlen = lengths[indices, np.arange(len(positions))]
200 return vmin, vlen
203def find_mic(v, cell, pbc=True):
204 """Finds the minimum-image representation of vector(s) v using either one
205 of two find mic algorithms depending on the given cell, v and pbc."""
207 cell = Cell(cell)
208 pbc = cell.any(1) & pbc2pbc(pbc)
209 dim = np.sum(pbc)
210 v = np.asarray(v)
211 single = v.ndim == 1
212 v = np.atleast_2d(v)
214 if dim > 0:
215 naive_find_mic_is_safe = False
216 if dim == 3:
217 vmin, vlen = naive_find_mic(v, cell)
218 # naive find mic is safe only for the following condition
219 if (vlen < 0.5 * min(cell.lengths())).all():
220 naive_find_mic_is_safe = True # hence skip Minkowski reduction
222 if not naive_find_mic_is_safe:
223 vmin, vlen = general_find_mic(v, cell, pbc=pbc)
224 else:
225 vmin = v.copy()
226 vlen = np.linalg.norm(vmin, axis=1)
228 if single:
229 return vmin[0], vlen[0]
230 else:
231 return vmin, vlen
234def conditional_find_mic(vectors, cell, pbc):
235 """Return vectors and their lengths considering cell and pbc.
237 The minimum image convention is applied if cell and pbc are set.
238 This can be used like a simple version of get_distances.
239 """
240 vectors = np.array(vectors)
241 if (cell is None) != (pbc is None):
242 raise ValueError("cell or pbc must be both set or both be None")
243 if cell is not None:
244 mics = [find_mic(v, cell, pbc) for v in vectors]
245 vectors, vector_lengths = zip(*mics)
246 else:
247 vector_lengths = np.sqrt(np.add.reduce(vectors**2, axis=-1))
248 return vectors, vector_lengths
251def get_angles(v0, v1, cell=None, pbc=None):
252 """Get angles formed by two lists of vectors.
254 Calculate angle in degrees between vectors v0 and v1
256 Set a cell and pbc to enable minimum image
257 convention, otherwise angles are taken as-is.
258 """
259 (v0, v1), (nv0, nv1) = conditional_find_mic([v0, v1], cell, pbc)
261 if (nv0 <= 0).any() or (nv1 <= 0).any():
262 raise ZeroDivisionError('Undefined angle')
263 v0n = v0 / nv0[:, np.newaxis]
264 v1n = v1 / nv1[:, np.newaxis]
265 # We just normalized the vectors, but in some cases we can get
266 # bad things like 1+2e-16. These we clip away:
267 angles = np.arccos(np.einsum('ij,ij->i', v0n, v1n).clip(-1.0, 1.0))
268 return np.degrees(angles)
271def get_angles_derivatives(v0, v1, cell=None, pbc=None):
272 """Get derivatives of angles formed by two lists of vectors (v0, v1) w.r.t.
273 Cartesian coordinates in degrees.
275 Set a cell and pbc to enable minimum image
276 convention, otherwise derivatives of angles are taken as-is.
278 There is a singularity in the derivatives for sin(angle) -> 0 for which
279 a ZeroDivisionError is raised.
281 Derivative output format: [[dx_a0, dy_a0, dz_a0], [...], [..., dz_a2].
282 """
283 (v0, v1), (nv0, nv1) = conditional_find_mic([v0, v1], cell, pbc)
285 angles = np.radians(get_angles(v0, v1, cell=cell, pbc=pbc))
286 sin_angles = np.sin(angles)
287 cos_angles = np.cos(angles)
288 if (sin_angles == 0.).any(): # identify singularities
289 raise ZeroDivisionError('Singularity for derivative of a planar angle')
291 product = nv0 * nv1
292 deriv_d0 = (-(v1 / product[:, np.newaxis] # derivatives by atom 0
293 - np.einsum('ij,i->ij', v0, cos_angles / nv0**2))
294 / sin_angles[:, np.newaxis])
295 deriv_d2 = (-(v0 / product[:, np.newaxis] # derivatives by atom 2
296 - np.einsum('ij,i->ij', v1, cos_angles / nv1**2))
297 / sin_angles[:, np.newaxis])
298 deriv_d1 = -(deriv_d0 + deriv_d2) # derivatives by atom 1
299 derivs = np.stack((deriv_d0, deriv_d1, deriv_d2), axis=1)
300 return np.degrees(derivs)
303def get_dihedrals(v0, v1, v2, cell=None, pbc=None):
304 """Get dihedral angles formed by three lists of vectors.
306 Calculate dihedral angle (in degrees) between the vectors a0->a1,
307 a1->a2 and a2->a3, written as v0, v1 and v2.
309 Set a cell and pbc to enable minimum image
310 convention, otherwise angles are taken as-is.
311 """
312 (v0, v1, v2), (_, nv1, _) = conditional_find_mic([v0, v1, v2], cell, pbc)
314 v1n = v1 / nv1[:, np.newaxis]
315 # v, w: projection of v0, v2 onto plane perpendicular to v1
316 v = -v0 - np.einsum('ij,ij,ik->ik', -v0, v1n, v1n)
317 w = v2 - np.einsum('ij,ij,ik->ik', v2, v1n, v1n)
319 # formula returns 0 for undefined dihedrals; prefer ZeroDivisionError
320 undefined_v = np.all(v == 0.0, axis=1)
321 undefined_w = np.all(w == 0.0, axis=1)
322 if np.any([undefined_v, undefined_w]):
323 raise ZeroDivisionError('Undefined dihedral for planar inner angle')
325 x = np.einsum('ij,ij->i', v, w)
326 y = np.einsum('ij,ij->i', np.cross(v1n, v, axis=1), w)
327 dihedrals = np.arctan2(y, x) # dihedral angle in [-pi, pi]
328 dihedrals[dihedrals < 0.] += 2 * np.pi # dihedral angle in [0, 2*pi]
329 return np.degrees(dihedrals)
332def get_dihedrals_derivatives(v0, v1, v2, cell=None, pbc=None):
333 """Get derivatives of dihedrals formed by three lists of vectors
334 (v0, v1, v2) w.r.t Cartesian coordinates in degrees.
336 Set a cell and pbc to enable minimum image
337 convention, otherwise dihedrals are taken as-is.
339 Derivative output format: [[dx_a0, dy_a0, dz_a0], ..., [..., dz_a3]].
340 """
341 (v0, v1, v2), (nv0, nv1, nv2) = conditional_find_mic([v0, v1, v2], cell,
342 pbc)
344 v0n = v0 / nv0[:, np.newaxis]
345 v1n = v1 / nv1[:, np.newaxis]
346 v2n = v2 / nv2[:, np.newaxis]
347 normal_v01 = np.cross(v0n, v1n, axis=1)
348 normal_v12 = np.cross(v1n, v2n, axis=1)
349 cos_psi01 = np.einsum('ij,ij->i', v0n, v1n) # == np.sum(v0 * v1, axis=1)
350 sin_psi01 = np.sin(np.arccos(cos_psi01))
351 cos_psi12 = np.einsum('ij,ij->i', v1n, v2n)
352 sin_psi12 = np.sin(np.arccos(cos_psi12))
353 if (sin_psi01 == 0.).any() or (sin_psi12 == 0.).any():
354 msg = ('Undefined derivative for undefined dihedral with planar inner '
355 'angle')
356 raise ZeroDivisionError(msg)
358 deriv_d0 = -normal_v01 / (nv0 * sin_psi01**2)[:, np.newaxis] # by atom 0
359 deriv_d3 = normal_v12 / (nv2 * sin_psi12**2)[:, np.newaxis] # by atom 3
360 deriv_d1 = (((nv1 + nv0 * cos_psi01) / nv1)[:, np.newaxis] * -deriv_d0
361 + (cos_psi12 * nv2 / nv1)[:, np.newaxis] * deriv_d3) # by a1
362 deriv_d2 = (-((nv1 + nv2 * cos_psi12) / nv1)[:, np.newaxis] * deriv_d3
363 - (cos_psi01 * nv0 / nv1)[:, np.newaxis] * -deriv_d0) # by a2
364 derivs = np.stack((deriv_d0, deriv_d1, deriv_d2, deriv_d3), axis=1)
365 return np.degrees(derivs)
368def get_distances(p1, p2=None, cell=None, pbc=None):
369 """Return distance matrix of every position in p1 with every position in p2
371 If p2 is not set, it is assumed that distances between all positions in p1
372 are desired. p2 will be set to p1 in this case.
374 Use set cell and pbc to use the minimum image convention.
375 """
376 p1 = np.atleast_2d(p1)
377 if p2 is None:
378 np1 = len(p1)
379 ind1, ind2 = np.triu_indices(np1, k=1)
380 D = p1[ind2] - p1[ind1]
381 else:
382 p2 = np.atleast_2d(p2)
383 D = (p2[np.newaxis, :, :] - p1[:, np.newaxis, :]).reshape((-1, 3))
385 (D, ), (D_len, ) = conditional_find_mic([D], cell=cell, pbc=pbc)
387 if p2 is None:
388 Dout = np.zeros((np1, np1, 3))
389 Dout[(ind1, ind2)] = D
390 Dout -= np.transpose(Dout, axes=(1, 0, 2))
392 Dout_len = np.zeros((np1, np1))
393 Dout_len[(ind1, ind2)] = D_len
394 Dout_len += Dout_len.T
395 return Dout, Dout_len
397 # Expand back to matrix indexing
398 D.shape = (-1, len(p2), 3)
399 D_len.shape = (-1, len(p2))
401 return D, D_len
404def get_distances_derivatives(v0, cell=None, pbc=None):
405 """Get derivatives of distances for all vectors in v0 w.r.t. Cartesian
406 coordinates in Angstrom.
408 Set cell and pbc to use the minimum image convention.
410 There is a singularity for distances -> 0 for which a ZeroDivisionError is
411 raised.
412 Derivative output format: [[dx_a0, dy_a0, dz_a0], [dx_a1, dy_a1, dz_a1]].
413 """
414 (v0, ), (dists, ) = conditional_find_mic([v0], cell, pbc)
416 if (dists <= 0.).any(): # identify singularities
417 raise ZeroDivisionError('Singularity for derivative of a zero distance')
419 derivs_d0 = np.einsum('i,ij->ij', -1. / dists, v0) # derivatives by atom 0
420 derivs_d1 = -derivs_d0 # derivatives by atom 1
421 derivs = np.stack((derivs_d0, derivs_d1), axis=1)
422 return derivs
425def get_duplicate_atoms(atoms, cutoff=0.1, delete=False):
426 """Get list of duplicate atoms and delete them if requested.
428 Identify all atoms which lie within the cutoff radius of each other.
429 Delete one set of them if delete == True.
430 """
431 from scipy.spatial.distance import pdist
432 dists = pdist(atoms.get_positions(), 'sqeuclidean')
433 dup = np.nonzero(dists < cutoff**2)
434 rem = np.array(_row_col_from_pdist(len(atoms), dup[0]))
435 if delete:
436 if rem.size != 0:
437 del atoms[rem[:, 0]]
438 else:
439 return rem
442def _row_col_from_pdist(dim, i):
443 """Calculate the i,j index in the square matrix for an index in a
444 condensed (triangular) matrix.
445 """
446 i = np.array(i)
447 b = 1 - 2 * dim
448 x = (np.floor((-b - np.sqrt(b**2 - 8 * i)) / 2)).astype(int)
449 y = (i + x * (b + x + 2) / 2 + 1).astype(int)
450 if i.shape:
451 return list(zip(x, y))
452 else:
453 return [(x, y)]
456def permute_axes(atoms, permutation):
457 """Permute axes of unit cell and atom positions. Considers only cell and
458 atomic positions. Other vector quantities such as momenta are not
459 modified."""
460 assert (np.sort(permutation) == np.arange(3)).all()
462 permuted = atoms.copy()
463 scaled = permuted.get_scaled_positions()
464 permuted.set_cell(permuted.cell.permute_axes(permutation),
465 scale_atoms=False)
466 permuted.set_scaled_positions(scaled[:, permutation])
467 permuted.set_pbc(permuted.pbc[permutation])
468 return permuted