Coverage for hiphive/input_output/shengBTE.py: 98%

124 statements  

« prev     ^ index     » next       coverage.py v7.10.1, created at 2025-08-01 17:04 +0000

1""" 

2The ``io.shengBTE`` module provides functions for reading and writing data 

3files in shengBTE format. 

4""" 

5 

6import numpy as np 

7from itertools import product 

8from collections import namedtuple 

9from ..core.structures import Supercell 

10import ase 

11import hiphive 

12 

13 

14def read_shengBTE_fc3( 

15 filename: str, 

16 prim: ase.Atoms, 

17 supercell: ase. Atoms, 

18 symprec: float = 1e-5, 

19): 

20 """Reads third-order force constants file in shengBTE format. 

21 

22 Parameters 

23 ---------- 

24 filename 

25 Input file name. 

26 prim 

27 Primitive structure for the force constants. 

28 supercell 

29 Supercell supercell onto which to map the force constants. 

30 symprec 

31 Structural symmetry tolerance. 

32 

33 Returns 

34 ------- 

35 Third-order force constants for the specified supercell 

36 as :class:`ForceConstants` object. 

37 """ 

38 

39 raw_sheng = _read_raw_sheng(filename) 

40 

41 sheng = _raw_to_fancy(raw_sheng, prim.cell) 

42 

43 fcs = _sheng_to_fcs(sheng, prim, supercell, symprec) 

44 

45 return fcs 

46 

47 

48def write_shengBTE_fc3( 

49 filename: str, 

50 fcs, 

51 prim: ase.Atoms, 

52 symprec: float = 1e-5, 

53 cutoff: float = np.inf, 

54 fc_tol: float = 1e-8, 

55): 

56 """Writes third-order force constants to file in shengBTE format. 

57 

58 Parameters 

59 ----------- 

60 filename 

61 Output file name. 

62 fcs 

63 Force constants in :class:`ForceConstants` format; the supercell 

64 associated with this object must be based on :attr:`prim`. 

65 prim 

66 Primitive structure (must be equivalent to structure used in the 

67 shengBTE calculation). 

68 symprec 

69 Structural symmetry tolerance. 

70 cutoff 

71 All atoms in cluster must be within this cutoff. 

72 fc_tol 

73 If the absolute value of the largest entry in a force constant is less 

74 than :attr:`fc_tol` it will not be included in the output. 

75 """ 

76 sheng = _fcs_to_sheng(fcs, prim, symprec, cutoff, fc_tol) 

77 raw_sheng = _fancy_to_raw(sheng) 

78 _write_raw_sheng(raw_sheng, filename) 

79 

80 

81def _read_raw_sheng(filename: str) -> list: 

82 """ Read shengBTE fc3 file. 

83 

84 Parameters 

85 ---------- 

86 filename 

87 Input file name. 

88 

89 Returns 

90 ------- 

91 list with shengBTE block entries, where an entry consist of 

92 `[i, j, k, cell_pos2, cell_pos3, fc3]`. 

93 """ 

94 lines_per_fc_block = 32 

95 

96 fc3_data_shengBTE = [] 

97 with open(filename, 'r') as f: 

98 lines = f.readlines() 

99 num_fcs = int(lines[0]) 

100 lind = 2 

101 for i in range(1, num_fcs+1): 

102 # sanity check 

103 assert int(lines[lind]) == i, (int(lines[lind]), i) 

104 

105 # read cell poses 

106 cell_pos2 = np.array([float(fld) for fld in lines[lind+1].split()]) 

107 cell_pos3 = np.array([float(fld) for fld in lines[lind+2].split()]) 

108 

109 # read basis indices 

110 i, j, k = (int(fld) for fld in lines[lind+3].split()) 

111 

112 # read fc 

113 fc3_ijk = np.zeros((3, 3, 3)) 

114 for n in range(27): 

115 x, y, z = [int(fld) - 1 for fld in lines[lind+4+n].split()[:3]] 

116 fc3_ijk[x, y, z] = float(lines[lind+4+n].split()[-1]) 

117 

118 # append entry 

119 entry = [i, j, k, cell_pos2, cell_pos3, fc3_ijk] 

120 fc3_data_shengBTE.append(entry) 

121 lind += lines_per_fc_block 

122 

123 return fc3_data_shengBTE 

124 

125 

126_ShengEntry = namedtuple('Entry', ['site_0', 'site_1', 'site_2', 'pos_1', 

127 'pos_2', 'fc', 'offset_1', 'offset_2']) 

128 

129 

130def _raw_to_fancy(raw_sheng, cell): 

131 """ 

132 Converts force constants as read from shengBTE file to the namedtuple 

133 format defined above (_ShengEntry). 

134 """ 

135 sheng = [] 

136 for raw_entry in raw_sheng: 

137 p1, p2 = raw_entry[3:5] 

138 offset_1 = np.linalg.solve(cell.T, p1).round(0).astype(int) 

139 offset_2 = np.linalg.solve(cell.T, p2).round(0).astype(int) 

140 entry = _ShengEntry(*(i - 1 for i in raw_entry[:3]), *raw_entry[3:], 

141 offset_1, offset_2) 

