Coverage for gpaw/response/integrators.py: 96%

1from abc import ABC, abstractmethod

2from functools import cached_property

3import numpy as np

4from gpaw.response import timer

5from scipy.spatial import Delaunay

6from scipy.linalg.blas import zher

8import gpaw.cgpaw as cgpaw

9from gpaw.utilities.blas import rk, mmm

10from gpaw.utilities.progressbar import ProgressBar

11from gpaw.response.pw_parallelization import Blocks1D

14class Integrand(ABC):

15 @abstractmethod

16 def matrix_element(self, point):

17 ...

19 @abstractmethod

20 def eigenvalues(self, point):

21 ...

24def czher(alpha: float, x, A) -> None:

25 """Hermetian rank-1 update of upper half of A.

27 A += alpha * np.outer(x.conj(), x)

29 """

30 AT = A.T

31 out = zher(alpha, x, 1, 1, 0, len(x), AT, 1)

32 assert out is AT

35class Integrator:

36 def __init__(self, cell_cv, context, blockcomm, kncomm):

37 """Baseclass for Brillouin zone integration and band summation.

39 Simple class to calculate integrals over Brilloun zones

40 and summation of bands.

42 context: ResponseContext

43 """

44 self.vol = abs(np.linalg.det(cell_cv))

45 self.context = context

46 self.blockcomm = blockcomm

47 self.kncomm = kncomm

49 def mydomain(self, domain):

50 from gpaw.response.pw_parallelization import Blocks1D

51 # This function does the same as distribute_domain

52 # but on a flat list and without all the fluff.

53 #

54 # In progress: Getting rid of distribute_domain(),

55 blocks = Blocks1D(self.kncomm, len(domain))

56 return [domain[i] for i in range(blocks.a, blocks.b)]

58 def integrate(self, **kwargs):

59 raise NotImplementedError

62class PointIntegrator(Integrator):

63 """Integrate brillouin zone using a broadening technique.

65 The broadening technique consists of smearing out the

66 delta functions appearing in many integrals by some factor

67 eta. In this code we use Lorentzians."""

69 def integrate(self, *, task, wd, domain, integrand, out_wxx):

70 """Integrate a response function over bands and kpoints."""

72 self.context.print('Integral kind:', task.kind)

74 mydomain = self.mydomain(domain)

76 prefactor = (2 * np.pi)**3 / self.vol / domain.nkpts

77 out_wxx /= prefactor * self.kncomm.size

79 # Sum kpoints

80 # Calculate integrations weight

81 pb = ProgressBar(self.context.fd)

82 for _, point in pb.enumerate(mydomain):

83 n_MG = integrand.matrix_element(point)

84 if n_MG is None:

85 continue

86 deps_M = integrand.eigenvalues(point)

88 task.run(wd, n_MG, deps_M, out_wxx)

90 # We have divided the broadcasted sum from previous update by

91 # kncomm.size, and thus here is will be back to its original value.

92 # This is to prevent allocating an extra large array.

93 self.kncomm.sum(out_wxx)

95 if self.blockcomm.size == 1 and task.symmetrizable_unless_blocked:

96 # Fill in upper/lower triangle also:

97 nx = out_wxx.shape[1]

98 il = np.tril_indices(nx, -1)

99 iu = il[::-1]

100

101 if isinstance(task, Hilbert):

102 # XXX special hack since one of them wants the other

103 # triangle.

104 for out_xx in out_wxx:

105 out_xx[il] = out_xx[iu].conj()

106 else:

107 for out_xx in out_wxx:

108 out_xx[iu] = out_xx[il].conj()

109

110 out_wxx *= prefactor

111

112

113class IntegralTask(ABC):

114 # Unique string for each kind of integral:

115 kind = '(unset)'

116

117 # Some integrals kinds like to calculate upper or lower half of the output

118 # when nblocks==1. In that case, this boolean signifies to the

119 # integrator that the output array should be symmetrized.

120 #

121 # Actually: We don't gain anything much by doing this boolean

122 # more systematically, since it's just Hermitian and Hilbert that need

123 # it, and then one of the Tetrahedron types which is not compatible

124 # anyway. We should probably not do this.

125 symmetrizable_unless_blocked = False

126

