Coverage for gpaw/response/chiks.py: 97%

1from __future__ import annotations

2from typing import Tuple

3import warnings

5import numpy as np

6from time import ctime

7from abc import abstractmethod

9from ase.units import Hartree

11from gpaw.utilities.blas import mmmx

13from gpaw.response import ResponseGroundStateAdapter, ResponseContext, timer

14from gpaw.response.symmetry import QSymmetryAnalyzer, QSymmetryInput

15from gpaw.response.symmetrize import BodySymmetryOperators

16from gpaw.response.frequencies import ComplexFrequencyDescriptor

17from gpaw.response.pw_parallelization import PlaneWaveBlockDistributor

18from gpaw.response.matrix_elements import (PlaneWaveMatrixElementCalculator,

19 NewPairDensityCalculator,

20 TransversePairPotentialCalculator)

21from gpaw.response.pair_integrator import PairFunctionIntegrator

22from gpaw.response.pair_transitions import PairTransitions

23from gpaw.response.qpd import SingleQPWDescriptor

24from gpaw.response.pair_functions import Chi

27class RealAxisWarning(UserWarning):

28 """For warning users to stay off the real axis."""

31class GeneralizedSuscetibilityCalculator(PairFunctionIntegrator):

32 r"""Abstract calculator for generalized Kohn-Sham susceptibilities.

34 For any combination of plane-wave matrix elements f(K) and g(K), one may

35 define a generalized Kohn-Sham susceptibility in the Lehmann representation

37 __ __ __

38 1 \ \ \

39 ̄x_KS,GG'^μν(q,ω+iη) = ‾ / / / σ^μ_ss' σ^ν_s's (f_nks - f_n'k+qs')

40 V ‾‾ ‾‾ ‾‾

41 k n,n' s,s'

42 f_(nks,n'k+qs')(G+q) g_(n'k+qs',nks)(-G'-q)

43 x ‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾

44 ħω - (ε_n'k+qs' - ε_nks) + iħη

46 where σ^μ and σ^ν are Pauli matrices and the plane-wave matrix elements are

47 defined in terms of some real local functional of the electron

48 (spin-)density f[n](r):

50 f_(nks,n'k+qs')(G+q) = <ψ_nks| e^-i(G+q)r f(r) |ψ_n'k+qs'>

51 """

53 def __init__(self, gs: ResponseGroundStateAdapter, context=None,

54 nblocks=1,

55 ecut=50, gammacentered=False,

56 nbands=None,

57 bandsummation='pairwise',

58 **kwargs):

59 """Contruct the GeneralizedSuscetibilityCalculator

61 Parameters

62 ----------

63 gs : ResponseGroundStateAdapter

64 context : ResponseContext

65 nblocks : int

66 Distribute the chiks_zGG array into nblocks (where nblocks is a

67 divisor of context.comm.size)

68 ecut : float (or None)

69 Plane-wave cutoff in eV

70 gammacentered : bool

71 Center the grid of plane waves around the Γ-point (or the q-vector)

72 nbands : int

73 Number of bands to include in the sum over states

74 bandsummation : str

75 Band summation strategy (does not change the result, but can affect

76 the run-time).

77 'pairwise': sum over pairs of bands

78 'double': double sum over band indices.

79 kwargs : see gpaw.response.pair_integrator.PairFunctionIntegrator

80 """

81 if context is None:

82 context = ResponseContext()

83 assert isinstance(context, ResponseContext)

85 super().__init__(gs, context, nblocks=nblocks, **kwargs)

87 self.ecut = None if ecut is None else ecut / Hartree # eV to Hartree

88 self.gammacentered = gammacentered

89 self.nbands = nbands

90 self.bandsummation = bandsummation

92 mecalc1, mecalc2 = self.create_matrix_element_calculators()

93 self.matrix_element_calc1 = mecalc1

94 self.matrix_element_calc2 = mecalc2

95 if mecalc2 is not mecalc1:

96 assert not self.qsymmetry.time_reversal, \

