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

1from pathlib import Path 

2import numpy as np 

3 

4from ase.units import Hartree 

5 

6from gpaw.response.qpd import SingleQPWDescriptor 

7 

8from gpaw.response.frequencies import ComplexFrequencyDescriptor 

9from gpaw.response.pair_integrator import DynamicPairFunction 

10from gpaw.response.pw_parallelization import (Blocks1D, 

11 PlaneWaveBlockDistributor) 

12 

13 

14class LatticePeriodicPairFunction(DynamicPairFunction): 

15 r"""Data object for lattice periodic pair functions. 

16 

17 Any spatial dependent pair function is considered to be lattice periodic, 

18 if it is invariant under translations of Bravais lattice vectors R: 

19 

20 pf(r, r', z) = pf(r + R, r' + R, z). 

21 

22 The Bloch lattice Fourier transform of a lattice periodic pair function, 

23 __ 

24 \ 

25 pf(r, r', q, z) = / e^(-iq.[r-r'-R']) pf(r, r' + R', z) 

26 ‾‾ 

27 R' 

28 

29 is then periodic in both r and r' independently and can be expressed in an 

30 arbitrary lattice periodic basis. 

31 

32 In the GPAW response code, lattice periodic pair functions are expanded in 

33 plane waves, 

34 

35 1 // 

36 pf_GG'(q, z) = ‾‾ || drdr' e^(-iG.r) pf(r, r', q, z) e^(iG'.r') 

37 V0 // 

38 V0 

39 

40 which are encoded in the SingleQPWDescriptor along with the wave vector q. 

41 """ 

42 

43 def __init__(self, 

44 qpd: SingleQPWDescriptor, 

45 zd: ComplexFrequencyDescriptor, 

46 blockdist: PlaneWaveBlockDistributor, 

47 distribution='ZgG'): 

48 """Contruct the LatticePeriodicPairFunction. 

49 

50 Parameters 

51 ---------- 

52 distribution : str 

53 Memory distribution of the pair function array. 

54 Choices: 'ZgG', 'GZg' and 'zGG'. 

55 """ 

56 self.qpd = qpd 

57 self.blockdist = blockdist 

58 self.distribution = distribution 

59 

60 self.blocks1d = None 

61 self.shape = None 

62 super().__init__(qpd.q_c, zd) 

63 

64 def zeros(self): 

65 if self.shape is None: 

66 self._initialize_block_distribution() 

67 return np.zeros(self.shape, complex) 

68 

69 def _initialize_block_distribution(self): 

70 """Initialize 1D block distribution and corresponding array shape.""" 

71 nz = len(self.zd) 

72 nG = self.qpd.ngmax 

73 blockdist = self.blockdist 

74 distribution = self.distribution 

75 

76 if distribution == 'ZgG': 

77 blocks1d = Blocks1D(blockdist.blockcomm, nG) 

78 shape = (nz, blocks1d.nlocal, nG) 

79 elif distribution == 'GZg': 

80 blocks1d = Blocks1D(blockdist.blockcomm, nG) 

81 shape = (nG, nz, blocks1d.nlocal) 

82 elif distribution == 'zGG': 

83 blocks1d = Blocks1D(blockdist.blockcomm, nz) 

84 shape = (blocks1d.nlocal, nG, nG) 

85 else: 

86 raise NotImplementedError(f'Distribution: {distribution}') 

87 

88 self.blocks1d = blocks1d 

89 self.shape = shape 

90 

91 def array_with_view(self, view): 

92 """Access a given view into the pair function array.""" 

93 if view == 'ZgG' and self.distribution in ['ZgG', 'GZg']: 

94 if self.distribution == 'GZg': 

95 pf_GZg = self.array 

96 pf_ZgG = pf_GZg.transpose((1, 2, 0)) 

97 else: 

98 pf_ZgG = self.array 

99 

100 pf_x = pf_ZgG 

101 else: 

102 raise ValueError(f'{view} is not a valid view, when array is of ' 

103 f'distribution {self.distribution}') 