127 @abstractmethod

128 def run(self, wd, n_mG, deps_m, out_wxx):

129 """Add contribution from one point to out_wxx."""

130

131

132class GenericUpdate(IntegralTask):

133 kind = 'response function'

134 symmetrizable_unless_blocked = False

135

136 def __init__(self, eta, blockcomm, eshift=None):

137 self.eta = eta

138 self.blockcomm = blockcomm

139 self.eshift = eshift or 0.0

140

141 # @timer('CHI_0 update')

142 def run(self, wd, n_mG, deps_m, chi0_wGG):

143 """Update chi."""

144

145 deps_m += self.eshift * np.sign(deps_m)

146 deps1_m = deps_m + 1j * self.eta

147 deps2_m = deps_m - 1j * self.eta

148

149 blocks1d = Blocks1D(self.blockcomm, chi0_wGG.shape[2])

150

151 for omega, chi0_GG in zip(wd.omega_w, chi0_wGG):

152 x_m = (1 / (omega + deps1_m) - 1 / (omega - deps2_m))

153 if blocks1d.blockcomm.size > 1:

154 nx_mG = n_mG[:, blocks1d.myslice] * x_m[:, np.newaxis]

155 else:

156 nx_mG = n_mG * x_m[:, np.newaxis]

157

158 mmm(1.0, np.ascontiguousarray(nx_mG.T), 'N', n_mG.conj(), 'N',

159 1.0, chi0_GG)

160

161

162class Hermitian(IntegralTask):

163 kind = 'hermitian response function'

164 symmetrizable_unless_blocked = True

165

166 def __init__(self, blockcomm, eshift=None):

167 self.blockcomm = blockcomm

168 self.eshift = eshift or 0.0

169

170 # @timer('CHI_0 hermetian update')

171 def run(self, wd, n_mG, deps_m, chi0_wGG):

172 """If eta=0 use hermitian update."""

173 deps_m += self.eshift * np.sign(deps_m)

174

175 blocks1d = Blocks1D(self.blockcomm, chi0_wGG.shape[2])

176

177 for w, omega in enumerate(wd.omega_w):

178 if blocks1d.blockcomm.size == 1:

179 x_m = np.abs(2 * deps_m / (omega.imag**2 + deps_m**2))**0.5

180 nx_mG = n_mG.conj() * x_m[:, np.newaxis]

181 rk(-1.0, nx_mG, 1.0, chi0_wGG[w], 'n')

182 else:

183 x_m = np.abs(2 * deps_m / (omega.imag**2 + deps_m**2))

184 mynx_mG = n_mG[:, blocks1d.myslice] * x_m[:, np.newaxis]

185 mmm(-1.0, mynx_mG, 'T', n_mG.conj(), 'N', 1.0, chi0_wGG[w])

186

187

188class Hilbert(IntegralTask):

189 kind = 'spectral function'

190 symmetrizable_unless_blocked = True

191

192 def __init__(self, blockcomm, eshift=None):

193 self.blockcomm = blockcomm

194 self.eshift = eshift or 0.0

195

196 # @timer('CHI_0 spectral function update (new)')

197 def run(self, wd, n_mG, deps_m, chi0_wGG):

198 """Update spectral function.

199

200 Updates spectral function A_wGG and saves it to chi0_wGG for

201 later hilbert-transform."""

202

203 deps_m += self.eshift * np.sign(deps_m)

204 o_m = abs(deps_m)

205 w_m = wd.get_floor_index(o_m)

206

207 blocks1d = Blocks1D(self.blockcomm, chi0_wGG.shape[2])

208

209 # Sort frequencies

210 argsw_m = np.argsort(w_m)

211 sortedo_m = o_m[argsw_m]

212 sortedw_m = w_m[argsw_m]

213 sortedn_mG = n_mG[argsw_m]

214

215 index = 0

216 while 1:

217 if index == len(sortedw_m):

218 break

219

220 w = sortedw_m[index]

221 startindex = index

222 while 1:

223 index += 1

224 if index == len(sortedw_m):

225 break

226 if w != sortedw_m[index]:

227 break

228

229 endindex = index

230

231 # Here, we have same frequency range w, for set of

232 # electron-hole excitations from startindex to endindex.