97 'Cannot make use of time-reversal symmetry for generalized ' \

98 'susceptibilities with two different matrix elements'

100 @abstractmethod

101 def create_matrix_element_calculators(self) -> Tuple[

102 PlaneWaveMatrixElementCalculator,

103 PlaneWaveMatrixElementCalculator]:

104 """Create the desired site matrix element calculators."""

105

106 def calculate(self, spincomponent, q_c, zd) -> Chi:

107 r"""Calculate ̄x_KS,GG'^μν(q,z), where z = ω + iη

108

109 Parameters

110 ----------

111 spincomponent : str

112 Spin component (μν) of the Kohn-Sham susceptibility.

113 Currently, '00', 'uu', 'dd', '+-' and '-+' are implemented.

114 q_c : list or np.array

115 Wave vector in relative coordinates

116 zd : ComplexFrequencyDescriptor

117 Complex frequencies z to evaluate ̄x_KS,GG'^μν(q,z) at.

118 """

119 if self.gs.metallic and not zd.upper_half_plane:

120 warnings.warn(

121 'The generalized susceptibility is not well behaved at the '

122 'real axis for metals. Please consider evaluating it in the '

123 'upper-half complex frequency plane.',

124 RealAxisWarning,

125 )

126 return self._calculate(*self._set_up_internals(spincomponent, q_c, zd))

127

128 def _set_up_internals(self, spincomponent, q_c, zd,

129 distribution='GZg'):

130 r"""Set up internal data objects."""

131 assert isinstance(zd, ComplexFrequencyDescriptor)

132

133 # Set up the internal plane-wave descriptor

134 qpdi = self.get_pw_descriptor(q_c, internal=True)

135

136 # Prepare to sum over bands and spins

137 transitions = self.get_band_and_spin_transitions(

138 spincomponent, nbands=self.nbands,

139 bandsummation=self.bandsummation)

140

141 self.context.print(self.get_info_string(

142 qpdi, len(zd), spincomponent, len(transitions)))

143

144 # Create data structure

145 chiks = self.create_chiks(spincomponent, qpdi, zd, distribution)

146

147 return chiks, transitions

148

149 def _calculate(self, chiks: Chi, transitions: PairTransitions):

150 r"""Integrate ̄x_KS according to the specified transitions."""

151 self.context.print('Initializing the matrix element PAW corrections')

152 self.matrix_element_calc1.initialize_paw_corrections(chiks.qpd)

153 if self.matrix_element_calc2 is not self.matrix_element_calc1:

154 self.matrix_element_calc2.initialize_paw_corrections(chiks.qpd)

155

156 # Perform the actual integration

157 symmetries = self._integrate(chiks, transitions)

158

159 # Symmetrize chiks according to the symmetries of the ground state

160 self.symmetrize(chiks, symmetries)

161

162 # Map to standard output format

163 chiks = self.post_process(chiks)

164

165 return chiks

166

167 def get_pw_descriptor(self, q_c, internal=False):

168 """Get plane-wave descriptor for the wave vector q_c.

169

170 Parameters

171 ----------

172 q_c : list or ndarray

173 Wave vector in relative coordinates

174 internal : bool

175 When using symmetries, the actual calculation of chiks must happen

176 using a q-centered plane wave basis. If internal==True, as it is by

177 default, the internal plane wave basis (used in the integration of

178 chiks.array) is returned, otherwise the external descriptor is

179 returned, corresponding to the requested chiks.

180 """

181 q_c = np.asarray(q_c, dtype=float)

182 gd = self.gs.gd

183

184 # Update to internal basis, if needed

185 if internal and self.gammacentered and not self.qsymmetry.disabled:

186 # In order to make use of the symmetries of the system to reduce

187 # the k-point integration, the internal code assumes a plane wave

188 # basis which is centered at q in reciprocal space.

189 gammacentered = False

190 # If we want to compute the pair function on a plane wave grid

191 # which is effectively centered in the gamma point instead of q, we

192 # need to extend the internal ecut such that the q-centered grid

