Coverage for gpaw/response/dyson.py: 98%

1from __future__ import annotations 

2 

3from abc import ABC, abstractmethod 

4from collections.abc import Sequence 

5from dataclasses import dataclass 

6 

7import numpy as np 

8 

9from gpaw.response import timer 

10from gpaw.response.pair_functions import Chi 

11 

12 

13class HXCScaling(ABC): 

14 """Helper for scaling the hxc contribution to Dyson equations.""" 

15 

16 def __init__(self, lambd=None): 

17 self._lambd = lambd 

18 

19 @property 

20 def lambd(self): 

21 return self._lambd 

22 

23 def calculate_scaling(self, dyson_equations): 

24 self._lambd = self._calculate_scaling(dyson_equations) 

25 

26 @abstractmethod 

27 def _calculate_scaling(self, dyson_equations: DysonEquations) -> float: 

28 """Calculate hxc scaling coefficient.""" 

29 

30 

31class PWKernel(ABC): 

32 """Hartree, exchange and/correleation kernels in a plane-wave basis.""" 

33 

34 @property 

35 def nG(self): 

36 return self.get_number_of_plane_waves() 

37 

38 def add_to(self, x_GG): 

39 assert x_GG.shape == (self.nG, self.nG) 

40 self._add_to(x_GG) 

41 

42 @abstractmethod 

43 def get_number_of_plane_waves(self): 

44 """Return the size of the plane-wave basis.""" 

45 

46 @abstractmethod 

47 def _add_to(self, x_GG): 

48 """Add plane-wave kernel to an input array.""" 

49 

50 

51class NoKernel(PWKernel): 

52 """A plane-wave kernel equal to zero.""" 

53 

54 def __init__(self, *, nG): 

55 self._nG = nG 

56 

57 @classmethod 

58 def from_qpd(cls, qpd): 

59 qG_Gv = qpd.get_reciprocal_vectors(add_q=True) 

60 return cls(nG=qG_Gv.shape[0]) 

61 

62 def get_number_of_plane_waves(self): 

63 return self._nG 

64 

65 def _add_to(self, x_GG): 

66 pass 

67 

68 

69@dataclass 

70class HXCKernel: 

71 """Hartree-exchange-correlation kernel in a plane-wave basis.""" 

72 hartree_kernel: PWKernel 

73 xc_kernel: PWKernel 

74 

75 def __post_init__(self): 

76 assert self.hartree_kernel.nG == self.xc_kernel.nG 

77 self.nG = self.hartree_kernel.nG 

78 

79 def get_Khxc_GG(self): 

80 """Hartree-exchange-correlation kernel.""" 

81 # Allocate array 

82 Khxc_GG = np.zeros((self.nG, self.nG), dtype=complex) 

83 self.hartree_kernel.add_to(Khxc_GG) 

84 self.xc_kernel.add_to(Khxc_GG) 

85 return Khxc_GG 

86 

87 

88class DysonSolver: 

89 """Class for inversion of Dyson-like equations.""" 

90 

91 def __init__(self, context): 

92 self.context = context 

93 

94 @timer('Invert Dyson-like equations') 

95 def __call__(self, chiks: Chi, self_interaction: HXCKernel | Chi, 

96 hxc_scaling: HXCScaling | None = None) -> Chi: 

97 """Solve the dyson equation and return the many-body susceptibility.""" 

98 dyson_equations = self.get_dyson_equations(chiks, self_interaction) 

99 if hxc_scaling: 

100 if not hxc_scaling.lambd: # calculate, if it hasn't been already 

101 hxc_scaling.calculate_scaling(dyson_equations) 

102 lambd = hxc_scaling.lambd 

103 self.context.print(r'Rescaling the self-enhancement function by a ' 

104 f'factor of λ={lambd}') 

105 self.context.print('Inverting Dyson-like equations') 

106 return dyson_equations.invert(hxc_scaling=hxc_scaling) 

107 

108 @staticmethod 

109 def get_dyson_equations(chiks, self_interaction): 

110 if isinstance(self_interaction, HXCKernel): 

111 return DysonEquationsWithKernel(chiks, hxc_kernel=self_interaction) 