233 o1 = wd.omega_w[w]

234 o2 = wd.omega_w[w + 1]

235 p = np.abs(1 / (o2 - o1)**2)

236 p1_m = np.array(p * (o2 - sortedo_m[startindex:endindex]))

237 p2_m = np.array(p * (sortedo_m[startindex:endindex] - o1))

238

239 if blocks1d.blockcomm.size > 1 and w + 1 < wd.wmax:

240 x_mG = sortedn_mG[startindex:endindex, blocks1d.myslice]

241 mmm(1.0,

242 np.concatenate((p1_m[:, None] * x_mG,

243 p2_m[:, None] * x_mG),

244 axis=1).T.copy(),

245 'N',

246 sortedn_mG[startindex:endindex].T.copy(),

247 'C',

248 1.0,

249 chi0_wGG[w:w + 2].reshape((2 * blocks1d.nlocal,

250 blocks1d.N)))

251

252 if blocks1d.blockcomm.size <= 1 and w + 1 < wd.wmax:

253 x_mG = sortedn_mG[startindex:endindex]

254 l_Gm = (p1_m[:, None] * x_mG).T.copy()

255 r_Gm = x_mG.T.copy()

256 mmm(1.0, r_Gm, 'N', l_Gm, 'C', 1.0, chi0_wGG[w])

257 l_Gm = (p2_m[:, None] * x_mG).T.copy()

258 mmm(1.0, r_Gm, 'N', l_Gm, 'C', 1.0, chi0_wGG[w + 1])

259

260

261class Intraband(IntegralTask):

262 kind = 'intraband'

263 symmetrizable_unless_blocked = False

264

265 # @timer('CHI_0 intraband update')

266 def run(self, wd, vel_mv, deps_M, chi0_wvv):

267 """Add intraband contributions"""

268 # Intraband is a little bit special, we use neither wd nor deps_M

269

270 for vel_v in vel_mv:

271 x_vv = np.outer(vel_v, vel_v)

272 chi0_wvv[0] += x_vv

273

274

275class OpticalLimit(IntegralTask):

276 kind = 'response function wings'

277 symmetrizable_unless_blocked = False

278

279 def __init__(self, eta, eshift=None):

280 self.eta = eta

281 self.eshift = eshift or 0.0

282

283 # @timer('CHI_0 optical limit update')

284 def run(self, wd, n_mG, deps_m, chi0_wxvG):

285 """Optical limit update of chi."""

286 deps_m += self.eshift * np.sign(deps_m)

287

288 deps1_m = deps_m + 1j * self.eta

289 deps2_m = deps_m - 1j * self.eta

290

291 for w, omega in enumerate(wd.omega_w):

292 x_m = (1 / (omega + deps1_m) - 1 / (omega - deps2_m))

293 chi0_wxvG[w, 0] += np.dot(x_m * n_mG[:, :3].T, n_mG.conj())

294 chi0_wxvG[w, 1] += np.dot(x_m * n_mG[:, :3].T.conj(), n_mG)

295

296

297class HermitianOpticalLimit(IntegralTask):

298 kind = 'hermitian response function wings'

299 symmetrizable_unless_blocked = False

300

301 def __init__(self, eshift=None):

302 self.eshift = eshift or 0.0

303

304 # @timer('CHI_0 hermitian optical limit update')

305 def run(self, wd, n_mG, deps_m, chi0_wxvG):

306 """Optical limit update of hermitian chi."""

307 deps_m += self.eshift * np.sign(deps_m)

308

309 for w, omega in enumerate(wd.omega_w):

310 x_m = - np.abs(2 * deps_m / (omega.imag**2 + deps_m**2))

311 chi0_wxvG[w, 0] += np.dot(x_m * n_mG[:, :3].T, n_mG.conj())

312 chi0_wxvG[w, 1] += np.dot(x_m * n_mG[:, :3].T.conj(), n_mG)

313

314

315class HilbertOpticalLimit(IntegralTask):

316 kind = 'spectral function wings'

317 symmetrizable_unless_blocked = False

318

319 def __init__(self, eshift=None):

320 self.eshift = eshift or 0.0

321

322 # @timer('CHI_0 optical limit hilbert-update')