193 # encompasses all reciprocal lattice points inside the gamma-

194 # centered sphere.

195 # The reduction to the global gamma-centered basis will then be

196 # carried out as a post processing step.

197

198 # Compute the extended internal ecut

199 B_cv = 2.0 * np.pi * gd.icell_cv # Reciprocal lattice vectors

200 q_v = q_c @ B_cv

201 ecut = get_ecut_to_encompass_centered_sphere(q_v, self.ecut)

202 else:

203 gammacentered = self.gammacentered

204 ecut = self.ecut

205

206 qpd = SingleQPWDescriptor.from_q(q_c, ecut, gd,

207 gammacentered=gammacentered)

208

209 return qpd

210

211 def create_chiks(self, spincomponent, qpd, zd, distribution):

212 """Create a new Chi object to be integrated."""

213 assert distribution in ['GZg', 'ZgG']

214 blockdist = PlaneWaveBlockDistributor(self.context.comm,

215 self.blockcomm,

216 self.intrablockcomm)

217 return Chi(spincomponent, qpd, zd,

218 blockdist, distribution=distribution)

219

220 @timer('Add integrand to ̄x_KS')

221 def add_integrand(self, kptpair, weight, chiks):

222 r"""Add generalized susceptibility integrand for a given k-point pair.

223

224 Calculates the relevant matrix elements and adds the susceptibility

225 integrand to the output data structure for all relevant band and spin

226 transitions of the given k-point pair, k -> k + q.

227

228 Depending on the bandsummation parameter, the integrand of the

229 generalized susceptibility is given by:

230

231 bandsummation: double

232

233 __

234 \ σ^μ_ss' σ^ν_s's (f_nks - f_n'k's')

235 (...)_k = / ‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾ f_kt(G+q) g_kt^*(G'+q)

236 ‾‾ ħz - (ε_n'k's' - ε_nks)

237 t

238

239 where f_kt(G+q) = f_nks,n'k's'(G+q) and k'=k+q up to a reciprocal wave

240 vector.

241

242 bandsummation: pairwise

243

244 __ /

245 \ | σ^μ_ss' σ^ν_s's (f_nks - f_n'k's')

246 (...)_k = / | ‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾

247 ‾‾ | ħz - (ε_n'k's' - ε_nks)

248 t \

249 \

250 σ^μ_s's σ^ν_ss' (f_nks - f_n'k's') |

251 - δ_n'>n ‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾ | f_kt(G+q) g_kt^*(G'+q)

252 ħz + (ε_n'k's' - ε_nks) |

253 /

254

255 The integrand is added to the output array chiks_x multiplied with the

256 supplied kptpair integral weight.

257 """

258 # Calculate the matrix elements f_kt(G+q) and g_kt(G+q)

259 matrix_element1 = self.matrix_element_calc1(kptpair, chiks.qpd)

260 if self.matrix_element_calc2 is self.matrix_element_calc1:

261 matrix_element2 = matrix_element1

262 else:

263 matrix_element2 = self.matrix_element_calc2(kptpair, chiks.qpd)

264 # Calculate the temporal part of the integrand

265 if chiks.spincomponent == '00' and self.gs.nspins == 1:

266 weight = 2 * weight

267 x_mytZ = get_temporal_part(chiks.spincomponent, chiks.zd.hz_z,

268 kptpair, self.bandsummation)

269 x_tZ = kptpair.tblocks.all_gather(x_mytZ)

270

271 self._add_integrand(

272 matrix_element1, matrix_element2, x_tZ, weight, chiks)

273

274 def _add_integrand(self, matrix_element1, matrix_element2, x_tZ,

275 weight, chiks):

276 r"""Add the generalized susceptibility integrand based on distribution.

277

278 This entail performing a sum of transition t and an outer product

279 in the plane-wave components G and G',

280 __

281 \

282 (...)_k = / x_t^μν(ħz) f_kt(G+q) g_kt^*(G'+q)

283 ‾‾

284 t

285

286 where x_t^μν(ħz) is the temporal part of ̄x_KS,GG'^μν(q,ω+iη).

287 """

