Coverage for gpaw/new/pwfd/etdm.py: 14%

1from __future__ import annotations 

2from functools import partial 

3 

4import numpy as np 

5from gpaw.core.arrays import DistributedArrays as XArray 

6from gpaw.core.atom_arrays import AtomArrays 

7from gpaw.new import zips 

8from gpaw.new.density import Density 

9from gpaw.new.eigensolver import Eigensolver 

10from gpaw.new.hamiltonian import Hamiltonian 

11from gpaw.new.potential import Potential 

12from gpaw.new.ibzwfs import IBZWaveFunctions 

13from gpaw.new.pwfd.lbfgs import LBFGS 

14from gpaw.new.energies import DFTEnergies 

15 

16 

17class ETDM(Eigensolver): 

18 def __init__(self, 

19 *, 

20 excited_state: bool = False, 

21 converge_unocc: bool = False): 

22 self.search_dir = LBFGS() 

23 self.grad_unX: list[XArray] = [] 

24 self.converge_unocc = converge_unocc 

25 self.dS_aii: AtomArrays 

26 self.nocc_s: list[int] = [] 

27 self.preconditioner 

28 

29 def new(self, **params) -> ETDM: 

30 return ETDM(**params) 

31 

32 def iterate(self, 

33 ibzwfs: IBZWaveFunctions, 

34 density: Density, 

35 potential: Potential, 

36 hamiltonian: Hamiltonian, 

37 pot_calc, 

38 energies) -> tuple[float, float, DFTEnergies]: 

39 

40 if len(self.nocc_s) == 0: 

41 xp = ibzwfs.xp 

42 self.nocc_s = find_number_of_ocupied_bands(ibzwfs) 

43 self.preconditioner = hamiltonian.create_preconditioner(10, xp=xp) 

44 self.dS_aii = pot_calc.setups.get_overlap_corrections( 

45 density.D_asii.layout.atomdist, xp) 

46 

47 dH = potential.dH 

48 Ht = partial(hamiltonian.apply, 

49 potential.vt_sR, 

50 potential.dedtaut_sR, 

51 ibzwfs, density.D_asii) 

52 

53 if len(self.grad_unX) == 0: 

54 

55 for wfs in ibzwfs: 

56 wfs._P_ani = None 

57 wfs.orthonormalized = False 

58 wfs.orthonormalize() 

59 wfs.subspace_diagonalize(Ht, dH) 

60 

61 energies, potential = update_density_and_potential( 

62 density, potential, pot_calc, ibzwfs, hamiltonian) 

63 Ht = partial(hamiltonian.apply, 

64 potential.vt_sR, 

65 potential.dedtaut_sR, 

66 ibzwfs, density.D_asii) 

67 

68 for wfs in ibzwfs: 

69 nocc = self.nocc_s[wfs.spin] 

70 psit_nX = wfs.psit_nX[:nocc] 

71 grad_nX = psit_nX.new() 

72 Ht(psit_nX, out=grad_nX, spin=wfs.spin) 

73 apply_non_local_hamiltonian(grad_nX, wfs, potential) 

74 project_gradient(grad_nX, wfs, self.dS_aii) 

75 weight_n = (wfs.weight * wfs.spin_degeneracy * 

76 wfs.myocc_n[:nocc]) 

77 grad_nX.data *= weight_n[:, np.newaxis] 

78 self.grad_unX.append(grad_nX) 

79 

80 psit_unX = [] 

81 for wfs in ibzwfs: 

82 nocc = self.nocc_s[wfs.spin] 

83 psit_nX = wfs.psit_nX[:nocc] 

84 psit_unX.append(psit_nX) 

85 

86 pg_unX = [] 

87 for psit_nX, grad_nX in zips(psit_unX, self.grad_unX): 

88 pg_nX = grad_nX.new() 

89 self.preconditioner(psit_nX, grad_nX, out=pg_nX) 

90 pg_nX.data *= -1.0 / (2 * (3 - len(self.nocc_s))) 

91 pg_unX.append(pg_nX) 

92 

93 p_unX = self.search_dir.update(psit_unX, pg_unX) 

94 for wfs, p_nX in zips(ibzwfs, p_unX): 

95 project_gradient(p_nX, wfs) 

96 

97 slength = sum(p_nX.norm2().sum() for p_nX in p_unX)**0.5 

98 max_step = 0.2 

99 alpha = max_step / slength if slength > max_step else 1.0 

100 

101 for psit_nX, p_nX in zips(psit_unX, p_unX): 

102 psit_nX.data += alpha * p_nX.data 

103 

104 for wfs in ibzwfs: 

105 wfs._P_ani = None 

106 wfs.orthonormalized = False 

107 wfs.orthonormalize() 

108 

109 energies, potential = update_density_and_potential( 

110 density, potential, pot_calc, ibzwfs, hamiltonian) 

111 

112 Ht = partial(hamiltonian.apply, 

113 potential.vt_sR, 

114 potential.dedtaut_sR, 

115 ibzwfs, density.D_asii) 

116 

117 error = 0.0 

118 for psit_nX, grad_nX, wfs in zips(psit_unX, self.grad_unX, ibzwfs): 

119 Ht(psit_nX, out=grad_nX, spin=wfs.spin) 

120 apply_non_local_hamiltonian(grad_nX, wfs, potential) 

121 project_gradient(grad_nX, wfs, self.dS_aii) 

