Coverage for gpaw/directmin/fdpw/er_localization.py: 10%

144 statements  

« prev     ^ index     » next       coverage.py v7.7.1, created at 2025-07-19 00:19 +0000

1""" 

2Potentials for orbital density dependent energy functionals 

3""" 

4 

5import numpy as np 

6from gpaw.utilities import pack_density, unpack_hermitian 

7from gpaw.lfc import LFC 

8from gpaw.transformers import Transformer 

9from gpaw.directmin.tools import d_matrix, get_n_occ 

10from gpaw.poisson import PoissonSolver 

11import gpaw.cgpaw as cgpaw 

12 

13 

14class ERLocalization: 

15 

16 """ 

17 Edmiston-Ruedenberg localisation functions. 

18 

19 This class return -1/2 sum_i f_i (2integrals rr') n_i n_i'/ (r-r') 

20 and it is gradients 

21 

22 """ 

23 def __init__(self, wfs, dens, 

24 sic_coarse_grid=True, 

25 poisson_solver='FPS'): 

26 

27 self.name = 'ER_SIC' 

28 # what we need from wfs 

29 self.setups = wfs.setups 

30 spos_ac = wfs.spos_ac 

31 self.cgd = wfs.gd 

32 

33 # what we need from dens 

34 self.finegd = dens.finegd 

35 self.sic_coarse_grid = sic_coarse_grid 

36 

37 if self.sic_coarse_grid: 

38 self.ghat = LFC(self.cgd, 

39 [setup.ghat_l for setup 

40 in self.setups], 

41 integral=np.sqrt(4 * np.pi), 

42 forces=True) 

43 self.ghat.set_positions(spos_ac) 

44 else: 

45 self.ghat = dens.ghat # we usually solve poiss. on finegd 

46 

47 if poisson_solver == 'FPS': 

48 self.poiss = PoissonSolver(use_charge_center=True, 

49 use_charged_periodic_corrections=True) 

50 elif poisson_solver == 'GS': 

51 self.poiss = PoissonSolver(name='fd', 

52 relax=poisson_solver, 

53 eps=1.0e-16, 

54 use_charge_center=True, 

55 use_charged_periodic_corrections=True) 

56 

57 if self.sic_coarse_grid is True: 

58 self.poiss.set_grid_descriptor(self.cgd) 

59 else: 

60 self.poiss.set_grid_descriptor(self.finegd) 

61 

62 self.interpolator = Transformer(self.cgd, self.finegd, 3) 

63 self.restrictor = Transformer(self.finegd, self.cgd, 3) 

64 self.dtype = wfs.dtype 

65 self.eigv_s = {} 

66 self.lagr_diag_s = {} 

67 self.e_sic_by_orbitals = {} 

68 self.counter = 0 # number of calls of this class 

69 

70 self.n_kps = wfs.kd.nibzkpts 

71 

72 def get_orbdens_compcharge_dm_kpt(self, wfs, kpt, n): 

73 

74 if wfs.mode == 'pw': 

75 nt_G = wfs.pd.gd.zeros(global_array=True) 

76 psit_G = wfs.pd.alltoall1(kpt.psit.array[n:n + 1], kpt.q) 

77 if psit_G is not None: 

78 psit_R = wfs.pd.ifft(psit_G, kpt.q, 

79 local=True, safe=False) 

80 cgpaw.add_to_density(1.0, psit_R, nt_G) 

81 wfs.pd.gd.comm.sum(nt_G) 

82 nt_G = wfs.pd.gd.distribute(nt_G) 

83 else: 

84 nt_G = np.absolute(kpt.psit_nG[n] ** 2) 

85 

86 # paw 

87 Q_aL = {} 

88 D_ap = {} 

89 for a, P_ni in kpt.P_ani.items(): 

90 P_i = P_ni[n] 

91 D_ii = np.outer(P_i, P_i.conj()).real 

92 D_ap[a] = D_p = pack_density(D_ii) 

93 Q_aL[a] = np.dot(D_p, self.setups[a].Delta_pL) 

94 

95 return nt_G, Q_aL, D_ap 

96 

97 def get_energy_and_gradients_kpt(self, wfs, kpt, grad_knG): 

98 

99 wfs.timer.start('SIC e/g grid calculations') 

100 k = self.n_kps * kpt.s + kpt.q 

101 n_occ = get_n_occ(kpt)[0] 

102 

103 e_total_sic = np.array([]) 

104 

105 for i in range(n_occ): 

106 nt_G, Q_aL, D_ap = \ 

107 self.get_orbdens_compcharge_dm_kpt(wfs, kpt, i) 

108 

109 # calculate - SI Hartree 

110 e_sic, v_ht_g = \ 

111 self.get_pseudo_pot(nt_G, Q_aL, i, kpoint=k) 

112 

113 # calculate PAW corrections 

114 e_sic_paw_m, dH_ap = \ 

115 self.get_paw_corrections(D_ap, v_ht_g) 

116 

117 if wfs.mode == 'pw': 

118 v_ht_g = wfs.pd.gd.collect(v_ht_g, broadcast=True) 

119 Q_G = wfs.pd.Q_qG[kpt.q] 

120 psit_G = wfs.pd.alltoall1(kpt.psit_nG[i: i + 1], kpt.q) 

121 if psit_G is not None: 

122 psit_R = wfs.pd.ifft( 

123 psit_G, kpt.q, local=True, safe=False) 