288 _add_integrand = self.get_add_integrand_method(chiks.distribution)

289 _add_integrand(matrix_element1, matrix_element2, x_tZ, weight, chiks)

290

291 def get_add_integrand_method(self, distribution):

292 """_add_integrand seletor."""

293 if distribution == 'ZgG':

294 _add_integrand = self._add_integrand_ZgG

295 elif distribution == 'GZg':

296 _add_integrand = self._add_integrand_GZg

297 else:

298 raise ValueError(f'Invalid distribution {distribution}')

299 return _add_integrand

300

301 def _add_integrand_ZgG(self, matrix_element1, matrix_element2, x_tZ,

302 weight, chiks):

303 """Add integrand in ZgG distribution.

304

305 Z = global complex frequency index

306 g = distributed G plane wave index

307 G = global G' plane wave index

308 """

309 chiks_ZgG = chiks.array

310 myslice = chiks.blocks1d.myslice

311

312 with self.context.timer('Set up gcc and xf'):

313 # Multiply the temporal part with the k-point integration weight

314 x_Zt = np.ascontiguousarray(weight * x_tZ.T)

315

316 # Set up f_kt(G+q) and g_kt^*(G'+q)

317 f_tG = matrix_element1.get_global_array()

318 if matrix_element2 is matrix_element1:

319 g_tG = f_tG

320 else:

321 g_tG = matrix_element2.get_global_array()

322 gcc_tG = g_tG.conj()

323

324 # Set up x_t^μν(ħz) f_kt(G+q)

325 f_gt = np.ascontiguousarray(f_tG[:, myslice].T)

326 xf_Zgt = x_Zt[:, np.newaxis, :] * f_gt[np.newaxis, :, :]

327

328 with self.context.timer('Perform sum over t-transitions of xf * gcc'):

329 for xf_gt, chiks_gG in zip(xf_Zgt, chiks_ZgG):

330 mmmx(1.0, xf_gt, 'N', gcc_tG, 'N', 1.0, chiks_gG) # slow step

331

332 def _add_integrand_GZg(self, matrix_element1, matrix_element2, x_tZ,

333 weight, chiks):

334 """Add integrand in GZg distribution.

335

336 G = global G' plane wave index

337 Z = global complex frequency index

338 g = distributed G plane wave index

339 """

340 chiks_GZg = chiks.array

341 myslice = chiks.blocks1d.myslice

342

343 with self.context.timer('Set up gcc and xf'):

344 # Multiply the temporal part with the k-point integration weight

345 x_tZ *= weight

346

347 # Set up f_kt(G+q) and g_kt^*(G'+q)

348 f_tG = matrix_element1.get_global_array()

349 if matrix_element2 is matrix_element1:

350 g_tG = f_tG

351 else:

352 g_tG = matrix_element2.get_global_array()

353 g_Gt = np.ascontiguousarray(g_tG.T)

354 gcc_Gt = g_Gt.conj()

355

356 # Set up x_t^μν(ħz) f_kt(G+q)

357 f_tg = f_tG[:, myslice]

358 xf_tZg = x_tZ[:, :, np.newaxis] * f_tg[:, np.newaxis, :]

359

360 with self.context.timer('Perform sum over t-transitions of gcc * xf'):

361 mmmx(1.0, gcc_Gt, 'N', xf_tZg, 'N', 1.0, chiks_GZg) # slow step

362

363 @timer('Symmetrizing chiks')

364 def symmetrize(self, chiks, symmetries):

365 """Symmetrize chiks_zGG."""

366 operators = BodySymmetryOperators(symmetries, chiks.qpd)

367 # Distribute over frequencies and symmetrize

368 nz = len(chiks.zd)

369 chiks_ZgG = chiks.array_with_view('ZgG')

370 tmp_zGG = chiks.blockdist.distribute_as(chiks_ZgG, nz, 'zGG')

371 operators.symmetrize_zGG(tmp_zGG)

372 # Distribute over plane waves

