Coverage for local_installation/dynasor/tools/structures.py: 94%
84 statements
« prev ^ index » next coverage.py v7.9.2, created at 2025-07-15 07:34 +0000
« prev ^ index » next coverage.py v7.9.2, created at 2025-07-15 07:34 +0000
1import numpy as np
2from ase import Atoms
3from dynasor.modes.atoms import Prim
6def align_structure(atoms: Atoms, atol: float = 1e-5):
7 """
9 Rotate and realign atoms object such that
10 * the first cell vector points along the x-directon
11 * the second cell vector lies in the xy-plane
13 Modifies input ``atoms`` object in place.
15 Parameters
16 ----------
17 atoms
18 input structure to be rotated aligned wtih the x,y,z coordinte system
19 atol
20 absolute tolerance used for sanity checking the cell
21 """
22 _align_a_onto_xy(atoms, atol)
23 _align_a_onto_x(atoms, atol)
24 _align_b_onto_xy(atoms, atol)
27def get_offset_index(primitive, supercell, tol=0.01, wrap=True):
28 """ Get the basis index and primitive cell offsets for a supercell
31 Simple iteration which should be fairly quick. If more stability is needed consider:
33 P = ideal.cell @ prim.cell_inv.T find the P-matrix
34 P *= len(ideal)/len(prim)/det(P) compensate for strain
35 ref_atoms = make_supercell(round(P), prim)
36 find assignment using ref_atoms and ideal via hungarian algorithm and mic distances
39 Parameters
40 ----------
41 primitive
42 primitive cell as either ASEAtoms
43 supercell
44 some ideal repetition of the prim cell and possible wrapping as either ASEAtoms
45 tol
46 tolerance length parameter. Increase to allow for slgihtly rattled or strained cells.
47 wrap
48 It might happen that the ideal cell boundary cuts through a unit cell
49 whos lattice points lie within the ideal cell. If there is a basis one
50 atom belonging to this unit cell might get wrapped while the other is
51 not. Then the wrapped atom now belong to a lattice point outside the P
52 matrix so to say. This would result in more lattice points than
53 expected from N_unit = len(ideal)/len(prim).
55 Returns
56 -------
57 offsets
58 The lattice points as integers in (N, 3)-array
59 index
60 The basis indices as integers in (N,)-array
61 """
63 if not isinstance(primitive, Atoms): 63 ↛ 64line 63 didn't jump to line 64 because the condition on line 63 was never true
64 raise ValueError
65 if not isinstance(supercell, Atoms): 65 ↛ 66line 65 didn't jump to line 66 because the condition on line 65 was never true
66 raise ValueError
68 prim = Prim(primitive)
70 from dynasor.tools.structures import get_P_matrix
71 from dynasor.modes.tools import inv
73 P = get_P_matrix(primitive.cell, supercell.cell) # P C = S
74 P_inv = inv(P)
76 lattice, basis = [], []
77 # Pick an atom in the supercell
78 for pos_ideal in supercell.positions:
79 # Does this atom perhaps belong to site "index"?
80 for index, pos_prim in enumerate(primitive.positions):
81 # if so we can remove the basis position vector and should end up on a lattice site
82 diff_pos = pos_ideal - pos_prim
83 # The lattice site has integer coordinates in reduced coordinates
84 prim_spos = diff_pos @ prim.inv_cell.T
85 # Rounding should not affect the coordinate much if it is integer
86 prim_spos_round = np.round(prim_spos).astype(int)
87 # If the rounded spos and unrounded spos are the same
88 if np.allclose(prim_spos, prim_spos_round, rtol=0, atol=tol):
89 # Since P_inv is represented using fractions we can neatly
90 # write the supercell spos of the lattice point using fractions
91 # and easily determine if it needs wrapping or not without
92 # worry about numerics
93 ideal_spos = prim_spos_round @ P_inv
94 # wrap if needed
95 ideal_spos_wrap = ideal_spos % 1 if wrap else ideal_spos
96 # This should be integer again
97 prim_spos_wrap = (ideal_spos_wrap @ P).astype(int)
98 # add the results and break out from the basis site loop
99 lattice.append(prim_spos_wrap)
100 basis.append(index)
101 break
102 else: # we get here by not breaking out from the basis site loop.
103 # This means that the candidate lattice site where not close to integers
105 raise ValueError(f' {prim_spos} {prim_spos_round} Supercell not compatible '
106 'with primitive cell.')
108 lattice = np.array(lattice)
109 basis = np.array(basis)
111 # We should have found len(ideal) unique positions
112 lattice_basis = [tuple((*lp, i)) for lp, i in zip(lattice, basis)]
113 assert len(set(lattice_basis)) == len(supercell)
115 return lattice, basis
118def get_P_matrix(c, S):
119 """ P c = S -> c.T P.T = S.T
121 The P matrix must be an integer matrix. c is the primitive cell metric and
122 S is the supercell metric as row-vectors
123 """
124 PT = np.linalg.solve(c.T, S.T)
125 P_float = PT.T
126 P = np.round(P_float).astype(int)
127 if not np.allclose(P_float, P) or not np.allclose(P @ c, S):
128 raise ValueError(
129 f'Please check that the supercell metric ({S}) is related to the'
130 f' the primitive cell {c} by an integer transformation matrix.')
131 return P
134def _align_a_onto_xy(atoms, atol):
135 """Rotate cell so that a is in the xy-plane"""
137 # get angle towards xy
138 # will break if a is along z
139 assert np.any(atoms.cell[0, :2])
141 cell = atoms.cell.array.copy()
143 a = cell[0]
144 a_xy = a.copy()
145 a_xy[2] = 0 # projection of a onto xy-plane
147 # angle between a and xy-plane
148 cosa = np.dot(a, a_xy) / np.linalg.norm(a) / np.linalg.norm(a_xy)
150 # cosa should be in the interval (0, 1]
151 assert not np.isclose(cosa, 0)
152 if cosa > 1: 152 ↛ 153line 152 didn't jump to line 153 because the condition on line 152 was never true
153 assert np.isclose(cosa, 1)
154 cosa = min(cosa, 1)
155 cosa = max(cosa, 0)
157 # angle between a and xy-plane in degs
158 angle_xy_deg = np.rad2deg(np.arccos(cosa))
160 # get unit vector to rotate around
161 vec = np.cross(a_xy, [0, 0, 1])
162 vec = vec / np.linalg.norm(vec)
163 assert vec[2] == 0
165 # Determine if the rotation should be positive or negative depending on
166 # whether a is pointing in the +z or -z direction
167 sign = -1 if a[2] > 0 else +1
169 # rotate
170 atoms.rotate(sign * angle_xy_deg, vec, rotate_cell=True)
172 assert np.isclose(atoms.cell[0, 2], 0, atol=atol, rtol=0), atoms.cell
175def _align_a_onto_x(atoms, atol):
176 assert np.isclose(atoms.cell[0, 2], 0, atol=atol, rtol=0) # make sure a is in xy-plane
178 a = atoms.cell[0]
179 a_x = a[0]
180 a_y = a[1]
182 # angle between a and x-axis (a is already in xy-plane)
184 # tan = y / x -> angle = arctan y / x "=" atan2(y, x)
185 angle_rad = np.arctan2(a_y, a_x)
186 angle_deg = np.rad2deg(angle_rad)
188 atoms.rotate(-angle_deg, [0, 0, 1], rotate_cell=True)
190 assert np.isclose(atoms.cell[0, 1], 0, atol=atol, rtol=0), atoms.cell
191 assert np.isclose(atoms.cell[0, 2], 0, atol=atol, rtol=0), atoms.cell
194def _align_b_onto_xy(atoms, atol):
195 assert np.isclose(atoms.cell[0, 1], 0, atol=atol, rtol=0) # make sure a is along x
196 assert np.isclose(atoms.cell[0, 2], 0, atol=atol, rtol=0) # make sure a is along x
198 # rotate so that b is in xy plane
199 # project b onto the yz-plane
200 b = atoms.cell[1]
201 b_y = b[1]
202 b_z = b[2]
203 angle_rad = np.arctan2(b_z, b_y)
204 angle_deg = np.rad2deg(angle_rad)
206 atoms.rotate(-angle_deg, [1, 0, 0], rotate_cell=True)
208 assert np.isclose(atoms.cell[0, 1], 0, atol=atol, rtol=0) # make sure a is in xy-plane
209 assert np.isclose(atoms.cell[0, 2], 0, atol=atol, rtol=0) # make sure a is in xy-plane
210 assert np.isclose(atoms.cell[1, 2], 0, atol=atol, rtol=0), atoms.cell