Coverage for /builds/debichem-team/python-ase/ase/calculators/singlepoint.py: 79.61%

206 statements  

« prev     ^ index     » next       coverage.py v7.5.3, created at 2025-03-06 04:00 +0000

1import numpy as np 

2 

3from ase.calculators.calculator import ( 

4 Calculator, 

5 PropertyNotImplementedError, 

6 PropertyNotPresent, 

7 all_properties, 

8) 

9from ase.outputs import Properties 

10from ase.utils import lazyproperty 

11 

12 

13class SinglePointCalculator(Calculator): 

14 """Special calculator for a single configuration. 

15 

16 Used to remember the energy, force and stress for a given 

17 configuration. If the positions, atomic numbers, unit cell, or 

18 boundary conditions are changed, then asking for 

19 energy/forces/stress will raise an exception.""" 

20 

21 name = 'unknown' 

22 

23 def __init__(self, atoms, **results): 

24 """Save energy, forces, stress, ... for the current configuration.""" 

25 Calculator.__init__(self) 

26 self.results = {} 

27 for property, value in results.items(): 

28 assert property in all_properties, property 

29 if value is None: 

30 continue 

31 if property in ['energy', 'magmom', 'free_energy']: 

32 self.results[property] = value 

33 else: 

34 self.results[property] = np.array(value, float) 

35 self.atoms = atoms.copy() 

36 

37 def __str__(self): 

38 tokens = [] 

39 for key, val in sorted(self.results.items()): 

40 if np.isscalar(val): 

41 txt = f'{key}={val}' 

42 else: 

43 txt = f'{key}=...' 

44 tokens.append(txt) 

45 return '{}({})'.format(self.__class__.__name__, ', '.join(tokens)) 

46 

47 def get_property(self, name, atoms=None, allow_calculation=True): 

48 if atoms is None: 

49 atoms = self.atoms 

50 if name not in self.results or self.check_state(atoms): 

51 if allow_calculation: 

52 raise PropertyNotImplementedError( 

53 f'The property "{name}" is not available.') 

54 return None 

55 

56 result = self.results[name] 

57 if isinstance(result, np.ndarray): 

58 result = result.copy() 

59 return result 

60 

61 

62class SinglePointKPoint: 

63 def __init__(self, weight, s, k, eps_n=None, f_n=None): 

64 self.weight = weight 

65 self.s = s # spin index 

66 self.k = k # k-point index 

67 if eps_n is None: 

68 eps_n = [] 

69 self.eps_n = eps_n 

70 if f_n is None: 

71 f_n = [] 

72 self.f_n = f_n 

73 

74 

75def arrays_to_kpoints(eigenvalues, occupations, weights): 

76 """Helper function for building SinglePointKPoints. 

77 

78 Convert eigenvalue, occupation, and weight arrays to list of 

79 SinglePointKPoint objects.""" 

80 nspins, nkpts, _nbands = eigenvalues.shape 

81 assert eigenvalues.shape == occupations.shape 

82 assert len(weights) == nkpts 

83 kpts = [] 

84 for s in range(nspins): 

85 for k in range(nkpts): 

86 kpt = SinglePointKPoint( 

87 weight=weights[k], s=s, k=k, 

88 eps_n=eigenvalues[s, k], f_n=occupations[s, k]) 

89 kpts.append(kpt) 

90 return kpts 

91 

92 

93class SinglePointDFTCalculator(SinglePointCalculator): 

94 def __init__(self, atoms, 

95 efermi=None, bzkpts=None, ibzkpts=None, bz2ibz=None, 

96 kpts=None, 

97 **results): 

98 self.bz_kpts = bzkpts 

99 self.ibz_kpts = ibzkpts 

100 self.bz2ibz = bz2ibz 

101 self.eFermi = efermi 

102 

103 SinglePointCalculator.__init__(self, atoms, **results) 

104 self.kpts = kpts 

105 

106 def get_fermi_level(self): 

107 """Return the Fermi-level(s).""" 

108 return self.eFermi 

109 

110 def get_bz_to_ibz_map(self): 

111 return self.bz2ibz 

112 

113 def get_bz_k_points(self): 

114 """Return the k-points.""" 

115 return self.bz_kpts 

116 

117 def get_number_of_spins(self): 

118 """Return the number of spins in the calculation. 

119 

120 Spin-paired calculations: 1, spin-polarized calculation: 2.""" 

121 if self.kpts is not None: 

122 nspin = set() 

123 for kpt in self.kpts: 

124 nspin.add(kpt.s) 

125 return len(nspin) 

126 return None 

127 

128 def get_number_of_bands(self): 

129 values = {len(kpt.eps_n) for kpt in self.kpts} 

130 if not values: 

131 return None 

132 elif len(values) == 1: 

133 return values.pop() 

134 else: 

135 raise RuntimeError('Multiple array sizes') 

136 

137 def get_spin_polarized(self): 

138 """Is it a spin-polarized calculation?""" 

139 nos = self.get_number_of_spins() 