373 chiks_ZgG[:] = chiks.blockdist.distribute_as(tmp_zGG, nz, 'ZgG')

374

375 @timer('Post processing')

376 def post_process(self, chiks):

377 """Cast a calculated chiks into a fixed output format."""

378 if chiks.distribution != 'ZgG':

379 # Always output chiks with distribution 'ZgG'

380 chiks = chiks.copy_with_distribution('ZgG')

381

382 if self.gammacentered and not self.qsymmetry.disabled:

383 # Reduce the q-centered plane-wave basis used internally to the

384 # gammacentered basis

385 assert not chiks.qpd.gammacentered # Internal qpd

386 qpd = self.get_pw_descriptor(chiks.q_c) # External qpd

387 chiks = chiks.copy_with_reduced_pd(qpd)

388

389 return chiks

390

391 def get_info_string(self, qpd, nz, spincomponent, nt):

392 r"""Get information about the ̄x_KS,GG'^μν(q,z) calculation"""

393 from gpaw.utilities.memory import maxrss

394

395 q_c = qpd.q_c

396 ecut = qpd.ecut * Hartree

397 Asize = nz * qpd.ngmax**2 * 16. / 1024**2 / self.blockcomm.size

398 cmem = maxrss() / 1024**2

399

400 isl = ['',

401 'Setting up a generalized Kohn-Sham susceptibility calculation '

402 'with:',

403 f' Spin component: {spincomponent}',

404 f' q_c: [{q_c[0]}, {q_c[1]}, {q_c[2]}]',

405 f' Number of frequency points: {nz}',

406 self.get_band_and_transitions_info_string(self.nbands, nt),

407 '',

408 self.get_basic_info_string(),

409 '',

410 'Plane-wave basis of the generalized Kohn-Sham susceptibility:',

411 f' Planewave cutoff: {ecut}',

412 f' Number of planewaves: {qpd.ngmax}',

413 ' Memory estimates:',

414 f' A_zGG: {Asize} M / cpu',

415 f' Memory usage before allocation: {cmem} M / cpu',

416 '',

417 f'{ctime()}']

418

419 return '\n'.join(isl)

420

421

422class ChiKSCalculator(GeneralizedSuscetibilityCalculator):

423 r"""Calculator class for the four-component Kohn-Sham susceptibility tensor

424

425 For collinear systems in absence of spin-orbit coupling,

426 see [PRB 103, 245110 (2021)],

427 __ __ __

428 1 \ \ \

429 χ_KS,GG'^μν(q,ω+iη) = ‾ / / / σ^μ_ss' σ^ν_s's (f_nks - f_n'k+qs')

430 V ‾‾ ‾‾ ‾‾

431 k n,n' s,s'

432 n_nks,n'k+qs'(G+q) n_n'k+qs',nks(-G'-q)

433 x ‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾

434 ħω - (ε_n'k+qs' - ε_nks) + iħη

435

436 where the matrix elements

437

438 n_nks,n'k+qs'(G+q) = <nks| e^-i(G+q)r |n'k+qs'>

439

440 are the plane-wave pair densities of each transition.

441 """

442

443 def create_matrix_element_calculators(self):

444 pair_density_calc = NewPairDensityCalculator(self.gs, self.context)

445 return pair_density_calc, pair_density_calc

446

447

448class SelfEnhancementCalculator(GeneralizedSuscetibilityCalculator):

449 r"""Calculator class for the transverse magnetic self-enhancement function.

450

451 For collinear systems in absence of spin-orbit coupling,

452 see [publication in preparation],

453 __ __ __

454 1 \ \ \

455 Ξ_GG'^++(q,ω+iη) = ‾ / / / σ^+_ss' σ^-_s's (f_nks - f_n'k+qs')

456 V ‾‾ ‾‾ ‾‾

457 k n,n' s,s'

458 n_nks,n'k+qs'(G+q) W^⟂_n'k+qs',nks(-G'-q)

459 x ‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾

460 ħω - (ε_n'k+qs' - ε_nks) + iħη

461

462 where the matrix elements

463

464 n_nks,n'k+qs'(G+q) = <nks| e^-i(G+q)r |n'k+qs'>

465

466 and

467

468 W^⟂_(nks,n'k+qs')(G+q) = <ψ_nks| e^-i(G+q)r f_LDA^-+(r) |ψ_n'k+qs'>

469

470 are the plane-wave pair densities and transverse magnetic pair potentials

471 respectively.

472 """