124 psit_R *= v_ht_g 

125 wfs.pd.fftplan.execute() 

126 vtpsit_G = wfs.pd.tmp_Q.ravel()[Q_G] 

127 else: 

128 vtpsit_G = wfs.pd.tmp_G 

129 wfs.pd.alltoall2(vtpsit_G, kpt.q, grad_knG[k][i: i + 1]) 

130 grad_knG[k][i] *= kpt.f_n[i] 

131 else: 

132 grad_knG[k][i] += kpt.psit_nG[i] * v_ht_g * kpt.f_n[i] 

133 

134 # total sic: 

135 e_sic += e_sic_paw_m 

136 c_axi = {} 

137 for a in kpt.P_ani.keys(): 

138 dH_ii = unpack_hermitian(dH_ap[a]) 

139 c_xi = np.dot(kpt.P_ani[a][i], dH_ii) 

140 c_axi[a] = c_xi * kpt.f_n[i] 

141 # add projectors to 

142 wfs.pt.add(grad_knG[k][i], c_axi, kpt.q) 

143 

144 e_total_sic = np.append(e_total_sic, 

145 kpt.f_n[i] * e_sic, axis=0) 

146 

147 self.e_sic_by_orbitals[k] = np.copy(e_total_sic) 

148 wfs.timer.stop('SIC e/g grid calculations') 

149 

150 return e_total_sic.sum() 

151 

152 def get_pseudo_pot(self, nt, Q_aL, i, kpoint=None): 

153 

154 if self.sic_coarse_grid is False: 

155 v_ht_g = self.finegd.zeros() 

156 nt_sg = self.finegd.zeros() 

157 else: 

158 v_ht_g = self.cgd.zeros() 

159 nt_sg = self.cgd.zeros() 

160 

161 if self.sic_coarse_grid is False: 

162 self.interpolator.apply(nt, nt_sg) 

163 nt_sg *= self.cgd.integrate(nt) / \ 

164 self.finegd.integrate(nt_sg) 

165 else: 

166 nt_sg = nt 

167 

168 self.ghat.add(nt_sg, Q_aL) 

169 

170 self.poiss.solve(v_ht_g, nt_sg, zero_initial_phi=False) 

171 

172 if self.sic_coarse_grid is False: 

173 ec = 0.5 * self.finegd.integrate(nt_sg * v_ht_g) 

174 else: 

175 ec = 0.5 * self.cgd.integrate(nt_sg * v_ht_g) 

176 

177 return np.array([-ec]), -1.0 * v_ht_g 

178 

179 def get_paw_corrections(self, D_ap, vHt_g): 

180 

181 # Hartree-PAW 

182 dH_ap = {} 

183 for a, D_p in D_ap.items(): 

184 dH_ap[a] = np.zeros(len(D_p)) 

185 

186 ec = 0.0 

187 W_aL = self.ghat.dict() 

188 self.ghat.integrate(vHt_g, W_aL) 

189 for a, D_p in D_ap.items(): 

190 setup = self.setups[a] 

191 M_p = np.dot(setup.M_pp, D_p) 

192 ec += np.dot(D_p, M_p) 

193 dH_ap[a] += (2.0 * M_p + np.dot(setup.Delta_pL, W_aL[a])) 

194 if self.sic_coarse_grid is False: 

195 ec = self.finegd.comm.sum_scalar(ec) 

196 else: 

197 ec = self.cgd.comm.sum_scalar(ec) 

198 

199 return np.array([-ec]), dH_ap 

200 

201 def get_energy_and_gradients_inner_loop( 

202 self, wfs, kpt, a_mat, evals, evec): 

203 

204 e_sic, l_odd = \ 

205 self.get_energy_and_hamiltonian_kpt(wfs, kpt) 

206 wfs.timer.start('Unitary gradients') 

207 f = np.ones(l_odd.shape[0]) 

208 

209 # calc_error: 

210 indz = np.absolute(l_odd) > 1.0e-4 

211 l_c = 2.0 * l_odd[indz] 

212 l_odd = f[:, np.newaxis] * l_odd.T.conj() - f * l_odd 

213 kappa = np.max(np.absolute(l_odd[indz]) / np.absolute(l_c)) 

214 

215 if a_mat is None: 

216 wfs.timer.stop('Unitary gradients') 

217 return l_odd.T.conj(), e_sic, kappa 

218 else: 

219 g_mat = evec.T.conj() @ l_odd.T.conj() @ evec 

220 g_mat = g_mat * d_matrix(evals) 

221 g_mat = evec @ g_mat @ evec.T.conj() 

222 

223 for i in range(g_mat.shape[0]): 

224 g_mat[i][i] *= 0.5 

225 

226 wfs.timer.stop('Unitary gradients') 

227 if a_mat.dtype == float: 

228 g_mat = g_mat.real 

229 return 2.0 * g_mat, e_sic, kappa 

230 

231 def get_energy_and_hamiltonian_kpt(self, wfs, kpt): 

232 

233 n_occ = get_n_occ(kpt)[0] 

234 k = self.n_kps * kpt.s + kpt.q 

235 grad = {k: np.zeros_like(kpt.psit_nG[:n_occ])} 

236 

237 e_sic = self.get_energy_and_gradients_kpt(wfs, kpt, grad) 

238 

239 l_odd = wfs.integrate(kpt.psit_nG[:n_occ], 

240 grad[k][:n_occ], True) 

241 

242 return e_sic, l_odd