104 

105 return pf_x 

106 

107 def copy_with_distribution(self, distribution='ZgG'): 

108 """Copy the pair function to a specified memory distribution.""" 

109 new_pf = self._new(*self.my_args(), distribution=distribution) 

110 new_pf.array[:] = self.array_with_view(distribution) 

111 

112 return new_pf 

113 

114 @classmethod 

115 def _new(cls, *args, **kwargs): 

116 return cls(*args, **kwargs) 

117 

118 def my_args(self, qpd=None, zd=None, blockdist=None): 

119 """Return the positional construction arguments of the 

120 LatticePeriodicPairFunction.""" 

121 if qpd is None: 

122 qpd = self.qpd 

123 if zd is None: 

124 zd = self.zd 

125 if blockdist is None: 

126 blockdist = self.blockdist 

127 

128 return qpd, zd, blockdist 

129 

130 def copy_with_reduced_pd(self, qpd): 

131 """Copy the pair function, but within a reduced plane-wave basis.""" 

132 new_pf = self._new(*self.my_args(qpd=qpd), 

133 distribution=self.distribution) 

134 if self.distribution == 'zGG': 

135 new_pf.array[:] = map_zGG_array_to_reduced_pd(self.qpd, qpd, 

136 self.array) 

137 elif self.distribution == 'ZgG': 

138 new_pf.array[:] = map_ZgG_array_to_reduced_pd(self.qpd, qpd, 

139 self.blockdist, 

140 self.array) 

141 else: 

142 raise NotImplementedError('Chi.copy_with_reduced_pd has not been ' 

143 'implemented for distribution ' 

144 f'{self.distribution}') 

145 return new_pf 

146 

147 def copy_with_global_frequency_distribution(self): 

148 """Copy the pair function, but with distribution zGG over world.""" 

149 # Make a copy, which is globally block distributed 

150 blockdist = self.blockdist.new_distributor(nblocks='max') 

151 new_pf = self._new(*self.my_args(blockdist=blockdist), 

152 distribution='zGG') 

153 

154 # Redistribute the data, distributing the frequencies over world 

155 assert self.distribution == 'ZgG' 

156 new_pf.array[:] = self.blockdist.distribute_frequencies(self.array, 

157 len(self.zd)) 

158 

159 return new_pf 

160 

161 

162def map_ZgG_array_to_reduced_pd(qpdi, qpd, blockdist, in_ZgG): 

163 """Map the array in_ZgG from the qpdi to the qpd plane-wave basis.""" 

164 # Distribute over frequencies 

165 nw = in_ZgG.shape[0] 

166 tmp_zGG = blockdist.distribute_as(in_ZgG, nw, 'zGG') 

167 

168 # Reduce the plane-wave basis 

169 tmp_zGG = map_zGG_array_to_reduced_pd(qpdi, qpd, tmp_zGG) 

170 

171 # Distribute over plane waves 

172 out_ZgG = blockdist.distribute_as(tmp_zGG, nw, 'ZgG') 

173 

174 return out_ZgG 

175 

176 

177def map_zGG_array_to_reduced_pd(qpdi, qpd, in_zGG): 

178 """Map the array in_zGG from the qpdi to the qpd plane-wave basis.""" 

179 from gpaw.pw.descriptor import PWMapping 

180 

181 # Initialize the basis mapping 

182 pwmapping = PWMapping(qpdi, qpd) 

183 G2_GG = tuple(np.meshgrid(pwmapping.G2_G1, pwmapping.G2_G1, 

184 indexing='ij')) 

185 G1_GG = tuple(np.meshgrid(pwmapping.G1, pwmapping.G1, 

186 indexing='ij')) 

187 

188 # Allocate array in the new basis 

189 nG = qpd.ngmax 

190 out_zGG_shape = (in_zGG.shape[0], nG, nG) 

191 out_zGG = np.zeros(out_zGG_shape, complex) 

192 

193 # Extract values 

194 for z, in_GG in enumerate(in_zGG): 