323 def run(self, wd, n_mG, deps_m, chi0_wxvG):

324 """Optical limit update of chi-head and -wings."""

325 deps_m += self.eshift * np.sign(deps_m)

326

327 for deps, n_G in zip(deps_m, n_mG):

328 o = abs(deps)

329 w = wd.get_floor_index(o)

330 if w + 1 >= wd.wmax:

331 continue

332 o1, o2 = wd.omega_w[w:w + 2]

333 if o > o2:

334 continue

335 else:

336 assert o1 <= o <= o2, (o1, o, o2)

337

338 p = 1 / (o2 - o1)**2

339 p1 = p * (o2 - o)

340 p2 = p * (o - o1)

341 x_vG = np.outer(n_G[:3], n_G.conj())

342 chi0_wxvG[w, 0, :, :] += p1 * x_vG

343 chi0_wxvG[w + 1, 0, :, :] += p2 * x_vG

344 chi0_wxvG[w, 1, :, :] += p1 * x_vG.conj()

345 chi0_wxvG[w + 1, 1, :, :] += p2 * x_vG.conj()

346

347

348class Point:

349 def __init__(self, kpt_c, K, spin):

350 self.kpt_c = kpt_c

351 self.K = K

352 self.spin = spin

353

354

355class Domain:

356 def __init__(self, kpts_kc, spins):

357 self.kpts_kc = kpts_kc

358 self.spins = spins

359

360 @property

361 def nkpts(self):

362 return len(self.kpts_kc)

363

364 @property

365 def nspins(self):

366 return len(self.spins)

367

368 def __len__(self):

369 return self.nkpts * self.nspins

370

371 def __getitem__(self, num) -> Point:

372 K = num // self.nspins

373 return Point(self.kpts_kc[K], K,

374 self.spins[num % self.nspins])

375

376 def tesselation(self):

377 tesselation = KPointTesselation(self.kpts_kc)

378 tesselated_domains = Domain(tesselation.bzk_kc, self.spins)

379 return tesselation, tesselated_domains

380

381

382class KPointTesselation:

383 def __init__(self, kpts):

384 self._td = Delaunay(kpts)

385

386 @property

387 def bzk_kc(self):

388 return self._td.points

389

390 @cached_property

391 def simplex_volumes(self):

392 volumes_s = np.zeros(self._td.nsimplex, float)

393 for s in range(self._td.nsimplex):

394 K_k = self._td.simplices[s]

395 k_kc = self._td.points[K_k]

396 volume = np.abs(np.linalg.det(k_kc[1:] - k_kc[0])) / 6.

397 volumes_s[s] = volume

398 return volumes_s

399

400 def tetrahedron_weight(self, K, deps_k, omega_w):

401 simplices_s = self.pts_k[K]

402 W_w = np.zeros(len(omega_w), float)

403 vol_s = self.simplex_volumes[simplices_s]

404 cgpaw.tetrahedron_weight(

405 deps_k, self._td.simplices, K, simplices_s, W_w, omega_w, vol_s)

406 return W_w

407

408 @cached_property

409 def pts_k(self):

410 pts_k = [[] for n in range(self.nkpts)]

411 for s, K_k in enumerate(self._td.simplices):

412 A_kv = np.append(self._td.points[K_k],

413 np.ones(4)[:, np.newaxis], axis=1)

414

415 D_kv = np.append((A_kv[:, :-1]**2).sum(1)[:, np.newaxis],

416 A_kv, axis=1)

417 a = np.linalg.det(D_kv[:, np.arange(5) != 0])

418

419 if np.abs(a) < 1e-10:

420 continue

421

422 for K in K_k:

423 pts_k[K].append(s)

424

425 return [np.array(pts_k[k], int) for k in range(self.nkpts)]

426

427 @property

428 def nkpts(self):

429 return self._td.npoints

430

431 @cached_property

432 def neighbours_k(self):

433 return [np.unique(self._td.simplices[self.pts_k[k]])

434 for k in range(self.nkpts)]

435

436

437class TetrahedronIntegrator(Integrator):

438 """Integrate brillouin zone using tetrahedron integration.

439

440 Tetrahedron integration uses linear interpolation of

441 the eigenenergies and of the matrix elements

442 between the vertices of the tetrahedron."""

