Coverage for /builds/ase/ase/ase/geometry/geometry.py: 97.95%
195 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-08-02 00:12 +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:
48 positions: float ndarray of shape (n, 3)
49 Positions of the atoms
50 cell: float ndarray of shape (3, 3)
51 Unit cell vectors.
52 pbc: one or 3 bool
53 For each axis in the unit cell decides whether the positions
54 will be moved along this axis.
55 center: three float
56 The positons in fractional coordinates that the new positions
57 will be nearest possible to.
58 pretty_translation: bool
59 Translates atoms such that fractional coordinates are minimized.
60 eps: float
61 Small number to prevent slightly negative coordinates from being
62 wrapped.
64 Example:
66 >>> from ase.geometry import wrap_positions
67 >>> wrap_positions([[-0.1, 1.01, -0.5]],
68 ... [[1, 0, 0], [0, 1, 0], [0, 0, 4]],
69 ... pbc=[1, 1, 0])
70 array([[ 0.9 , 0.01, -0.5 ]])
71 """
73 if not hasattr(center, '__len__'):
74 center = (center,) * 3
76 pbc = pbc2pbc(pbc)
77 shift = np.asarray(center) - 0.5 - eps
79 # Don't change coordinates when pbc is False
80 shift[np.logical_not(pbc)] = 0.0
82 assert np.asarray(cell)[np.asarray(pbc)].any(axis=1).all(), (cell, pbc)
84 cell = complete_cell(cell)
85 fractional = np.linalg.solve(cell.T,
86 np.asarray(positions).T).T - shift
88 if pretty_translation:
89 fractional = translate_pretty(fractional, pbc)
90 shift = np.asarray(center) - 0.5
91 shift[np.logical_not(pbc)] = 0.0
92 fractional += shift
93 else:
94 for i, periodic in enumerate(pbc):
95 if periodic:
96 fractional[:, i] %= 1.0
97 fractional[:, i] += shift[i]
99 return np.dot(fractional, cell)
102def get_layers(atoms, miller, tolerance=0.001):
103 """Returns two arrays describing which layer each atom belongs
104 to and the distance between the layers and origo.
106 Parameters:
108 miller: 3 integers
109 The Miller indices of the planes. Actually, any direction
110 in reciprocal space works, so if a and b are two float
111 vectors spanning an atomic plane, you can get all layers
112 parallel to this with miller=np.cross(a,b).
113 tolerance: float
114 The maximum distance in Angstrom along the plane normal for
115 counting two atoms as belonging to the same plane.
117 Returns:
119 tags: array of integres
120 Array of layer indices for each atom.
121 levels: array of floats
122 Array of distances in Angstrom from each layer to origo.
124 Example:
126 >>> import numpy as np
128 >>> from ase.spacegroup import crystal
129 >>> from ase.geometry.geometry import get_layers
131 >>> atoms = crystal('Al', [(0,0,0)], spacegroup=225, cellpar=4.05)
132 >>> np.round(atoms.positions, decimals=5) # doctest: +NORMALIZE_WHITESPACE
133 array([[ 0. , 0. , 0. ],
134 [ 0. , 2.025, 2.025],
135 [ 2.025, 0. , 2.025],
136 [ 2.025, 2.025, 0. ]])
137 >>> get_layers(atoms, (0,0,1)) # doctest: +ELLIPSIS
138 (array([0, 1, 1, 0]...), array([ 0. , 2.025]))
139 """
140 miller = np.asarray(miller)
142 metric = np.dot(atoms.cell, atoms.cell.T)
143 c = np.linalg.solve(metric.T, miller.T).T
144 miller_norm = np.sqrt(np.dot(c, miller))
145 d = np.dot(atoms.get_scaled_positions(), miller) / miller_norm
147 keys = np.argsort(d)
148 ikeys = np.argsort(keys)
149 mask = np.concatenate(([True], np.diff(d[keys]) > tolerance))
150 tags = np.cumsum(mask)[ikeys]
151 if tags.min() == 1:
152 tags -= 1
154 levels = d[keys][mask]
155 return tags, levels
158def naive_find_mic(v, cell):
159 """Finds the minimum-image representation of vector(s) v.
160 Safe to use for (pbc.all() and (norm(v_mic) < 0.5 * min(cell.lengths()))).
161 Can otherwise fail for non-orthorhombic cells.
162 Described in:
163 W. Smith, "The Minimum Image Convention in Non-Cubic MD Cells", 1989,
164 http://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.57.1696."""
165 f = Cell(cell).scaled_positions(v)
166 f -= np.floor(f + 0.5)
167 vmin = f @ cell
168 vlen = np.linalg.norm(vmin, axis=1)
169 return vmin, vlen
172def general_find_mic(v, cell, pbc=True):
173 """Finds the minimum-image representation of vector(s) v. Using the
174 Minkowski reduction the algorithm is relatively slow but safe for any cell.
175 """
177 cell = complete_cell(cell)
178 rcell, _ = minkowski_reduce(cell, pbc=pbc)
179 positions = wrap_positions(v, rcell, pbc=pbc, eps=0)
181 # In a Minkowski-reduced cell we only need to test nearest neighbors,
182 # or "Voronoi-relevant" vectors. These are a subset of combinations of
183 # [-1, 0, 1] of the reduced cell vectors.
185 # Define ranges [-1, 0, 1] for periodic directions and [0] for aperiodic
186 # directions.
187 ranges = [np.arange(-1 * p, p + 1) for p in pbc]
189 # Get Voronoi-relevant vectors.
190 # Pre-pend (0, 0, 0) to resolve issue #772
191 hkls = np.array([(0, 0, 0)] + list(itertools.product(*ranges)))
192 vrvecs = hkls @ rcell
194 # Map positions into neighbouring cells.
195 x = positions + vrvecs[:, None]
197 # Find minimum images
198 lengths = np.linalg.norm(x, axis=2)
199 indices = np.argmin(lengths, axis=0)
200 vmin = x[indices, np.arange(len(positions)), :]
201 vlen = lengths[indices, np.arange(len(positions))]
202 return vmin, vlen
205def find_mic(v, cell, pbc=True):
206 """Finds the minimum-image representation of vector(s) v using either one
207 of two find mic algorithms depending on the given cell, v and pbc."""
209 cell = Cell(cell)
210 pbc = cell.any(1) & pbc2pbc(pbc)
211 dim = np.sum(pbc)
212 v = np.asarray(v)
213 single = v.ndim == 1
214 v = np.atleast_2d(v)
216 if dim > 0:
217 naive_find_mic_is_safe = False
218 if dim == 3:
219 vmin, vlen = naive_find_mic(v, cell)
220 # naive find mic is safe only for the following condition
221 if (vlen < 0.5 * min(cell.lengths())).all():
222 naive_find_mic_is_safe = True # hence skip Minkowski reduction
224 if not naive_find_mic_is_safe:
225 vmin, vlen = general_find_mic(v, cell, pbc=pbc)
226 else:
227 vmin = v.copy()
228 vlen = np.linalg.norm(vmin, axis=1)
230 if single:
231 return vmin[0], vlen[0]
232 else:
233 return vmin, vlen
236def conditional_find_mic(vectors, cell, pbc):
237 """Return vectors and their lengths considering cell and pbc.
239 The minimum image convention is applied if cell and pbc are set.
240 This can be used like a simple version of get_distances.
241 """
242 vectors = np.array(vectors)
243 if (cell is None) != (pbc is None):
244 raise ValueError("cell or pbc must be both set or both be None")
245 if cell is not None:
246 mics = [find_mic(v, cell, pbc) for v in vectors]
247 vectors, vector_lengths = zip(*mics)
248 else:
249 vector_lengths = np.sqrt(np.add.reduce(vectors**2, axis=-1))
250 return vectors, vector_lengths
253def get_angles(v0, v1, cell=None, pbc=None):
254 """Get angles formed by two lists of vectors.
256 Calculate angle in degrees between vectors v0 and v1
258 Set a cell and pbc to enable minimum image
259 convention, otherwise angles are taken as-is.
260 """
261 (v0, v1), (nv0, nv1) = conditional_find_mic([v0, v1], cell, pbc)
263 if (nv0 <= 0).any() or (nv1 <= 0).any():
264 raise ZeroDivisionError('Undefined angle')
265 v0n = v0 / nv0[:, np.newaxis]
266 v1n = v1 / nv1[:, np.newaxis]
267 # We just normalized the vectors, but in some cases we can get
268 # bad things like 1+2e-16. These we clip away:
269 angles = np.arccos(np.einsum('ij,ij->i', v0n, v1n).clip(-1.0, 1.0))
270 return np.degrees(angles)
273def get_angles_derivatives(v0, v1, cell=None, pbc=None):
274 """Get derivatives of angles formed by two lists of vectors (v0, v1) w.r.t.
275 Cartesian coordinates in degrees.
277 Set a cell and pbc to enable minimum image
278 convention, otherwise derivatives of angles are taken as-is.
280 There is a singularity in the derivatives for sin(angle) -> 0 for which
281 a ZeroDivisionError is raised.
283 Derivative output format: [[dx_a0, dy_a0, dz_a0], [...], [..., dz_a2].
284 """
285 (v0, v1), (nv0, nv1) = conditional_find_mic([v0, v1], cell, pbc)
287 angles = np.radians(get_angles(v0, v1, cell=cell, pbc=pbc))
288 sin_angles = np.sin(angles)
289 cos_angles = np.cos(angles)
290 if (sin_angles == 0.).any(): # identify singularities
291 raise ZeroDivisionError('Singularity for derivative of a planar angle')
293 product = nv0 * nv1
294 deriv_d0 = (-(v1 / product[:, np.newaxis] # derivatives by atom 0
295 - np.einsum('ij,i->ij', v0, cos_angles / nv0**2))
296 / sin_angles[:, np.newaxis])
297 deriv_d2 = (-(v0 / product[:, np.newaxis] # derivatives by atom 2
298 - np.einsum('ij,i->ij', v1, cos_angles / nv1**2))
299 / sin_angles[:, np.newaxis])
300 deriv_d1 = -(deriv_d0 + deriv_d2) # derivatives by atom 1
301 derivs = np.stack((deriv_d0, deriv_d1, deriv_d2), axis=1)
302 return np.degrees(derivs)
305def get_dihedrals(v0, v1, v2, cell=None, pbc=None):
306 """Get dihedral angles formed by three lists of vectors.
308 Calculate dihedral angle (in degrees) between the vectors a0->a1,
309 a1->a2 and a2->a3, written as v0, v1 and v2.
311 Set a cell and pbc to enable minimum image
312 convention, otherwise angles are taken as-is.
313 """
314 (v0, v1, v2), (_, nv1, _) = conditional_find_mic([v0, v1, v2], cell, pbc)
316 v1n = v1 / nv1[:, np.newaxis]
317 # v, w: projection of v0, v2 onto plane perpendicular to v1
318 v = -v0 - np.einsum('ij,ij,ik->ik', -v0, v1n, v1n)
319 w = v2 - np.einsum('ij,ij,ik->ik', v2, v1n, v1n)
321 # formula returns 0 for undefined dihedrals; prefer ZeroDivisionError
322 undefined_v = np.all(v == 0.0, axis=1)
323 undefined_w = np.all(w == 0.0, axis=1)
324 if np.any([undefined_v, undefined_w]):
325 raise ZeroDivisionError('Undefined dihedral for planar inner angle')
327 x = np.einsum('ij,ij->i', v, w)
328 y = np.einsum('ij,ij->i', np.cross(v1n, v, axis=1), w)
329 dihedrals = np.arctan2(y, x) # dihedral angle in [-pi, pi]
330 dihedrals[dihedrals < 0.] += 2 * np.pi # dihedral angle in [0, 2*pi]
331 return np.degrees(dihedrals)
334def get_dihedrals_derivatives(v0, v1, v2, cell=None, pbc=None):
335 """Get derivatives of dihedrals formed by three lists of vectors
336 (v0, v1, v2) w.r.t Cartesian coordinates in degrees.
338 Set a cell and pbc to enable minimum image
339 convention, otherwise dihedrals are taken as-is.
341 Derivative output format: [[dx_a0, dy_a0, dz_a0], ..., [..., dz_a3]].
342 """
343 (v0, v1, v2), (nv0, nv1, nv2) = conditional_find_mic([v0, v1, v2], cell,
344 pbc)
346 v0n = v0 / nv0[:, np.newaxis]
347 v1n = v1 / nv1[:, np.newaxis]
348 v2n = v2 / nv2[:, np.newaxis]
349 normal_v01 = np.cross(v0n, v1n, axis=1)
350 normal_v12 = np.cross(v1n, v2n, axis=1)
351 cos_psi01 = np.einsum('ij,ij->i', v0n, v1n) # == np.sum(v0 * v1, axis=1)
352 sin_psi01 = np.sin(np.arccos(cos_psi01))
353 cos_psi12 = np.einsum('ij,ij->i', v1n, v2n)
354 sin_psi12 = np.sin(np.arccos(cos_psi12))
355 if (sin_psi01 == 0.).any() or (sin_psi12 == 0.).any():
356 msg = ('Undefined derivative for undefined dihedral with planar inner '
357 'angle')
358 raise ZeroDivisionError(msg)
360 deriv_d0 = -normal_v01 / (nv0 * sin_psi01**2)[:, np.newaxis] # by atom 0
361 deriv_d3 = normal_v12 / (nv2 * sin_psi12**2)[:, np.newaxis] # by atom 3
362 deriv_d1 = (((nv1 + nv0 * cos_psi01) / nv1)[:, np.newaxis] * -deriv_d0
363 + (cos_psi12 * nv2 / nv1)[:, np.newaxis] * deriv_d3) # by a1
364 deriv_d2 = (-((nv1 + nv2 * cos_psi12) / nv1)[:, np.newaxis] * deriv_d3
365 - (cos_psi01 * nv0 / nv1)[:, np.newaxis] * -deriv_d0) # by a2
366 derivs = np.stack((deriv_d0, deriv_d1, deriv_d2, deriv_d3), axis=1)
367 return np.degrees(derivs)
370def get_distances(p1, p2=None, cell=None, pbc=None):
371 """Return distance matrix of every position in p1 with every position in p2
373 If p2 is not set, it is assumed that distances between all positions in p1
374 are desired. p2 will be set to p1 in this case.
376 Use set cell and pbc to use the minimum image convention.
377 """
378 p1 = np.atleast_2d(p1)
379 if p2 is None:
380 np1 = len(p1)
381 ind1, ind2 = np.triu_indices(np1, k=1)
382 D = p1[ind2] - p1[ind1]
383 else:
384 p2 = np.atleast_2d(p2)
385 D = (p2[np.newaxis, :, :] - p1[:, np.newaxis, :]).reshape((-1, 3))
387 (D, ), (D_len, ) = conditional_find_mic([D], cell=cell, pbc=pbc)
389 if p2 is None:
390 Dout = np.zeros((np1, np1, 3))
391 Dout[(ind1, ind2)] = D
392 Dout -= np.transpose(Dout, axes=(1, 0, 2))
394 Dout_len = np.zeros((np1, np1))
395 Dout_len[(ind1, ind2)] = D_len
396 Dout_len += Dout_len.T
397 return Dout, Dout_len
399 # Expand back to matrix indexing
400 D.shape = (-1, len(p2), 3)
401 D_len.shape = (-1, len(p2))
403 return D, D_len
406def get_distances_derivatives(v0, cell=None, pbc=None):
407 """Get derivatives of distances for all vectors in v0 w.r.t. Cartesian
408 coordinates in Angstrom.
410 Set cell and pbc to use the minimum image convention.
412 There is a singularity for distances -> 0 for which a ZeroDivisionError is
413 raised.
414 Derivative output format: [[dx_a0, dy_a0, dz_a0], [dx_a1, dy_a1, dz_a1]].
415 """
416 (v0, ), (dists, ) = conditional_find_mic([v0], cell, pbc)
418 if (dists <= 0.).any(): # identify singularities
419 raise ZeroDivisionError('Singularity for derivative of a zero distance')
421 derivs_d0 = np.einsum('i,ij->ij', -1. / dists, v0) # derivatives by atom 0
422 derivs_d1 = -derivs_d0 # derivatives by atom 1
423 derivs = np.stack((derivs_d0, derivs_d1), axis=1)
424 return derivs
427def get_duplicate_atoms(atoms, cutoff=0.1, delete=False):
428 """Get list of duplicate atoms and delete them if requested.
430 Identify all atoms which lie within the cutoff radius of each other.
431 Delete one set of them if delete == True.
432 """
433 dists = get_distances(atoms.positions, cell=atoms.cell, pbc=atoms.pbc)[1]
434 dup = np.argwhere(dists < cutoff)
435 dup = dup[dup[:, 0] < dup[:, 1]] # indices at upper triangle
436 if delete and dup.size != 0:
437 del atoms[dup[:, 0]]
438 return dup
441def permute_axes(atoms, permutation):
442 """Permute axes of unit cell and atom positions. Considers only cell and
443 atomic positions. Other vector quantities such as momenta are not
444 modified."""
445 assert (np.sort(permutation) == np.arange(3)).all()
447 permuted = atoms.copy()
448 scaled = permuted.get_scaled_positions()
449 permuted.set_cell(permuted.cell.permute_axes(permutation),
450 scale_atoms=False)
451 permuted.set_scaled_positions(scaled[:, permutation])
452 permuted.set_pbc(permuted.pbc[permutation])
453 return permuted