Coverage for gpaw/hybrids/wrapper.py: 96%

1from typing import Tuple, Dict 

2 

3import numpy as np 

4 

5from gpaw.xc import XC 

6from .coulomb import coulomb_interaction 

7from .forces import calculate_forces 

8from .paw import calculate_paw_stuff 

9from .scf import apply1, apply2 

10from .symmetry import Symmetry 

11 

12 

13class HybridXC: 

14 orbital_dependent = True 

15 type = 'HYB' 

16 

17 def __init__(self, 

18 xcname: str, 

19 fraction: float = None, 

20 omega: float = None, 

21 yukawa=None): 

22 from . import parse_name 

23 if xcname in ['EXX', 'PBE0', 'HSE03', 'HSE06', 'B3LYP', 'YS-PBE0']: 

24 if fraction is not None or omega is not None or yukawa is not None: 

25 raise ValueError 

26 self.name = xcname 

27 xcname, fraction, omega, yukawa = parse_name(xcname) 

28 else: 

29 if fraction is None or omega is None: 

30 raise ValueError 

31 self.name = f'{xcname}-{fraction:.3f}-{omega:.3f}' 

32 

33 self.xc = XC(xcname) 

34 self.exx_fraction = fraction 

35 self.omega = omega 

36 self.yukawa = yukawa 

37 

38 if xcname == 'null': 

39 self.description = '' 

40 else: 

41 self.description = f'{xcname} + ' 

42 self.description += f'{fraction} * EXX(omega = {omega} bohr^-1)' 

43 

44 self.vlda_sR = None 

45 self.v_sknG: Dict[Tuple[int, int], np.ndarray] = {} 

46 

47 self.ecc = np.nan 

48 self.evc = np.nan 

49 self.evv = np.nan 

50 self.ekin = np.nan 

51 

52 self.sym = None 

53 self.coulomb = None 

54 

55 def get_setup_name(self): 

56 return 'PBE' 

57 

58 def initialize(self, dens, ham, wfs): 

59 self.dens = dens 

60 self.wfs = wfs 

61 self.ecc = sum(setup.ExxC for setup in wfs.setups) * self.exx_fraction 

62 assert wfs.world.size == wfs.gd.comm.size 

63 

64 def get_description(self): 

65 return self.description 

66 

67 def set_positions(self, spos_ac): 

68 self.spos_ac = spos_ac 

69 

70 def calculate(self, gd, nt_sr, vt_sr): 

71 energy = self.ecc + self.evv + self.evc 

72 energy += self.xc.calculate(gd, nt_sr, vt_sr) 

73 return energy 

74 

75 def calculate_paw_correction(self, setup, D_sp, dH_sp=None, a=None): 

76 return self.xc.calculate_paw_correction(setup, D_sp, dH_sp, a=a) 

77 

78 def get_kinetic_energy_correction(self): 

79 return self.ekin 

80 

81 def apply_orbital_dependent_hamiltonian(self, kpt, psit_xG, 

82 Htpsit_xG=None, dH_asp=None): 

83 wfs = self.wfs 

84 if self.coulomb is None: 

85 self.coulomb = coulomb_interaction( 

86 self.omega, wfs.gd, wfs.kd, yukawa=self.yukawa) 

87 self.description += '\n' + self.coulomb.get_description() 

88 self.sym = Symmetry(wfs.kd) 

89 

90 paw_s = calculate_paw_stuff(wfs, self.dens) # ??????? 

91 

92 if kpt.f_n is None: 

93 # Just use LDA_X for first step: 

94 if self.vlda_sR is None: 

95 # First time: 

96 self.vlda_sR = self.calculate_lda_potential() 

97 pd = kpt.psit.pd 

98 for psit_G, Htpsit_G in zip(psit_xG, Htpsit_xG): 

99 Htpsit_G += pd.fft(self.vlda_sR[kpt.s] * 

100 pd.ifft(psit_G, kpt.k), kpt.q) 

101 else: 

102 self.vlda_sR = None 

103 if kpt.psit.array.base is psit_xG.base: 

104 if (kpt.s, kpt.k) not in self.v_sknG: 

105 assert not any(s == kpt.s for s, k in self.v_sknG) 

106 evc, evv, ekin, v_knG = apply1( 

107 kpt, Htpsit_xG, 

108 wfs, 

109 self.coulomb, self.sym, 

110 paw_s[kpt.s]) 

111 if kpt.s == 0: 

112 self.evc = 0.0 

113 self.evv = 0.0 

114 self.ekin = 0.0 

115 scale = 2 / wfs.nspins * self.exx_fraction 

116 self.evc += evc * scale 

117 self.evv += evv * scale 

118 self.ekin += ekin * scale 

119 self.v_sknG = {(kpt.s, k): v_nG 

120 for k, v_nG in v_knG.items()} 

121 v_nG = self.v_sknG.pop((kpt.s, kpt.k)) 

122 else: 

123 v_nG = apply2(kpt, psit_xG, Htpsit_xG, wfs, 

124 self.coulomb, self.sym, 

125 paw_s[kpt.s]) 

126 Htpsit_xG += v_nG * self.exx_fraction 

127 

128 def calculate_lda_potential(self): 

129 from gpaw.xc import XC 

130 lda = XC('LDA_X') 

131 nt_sr = self.dens.nt_sg 

132 vt_sr = np.zeros_like(nt_sr) 

133 vlda_sR = self.dens.gd.zeros(self.wfs.nspins) 

134 lda.calculate(self.dens.finegd, nt_sr, vt_sr) 

135 for vt_R, vt_r in zip(vlda_sR, vt_sr): 

136 vt_R[:], _ = self.dens.pd3.restrict(vt_r, self.dens.pd2) 

137 return vlda_sR * self.exx_fraction 

138 

139 def summary(self, log): 

140 log(self.description) 

141 

142 def add_forces(self, F_av): 

143 paw_s = calculate_paw_stuff(self.wfs, self.dens) 

144 F_av += calculate_forces(self.wfs, 

145 self.coulomb, 

146 self.sym, 

147 paw_s) * self.exx_fraction 

148 

149 def correct_hamiltonian_matrix(self, kpt, H_nn): 

150 return 

151 

152 def rotate(self, kpt, U_nn): 

153 pass # 1 / 0 

154 

155 def add_correction(self, kpt, psit_xG, Htpsit_xG, P_axi, c_axi, n_x, 

156 calculate_change=False): 

157 pass # 1 / 0 

158 

159 def read(self, reader): 

160 pass 

161 

162 def write(self, writer): 

163 pass 

164 

165 def set_grid_descriptor(self, gd): 

166 pass