195 out_zGG[z][G2_GG] = in_GG[G1_GG] 

196 

197 return out_zGG 

198 

199 

200class Chi(LatticePeriodicPairFunction): 

201 r"""Data object for the four-component susceptibility tensor χ_GG'^μν(q,z). 

202 """ 

203 

204 def __init__(self, spincomponent, qpd, zd, 

205 blockdist, distribution='ZgG'): 

206 r"""Construct a susceptibility of a given spin-component (μν).""" 

207 self.spincomponent = spincomponent 

208 super().__init__(qpd, zd, blockdist, distribution=distribution) 

209 

210 def new(self, **kwargs): 

211 return self._new(*self.my_args_and_kwargs(**kwargs)) 

212 

213 def my_args(self, spincomponent=None, qpd=None, zd=None, blockdist=None): 

214 """Return positional construction arguments of the Chi object.""" 

215 if spincomponent is None: 

216 spincomponent = self.spincomponent 

217 qpd, zd, blockdist = super().my_args(qpd=qpd, zd=zd, 

218 blockdist=blockdist) 

219 

220 return spincomponent, qpd, zd, blockdist 

221 

222 def my_args_and_kwargs(self, distribution=None, **args): 

223 """Return all the construction arguments of Chi, in order.""" 

224 args = self.my_args(**args) 

225 if distribution is None: 

226 distribution = self.distribution 

227 return args + (distribution,) 

228 

229 def copy_with_reduced_ecut(self, ecut): 

230 """Copy the susceptibility, but with a reduced ecut.""" 

231 ecut = ecut / Hartree # eV -> Hartree 

232 assert ecut <= self.qpd.ecut 

233 qpd = self.qpd.copy_with(ecut=ecut) 

234 return self.copy_with_reduced_pd(qpd) 

235 

236 def copy_reactive_part(self): 

237 r"""Return a copy of the reactive part of the susceptibility. 

238 

239 The reactive part of the susceptibility is defined as (see 

240 [PRB 103, 245110 (2021)]): 

241 

242 1 

243 χ_GG'^(μν')(q,z) = ‾ [χ_GG'^μν(q,z) + χ_(-G'-G)^νμ(-q,-z*)]. 

244 2 

245 

246 However if the density operators n^μ(r) and n^ν(r) are each others 

247 Hermitian conjugates, the reactive part simply becomes the Hermitian 

248 part in terms of the plane-wave basis: 

249 

250 1 

251 χ_GG'^(μν')(q,z) = ‾ [χ_GG'^μν(q,z) + χ_G'G^(μν*)(q,z)], 

252 2 

253 

254 which is trivial to evaluate. 

255 """ 

256 assert self.distribution == 'zGG' or \ 

257 (self.distribution == 'ZgG' and self.blockdist.blockcomm.size == 1) 

258 assert self.spincomponent in ['00', 'uu', 'dd', '+-', '-+'], \ 

259 'Spin-density operators has to be each others hermitian conjugates' 

260 chiksr = self.new(distribution='zGG') 

261 chiks_zGG = self.array 

262 chiksr.array += chiks_zGG 

263 chiksr.array += np.conj(np.transpose(chiks_zGG, (0, 2, 1))) 

264 chiksr.array /= 2. 

265 

266 return chiksr 

267 

268 def copy_dissipative_part(self): 

269 r"""Return a copy of the dissipative part of the susceptibility. 

270 

271 The dissipative part of the susceptibility is defined as (see 

272 [PRB 103, 245110 (2021)]): 

273 

274 1 

275 χ_GG'^(μν")(q,z) = ‾‾ [χ_GG'^μν(q,z) - χ_(-G'-G)^νμ(-q,-z*)]. 

276 2i 

277 

278 Similar to the reactive part, this expression reduces to the 

279 anti-Hermitian part of the susceptibility in terms of the plane-wave 

280 basis, whenever the density operators n^μ(r) and n^ν(r) are each others 

281 Hermitian conjugates: 

282 

283 1 

284 χ_GG'^(μν")(q,z) = ‾‾ [χ_GG'^μν(q,z) - χ_G'G^(μν*)(q,z)]. 

285 2i 

286 """ 