112 elif isinstance(self_interaction, Chi): 

113 return DysonEquationsWithXi(chiks, xi=self_interaction) 

114 else: 

115 raise ValueError( 

116 f'Invalid encoding of the self-interaction {self_interaction}') 

117 

118 

119class DysonEquations(Sequence): 

120 """Sequence of Dyson-like equations at different complex frequencies z.""" 

121 

122 def __init__(self, chiks: Chi): 

123 assert chiks.distribution == 'zGG' and\ 

124 chiks.blockdist.fully_block_distributed, \ 

125 "chiks' frequencies need to be fully distributed over world" 

126 self.chiks = chiks 

127 # Inherit basic descriptors from chiks 

128 self.qpd = chiks.qpd 

129 self.zd = chiks.zd 

130 self.zblocks = chiks.blocks1d 

131 self.spincomponent = chiks.spincomponent 

132 

133 def __len__(self): 

134 return self.zblocks.nlocal 

135 

136 def invert(self, hxc_scaling: HXCScaling | None = None) -> Chi: 

137 """Invert Dyson equations to obtain χ(z).""" 

138 # Scaling coefficient of the self-enhancement function 

139 lambd = hxc_scaling.lambd if hxc_scaling else None 

140 chi = self.chiks.new() 

141 for z, dyson_equation in enumerate(self): 

142 chi.array[z] = dyson_equation.invert(lambd=lambd) 

143 return chi 

144 

145 

146class DysonEquationsWithKernel(DysonEquations): 

147 def __init__(self, chiks: Chi, *, hxc_kernel: HXCKernel): 

148 # Check compatibility 

149 nG = hxc_kernel.nG 

150 assert chiks.array.shape[1:] == (nG, nG) 

151 # Initialize 

152 super().__init__(chiks) 

153 self.Khxc_GG = hxc_kernel.get_Khxc_GG() 

154 

155 def __getitem__(self, z): 

156 chiks_GG = self.chiks.array[z] 

157 xi_GG = chiks_GG @ self.Khxc_GG 

158 return DysonEquation(chiks_GG, xi_GG) 

159 

160 

161class DysonEquationsWithXi(DysonEquations): 

162 def __init__(self, chiks: Chi, *, xi: Chi): 

163 # Check compatibility 

164 assert xi.distribution == 'zGG' and \ 

165 xi.blockdist.fully_block_distributed 

166 assert chiks.spincomponent == xi.spincomponent 

167 assert np.allclose(chiks.zd.hz_z, xi.zd.hz_z) 

168 assert np.allclose(chiks.qpd.q_c, xi.qpd.q_c) 

169 # Initialize 

170 super().__init__(chiks) 

171 self.xi = xi 

172 

173 def __getitem__(self, z): 

174 return DysonEquation(self.chiks.array[z], self.xi.array[z]) 

175 

176 

177class DysonEquation: 

178 """Dyson equation at wave vector q and frequency z. 

179 

180 The Dyson equation is given in plane-wave components as 

181 

182 χ(q,z) = χ_KS(q,z) + Ξ(q,z) χ(q,z), 

183 

184 where the self-enhancement function encodes the electron correlations 

185 induced by by the effective (Hartree-exchange-correlation) interaction: 

186 

187 Ξ(q,z) = χ_KS(q,z) K_hxc(q,z) 

188 

189 See [to be published] for more information. 

190 """ 

191 

192 def __init__(self, chiks_GG, xi_GG): 

193 self.nG = chiks_GG.shape[0] 

194 self.chiks_GG = chiks_GG 

195 self.xi_GG = xi_GG 

196 

197 def invert(self, lambd: float | None = None): 

198 """Invert the Dyson equation (with or without a rescaling of Ξ). 

199 

200 χ(q,z) = [1 - λ Ξ(q,z)]^(-1) χ_KS(q,z) 

201 """ 

202 if lambd is None: 

203 lambd = 1. # no rescaling 

204 leftside_GG = np.eye(self.nG) - lambd * self.xi_GG 

205 return np.linalg.solve(leftside_GG, self.chiks_GG)