Coverage for gpaw/response/chi0_data.py: 100%

1from __future__ import annotations 

2 

3import numpy as np 

4 

5from ase.units import Ha 

6 

7from gpaw.pw.descriptor import PWMapping 

8 

9from gpaw.response.pw_parallelization import (Blocks1D, 

10 PlaneWaveBlockDistributor) 

11from gpaw.response.frequencies import (FrequencyDescriptor, 

12 ComplexFrequencyDescriptor) 

13from gpaw.response.pair_functions import (SingleQPWDescriptor, 

14 map_ZgG_array_to_reduced_pd) 

15 

16 

17class Chi0RelatedData: 

18 """Base class for chi0 related data objects. 

19 

20 Right now, all we do is to limit boiler plate code...""" 

21 

22 def __init__(self, 

23 wd: FrequencyDescriptor, 

24 qpd: SingleQPWDescriptor): 

25 self.wd = wd 

26 self.qpd = qpd 

27 self.q_c = qpd.q_c 

28 

29 # Basis set size 

30 self.nG = qpd.ngmax 

31 self.nw = len(wd) 

32 

33 

34class Chi0BodyData(Chi0RelatedData): 

35 """Data object containing the response body data arrays 

36 for a single q-point, while holding also the corresponding 

37 basis descriptors and block distributor.""" 

38 

39 def __init__(self, wd, qpd, 

40 blockdist: PlaneWaveBlockDistributor): 

41 super().__init__(wd, qpd) 

42 

43 # Initialize block distibution of plane wave basis 

44 self.blockdist = blockdist 

45 self.blocks1d = Blocks1D(blockdist.blockcomm, self.nG) 

46 

47 # Data array 

48 self.data_WgG = self.zeros() 

49 

50 def zeros(self): 

51 return np.zeros(self.WgG_shape, complex) 

52 

53 @property 

54 def mynG(self): 

55 return self.blocks1d.nlocal 

56 

57 @property 

58 def WgG_shape(self): 

59 return (self.nw, self.mynG, self.nG) 

60 

61 def get_distributed_frequencies_array(self): 

62 """Copy data to a 'wGG'-like array, distributed over the entire world. 

63 

64 This differs from copy_array_with_distribution('wGG'), in that the 

65 frequencies are distributed over world, instead of among the block 

66 communicator.""" 

67 return self.blockdist.distribute_frequencies(self.data_WgG, self.nw) 

68 

69 def get_distributed_frequencies_blocks1d(self): 

70 """Get Blocks1D for the global frequency distribution.""" 

71 return Blocks1D(self.blockdist.world, len(self.wd)) 

72 

73 def copy_array_with_distribution(self, distribution): 

74 """Copy data to a new array of a desired distribution. 

75 

76 Parameters 

77 ---------- 

78 distribution: str 

79 Array distribution. Choices: 'wGG' and 'WgG' 

80 """ 

81 data_x = self.blockdist.distribute_as(self.data_WgG, self.nw, 

82 distribution) 

83 

84 if data_x is self.data_WgG: 

85 # When asking for 'WgG' distribution or when there is no block 

86 # distribution at all, we may still be pointing to the original 

87 # array, but we want strictly to return a copy 

88 assert distribution == 'WgG' or \ 

89 self.blockdist.blockcomm.size == 1 

90 data_x = self.data_WgG.copy() 

91 

92 return data_x 

93 

94 def copy_with_reduced_pd(self, qpd): 

95 """Make a copy corresponding to a new plane-wave description.""" 

96 new_chi0_body = Chi0BodyData(self.wd, qpd, self.blockdist) 

97 

98 # Map data to reduced plane-wave representation 

99 new_chi0_body.data_WgG[:] = map_ZgG_array_to_reduced_pd( 

100 self.qpd, qpd, self.blockdist, self.data_WgG) 

101 

102 return new_chi0_body 

103 

104 

105class Chi0DrudeData: 

106 def __init__(self, zd: ComplexFrequencyDescriptor): 

107 self.zd = zd 

108 self.plasmafreq_vv, self.chi_Zvv = self.zeros() 

109 

110 def zeros(self): 

111 return (np.zeros(self.vv_shape, complex), # plasmafreq 

112 np.zeros(self.Zvv_shape, complex)) # chi0_drude 

113 

114 @staticmethod 