287 assert self.distribution == 'zGG' or \ 

288 (self.distribution == 'ZgG' and self.blockdist.blockcomm.size == 1) 

289 assert self.spincomponent in ['00', 'uu', 'dd', '+-', '-+'], \ 

290 'Spin-density operators has to be each others hermitian conjugates' 

291 chiksd = self.new(distribution='zGG') 

292 chiks_zGG = self.array 

293 chiksd.array += chiks_zGG 

294 chiksd.array -= np.conj(np.transpose(chiks_zGG, (0, 2, 1))) 

295 chiksd.array /= 2.j 

296 

297 return chiksd 

298 

299 def symmetrize_reciprocity(self): 

300 r"""Symmetrize the reciprocity of the susceptibility (for q=0). 

301 

302 In collinear systems without spin-orbit coupling, the plane-wave 

303 susceptibility is reciprocal in the sense that 

304 

305 χ_GG'^(μν)(q, ω) = χ_(-G'-G)^(μν)(-q, ω) 

306 

307 for all μν ∈ {00, uu, dd, +-, -+}, see [PRB 106, 085131 (2022)]. 

308 For q=0, we may symmetrize the susceptibility in this sense for free. 

309 """ 

310 assert np.allclose(self.q_c, 0.) 

311 assert self.distribution == 'zGG' or \ 

312 (self.distribution == 'ZgG' and self.blockdist.blockcomm.size == 1) 

313 assert self.spincomponent in ['00', 'uu', 'dd', '+-', '-+'], \ 

314 'Spin-density operators has to be each others hermitian conjugates' 

315 invmap_GG = get_inverted_pw_mapping(self.qpd, self.qpd) 

316 for chi_GG in self.array: 

317 # Symmetrize [χ_(GG')(q, ω) + χ_(-G'-G)(-q, ω)] / 2 

318 chi_GG[:] = (chi_GG + chi_GG[invmap_GG].T) / 2. 

319 

320 def write_macroscopic_component(self, filename): 

321 """Write the spatially averaged (macroscopic) component of the 

322 susceptibility to a file along with the frequency grid.""" 

323 chi_Z = self.get_macroscopic_component() 

324 if self.blocks1d.blockcomm.rank == 0: 

325 write_pair_function(filename, self.zd, chi_Z) 

326 

327 def get_macroscopic_component(self): 

328 """Get the macroscopic (G=0) component, all-gathered.""" 

329 assert self.distribution == 'zGG' 

330 chi_zGG = self.array 

331 chi_z = chi_zGG[:, 0, 0] # Macroscopic component 

332 chi_Z = self.blocks1d.all_gather(chi_z) 

333 return chi_Z 

334 

335 def write_array(self, filename): 

336 """Write the full susceptibility array to a file along with the 

337 frequency grid and plane-wave components.""" 

338 assert self.distribution == 'zGG' 

339 chi_ZGG = self.blocks1d.gather(self.array) 

340 if self.blocks1d.blockcomm.rank == 0: 

341 write_susceptibility_array(filename, self.zd, self.qpd, chi_ZGG) 

342 

343 def write_diagonal(self, filename): 

344 """Write the diagonal of the many-body susceptibility within a reduced 

345 plane-wave basis to a file along with the frequency grid.""" 

346 assert self.distribution == 'zGG' 

347 chi_zGG = self.array 

348 chi_zG = np.diagonal(chi_zGG, axis1=1, axis2=2) 

349 chi_ZG = self.blocks1d.gather(chi_zG) 

350 if self.blocks1d.blockcomm.rank == 0: 

351 write_susceptibility_array(filename, self.zd, self.qpd, chi_ZG) 

352 

353 

354def get_inverted_pw_mapping(qpd1, qpd2): 

355 """Get the planewave coefficients mapping GG' of qpd1 into -G-G' of qpd2""" 