140 if nos is not None: 

141 return nos == 2 

142 return None 

143 

144 def get_ibz_k_points(self): 

145 """Return k-points in the irreducible part of the Brillouin zone.""" 

146 return self.ibz_kpts 

147 

148 def get_kpt(self, kpt=0, spin=0): 

149 if self.kpts is not None: 

150 counter = 0 

151 for kpoint in self.kpts: 

152 if kpoint.s == spin: 

153 if kpt == counter: 

154 return kpoint 

155 counter += 1 

156 return None 

157 

158 def get_k_point_weights(self): 

159 """ Retunrs the weights of the k points """ 

160 if self.kpts is not None: 

161 weights = [] 

162 for kpoint in self.kpts: 

163 if kpoint.s == 0: 

164 weights.append(kpoint.weight) 

165 return np.array(weights) 

166 return None 

167 

168 def get_occupation_numbers(self, kpt=0, spin=0): 

169 """Return occupation number array.""" 

170 kpoint = self.get_kpt(kpt, spin) 

171 if kpoint is not None: 

172 if len(kpoint.f_n): 

173 return kpoint.f_n 

174 return None 

175 

176 def get_eigenvalues(self, kpt=0, spin=0): 

177 """Return eigenvalue array.""" 

178 kpoint = self.get_kpt(kpt, spin) 

179 if kpoint is not None: 

180 return kpoint.eps_n 

181 return None 

182 

183 def get_homo_lumo(self): 

184 """Return HOMO and LUMO energies.""" 

185 if self.kpts is None: 

186 raise RuntimeError('No kpts') 

187 eH = -np.inf 

188 eL = np.inf 

189 for spin in range(self.get_number_of_spins()): 

190 homo, lumo = self.get_homo_lumo_by_spin(spin) 

191 eH = max(eH, homo) 

192 eL = min(eL, lumo) 

193 return eH, eL 

194 

195 def get_homo_lumo_by_spin(self, spin=0): 

196 """Return HOMO and LUMO energies for a given spin.""" 

197 if self.kpts is None: 

198 raise RuntimeError('No kpts') 

199 for kpt in self.kpts: 

200 if kpt.s == spin: 

201 break 

202 else: 

203 raise RuntimeError(f'No k-point with spin {spin}') 

204 if self.eFermi is None: 

205 raise RuntimeError('Fermi level is not available') 

206 eH = -1.e32 

207 eL = 1.e32 

208 for kpt in self.kpts: 

209 if kpt.s == spin: 

210 for e in kpt.eps_n: 

211 if e <= self.eFermi: 

212 eH = max(eH, e) 

213 else: 

214 eL = min(eL, e) 

215 return eH, eL 

216 

217 def properties(self) -> Properties: 

218 return OutputPropertyWrapper(self).properties() 

219 

220 

221def propertygetter(func): 

222 from functools import wraps 

223 

224 @wraps(func) 

225 def getter(self): 

226 value = func(self) 

227 if value is None: 

228 raise PropertyNotPresent(func.__name__) 

229 return value 

230 return lazyproperty(getter) 

231 

232 

233class OutputPropertyWrapper: 

234 def __init__(self, calc): 

235 self.calc = calc 

236 

237 @propertygetter 

238 def nspins(self): 

239 return self.calc.get_number_of_spins() 

240 

241 @propertygetter 

242 def nbands(self): 

243 return self.calc.get_number_of_bands() 

244 

245 @propertygetter 

246 def nkpts(self): 

247 return len(self.calc.kpts) // self.nspins 

248 

249 def _build_eig_occ_array(self, getter): 

250 arr = np.empty((self.nspins, self.nkpts, self.nbands)) 

251 for s in range(self.nspins): 

252 for k in range(self.nkpts): 

253 value = getter(spin=s, kpt=k) 

254 if value is None: 

255 return None 

256 arr[s, k, :] = value 

257 return arr 

258 

259 @propertygetter 

260 def eigenvalues(self): 

261 return self._build_eig_occ_array(self.calc.get_eigenvalues) 

262 

263 @propertygetter 

264 def occupations(self): 

265 return self._build_eig_occ_array(self.calc.get_occupation_numbers) 

266 

267 @propertygetter 

268 def fermi_level(self): 

269 return self.calc.get_fermi_level() 

270 

271 @propertygetter 

272 def kpoint_weights(self): 

273 return self.calc.get_k_point_weights() 

274 

275 @propertygetter 

276 def ibz_kpoints(self): 

277 return self.calc.get_ibz_k_points() 

278 

279 def properties(self) -> Properties: 

280 dct = {} 

281 for name in ['eigenvalues', 'occupations', 'fermi_level', 

282 'kpoint_weights', 'ibz_kpoints']: 

283 try: 

284 value = getattr(self, name) 

285 except PropertyNotPresent: 

286 pass 

287 else: 

288 dct[name] = value 

289 

290 for name, value in self.calc.results.items(): 

291 dct[name] = value 

292 

293 return Properties(dct)