443

444 @timer('Spectral function integration')

445 def integrate(self, *, domain, integrand, wd, out_wxx, task):

446 """Integrate response function.

447

448 Assume that the integral has the

449 form of a response function. For the linear tetrahedron

450 method it is possible calculate frequency dependent weights

451 and do a point summation using these weights."""

452

453 tesselation, alldomains = domain.tesselation()

454 mydomain = self.mydomain(alldomains)

455

456 with self.context.timer('eigenvalues'):

457 deps_tMk = None # t for term

458

459 for point in alldomains:

460 deps_M = -integrand.eigenvalues(point)

461 if deps_tMk is None:

462 deps_tMk = np.zeros([alldomains.nspins, *deps_M.shape,

463 tesselation.nkpts], float)

464 deps_tMk[point.spin, :, point.K] = deps_M

465

466 # Calculate integrations weight

467 pb = ProgressBar(self.context.fd)

468 for _, point in pb.enumerate(mydomain):

469 deps_Mk = deps_tMk[point.spin]

470 teteps_Mk = deps_Mk[:, tesselation.neighbours_k[point.K]]

471 n_MG = integrand.matrix_element(point)

472

473 # Generate frequency weights

474 i0_M, i1_M = wd.get_index_range(teteps_Mk.min(1), teteps_Mk.max(1))

475 with self.context.timer('tetrahedron weight'):

476 W_Mw = []

477 for deps_k, i0, i1 in zip(deps_Mk, i0_M, i1_M):

478 W_w = tesselation.tetrahedron_weight(

479 point.K, deps_k, wd.omega_w[i0:i1])

480 W_Mw.append(W_w)

481

482 task.run(n_MG, deps_Mk, W_Mw, i0_M, i1_M, out_wxx)

483

484 self.kncomm.sum(out_wxx)

485

486 if self.blockcomm.size == 1 and task.symmetrizable_unless_blocked:

487 # Fill in upper/lower triangle also:

488 nx = out_wxx.shape[1]

489 il = np.tril_indices(nx, -1)

490 iu = il[::-1]

491 for out_xx in out_wxx:

492 out_xx[il] = out_xx[iu].conj()

493

494

495class HilbertTetrahedron:

496 kind = 'spectral function'

497 symmetrizable_unless_blocked = True

498

499 def __init__(self, blockcomm):

500 self.blockcomm = blockcomm

501

502 def run(self, n_MG, deps_Mk, W_Mw, i0_M, i1_M, out_wxx):

503 """Update output array with dissipative part."""

504 blocks1d = Blocks1D(self.blockcomm, out_wxx.shape[2])

505

506 for n_G, deps_k, W_w, i0, i1 in zip(n_MG, deps_Mk, W_Mw,

507 i0_M, i1_M):

508 if i0 == i1:

509 continue

510

511 for iw, weight in enumerate(W_w):

512 if blocks1d.blockcomm.size > 1:

513 myn_G = n_G[blocks1d.myslice].reshape((-1, 1))

514 # gemm(weight, n_G.reshape((-1, 1)), myn_G,

515 # 1.0, out_wxx[i0 + iw], 'c')

516 mmm(weight, myn_G, 'N', n_G.reshape((-1, 1)), 'C',

517 1.0, out_wxx[i0 + iw])

518 else:

519 czher(weight, n_G.conj(), out_wxx[i0 + iw])

520

521

522class HilbertOpticalLimitTetrahedron:

523 kind = 'spectral function wings'

524 symmetrizable_unless_blocked = False

525

526 def run(self, n_MG, deps_Mk, W_Mw, i0_M, i1_M, out_wxvG):

527 """Update optical limit output array with dissipative part of the head

528 and wings."""

529 for n_G, deps_k, W_w, i0, i1 in zip(n_MG, deps_Mk, W_Mw,

530 i0_M, i1_M):

531 if i0 == i1:

532 continue

533

534 for iw, weight in enumerate(W_w):

535 x_vG = np.outer(n_G[:3], n_G.conj())

536 out_wxvG[i0 + iw, 0, :, :] += weight * x_vG

537 out_wxvG[i0 + iw, 1, :, :] += weight * x_vG.conj()