Coverage for dynasor/tools/structures.py: 100%
110 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-18 09:03 +0000
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-18 09:03 +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 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 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 cell missmatch more than tol value error 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=True)
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, atoms_ref: Atoms) -> list[int]:
51 """ Returns the best permutation of atoms for mapping one
52 configuration onto another.
54 Parameters
55 ----------
56 atoms
57 Configuration to be permuted.
58 atoms_ref
59 Configuration onto which to map.
61 Example
62 -------
63 After obtaining the permutation via ``p = find_permutation(atoms1, atoms2)``
64 the reordered structure ``atoms1[p]`` will give the closest match
65 to ``atoms2``.
66 """
67 if np.linalg.norm(atoms.cell - atoms_ref.cell) > 1e-4:
68 raise ValueError('Cells do not match.')
70 permutation = []
71 for i in range(len(atoms_ref)):
72 dist_row = get_distances(
73 atoms.positions, atoms_ref.positions[i], cell=atoms_ref.cell, pbc=True)[1][:, 0]
74 permutation.append(np.argmin(dist_row))
76 if len(set(permutation)) != len(permutation):
77 raise Exception('Duplicates in permutation.')
78 for i, p in enumerate(permutation):
79 if atoms[p].symbol != atoms_ref[i].symbol:
80 raise Exception('Matching lattice sites have different occupation.')
81 return permutation
84def align_structure(atoms: Atoms, atol: Optional[float] = 1e-5) -> None:
85 """
86 Rotates and realigns a structure such that
87 * the first cell vector points along the x-directon
88 * the second cell vector lies in the xy-plane
90 Note that this function modifies the :attr:`atoms` object in place.
92 Parameters
93 ----------
94 atoms
95 Input structure to be rotated aligned with the x,y,z coordinte system.
96 atol
97 Absolute tolerance used for sanity checking the cell.
98 """
99 _align_a_onto_xy(atoms, atol)
100 _align_a_onto_x(atoms, atol)
101 _align_b_onto_xy(atoms, atol)
104def get_offset_index(
105 primitive: Atoms,
106 supercell: Atoms,
107 tol: Optional[float] = 0.01,
108 wrap: Optional[bool] = True,
109) -> tuple[NDArray[float], NDArray[float]]:
110 """ Returns the basis index and primitive cell offsets for a supercell.
112 This implementation uses a simple iteration procedure that should be fairly quick.
113 If more stability is needed consider the following approach:
115 * find the P-matrix: `P = ideal.cell @ prim.cell_inv.T`
116 * compensate for strain: `P *= len(ideal)/len(prim)/det(P)`
117 * generate the reference structure: `ref_atoms = make_supercell(round(P), prim)`
118 * find the assignment using `ref_atoms` via the Hungarian algorithm using the mic distances
120 Parameters
121 ----------
122 primitive
123 Primitive cell.
124 supercell
125 Some ideal repetition of the primitive cell.
126 tol
127 Tolerance length parameter. Increase to allow for slgihtly rattled or strained cells.
128 wrap
129 It might happen that the ideal cell boundary cuts through a unit cell
130 whose lattice points lie inside the ideal cell. If there is a basis, an
131 atom belonging to this unit cell might get wrapped while another is
132 not. Then the wrapped atom now belongs to a lattice point outside the P
133 matrix so to say. This would result in more lattice points than
134 expected from `N_unit = len(ideal)/len(prim)`.
136 Returns
137 -------
138 offsets
139 The lattice points as integers in `(N, 3)`-array.
140 index
141 The basis indices as integers in `(N,)`-array.
142 """
144 if not isinstance(primitive, Atoms):
145 raise ValueError
146 if not isinstance(supercell, Atoms):
147 raise ValueError
149 prim = Prim(primitive)
151 from dynasor.modes.tools import inv
153 P = get_P_matrix(primitive.cell, supercell.cell) # P C = S
154 P_inv = inv(P)
156 lattice, basis = [], []
157 # Pick an atom in the supercell
158 for pos_ideal in supercell.positions:
159 # Does this atom perhaps belong to site "index"?
160 for index, pos_prim in enumerate(primitive.positions):
161 # if so we can remove the basis position vector and should end up on a lattice site
162 diff_pos = pos_ideal - pos_prim
163 # The lattice site has integer coordinates in reduced coordinates
164 prim_spos = diff_pos @ prim.inv_cell.T
165 # Rounding should not affect the coordinate much if it is integer
166 prim_spos_round = np.round(prim_spos).astype(int)
167 # If the rounded spos and unrounded spos are the same
168 if np.allclose(prim_spos, prim_spos_round, rtol=0, atol=tol):
169 # Since P_inv is represented using fractions we can neatly
170 # write the supercell spos of the lattice point using fractions
171 # and easily determine if it needs wrapping or not without
172 # worry about numerics
173 ideal_spos = prim_spos_round @ P_inv
174 # wrap if needed
175 ideal_spos_wrap = ideal_spos % 1 if wrap else ideal_spos
176 # This should be integer again
177 prim_spos_wrap = (ideal_spos_wrap @ P).astype(int)
178 # add the results and break out from the basis site loop
179 lattice.append(prim_spos_wrap)
180 basis.append(index)
181 break
182 else: # we get here by not breaking out from the basis site loop.
183 # This means that the candidate lattice site where not close to integers
185 raise ValueError(f' {prim_spos} {prim_spos_round} Supercell not compatible '
186 'with primitive cell.')
188 lattice = np.array(lattice)
189 basis = np.array(basis)
191 # We should have found len(ideal) unique positions
192 lattice_basis = [tuple((*lp, i)) for lp, i in zip(lattice, basis)]
193 assert len(set(lattice_basis)) == len(supercell)
195 return lattice, basis
198def get_P_matrix(
199 c: NDArray[float],
200 S: NDArray[float],
201) -> NDArray[float]:
202 """Returns the P matrix, i.e., the `3x3` integer matrix :math:`P` that satisfies
204 .. math::
206 P c = S
208 Here, :math:`c` is the primitive cell metric and :math:`S` is the
209 supercell metric as row vectors. Note that the above condition is
210 equivalent to:
212 .. math::
214 c^T P^T = S^T
216 Parameters
217 ----------
218 c
219 Cell metric of the primitive structure.
220 S
221 Cell metric of the supercell.
222 """
223 PT = np.linalg.solve(c.T, S.T)
224 P_float = PT.T
225 P = np.round(P_float).astype(int)
226 if not np.allclose(P_float, P) or not np.allclose(P @ c, S):
227 raise ValueError(
228 f'Please check that the supercell metric ({S}) is related to the'
229 f' the primitive cell {c} by an integer transformation matrix.')
230 return P
233def _align_a_onto_xy(atoms: Atoms, atol: float) -> None:
234 """ Rotate cell so that a is in the xy-plane. """
236 cell = atoms.cell.array.copy()
238 a = cell[0]
239 a_xy = a.copy()
240 a_xy[2] = 0 # projection of a onto xy-plane
242 norm_a = np.linalg.norm(a)
243 norm_a_xy = np.linalg.norm(a_xy)
245 assert norm_a >= atol, \
246 f'First cell vector a has near-zero length ({norm_a:.3e}); cannot align cell.'
247 assert norm_a_xy >= atol, \
248 f'First cell vector a is nearly parallel to the z-axis ' \
249 f'(|a_xy| = {norm_a_xy:.3e}); cannot align a onto the xy-plane.'
251 # cosine of the angle between a and its xy-projection
252 cosa = norm_a_xy / norm_a
253 cosa = min(cosa, 1.0) # clamp for floating-point safety
255 # angle between a and xy-plane in degs
256 angle_xy_deg = np.rad2deg(np.arccos(cosa))
258 # get unit vector to rotate around
259 vec = np.cross(a_xy, [0, 0, 1])
260 vec = vec / np.linalg.norm(vec)
261 assert vec[2] == 0
263 # Determine if the rotation should be positive or negative depending on
264 # whether a is pointing in the +z or -z direction
265 sign = -1 if a[2] > 0 else +1
267 # rotate
268 atoms.rotate(sign * angle_xy_deg, vec, rotate_cell=True)
270 assert np.isclose(atoms.cell[0, 2], 0, atol=atol, rtol=0), atoms.cell
273def _align_a_onto_x(atoms: Atoms, atol: float) -> None:
274 assert np.isclose(atoms.cell[0, 2], 0, atol=atol, rtol=0) # make sure a is in xy-plane
276 a = atoms.cell[0]
277 a_x = a[0]
278 a_y = a[1]
280 # angle between a and x-axis (a is already in xy-plane)
282 # tan = y / x -> angle = arctan y / x "=" atan2(y, x)
283 angle_rad = np.arctan2(a_y, a_x)
284 angle_deg = np.rad2deg(angle_rad)
286 atoms.rotate(-angle_deg, [0, 0, 1], rotate_cell=True)
288 assert np.isclose(atoms.cell[0, 1], 0, atol=atol, rtol=0), atoms.cell
289 assert np.isclose(atoms.cell[0, 2], 0, atol=atol, rtol=0), atoms.cell
292def _align_b_onto_xy(atoms: Atoms, atol: float) -> None:
293 assert np.isclose(atoms.cell[0, 1], 0, atol=atol, rtol=0) # make sure a is along x
294 assert np.isclose(atoms.cell[0, 2], 0, atol=atol, rtol=0) # make sure a is along x
296 # rotate so that b is in xy plane
297 # project b onto the yz-plane
298 b = atoms.cell[1]
299 b_y = b[1]
300 b_z = b[2]
301 angle_rad = np.arctan2(b_z, b_y)
302 angle_deg = np.rad2deg(angle_rad)
304 atoms.rotate(-angle_deg, [1, 0, 0], rotate_cell=True)
306 assert np.isclose(atoms.cell[0, 1], 0, atol=atol, rtol=0) # make sure a is in xy-plane
307 assert np.isclose(atoms.cell[0, 2], 0, atol=atol, rtol=0) # make sure a is in xy-plane
308 assert np.isclose(atoms.cell[1, 2], 0, atol=atol, rtol=0), atoms.cell