Coverage for gpaw/new/fd/pot_calc.py: 87%

1from __future__ import annotations 

2 

3from math import pi 

4 

5import numpy as np 

6 

7from gpaw.core import UGDesc 

8from gpaw.new import spinsum, trace, zips 

9from gpaw.new.pot_calc import PotentialCalculator 

10 

11 

12class FDPotentialCalculator(PotentialCalculator): 

13 def __init__(self, 

14 wf_grid: UGDesc, 

15 fine_grid: UGDesc, 

16 setups, 

17 xc, 

18 poisson_solver, 

19 *, 

20 relpos_ac, 

21 atomdist, 

22 interpolation_stencil_range=3, 

23 environment=None, 

24 extensions=None, 

25 xp=np): 

26 self.fine_grid = fine_grid 

27 self.grid = wf_grid 

28 

29 self.vbar_ar = setups.create_local_potentials(fine_grid, relpos_ac, 

30 atomdist, xp=xp) 

31 self.ghat_aLr = setups.create_compensation_charges(fine_grid, 

32 relpos_ac, 

33 atomdist, 

34 xp=xp) 

35 

36 self.vbar_r = fine_grid.empty(xp=xp) 

37 self.vbar_ar.to_uniform_grid(out=self.vbar_r) 

38 

39 n = interpolation_stencil_range 

40 self.interpolation_stencil_range = n 

41 self._interpolate = wf_grid.transformer(fine_grid, n, xp=xp) 

42 self._restrict = fine_grid.transformer(wf_grid, n, xp=xp) 

43 

44 self.xp = xp 

45 

46 super().__init__(xc, poisson_solver, setups, 

47 relpos_ac=relpos_ac, 

48 environment=environment, 

49 extensions=extensions) 

50 

51 def __str__(self): 

52 txt = super().__str__() 

53 degree = self.interpolation_stencil_range * 2 - 1 

54 name = ['linear', 'cubic', 'quintic', 'heptic'][degree // 2] 

55 txt += (f'interpolation: tri-{name}' + 

56 f' # {degree}. degree polynomial\n') 

57 return txt 

58 

59 def interpolate(self, a_xR, a_xr=None): 

60 return self._interpolate(a_xR, a_xr) 

61 

62 def restrict(self, a_xr, a_xR=None): 

63 return self._restrict(a_xr, a_xR) 

64 

65 def calculate_non_selfconsistent_exc(self, xc, density): 

66 nt_sr, _, _ = self._interpolate_density(density.nt_sR) 

67 if density.taut_sR is not None: 

68 taut_sr = self.interpolate(density.taut_sR) 

69 else: 

70 taut_sr = None 

71 e_xc, _, _ = xc.calculate(nt_sr, taut_sr) 

72 return e_xc 

73 

74 def _interpolate_density(self, nt_sR): 

75 nt_sr = self.interpolate(nt_sR) 

76 if not nt_sR.desc.pbc_c.all(): 

77 Nt1_s = nt_sR.integrate() 

78 Nt2_s = nt_sr.integrate() 

79 for Nt1, Nt2, nt_r in zips(Nt1_s, Nt2_s, nt_sr): 

80 if float(Nt2) > 1e-14: 

81 nt_r.data *= Nt1 / Nt2 

82 return nt_sr, None, None 

83 

84 @trace 

85 def calculate_pseudo_potential(self, density, ibzwfs, vHt_r): 

86 nt_sr, _, _ = self._interpolate_density(density.nt_sR) 

87 grid2 = nt_sr.desc 

88 

89 if density.taut_sR is not None: 

90 taut_sr = self.interpolate(density.taut_sR) 

91 else: 

92 taut_sr = None 

93 

94 e_xc, vxct_sr, dedtaut_sr = self.xc.calculate(nt_sr, taut_sr) 

95 

96 charge_r = grid2.empty(xp=self.xp) 

97 charge_r.data[:] = nt_sr.data[:density.ndensities].sum(axis=0) 

98 nt_r = charge_r.copy() 

99 e_zero = self.vbar_r.integrate(nt_r) 

100 

101 ccc_aL = density.calculate_compensation_charge_coefficients() 

102 

103 # Normalize: (LCAO basis functions may extend outside box) 

104 comp_charge = (4 * pi)**0.5 * sum(float(ccc_L[0]) 

105 for ccc_L in ccc_aL.values()) 

106 comp_charge = ccc_aL.layout.atomdist.comm.sum_scalar(comp_charge) 

107 pseudo_charge = charge_r.integrate() 

108 if abs(pseudo_charge) > 1e-10: 

109 pc = -comp_charge - density.charge + self.environment.charge 

110 charge_r.data *= pc / pseudo_charge 

111 

112 self.environment.update1(charge_r) 

113 

114 self.ghat_aLr.add_to(charge_r, ccc_aL) 

115 

116 if vHt_r is None: 

117 vHt_r = grid2.zeros(xp=self.xp) 

118 self.poisson_solver.solve(vHt_r, charge_r) 

119 e_coulomb = 0.5 * vHt_r.integrate(charge_r) 

120 

121 vt_sr = vxct_sr 

122 vt_sr.data += vHt_r.data + self.vbar_r.data 

123 

124 e_env = self.environment.update2(nt_r, vHt_r, vt_sr) 

125 

126 vt_sR = self.restrict(vt_sr) 

127 

128 e_external = e_env 

129 

130 V_aL = self.ghat_aLr.integrate(vHt_r) 

131 

132 return ({'coulomb': e_coulomb, 

133 'zero': e_zero, 

134 'xc': e_xc, 

135 'external': e_external}, 

136 vt_sR, 

137 dedtaut_sr, 

138 vHt_r, 

139 V_aL, 

140 np.nan) 

141 

142 def move(self, relpos_ac, atomdist): 

143 super().move(relpos_ac, atomdist) 

144 self.ghat_aLr.move(relpos_ac, atomdist) 

145 self.vbar_ar.move(relpos_ac, atomdist) 

146 self.vbar_ar.to_uniform_grid(out=self.vbar_r) 

147 

148 def force_contributions(self, Q_aL, density, potential): 

149 nt_R = spinsum(density.nt_sR) 

150 vt_R = spinsum(potential.vt_sR, mean=True) 

151 dedtaut_sR = potential.dedtaut_sR 

152 if dedtaut_sR is not None: 

153 dedtaut_R = spinsum(dedtaut_sR, mean=True) 

154 Ftauct_av = density.tauct_aX.derivative(dedtaut_R) 

155 else: 

156 Ftauct_av = None 

157 

158 nt_r = self.interpolate(nt_R) 

159 if not nt_r.desc.pbc_c.all(): 

160 scale = nt_R.integrate() / nt_r.integrate() 

161 nt_r.data *= scale 

162 

163 F_avL = self.ghat_aLr.derivative(potential.vHt_x) 

164 force_av = np.zeros((len(Q_aL), 3)) 

165 for a, dF_vL in F_avL.items(): 

166 force_av[a] += dF_vL @ Q_aL[a] 

167 

168 force_av += self.environment.forces(nt_r, potential.vHt_x) 

169 

170 return (force_av, 

171 density.nct_aX.derivative(vt_R), 

172 Ftauct_av, 

173 self.vbar_ar.derivative(nt_r), 

174 self.extensions_force_av)