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

1import numpy as np 

2from ase import Atoms 

3from dynasor.modes.atoms import Prim 

4 

5 

6def align_structure(atoms: Atoms, atol: float = 1e-5): 

7 """ 

8 

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 

12 

13 Modifies input ``atoms`` object in place. 

14 

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) 

25 

26 

27def get_offset_index(primitive, supercell, tol=0.01, wrap=True): 

28 """ Get the basis index and primitive cell offsets for a supercell 

29 

30 

31 Simple iteration which should be fairly quick. If more stability is needed consider: 

32 

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 

37 

38 

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

54 

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

62 

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 

67 

68 prim = Prim(primitive) 

69 

70 from dynasor.tools.structures import get_P_matrix 

71 from dynasor.modes.tools import inv 

72 

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

74 P_inv = inv(P) 

75 

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 

104 

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

106 'with primitive cell.') 

107 

108 lattice = np.array(lattice) 

109 basis = np.array(basis) 

110 

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) 

114 

115 return lattice, basis 

116 

117 

118def get_P_matrix(c, S): 

119 """ P c = S -> c.T P.T = S.T 

120 

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 

132 

133 

134def _align_a_onto_xy(atoms, atol): 

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

136 

137 # get angle towards xy 

138 # will break if a is along z 

139 assert np.any(atoms.cell[0, :2]) 

140 

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

142 

143 a = cell[0] 

144 a_xy = a.copy() 

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

146 

147 # angle between a and xy-plane 

148 cosa = np.dot(a, a_xy) / np.linalg.norm(a) / np.linalg.norm(a_xy) 

149 

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) 

156 

157 # angle between a and xy-plane in degs 

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

159 

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 

164 

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 

168 

169 # rotate 

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

171 

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

173 

174 

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 

177 

178 a = atoms.cell[0] 

179 a_x = a[0] 

180 a_y = a[1] 

181 

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

183 

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) 

187 

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

189 

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 

192 

193 

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 

197 

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) 

205 

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

207 

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