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

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 

8 

9 

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. 

17 

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 """ 

29 

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.') 

34 

35 u = atoms.positions - atoms_ideal.positions 

36 return get_displacements_from_u(u, atoms_ideal.cell, check_mic=True) 

37 

38 

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 

48 

49 

50def find_permutation(atoms: Atoms, atoms_ref: Atoms) -> list[int]: 

51 """ Returns the best permutation of atoms for mapping one 

52 configuration onto another. 

53 

54 Parameters 

55 ---------- 

56 atoms 

57 Configuration to be permuted. 

58 atoms_ref 

59 Configuration onto which to map. 

60 

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.') 

69 

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)) 

75 

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 

82 

83 

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 

89 

90 Note that this function modifies the :attr:`atoms` object in place. 

91 

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) 

102 

103 

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. 

111 

112 This implementation uses a simple iteration procedure that should be fairly quick. 

113 If more stability is needed consider the following approach: 

114 

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 

119 

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)`. 

135 

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 """ 

143 

144 if not isinstance(primitive, Atoms): 

145 raise ValueError 

146 if not isinstance(supercell, Atoms): 

147 raise ValueError 

148 

149 prim = Prim(primitive) 

150 

151 from dynasor.modes.tools import inv 

152 

153 P = get_P_matrix(primitive.cell, supercell.cell) # P C = S 

154 P_inv = inv(P) 

155 

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 

184 

185 raise ValueError(f' {prim_spos} {prim_spos_round} Supercell not compatible ' 

186 'with primitive cell.') 

187 

188 lattice = np.array(lattice) 

189 basis = np.array(basis) 

190 

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) 

194 

195 return lattice, basis 

196 

197 

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 

203 

204 .. math:: 

205 

206 P c = S 

207 

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: 

211 

212 .. math:: 

213 

214 c^T P^T = S^T 

215 

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 

231 

232 

233def _align_a_onto_xy(atoms: Atoms, atol: float) -> None: 

234 """ Rotate cell so that a is in the xy-plane. """ 

235 

236 cell = atoms.cell.array.copy() 

237 

238 a = cell[0] 

239 a_xy = a.copy() 

240 a_xy[2] = 0 # projection of a onto xy-plane 

241 

242 norm_a = np.linalg.norm(a) 

243 norm_a_xy = np.linalg.norm(a_xy) 

244 

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.' 

250 

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 

254 

255 # angle between a and xy-plane in degs 

256 angle_xy_deg = np.rad2deg(np.arccos(cosa)) 

257 

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 

262 

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 

266 

267 # rotate 

268 atoms.rotate(sign * angle_xy_deg, vec, rotate_cell=True) 

269 

270 assert np.isclose(atoms.cell[0, 2], 0, atol=atol, rtol=0), atoms.cell 

271 

272 

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 

275 

276 a = atoms.cell[0] 

277 a_x = a[0] 

278 a_y = a[1] 

279 

280 # angle between a and x-axis (a is already in xy-plane) 

281 

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) 

285 

286 atoms.rotate(-angle_deg, [0, 0, 1], rotate_cell=True) 

287 

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 

290 

291 

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 

295 

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) 

303 

304 atoms.rotate(-angle_deg, [1, 0, 0], rotate_cell=True) 

305 

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