473

474 def __init__(self, gs: ResponseGroundStateAdapter, context=None,

475 rshelmax: int = -1,

476 rshewmin: float | None = None,

477 qsymmetry: QSymmetryInput = True,

478 **kwargs):

479 """Construct the SelfEnhancementCalculator.

480

481 Parameters

482 ----------

483 rshelmax : int

484 The maximum index l (l < 6) to use in the expansion of the

485 xc-kernel f_LDA^-+(r) into real spherical harmonics for the PAW

486 correction.

487 rshewmin : float or None

488 If None, f_LDA^-+(r) will be fully expanded up to the chosen lmax.

489 Given as a float (0 < rshewmin < 1), rshewmin indicates what

490 coefficients to use in the expansion. If any (l,m) coefficient

491 contributes with less than a fraction of rshewmin on average, it

492 will not be included.

493 """

494 self.rshelmax = rshelmax

495 self.rshewmin = rshewmin

496

497 qsymmetry = QSymmetryAnalyzer.from_input(qsymmetry)

498 super().__init__(gs, context=context,

499 qsymmetry=QSymmetryAnalyzer(

500 point_group=qsymmetry.point_group,

501 time_reversal=False),

502 **kwargs)

503

504 def create_matrix_element_calculators(self):

505 pair_density_calc = NewPairDensityCalculator(self.gs, self.context)

506 pair_potential_calc = TransversePairPotentialCalculator(

507 self.gs, self.context,

508 rshelmax=self.rshelmax,

509 rshewmin=self.rshewmin)

510 return pair_density_calc, pair_potential_calc

511

512 def _set_up_internals(self, spincomponent, *args, **kwargs):

513 # For now, we are hardcoded to use the transverse pair potential,

514 # calculating Ξ^++ corresponding to χ^+-

515 assert spincomponent == '+-'

516 return super()._set_up_internals(spincomponent, *args, **kwargs)

517

518

519def get_ecut_to_encompass_centered_sphere(q_v, ecut):

520 """Calculate the minimal ecut which results in a q-centered plane wave

521 basis containing all the reciprocal lattice vectors G, which lie inside a

522 specific gamma-centered sphere:

523

524 |G|^2 < 2 * ecut

525 """

526 q = np.linalg.norm(q_v)

527 ecut += q * (np.sqrt(2 * ecut) + q / 2)

528

529 return ecut

530

531

532def get_temporal_part(spincomponent, hz_z, kptpair, bandsummation):

533 """Get the temporal part of a (causal linear) susceptibility, x_t^μν(ħz).

534 """

535 _get_temporal_part = create_get_temporal_part(bandsummation)

536 return _get_temporal_part(spincomponent, hz_z, kptpair)

537

538

539def create_get_temporal_part(bandsummation):

540 """Creator component, deciding how to calculate the temporal part"""

541 if bandsummation == 'double':

542 return get_double_temporal_part

543 elif bandsummation == 'pairwise':

544 return get_pairwise_temporal_part

545 raise ValueError(bandsummation)

546

547

548def get_double_temporal_part(spincomponent, hz_z, kptpair):

549 r"""Get:

550

551 σ^μ_ss' σ^ν_s's (f_nks - f_n'k's')

552 x_t^μν(ħz) = ‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾

553 ħz - (ε_n'k's' - ε_nks)

554 """

555 # Calculate σ^μ_ss' σ^ν_s's

556 s1_myt, s2_myt = kptpair.get_local_spin_indices()

557 scomps_myt = get_smat_components(spincomponent, s1_myt, s2_myt)

558 # Calculate nominator