122 weight_n = (wfs.weight * wfs.spin_degeneracy * 

123 wfs.myocc_n[:nocc]) 

124 error += grad_nX.norm2() @ weight_n 

125 grad_nX.data *= weight_n[:, np.newaxis] 

126 

127 return 0.0, error, energies 

128 

129 def postprocess(self, ibzwfs, density, potential, hamiltonian): 

130 if not self.converge_unocc: 

131 return 

132 

133 # dH = potential.dH 

134 Ht = partial(hamiltonian.apply, 

135 potential.vt_sR, 

136 potential.dedtaut_sR, 

137 ibzwfs, density.D_asii) 

138 

139 grad_unX = [] 

140 psit_unX = [] 

141 for wfs in ibzwfs: 

142 nocc = self.nocc_s[wfs.spin] 

143 psit_nX = wfs.psit_nX[nocc:] 

144 psit_unX.append(psit_nX) 

145 grad_nX = psit_nX.new() 

146 Ht(psit_nX, out=grad_nX, spin=wfs.spin) 

147 apply_non_local_hamiltonian(grad_nX, wfs, potential, 

148 slice(nocc, None)) 

149 project_gradient(grad_nX, wfs, self.dS_aii) 

150 weight = wfs.weight * wfs.spin_degeneracy 

151 grad_nX.data *= weight 

152 grad_unX.append(grad_nX) 

153 

154 while 1: 

155 pg_unX = [] 

156 for psit_nX, grad_nX in zips(psit_unX, grad_unX): 

157 pg_nX = grad_nX.new() 

158 self.preconditioner(psit_nX, grad_nX, out=pg_nX) 

159 pg_nX.data *= -1.0 / (2 * (3 - len(self.nocc_s))) 

160 pg_unX.append(pg_nX) 

161 

162 p_unX = self.search_dir.update(psit_unX, pg_unX) 

163 for wfs, p_nX in zips(ibzwfs, p_unX): 

164 project_gradient(p_nX, wfs) 

165 

166 slength = sum(p_nX.norm2().sum() for p_nX in p_unX)**0.5 

167 max_step = 0.2 

168 alpha = max_step / slength if slength > max_step else 1.0 

169 

170 for psit_nX, p_nX in zips(psit_unX, p_unX): 

171 psit_nX.data += alpha * p_nX.data 

172 

173 for wfs in ibzwfs: 

174 wfs._P_ani = None 

175 wfs.orthonormalized = False 

176 wfs.orthonormalize() 

177 

178 error = 0.0 

179 for psit_nX, grad_nX, wfs in zips(psit_unX, grad_unX, ibzwfs): 

180 Ht(psit_nX, out=grad_nX, spin=wfs.spin) 

181 apply_non_local_hamiltonian(grad_nX, wfs, potential) 

182 project_gradient(grad_nX, wfs, self.dS_aii) 

183 weight = wfs.weight * wfs.spin_degeneracy 

184 error += grad_nX.norm2().sum() * weight 

185 grad_nX.data *= weight 

186 print(error) 

187 

188 

189def apply_non_local_hamiltonian(Htpsit_nX, 

190 wfs, 

191 potential: Potential, 

192 bands: slice | None = None) -> None: 

193 bands = bands or slice(len(Htpsit_nX)) 

194 c_ani = {} 

195 dH_asii = potential.dH_asii 

196 for a, P_ni in wfs.P_ani.items(): 

197 dH_ii = dH_asii[a][wfs.spin] 

198 c_ani[a] = P_ni[bands] @ dH_ii 

199 wfs.pt_aiX.add_to(Htpsit_nX, c_ani) 

200 

201 

202def project_gradient(grad_nX: XArray, 

203 wfs, 

204 dS_aii=None): 

205 nocc = len(grad_nX) 

206 psit_nX = wfs.psit_nX[:nocc] 

207 

208 M_nn = grad_nX.integrate(psit_nX) 

209 M_nn += M_nn.T.conj() 

210 M_nn *= 0.5 

211 grad_nX.data -= M_nn @ psit_nX.data 

212 if dS_aii: 

213 c_ani = {} 

214 for a, P_ni in wfs.P_ani.items(): 

215 c_ani[a] = M_nn @ P_ni[:nocc] @ -dS_aii[a] 

216 wfs.pt_aiX.add_to(grad_nX, c_ani) 

217 

218 

219def update_density_and_potential(density, 

220 potential, 

221 pot_calc, 

222 ibzwfs, 

223 hamiltonian) -> tuple[float, Potential]: 

224 density.update(ibzwfs, ked=pot_calc.xc.type == 'MGGA') 

225 potential, energies, _ = pot_calc.calculate(density, 

226 ibzwfs, 

227 potential.vHt_x) 

228 energies.set(kinetic=ibzwfs.calculate_kinetic_energy(hamiltonian, density), 

229 band=0.0) 

230 return energies, potential 

231 

232 

233def find_number_of_ocupied_bands(ibzwfs: IBZWaveFunctions) -> list[int]: 

234 nocc_s = [-1] * ibzwfs.nspins 

235 for wfs in ibzwfs: 

236 nocc = (wfs.occ_n > 0.5).sum() 

237 n = nocc_s[wfs.spin] 

238 if n != -1: 

239 assert nocc == n 

240 else: 

241 nocc_s[wfs.spin] = nocc 

242 return nocc_s