Coverage for ase / geometry / geometry.py: 97.95%
195 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-03-30 08:22 +0000
1# fmt: off
3# Copyright (C) 2010, Jesper Friis
4# (see accompanying license files for details).
6"""Utility tools for atoms/geometry manipulations.
7 - convenient creation of slabs and interfaces of
8different orientations.
9 - detection of duplicate atoms / atoms within cutoff radius
10"""
12import itertools
14import numpy as np
16from ase.cell import Cell
17from ase.geometry import complete_cell
18from ase.geometry.minkowski_reduction import minkowski_reduce
19from ase.utils import pbc2pbc
22def translate_pretty(fractional, pbc):
23 """Translates atoms such that fractional positions are minimized."""
25 for i in range(3):
26 if not pbc[i]:
27 continue
29 indices = np.argsort(fractional[:, i])
30 sp = fractional[indices, i]
32 widths = (np.roll(sp, 1) - sp) % 1.0
33 fractional[:, i] -= sp[np.argmin(widths)]
34 fractional[:, i] %= 1.0
35 return fractional
38def wrap_positions(positions, cell, pbc=True, center=(0.5, 0.5, 0.5),
39 pretty_translation=False, eps=1e-7):
40 """Wrap positions to unit cell.
42 Returns positions changed by a multiple of the unit cell vectors to
43 fit inside the space spanned by these vectors. See also the
44 :meth:`ase.Atoms.wrap` method.
46 Parameters
47 ----------
49 positions: float ndarray of shape (n, 3)
50 Positions of the atoms
51 cell: float ndarray of shape (3, 3)
52 Unit cell vectors.
53 pbc: one or 3 bool
54 For each axis in the unit cell decides whether the positions
55 will be moved along this axis.
56 center: three float
57 The positons in fractional coordinates that the new positions
58 will be nearest possible to.
59 pretty_translation: bool
60 Translates atoms such that fractional coordinates are minimized.
61 eps: float
62 Small number to prevent slightly negative coordinates from being
63 wrapped.
65 Examples
66 --------
67 >>> from ase.geometry import wrap_positions
68 >>> wrap_positions([[-0.1, 1.01, -0.5]],
69 ... [[1, 0, 0], [0, 1, 0], [0, 0, 4]],
70 ... pbc=[1, 1, 0])
71 array([[ 0.9 , 0.01, -0.5 ]])
72 """
74 if not hasattr(center, '__len__'):
75 center = (center,) * 3
77 pbc = pbc2pbc(pbc)
78 shift = np.asarray(center) - 0.5 - eps
80 # Don't change coordinates when pbc is False
81 shift[np.logical_not(pbc)] = 0.0
83 assert np.asarray(cell)[np.asarray(pbc)].any(axis=1).all(), (cell, pbc)
85 cell = complete_cell(cell)
86 fractional = np.linalg.solve(cell.T,
87 np.asarray(positions).T).T - shift
89 if pretty_translation:
90 fractional = translate_pretty(fractional, pbc)
91 shift = np.asarray(center) - 0.5
92 shift[np.logical_not(pbc)] = 0.0
93 fractional += shift
94 else:
95 for i, periodic in enumerate(pbc):
96 if periodic:
97 fractional[:, i] %= 1.0
98 fractional[:, i] += shift[i]
100 return np.dot(fractional, cell)
103def get_layers(atoms, miller, tolerance=0.001):
104 """Returns two arrays describing which layer each atom belongs
105 to and the distance between the layers and origo.
107 Parameters
108 ----------
110 miller: 3 integers
111 The Miller indices of the planes. Actually, any direction
112 in reciprocal space works, so if a and b are two float
113 vectors spanning an atomic plane, you can get all layers
114 parallel to this with miller=np.cross(a,b).
115 tolerance: float
116 The maximum distance in Angstrom along the plane normal for
117 counting two atoms as belonging to the same plane.
119 Returns
120 -------
122 tags: array of integres
123 Array of layer indices for each atom.
124 levels: array of floats
125 Array of distances in Angstrom from each layer to origo.
127 Example:
129 >>> import numpy as np
131 >>> from ase.spacegroup import crystal
132 >>> from ase.geometry.geometry import get_layers
134 >>> atoms = crystal('Al', [(0,0,0)], spacegroup=225, cellpar=4.05)
135 >>> np.round(atoms.positions, decimals=5) # doctest: +NORMALIZE_WHITESPACE
136 array([[ 0. , 0. , 0. ],
137 [ 0. , 2.025, 2.025],
138 [ 2.025, 0. , 2.025],
139 [ 2.025, 2.025, 0. ]])
140 >>> get_layers(atoms, (0,0,1)) # doctest: +ELLIPSIS
141 (array([0, 1, 1, 0]...), array([ 0. , 2.025]))
142 """
143 miller = np.asarray(miller)
145 metric = np.dot(atoms.cell, atoms.cell.T)
146 c = np.linalg.solve(metric.T, miller.T).T
147 miller_norm = np.sqrt(np.dot(c, miller))
148 d = np.dot(atoms.get_scaled_positions(), miller) / miller_norm
150 keys = np.argsort(d)
151 ikeys = np.argsort(keys)
152 mask = np.concatenate(([True], np.diff(d[keys]) > tolerance))
153 tags = np.cumsum(mask)[ikeys]
154 if tags.min() == 1:
155 tags -= 1
157 levels = d[keys][mask]
158 return tags, levels
161def naive_find_mic(v, cell):
162 """Finds the minimum-image representation of vector(s) v.
163 Safe to use for (pbc.all() and (norm(v_mic) < 0.5 * min(cell.lengths()))).
164 Can otherwise fail for non-orthorhombic cells.
165 Described in:
166 W. Smith, "The Minimum Image Convention in Non-Cubic MD Cells", 1989,
167 http://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.57.1696."""
168 f = Cell(cell).scaled_positions(v)
169 f -= np.floor(f + 0.5)
170 vmin = f @ cell
171 vlen = np.linalg.norm(vmin, axis=1)
172 return vmin, vlen
175def general_find_mic(v, cell, pbc=True):
176 """Finds the minimum-image representation of vector(s) v. Using the
177 Minkowski reduction the algorithm is relatively slow but safe for any cell.
178 """
180 cell = complete_cell(cell)
181 rcell, _ = minkowski_reduce(cell, pbc=pbc)
182 positions = wrap_positions(v, rcell, pbc=pbc, eps=0)
184 # In a Minkowski-reduced cell we only need to test nearest neighbors,
185 # or "Voronoi-relevant" vectors. These are a subset of combinations of
186 # [-1, 0, 1] of the reduced cell vectors.
188 # Define ranges [-1, 0, 1] for periodic directions and [0] for aperiodic
189 # directions.
190 ranges = [np.arange(-1 * p, p + 1) for p in pbc]
192 # Get Voronoi-relevant vectors.
193 # Pre-pend (0, 0, 0) to resolve issue #772
194 hkls = np.array([(0, 0, 0)] + list(itertools.product(*ranges)))
195 vrvecs = hkls @ rcell
197 # Map positions into neighbouring cells.
198 x = positions + vrvecs[:, None]
200 # Find minimum images
201 lengths = np.linalg.norm(x, axis=2)
202 indices = np.argmin(lengths, axis=0)
203 vmin = x[indices, np.arange(len(positions)), :]
204 vlen = lengths[indices, np.arange(len(positions))]
205 return vmin, vlen
208def find_mic(v, cell, pbc=True):
209 """Finds the minimum-image representation of vector(s) v using either one
210 of two find mic algorithms depending on the given cell, v and pbc."""
212 cell = Cell(cell)
213 pbc = cell.any(1) & pbc2pbc(pbc)
214 dim = np.sum(pbc)
215 v = np.asarray(v)
216 single = v.ndim == 1
217 v = np.atleast_2d(v)
219 if dim > 0:
220 naive_find_mic_is_safe = False
221 if dim == 3:
222 vmin, vlen = naive_find_mic(v, cell)
223 # naive find mic is safe only for the following condition
224 if (vlen < 0.5 * min(cell.lengths())).all():
225 naive_find_mic_is_safe = True # hence skip Minkowski reduction
227 if not naive_find_mic_is_safe:
228 vmin, vlen = general_find_mic(v, cell, pbc=pbc)
229 else:
230 vmin = v.copy()
231 vlen = np.linalg.norm(vmin, axis=1)
233 if single:
234 return vmin[0], vlen[0]
235 else:
236 return vmin, vlen
239def conditional_find_mic(vectors, cell, pbc):
240 """Return vectors and their lengths considering cell and pbc.
242 The minimum image convention is applied if cell and pbc are set.
243 This can be used like a simple version of get_distances.
244 """
245 vectors = np.array(vectors)
246 if (cell is None) != (pbc is None):
247 raise ValueError("cell or pbc must be both set or both be None")
248 if cell is not None:
249 mics = [find_mic(v, cell, pbc) for v in vectors]
250 vectors, vector_lengths = zip(*mics)
251 else:
252 vector_lengths = np.sqrt(np.add.reduce(vectors**2, axis=-1))
253 return vectors, vector_lengths
256def get_angles(v0, v1, cell=None, pbc=None):
257 """Get angles formed by two lists of vectors.
259 Calculate angle in degrees between vectors v0 and v1
261 Set a cell and pbc to enable minimum image
262 convention, otherwise angles are taken as-is.
263 """
264 (v0, v1), (nv0, nv1) = conditional_find_mic([v0, v1], cell, pbc)
266 if (nv0 <= 0).any() or (nv1 <= 0).any():
267 raise ZeroDivisionError('Undefined angle')
268 v0n = v0 / nv0[:, np.newaxis]
269 v1n = v1 / nv1[:, np.newaxis]
270 # We just normalized the vectors, but in some cases we can get
271 # bad things like 1+2e-16. These we clip away:
272 angles = np.arccos(np.einsum('ij,ij->i', v0n, v1n).clip(-1.0, 1.0))
273 return np.degrees(angles)
276def get_angles_derivatives(v0, v1, cell=None, pbc=None):
277 """Get derivatives of angles formed by two lists of vectors (v0, v1) w.r.t.
278 Cartesian coordinates in degrees.
280 Set a cell and pbc to enable minimum image
281 convention, otherwise derivatives of angles are taken as-is.
283 There is a singularity in the derivatives for sin(angle) -> 0 for which
284 a ZeroDivisionError is raised.
286 Derivative output format: [[dx_a0, dy_a0, dz_a0], [...], [..., dz_a2].
287 """
288 (v0, v1), (nv0, nv1) = conditional_find_mic([v0, v1], cell, pbc)
290 angles = np.radians(get_angles(v0, v1, cell=cell, pbc=pbc))
291 sin_angles = np.sin(angles)
292 cos_angles = np.cos(angles)
293 if (sin_angles == 0.).any(): # identify singularities
294 raise ZeroDivisionError('Singularity for derivative of a planar angle')
296 product = nv0 * nv1
297 deriv_d0 = (-(v1 / product[:, np.newaxis] # derivatives by atom 0
298 - np.einsum('ij,i->ij', v0, cos_angles / nv0**2))
299 / sin_angles[:, np.newaxis])
300 deriv_d2 = (-(v0 / product[:, np.newaxis] # derivatives by atom 2
301 - np.einsum('ij,i->ij', v1, cos_angles / nv1**2))
302 / sin_angles[:, np.newaxis])
303 deriv_d1 = -(deriv_d0 + deriv_d2) # derivatives by atom 1
304 derivs = np.stack((deriv_d0, deriv_d1, deriv_d2), axis=1)
305 return np.degrees(derivs)
308def get_dihedrals(v0, v1, v2, cell=None, pbc=None):
309 """Get dihedral angles formed by three lists of vectors.
311 Calculate dihedral angle (in degrees) between the vectors a0->a1,
312 a1->a2 and a2->a3, written as v0, v1 and v2.
314 Set a cell and pbc to enable minimum image
315 convention, otherwise angles are taken as-is.
316 """
317 (v0, v1, v2), (_, nv1, _) = conditional_find_mic([v0, v1, v2], cell, pbc)
319 v1n = v1 / nv1[:, np.newaxis]
320 # v, w: projection of v0, v2 onto plane perpendicular to v1
321 v = -v0 - np.einsum('ij,ij,ik->ik', -v0, v1n, v1n)
322 w = v2 - np.einsum('ij,ij,ik->ik', v2, v1n, v1n)
324 # formula returns 0 for undefined dihedrals; prefer ZeroDivisionError
325 undefined_v = np.all(v == 0.0, axis=1)
326 undefined_w = np.all(w == 0.0, axis=1)
327 if np.any([undefined_v, undefined_w]):
328 raise ZeroDivisionError('Undefined dihedral for planar inner angle')
330 x = np.einsum('ij,ij->i', v, w)
331 y = np.einsum('ij,ij->i', np.cross(v1n, v, axis=1), w)
332 dihedrals = np.arctan2(y, x) # dihedral angle in [-pi, pi]
333 dihedrals[dihedrals < 0.] += 2 * np.pi # dihedral angle in [0, 2*pi]
334 return np.degrees(dihedrals)
337def get_dihedrals_derivatives(v0, v1, v2, cell=None, pbc=None):
338 """Get derivatives of dihedrals formed by three lists of vectors
339 (v0, v1, v2) w.r.t Cartesian coordinates in degrees.
341 Set a cell and pbc to enable minimum image
342 convention, otherwise dihedrals are taken as-is.
344 Derivative output format: [[dx_a0, dy_a0, dz_a0], ..., [..., dz_a3]].
345 """
346 (v0, v1, v2), (nv0, nv1, nv2) = conditional_find_mic([v0, v1, v2], cell,
347 pbc)
349 v0n = v0 / nv0[:, np.newaxis]
350 v1n = v1 / nv1[:, np.newaxis]
351 v2n = v2 / nv2[:, np.newaxis]
352 normal_v01 = np.cross(v0n, v1n, axis=1)
353 normal_v12 = np.cross(v1n, v2n, axis=1)
354 cos_psi01 = np.einsum('ij,ij->i', v0n, v1n) # == np.sum(v0 * v1, axis=1)
355 sin_psi01 = np.sin(np.arccos(cos_psi01))
356 cos_psi12 = np.einsum('ij,ij->i', v1n, v2n)
357 sin_psi12 = np.sin(np.arccos(cos_psi12))
358 if (sin_psi01 == 0.).any() or (sin_psi12 == 0.).any():
359 msg = ('Undefined derivative for undefined dihedral with planar inner '
360 'angle')
361 raise ZeroDivisionError(msg)
363 deriv_d0 = -normal_v01 / (nv0 * sin_psi01**2)[:, np.newaxis] # by atom 0
364 deriv_d3 = normal_v12 / (nv2 * sin_psi12**2)[:, np.newaxis] # by atom 3
365 deriv_d1 = (((nv1 + nv0 * cos_psi01) / nv1)[:, np.newaxis] * -deriv_d0
366 + (cos_psi12 * nv2 / nv1)[:, np.newaxis] * deriv_d3) # by a1
367 deriv_d2 = (-((nv1 + nv2 * cos_psi12) / nv1)[:, np.newaxis] * deriv_d3
368 - (cos_psi01 * nv0 / nv1)[:, np.newaxis] * -deriv_d0) # by a2
369 derivs = np.stack((deriv_d0, deriv_d1, deriv_d2, deriv_d3), axis=1)
370 return np.degrees(derivs)
373def get_distances(p1, p2=None, cell=None, pbc=None):
374 """Return distance matrix of every position in p1 with every position in p2
376 If p2 is not set, it is assumed that distances between all positions in p1
377 are desired. p2 will be set to p1 in this case.
379 Use set cell and pbc to use the minimum image convention.
380 """
381 p1 = np.atleast_2d(p1)
382 if p2 is None:
383 np1 = len(p1)
384 ind1, ind2 = np.triu_indices(np1, k=1)
385 D = p1[ind2] - p1[ind1]
386 else:
387 p2 = np.atleast_2d(p2)
388 D = (p2[np.newaxis, :, :] - p1[:, np.newaxis, :]).reshape((-1, 3))
390 (D, ), (D_len, ) = conditional_find_mic([D], cell=cell, pbc=pbc)
392 if p2 is None:
393 Dout = np.zeros((np1, np1, 3))
394 Dout[(ind1, ind2)] = D
395 Dout -= np.transpose(Dout, axes=(1, 0, 2))
397 Dout_len = np.zeros((np1, np1))
398 Dout_len[(ind1, ind2)] = D_len
399 Dout_len += Dout_len.T
400 return Dout, Dout_len
402 # Expand back to matrix indexing
403 D.shape = (-1, len(p2), 3)
404 D_len.shape = (-1, len(p2))
406 return D, D_len
409def get_distances_derivatives(v0, cell=None, pbc=None):
410 """Get derivatives of distances for all vectors in v0 w.r.t. Cartesian
411 coordinates in Angstrom.
413 Set cell and pbc to use the minimum image convention.
415 There is a singularity for distances -> 0 for which a ZeroDivisionError is
416 raised.
417 Derivative output format: [[dx_a0, dy_a0, dz_a0], [dx_a1, dy_a1, dz_a1]].
418 """
419 (v0, ), (dists, ) = conditional_find_mic([v0], cell, pbc)
421 if (dists <= 0.).any(): # identify singularities
422 raise ZeroDivisionError('Singularity for derivative of a zero distance')
424 derivs_d0 = np.einsum('i,ij->ij', -1. / dists, v0) # derivatives by atom 0
425 derivs_d1 = -derivs_d0 # derivatives by atom 1
426 derivs = np.stack((derivs_d0, derivs_d1), axis=1)
427 return derivs
430def get_duplicate_atoms(atoms, cutoff=0.1, delete=False):
431 """Get list of duplicate atoms and delete them if requested.
433 Identify all atoms which lie within the cutoff radius of each other.
434 Delete one set of them if delete == True.
435 """
436 dists = get_distances(atoms.positions, cell=atoms.cell, pbc=atoms.pbc)[1]
437 dup = np.argwhere(dists < cutoff)
438 dup = dup[dup[:, 0] < dup[:, 1]] # indices at upper triangle
439 if delete and dup.size != 0:
440 del atoms[dup[:, 0]]
441 return dup
444def permute_axes(atoms, permutation):
445 """Permute axes of unit cell and atom positions. Considers only cell and
446 atomic positions. Other vector quantities such as momenta are not
447 modified."""
448 assert (np.sort(permutation) == np.arange(3)).all()
450 permuted = atoms.copy()
451 scaled = permuted.get_scaled_positions()
452 permuted.set_cell(permuted.cell.permute_axes(permutation),
453 scale_atoms=False)
454 permuted.set_scaled_positions(scaled[:, permutation])
455 permuted.set_pbc(permuted.pbc[permutation])
456 return permuted