559 nom_myt = - scomps_myt * kptpair.df_myt # df = (f_n'k's' - f_nks)

560 # Calculate denominator

561 deps_myt = kptpair.deps_myt # dε = (ε_n'k's' - ε_nks)

562 denom_mytz = hz_z[np.newaxis] - deps_myt[:, np.newaxis]

563 regularize_intraband_transitions(denom_mytz, kptpair)

564

565 return nom_myt[:, np.newaxis] / denom_mytz

566

567

568def get_pairwise_temporal_part(spincomponent, hz_z, kptpair):

569 r"""Get:

570

571 /

572 | σ^μ_ss' σ^ν_s's (f_nks - f_n'k's')

573 x_t^μν(ħz) = | ‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾

574 | ħz - (ε_n'k's' - ε_nks)

575 \

576 \

577 σ^μ_s's σ^ν_ss' (f_nks - f_n'k's') |

578 - δ_n'>n ‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾ |

579 ħz + (ε_n'k's' - ε_nks) |

580 /

581 """

582 # Kroenecker delta

583 n1_myt, n2_myt = kptpair.get_local_band_indices()

584 delta_myt = np.ones(len(n1_myt))

585 delta_myt[n2_myt <= n1_myt] = 0

586 # Calculate σ^μ_ss' σ^ν_s's and σ^μ_s's σ^ν_ss'

587 s1_myt, s2_myt = kptpair.get_local_spin_indices()

588 scomps1_myt = get_smat_components(spincomponent, s1_myt, s2_myt)

589 scomps2_myt = get_smat_components(spincomponent, s2_myt, s1_myt)

590 # Calculate nominators

591 df_myt = kptpair.df_myt # df = (f_n'k's' - f_nks)

592 nom1_myt = - scomps1_myt * df_myt

593 nom2_myt = - delta_myt * scomps2_myt * df_myt

594 # Calculate denominators

595 deps_myt = kptpair.deps_myt # dε = (ε_n'k's' - ε_nks)

596 denom1_mytz = hz_z[np.newaxis] - deps_myt[:, np.newaxis]

597 denom2_mytz = hz_z[np.newaxis] + deps_myt[:, np.newaxis]

598 regularize_intraband_transitions(denom1_mytz, kptpair)

599 regularize_intraband_transitions(denom2_mytz, kptpair)

600

601 return nom1_myt[:, np.newaxis] / denom1_mytz\

602 - nom2_myt[:, np.newaxis] / denom2_mytz

603

604

605def regularize_intraband_transitions(denom_mytx, kptpair):

606 """Regularize the denominator of the temporal part in case of degeneracy.

607

608 If the q-vector connects two symmetrically equivalent k-points inside a

609 band, the occupation differences vanish and we regularize the denominator.

610

611 NB: In principle there *should* be a contribution from the intraband

612 transitions, but this is left for future work for now."""

613 intraband_myt = kptpair.get_local_intraband_mask()

614 degenerate_myt = np.abs(kptpair.deps_myt) < 1e-8

615 denom_mytx[intraband_myt & degenerate_myt, ...] = 1.

616

617

618def get_smat_components(spincomponent, s1_t, s2_t):

619 """Calculate σ^μ_ss' σ^ν_s's for spincomponent (μν)."""

620 smatmu = smat(spincomponent[0])

621 smatnu = smat(spincomponent[1])

622 return smatmu[s1_t, s2_t] * smatnu[s2_t, s1_t]

623

624

625def smat(spinrot):

626 if spinrot == '0':

627 return np.array([[1, 0], [0, 1]])

628 elif spinrot == 'u':

629 return np.array([[1, 0], [0, 0]])

630 elif spinrot == 'd':

631 return np.array([[0, 0], [0, 1]])

632 elif spinrot == '-':

633 return np.array([[0, 0], [1, 0]])

634 elif spinrot == '+':

635 return np.array([[0, 1], [0, 0]])

636 elif spinrot == 'z':

637 return np.array([[1, 0], [0, -1]])

638 else:

639 raise ValueError(spinrot)