142 sheng.append(entry) 

143 return sheng 

144 

145 

146def _fancy_to_raw(sheng): 

147 """ 

148 Converts force constants namedtuple format defined above (_ShengEntry) to 

149 format used for writing shengBTE files. 

150 """ 

151 raw_sheng = [] 

152 for entry in sheng: 

153 raw_entry = list(entry[:6]) 

154 raw_entry[0] += 1 

155 raw_entry[1] += 1 

156 raw_entry[2] += 1 

157 raw_sheng.append(raw_entry) 

158 

159 return raw_sheng 

160 

161 

162def _write_raw_sheng(raw_sheng, filename): 

163 """ See corresponding read function. """ 

164 

165 with open(filename, 'w') as f: 

166 f.write('{}\n\n'.format(len(raw_sheng))) 

167 

168 for index, fc3_row in enumerate(raw_sheng, start=1): 

169 i, j, k, cell_pos2, cell_pos3, fc3_ijk = fc3_row 

170 

171 f.write('{:5d}\n'.format(index)) 

172 

173 f.write((3*'{:14.10f} '+'\n').format(*cell_pos2)) 

174 f.write((3*'{:14.10f} '+'\n').format(*cell_pos3)) 

175 f.write((3*'{:5d}'+'\n').format(i, j, k)) 

176 

177 for x, y, z in product(range(3), repeat=3): 

178 f.write((3*' {:}').format(x+1, y+1, z+1)) 

179 f.write(' {:14.10f}\n'.format(fc3_ijk[x, y, z])) 

180 f.write('\n') 

181 

182 

183def _fcs_to_sheng(fcs, prim, symprec, cutoff, fc_tol): 

184 """ Produces a list containing fcs in shengBTE format 

185 

186 prim and fcs.supercell must be aligned. 

187 """ 

188 

189 supercell = Supercell(fcs.supercell, prim, symprec) 

190 assert all(fcs.supercell.pbc) and all(prim.pbc) 

191 

192 n_atoms = len(supercell) 

193 

194 D = fcs.supercell.get_all_distances(mic=False, vector=True) 

195 D_mic = fcs.supercell.get_all_distances(mic=True, vector=True) 

196 M = np.eye(n_atoms, dtype=bool) 

197 for i in range(n_atoms): 

198 for j in range(i + 1, n_atoms): 

199 M[i, j] = (np.allclose(D[i, j], D_mic[i, j], atol=symprec, rtol=0) 

200 and np.linalg.norm(D[i, j]) < cutoff) 

201 M[j, i] = M[i, j] 

202 

203 data = {} 

204 for a0 in supercell: 

205 for a1 in supercell: 

206 if not M[a0.index, a1.index]: 

207 continue 

208 for a2 in supercell: 

209 if not (M[a0.index, a2.index] and M[a1.index, a2.index]): 

210 continue 

211 

212 offset_1 = np.subtract(a1.offset, a0.offset) 

213 offset_2 = np.subtract(a2.offset, a0.offset) 

214 

215 sites = (a0.site, a1.site, a2.site) 

216 

217 key = sites + tuple(offset_1) + tuple(offset_2) 

218 

219 ijk = (a0.index, a1.index, a2.index) 

220 

221 fc = fcs[ijk] 

222 

223 if key in data: 

224 assert np.allclose(data[key], fc, atol=fc_tol) 

225 else: 

226 data[key] = fc 

227 

228 sheng = [] 

229 for k, fc in data.items(): 

230 if np.max(np.abs(fc)) < fc_tol: 

231 continue 

232 offset_1 = k[3:6] 

233 pos_1 = np.dot(offset_1, prim.cell) 

234 offset_2 = k[6:9] 

235 pos_2 = np.dot(offset_2, prim.cell) 

236 entry = _ShengEntry(*k[:3], pos_1, pos_2, fc, offset_1, offset_2) 

237 sheng.append(entry) 

238 

239 return sheng 

240 

241 

242def _sheng_to_fcs(sheng, prim, supercell, symprec): 

243 supercell_map = Supercell(supercell, prim, symprec) 

244 fc_array = np.zeros((len(supercell),) * 3 + (3,) * 3) 

245 mapped_clusters = np.zeros((len(supercell),) * 3, dtype=bool) 

246 

247 for atom in supercell_map: 

248 i = atom.index 

249 for entry in sheng: 

250 if atom.site != entry.site_0: 250 ↛ 251line 250 didn't jump to line 251 because the condition on line 250 was never true

251 continue 

252 j = supercell_map.index(entry.site_1, entry.offset_1 + atom.offset) 

253 k = supercell_map.index(entry.site_2, entry.offset_2 + atom.offset) 

254 ijk = (i, j, k) 

255 if mapped_clusters[ijk]: 255 ↛ 256line 255 didn't jump to line 256 because the condition on line 255 was never true

256 raise Exception('Ambiguous force constant.' 

257 ' Supercell is too small') 

258 fc_array[ijk] = entry.fc 

259 mapped_clusters[ijk] = True 

260 

261 fcs = hiphive.force_constants.ForceConstants.from_arrays( 

262 supercell, fc3_array=fc_array) 

263 return fcs