115 def from_frequency_descriptor(wd, rate): 

116 """Construct the Chi0DrudeData object from a frequency descriptor and 

117 the imaginary part (in eV) of the resulting horizontal frequency 

118 contour""" 

119 rate = rate / Ha # eV -> Hartree 

120 zd = ComplexFrequencyDescriptor(wd.omega_w + 1.j * rate) 

121 

122 return Chi0DrudeData(zd) 

123 

124 @property 

125 def nz(self): 

126 return len(self.zd) 

127 

128 @property 

129 def vv_shape(self): 

130 return (3, 3) 

131 

132 @property 

133 def Zvv_shape(self): 

134 return (self.nz,) + self.vv_shape 

135 

136 

137class Chi0OpticalExtensionData(Chi0RelatedData): 

138 def __init__(self, wd, qpd): 

139 assert qpd.optical_limit 

140 super().__init__(wd, qpd) 

141 

142 self.head_Wvv, self.wings_WxvG = self.zeros() 

143 

144 def zeros(self): 

145 return (np.zeros(self.Wvv_shape, complex), # head 

146 np.zeros(self.WxvG_shape, complex)) # wings 

147 

148 @property 

149 def Wvv_shape(self): 

150 return (self.nw, 3, 3) 

151 

152 @property 

153 def WxvG_shape(self): 

154 return (self.nw, 2, 3, self.nG) 

155 

156 def copy_with_reduced_pd(self, qpd): 

157 """Make a copy corresponding to a new plane-wave description.""" 

158 new_chi0_optical_extension = Chi0OpticalExtensionData(self.wd, qpd) 

159 

160 # Copy the head (present in any plane-wave representation) 

161 new_chi0_optical_extension.head_Wvv[:] = self.head_Wvv 

162 

163 # Map the wings to the reduced plane-wave description 

164 G2_G1 = PWMapping(qpd, self.qpd).G2_G1 

165 new_chi0_optical_extension.wings_WxvG[:] \ 

166 = self.wings_WxvG[..., G2_G1] 

167 

168 return new_chi0_optical_extension 

169 

170 

171class Chi0Data(Chi0RelatedData): 

172 """Container object for the chi0 data objects for a single q-point, 

173 while holding also the corresponding basis descriptors and block 

174 distributor.""" 

175 

176 def __init__(self, 

177 chi0_body: Chi0BodyData, 

178 chi0_opt_ext: Chi0OpticalExtensionData | None = None): 

179 super().__init__(chi0_body.wd, chi0_body.qpd) 

180 self.body = chi0_body 

181 

182 self.optical_limit = self.qpd.optical_limit 

183 if self.optical_limit: 

184 assert isinstance(chi0_opt_ext, Chi0OpticalExtensionData) 

185 assert chi0_opt_ext.wd is self.wd 

186 assert chi0_opt_ext.qpd is self.qpd 

187 else: 

188 assert chi0_opt_ext is None 

189 self.optical_extension = chi0_opt_ext 

190 

191 @staticmethod 

192 def from_chi0_body_data(chi0_body): 

193 """Construct the container from a chi0 body data instance.""" 

194 qpd = chi0_body.qpd 

195 if qpd.optical_limit: 

196 wd = chi0_body.wd 

197 chi0_optical_extension = Chi0OpticalExtensionData(wd, qpd) 

198 else: 

199 chi0_optical_extension = None 

200 

201 return Chi0Data(chi0_body, chi0_optical_extension) 

202 

203 def copy_with_reduced_pd(self, qpd): 

204 """Make a copy of the data object, reducing the plane wave basis.""" 

205 new_body = self.body.copy_with_reduced_pd(qpd) 

206 if self.optical_limit: 

207 new_optical_extension = \ 

208 self.optical_extension.copy_with_reduced_pd(qpd) 

209 else: 

210 new_optical_extension = None 

211 

212 return Chi0Data(new_body, new_optical_extension) 

213 

214 @property 

215 def chi0_WgG(self): 

216 return self.body.data_WgG 

217 

218 @property 

219 def chi0_Wvv(self): 

220 if self.optical_limit: 

221 return self.optical_extension.head_Wvv 

222 

223 @property 

224 def chi0_WxvG(self): 

225 if self.optical_limit: 

226 return self.optical_extension.wings_WxvG