Coverage for gpaw/pw/hamiltonian.py: 100%

1import numpy as np 

2from gpaw.arraydict import ArrayDict 

3from gpaw.density import Density 

4from gpaw.external import NoExternalPotential 

5from gpaw.hamiltonian import Hamiltonian 

6from gpaw.pw.lfc import PWLFC 

7from gpaw.pw.poisson import (ChargedReciprocalSpacePoissonSolver, 

8 ReciprocalSpacePoissonSolver) 

9from gpaw.typing import Array3D 

10 

11 

12class ReciprocalSpaceHamiltonian(Hamiltonian): 

13 def __init__(self, gd, finegd, pd2, pd3, nspins, collinear, 

14 setups, timer, xc, world, xc_redistributor, 

15 vext=None, 

16 psolver=None, redistributor=None, realpbc_c=None, 

17 charge=0.0): 

18 

19 assert redistributor is not None # XXX should not be like this 

20 

21 if vext is None: 

22 vext = NoExternalPotential() 

23 

24 Hamiltonian.__init__(self, gd, finegd, nspins, collinear, setups, 

25 timer, xc, world, vext=vext, 

26 redistributor=redistributor) 

27 

28 self.vbar = PWLFC([[setup.vbar] for setup in setups], pd2) 

29 self.pd2 = pd2 

30 self.pd3 = pd3 

31 self.xc_redistributor = xc_redistributor 

32 self.charge = charge 

33 

34 self.vHt_q = pd3.empty() 

35 

36 if psolver is None: 

37 psolver = {} 

38 elif not isinstance(psolver, dict): 

39 psolver = psolver.todict() 

40 

41 if psolver.get('name') == 'nointeraction': 

42 self.poisson = ReciprocalSpacePoissonSolver(pd3, charge, 0.0) 

43 else: 

44 if charge == 0.0 or realpbc_c.any(): 

45 self.poisson = ReciprocalSpacePoissonSolver(pd3, charge) 

46 else: 

47 self.poisson = ChargedReciprocalSpacePoissonSolver(pd3, charge) 

48 

49 if 'dipolelayer' in psolver: 

50 direction = psolver['dipolelayer'] 

51 assert len(psolver) == 1 

52 from gpaw.dipole_correction import DipoleCorrection 

53 self.poisson = DipoleCorrection(self.poisson, direction) 

54 self.poisson.check_direction(gd, realpbc_c) 

55 else: 

56 assert not psolver 

57 

58 self.npoisson = 0 

59 

60 self.vbar_Q = None 

61 self.vt_Q = None 

62 self.estress = None 

63 

64 def __str__(self): 

65 s = Hamiltonian.__str__(self) 

66 if self.charge != 0.0: 

67 s += f'Poisson solver:\n {self.poisson}\n' 

68 return s 

69 

70 @property 

71 def xc_gd(self): 

72 if self.xc_redistributor is None: 

73 return self.finegd 

74 return self.xc_redistributor.aux_gd 

75 

76 def set_positions(self, spos_ac, atom_partition): 

77 Hamiltonian.set_positions(self, spos_ac, atom_partition) 

78 self.vbar_Q = self.pd2.zeros() 

79 self.vbar.add(self.vbar_Q) 

80 

81 def update_pseudo_potential(self, dens): 

82 ebar = self.pd2.integrate(self.vbar_Q, dens.nt_Q, 

83 global_integral=False) 

84 with self.timer('Poisson'): 

85 epot = self.poisson.solve(self.vHt_q, dens) 

86 epot /= self.finegd.comm.size 

87 

88 self.vt_Q = self.vbar_Q.copy() 

89 

90 dens.map23.add_to1(self.vt_Q, self.vHt_q) 

91 

92 # vt_sG[:] = pd2.ifft(vt_Q) 

93 eext = self.vext.update_potential_pw(self, dens) 

94 

95 self.timer.start('XC 3D grid') 

96 

97 nt_xg = dens.nt_xg 

98 

99 # If we have a redistributor, we want to do the 

100 # good old distribute-calculate-collect: 

101 redist = self.xc_redistributor 

102 if redist is not None: 

103 nt_xg = redist.distribute(nt_xg) 

104 

105 vxct_xg = np.zeros_like(nt_xg) 

106 exc = self.xc.calculate(self.xc_gd, nt_xg, vxct_xg) 

107 exc /= self.finegd.comm.size 

108 if redist is not None: 

109 vxct_xg = redist.collect(vxct_xg) 

110 

111 for x, (vt_G, vxct_g) in enumerate(zip(self.vt_xG, vxct_xg)): 

112 vxc_G, vxc_Q = self.pd3.restrict(vxct_g, self.pd2) 

113 if x < self.nspins: 

114 vt_G += vxc_G 

115 self.vt_Q += vxc_Q / self.nspins 

116 else: 

117 vt_G += vxc_G 

118 

119 self.timer.stop('XC 3D grid') 

120 

121 energies = np.array([epot, ebar, eext, exc]) 

122 self.estress = self.gd.comm.sum_scalar(epot + ebar) 

123 return energies 

124 

125 def calculate_atomic_hamiltonians(self, density): 

126 def getshape(a): 

127 return sum(2 * l + 1 

128 for l, _ in enumerate(self.setups[a].ghat_l)), 

129 W_aL = ArrayDict(self.atomdist.aux_partition, getshape, float) 

130 

131 self.vext.update_atomic_hamiltonians_pw(self, W_aL, density) 

132 return self.atomdist.to_work(self.atomdist.from_aux(W_aL)) 

133 

134 def calculate_kinetic_energy(self, density): 

135 ekin = 0.0 

136 for vt_G, nt_G in zip(self.vt_xG, density.nt_xG): 

137 ekin -= self.gd.integrate(vt_G, nt_G, global_integral=False) 

138 ekin += self.gd.integrate(self.vt_sG, density.nct_G, 

139 global_integral=False).sum() 

140 return ekin 

141 

142 def restrict(self, in_xR, out_xR=None): 

143 """Restrict array.""" 

144 if out_xR is None: 

145 out_xR = self.gd.empty(in_xR.shape[:-3]) 

146 

147 a_xR = in_xR.reshape((-1,) + in_xR.shape[-3:]) 

148 b_xR = out_xR.reshape((-1,) + out_xR.shape[-3:]) 

149 

150 for in_R, out_R in zip(a_xR, b_xR): 

151 out_R[:] = self.pd3.restrict(in_R, self.pd2)[0] 

152 

153 return out_xR 

154 

155 restrict_and_collect = restrict 

156 

157 def calculate_forces2(self, dens, ghat_aLv, nct_av, vbar_av): 

158 self.vext.derivative_pw(self, ghat_aLv, dens) 

159 dens.nct.derivative(self.vt_Q, nct_av) 

160 self.vbar.derivative(dens.nt_Q, vbar_av) 

161 

162 def get_electrostatic_potential(self, dens: Density) -> Array3D: 

163 self.poisson.solve(self.vHt_q, dens) 

164 vHt_R = self.pd3.ifft(self.vHt_q, distribute=False) 

165 self.pd3.comm.broadcast(vHt_R, 0) 

166 return vHt_R