356 G1_Gc = get_pw_coordinates(qpd1) 

357 G2_Gc = get_pw_coordinates(qpd2) 

358 

359 mG2_G1 = [] 

360 for G1_c in G1_Gc: 

361 found_match = False 

362 for G2, G2_c in enumerate(G2_Gc): 

363 if np.all(G2_c == -G1_c): 

364 mG2_G1.append(G2) 

365 found_match = True 

366 break 

367 if not found_match: 

368 raise ValueError('Could not match qpd1 and qpd2') 

369 

370 # Set up mapping from GG' to -G-G' 

371 invmap_GG = tuple(np.meshgrid(mG2_G1, mG2_G1, indexing='ij')) 

372 

373 return invmap_GG 

374 

375 

376def get_pw_coordinates(qpd): 

377 """Get the reciprocal lattice vector coordinates corresponding to a 

378 givne plane wave basis. 

379 

380 Please remark, that the response code currently works with one q-vector 

381 at a time, at thus only a single plane wave representation at a time. 

382 

383 Returns 

384 ------- 

385 G_Gc : nd.array (dtype=int) 

386 Coordinates on the reciprocal lattice 

387 """ 

388 # List of all plane waves 

389 G_Gv = np.array([qpd.G_Qv[Q] for Q in qpd.Q_qG[0]]) 

390 

391 # Use cell to get coordinates 

392 B_cv = 2.0 * np.pi * qpd.gd.icell_cv 

393 return np.round(np.dot(G_Gv, np.linalg.inv(B_cv))).astype(int) 

394 

395 

396def write_pair_function(filename, zd, pf_z): 

397 """Write a pair function pf(q,z) for a specific q.""" 

398 # For now, we assume that the complex frequencies lie on a horizontal 

399 # contour and that the pair function is complex. This could be easily 

400 # generalized at a later stage. 

401 assert zd.horizontal_contour 

402 omega_w = zd.omega_w * Hartree # Ha -> eV 

403 pf_w = pf_z # Atomic units 

404 assert pf_w.dtype == complex 

405 

406 # Write results file 

407 with open(filename, 'w') as fd: 

408 print('# {:>11}, {:>11}, {:>11}'.format( 

409 'omega [eV]', 'pf_w.real', 'pf_w.imag'), file=fd) 

410 for omega, pf in zip(omega_w, pf_w): 

411 print(' {:11.6f}, {:11.6f}, {:11.6f}'.format( 

412 omega, pf.real, pf.imag), file=fd) 

413 

414 

415def read_pair_function(filename): 

416 """Read a stored pair function file.""" 

417 d = np.loadtxt(filename, delimiter=',') 

418 

419 if d.shape[1] == 3: 

420 # Complex pair function on a horizontal contour 

421 omega_w = np.array(d[:, 0], float) 

422 pf_w = np.array(d[:, 1], complex) 

423 pf_w.imag = d[:, 2] 

424 else: 

425 raise ValueError(f'Unexpected array dimension {d.shape}') 

426 

427 return omega_w, pf_w 

428 

429 

430def write_susceptibility_array(filename, zd, qpd, chi_zx): 

431 """Write dynamic susceptibility array to a .npz file.""" 

432 assert Path(filename).suffix == '.npz', filename 

433 # For now, we assume that the complex frequencies lie on a horizontal 

434 # contour 

435 assert zd.horizontal_contour 

436 omega_w = zd.omega_w * Hartree # Ha -> eV 

437 G_Gc = get_pw_coordinates(qpd) 

438 chi_wx = chi_zx 

439 np.savez(filename, omega_w=omega_w, G_Gc=G_Gc, chi_wx=chi_wx) 

440 

441 

442def read_susceptibility_array(filename): 

443 """Read a stored dynamic susceptibility array file.""" 

444 assert Path(filename).suffix == '.npz', filename 

445 npzfile = np.load(filename) 

446 return npzfile['omega_w'], npzfile['G_Gc'], npzfile['chi_wx']