Coverage for dynasor/tools/structures.py: 100%
123 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-10 08:27 +0000
1from typing import Optional
2import numpy as np
3from ase import Atoms
4from ase.geometry import get_distances
5from ase.geometry import find_mic
6from dynasor.modes.atoms import Prim
7from numpy.typing import NDArray
10def get_displacements(atoms: Atoms,
11 atoms_ideal: Atoms,
12 check_mic: Optional[bool] = True,
13 cell_tol: Optional[float] = 1e-4) -> NDArray[float]:
14 """Returns the smallest possible displacements between a
15 displaced configuration relative to an ideal (reference)
16 configuration.
18 Parameters
19 ----------
20 atoms
21 Structure with displaced atoms.
22 atoms_ideal
23 Ideal configuration relative to which displacements are computed.
24 check_mic
25 Whether to check minimum image convention.
26 cell_tol
27 Cell tolerance; if the cell mismatch exceeds this value, a `ValueError` is raised.
28 """
30 if not np.array_equal(atoms.numbers, atoms_ideal.numbers):
31 raise ValueError('Atomic numbers do not match.')
32 if np.linalg.norm(atoms.cell - atoms_ideal.cell) > cell_tol:
33 raise ValueError('Cells do not match.')
35 u = atoms.positions - atoms_ideal.positions
36 return get_displacements_from_u(u, atoms_ideal.cell, check_mic=check_mic)
39def get_displacements_from_u(
40 u: NDArray[float],
41 cell: NDArray[float],
42 check_mic: Optional[bool] = True,
43) -> NDArray[float]:
44 """wraps displacements using mic"""
45 if check_mic:
46 u, _ = find_mic(u, cell)
47 return u
50def find_permutation(atoms: Atoms,
51 atoms_ref: Atoms,
52 cell_tol: Optional[float] = 1e-4) -> list[int]:
53 """ Returns the best permutation of atoms for mapping one
54 configuration onto another.
56 Parameters
57 ----------
58 atoms
59 Configuration to be permuted.
60 atoms_ref
61 Configuration onto which to map.
62 cell_tol
63 Cell tolerance; if the cell mismatch exceeds this value, a `ValueError` is raised.
65 Example
66 -------
67 After obtaining the permutation via ``p = find_permutation(atoms1, atoms2)``
68 the reordered structure ``atoms1[p]`` will give the closest match
69 to ``atoms2``.
70 """
71 if np.linalg.norm(atoms.cell - atoms_ref.cell) > cell_tol:
72 raise ValueError('Cells do not match.')
74 permutation = []
75 for i in range(len(atoms_ref)):
76 dist_row = get_distances(
77 atoms.positions, atoms_ref.positions[i], cell=atoms_ref.cell, pbc=True)[1][:, 0]
78 permutation.append(np.argmin(dist_row))
80 if len(set(permutation)) != len(permutation):
81 raise ValueError('Duplicates in permutation.')
82 for i, p in enumerate(permutation):
83 if atoms[p].symbol != atoms_ref[i].symbol:
84 raise ValueError('Matching lattice sites have different occupation.')
85 return permutation
88def align_structure(
89 atoms: Atoms,
90 atol: Optional[float] = 1e-5,
91 method: Optional[str] = 'sequential',
92) -> None:
93 """
94 Rotates a structure into alignment with the x,y,z coordinate system.
96 With :attr:`method` set to ``'sequential'`` the structure is rotated such that
98 * the first cell vector points along the x-direction
99 * the second cell vector lies in the xy-plane
101 which yields a lower triangular cell metric. With :attr:`method` set to
102 ``'best_fit'`` a single rotation is applied that maximizes the sum of the cosines of
103 the angles between the three cell vectors and the x, y and z axes, respectively.
104 All three vectors may then deviate from their respective axis, but the total
105 misalignment is as small as possible; only the directions of the cell vectors enter,
106 not their lengths.
108 Note that this function modifies the :attr:`atoms` object in place, and that only
109 the positions and the cell are rotated; velocities and momenta are left unchanged.
111 Parameters
112 ----------
113 atoms
114 Input structure to be aligned with the x,y,z coordinate system.
115 atol
116 Absolute tolerance used for sanity checking the cell.
117 method
118 Alignment method to use; must be either ``'sequential'`` or ``'best_fit'``.
119 """
120 if method == 'sequential':
121 _align_a_onto_xy(atoms, atol)
122 _align_a_onto_x(atoms, atol)
123 _align_b_onto_xy(atoms, atol)
124 elif method == 'best_fit':
125 _align_best_fit(atoms, atol)
126 else:
127 raise ValueError(f'Unknown method {method}; must be "sequential" or "best_fit".')
130def get_offset_index(
131 primitive: Atoms,
132 supercell: Atoms,
133 tol: Optional[float] = 0.01,
134 wrap: Optional[bool] = True,
135) -> tuple[NDArray[float], NDArray[float]]:
136 """ Returns the basis index and primitive cell offsets for a supercell.
138 This implementation uses a simple iteration procedure that should be fairly quick.
139 If more stability is needed consider the following approach:
141 * find the P-matrix: `P = ideal.cell @ prim.cell_inv.T`
142 * compensate for strain: `P *= len(ideal)/len(prim)/det(P)`
143 * generate the reference structure: `ref_atoms = make_supercell(round(P), prim)`
144 * find the assignment using `ref_atoms` via the Hungarian algorithm using the mic distances
146 Parameters
147 ----------
148 primitive
149 Primitive cell.
150 supercell
151 Some ideal repetition of the primitive cell.
152 tol
153 Tolerance length parameter. Increase to allow for slightly rattled or strained cells.
154 wrap
155 It might happen that the ideal cell boundary cuts through a unit cell
156 whose lattice points lie inside the ideal cell. If there is a basis, an
157 atom belonging to this unit cell might get wrapped while another is
158 not. Then the wrapped atom now belongs to a lattice point outside the P
159 matrix so to say. This would result in more lattice points than
160 expected from `N_unit = len(ideal)/len(prim)`.
162 Returns
163 -------
164 offsets
165 The lattice points as integers in `(N, 3)`-array.
166 index
167 The basis indices as integers in `(N,)`-array.
168 """
170 if not isinstance(primitive, Atoms):
171 raise ValueError('primitive must be an ASE Atoms object.')
172 if not isinstance(supercell, Atoms):
173 raise ValueError('supercell must be an ASE Atoms object.')
175 prim = Prim(primitive)
177 from dynasor.modes.tools import inv
179 P = get_P_matrix(primitive.cell, supercell.cell) # P C = S
180 P_inv = inv(P)
182 lattice, basis = [], []
183 # Pick an atom in the supercell
184 for pos_ideal in supercell.positions:
185 # Does this atom perhaps belong to site "index"?
186 for index, pos_prim in enumerate(primitive.positions):
187 # if so we can remove the basis position vector and should end up on a lattice site
188 diff_pos = pos_ideal - pos_prim
189 # The lattice site has integer coordinates in reduced coordinates
190 prim_spos = diff_pos @ prim.inv_cell.T
191 # Rounding should not affect the coordinate much if it is integer
192 prim_spos_round = np.round(prim_spos).astype(int)
193 # If the rounded spos and unrounded spos are the same
194 if np.allclose(prim_spos, prim_spos_round, rtol=0, atol=tol):
195 # Since P_inv is represented using fractions we can neatly
196 # write the supercell spos of the lattice point using fractions
197 # and easily determine if it needs wrapping or not without
198 # worry about numerics
199 ideal_spos = prim_spos_round @ P_inv
200 # wrap if needed
201 ideal_spos_wrap = ideal_spos % 1 if wrap else ideal_spos
202 # This should be integer again
203 prim_spos_wrap = (ideal_spos_wrap @ P).astype(int)
204 # add the results and break out from the basis site loop
205 lattice.append(prim_spos_wrap)
206 basis.append(index)
207 break
208 else: # we get here by not breaking out from the basis site loop.
209 # This means that the candidate lattice site where not close to integers
211 raise ValueError('Supercell not compatible with primitive cell; the atom at '
212 f'{pos_ideal} does not map onto a lattice point (reduced '
213 f'coordinates {prim_spos}, nearest integer {prim_spos_round}).')
215 lattice = np.array(lattice)
216 basis = np.array(basis)
218 # We should have found len(ideal) unique positions
219 lattice_basis = [tuple((*lp, i)) for lp, i in zip(lattice, basis)]
220 assert len(set(lattice_basis)) == len(supercell)
222 return lattice, basis
225def get_P_matrix(
226 c: NDArray[float],
227 S: NDArray[float],
228) -> NDArray[float]:
229 """Returns the P matrix, i.e., the `3x3` integer matrix :math:`P` that satisfies
231 .. math::
233 P c = S
235 Here, :math:`c` is the primitive cell metric and :math:`S` is the
236 supercell metric as row vectors. Note that the above condition is
237 equivalent to:
239 .. math::
241 c^T P^T = S^T
243 Parameters
244 ----------
245 c
246 Cell metric of the primitive structure.
247 S
248 Cell metric of the supercell.
249 """
250 PT = np.linalg.solve(c.T, S.T)
251 P_float = PT.T
252 P = np.round(P_float).astype(int)
253 if not np.allclose(P_float, P) or not np.allclose(P @ c, S):
254 raise ValueError(
255 f'Please check that the supercell metric ({S}) is related to the'
256 f' the primitive cell {c} by an integer transformation matrix.')
257 return P
260def _align_a_onto_xy(atoms: Atoms, atol: float) -> None:
261 """ Rotate cell so that a is in the xy-plane. """
263 cell = atoms.cell.array.copy()
265 a = cell[0]
266 a_xy = a.copy()
267 a_xy[2] = 0 # projection of a onto xy-plane
269 norm_a = np.linalg.norm(a)
270 norm_a_xy = np.linalg.norm(a_xy)
272 assert norm_a >= atol, \
273 f'First cell vector a has near-zero length ({norm_a:.3e}); cannot align cell.'
274 assert norm_a_xy >= atol, \
275 f'First cell vector a is nearly parallel to the z-axis ' \
276 f'(|a_xy| = {norm_a_xy:.3e}); cannot align a onto the xy-plane.'
278 # cosine of the angle between a and its xy-projection
279 cosa = norm_a_xy / norm_a
280 cosa = min(cosa, 1.0) # clamp for floating-point safety
282 # angle between a and xy-plane in degs
283 angle_xy_deg = np.rad2deg(np.arccos(cosa))
285 # get unit vector to rotate around
286 vec = np.cross(a_xy, [0, 0, 1])
287 vec = vec / np.linalg.norm(vec)
288 assert vec[2] == 0
290 # Determine if the rotation should be positive or negative depending on
291 # whether a is pointing in the +z or -z direction
292 sign = -1 if a[2] > 0 else +1
294 # rotate
295 atoms.rotate(sign * angle_xy_deg, vec, rotate_cell=True)
297 assert np.isclose(atoms.cell[0, 2], 0, atol=atol, rtol=0), atoms.cell
300def _align_a_onto_x(atoms: Atoms, atol: float) -> None:
301 assert np.isclose(atoms.cell[0, 2], 0, atol=atol, rtol=0) # make sure a is in xy-plane
303 a = atoms.cell[0]
304 a_x = a[0]
305 a_y = a[1]
307 # angle between a and x-axis (a is already in xy-plane)
309 # tan = y / x -> angle = arctan y / x "=" atan2(y, x)
310 angle_rad = np.arctan2(a_y, a_x)
311 angle_deg = np.rad2deg(angle_rad)
313 atoms.rotate(-angle_deg, [0, 0, 1], rotate_cell=True)
315 assert np.isclose(atoms.cell[0, 1], 0, atol=atol, rtol=0), atoms.cell
316 assert np.isclose(atoms.cell[0, 2], 0, atol=atol, rtol=0), atoms.cell
319def _align_b_onto_xy(atoms: Atoms, atol: float) -> None:
320 assert np.isclose(atoms.cell[0, 1], 0, atol=atol, rtol=0) # make sure a is along x
321 assert np.isclose(atoms.cell[0, 2], 0, atol=atol, rtol=0) # make sure a is along x
323 # rotate so that b is in xy plane
324 # project b onto the yz-plane
325 b = atoms.cell[1]
326 b_y = b[1]
327 b_z = b[2]
328 angle_rad = np.arctan2(b_z, b_y)
329 angle_deg = np.rad2deg(angle_rad)
331 atoms.rotate(-angle_deg, [1, 0, 0], rotate_cell=True)
333 assert np.isclose(atoms.cell[0, 1], 0, atol=atol, rtol=0) # make sure a is in xy-plane
334 assert np.isclose(atoms.cell[0, 2], 0, atol=atol, rtol=0) # make sure a is in xy-plane
335 assert np.isclose(atoms.cell[1, 2], 0, atol=atol, rtol=0), atoms.cell
338def _align_best_fit(atoms: Atoms, atol: float) -> None:
339 """ Rotate cell such that a, b and c match x, y and z as closely as possible. """
341 cell = atoms.cell.array.copy()
343 norms = np.linalg.norm(cell, axis=1)
344 assert np.all(norms >= atol), \
345 f'Cell vector with near-zero length ({np.min(norms):.3e}); cannot align cell.'
347 # rotation that maximizes the sum of e_i . R a_i over proper rotations,
348 # i.e., the orthogonal Procrustes solution for the normalized cell vectors
349 u, _, vh = np.linalg.svd(cell / norms[:, None])
350 d = np.diag([1, 1, np.sign(np.linalg.det(u @ vh))])
351 R = u @ d @ vh
353 atoms.positions[:] = atoms.positions @ R.T
354 atoms.set_cell(cell @ R.T)