Coverage for gpaw/lcao/dipoletransition.py: 100%

1import numpy as np 

2 

3from gpaw.mpi import world 

4from gpaw.typing import ArrayND 

5 

6 

7# NOTE: This routine is not specific to Raman per se. Maybe it should go 

8# somewhere else? 

9def get_dipole_transitions(wfs) -> ArrayND: 

10 r""" 

11 Finds dipole transition matrix elements based on the velocity form. 

12 

13 Dipole and momentum matrix elements are related by the expression: 

14 <nk|r|mk> = -i hbar/m <nk|p|mk> / (E_nk-E_mk) 

15 

16 For n=m this is ill defined, and element will be set to zero. 

17 

18 Parameters 

19 ---------- 

20 wfs 

21 LCAO WaveFunctions object 

22 """ 

23 p_skvnm = get_momentum_transitions(wfs, False) 

24 r_skvnm = np.zeros_like(p_skvnm) 

25 for kpt in wfs.kpt_u: 

26 # NOTE: Check whether it's this way or other way around. 

27 deltaE = kpt.eps_n[:, None] - kpt.eps_n[None, :] 

28 # The treatment of the energy difference in the denominator is 

29 # from https://doi.org/10.1103/PhysRevB.52.14636 eq 7 and note 18 

30 # therein 

31 r_skvnm[kpt.s, kpt.k, :] = -1j * p_skvnm[kpt.s, kpt.k, :] * \ 

32 np.reciprocal(deltaE, where=~np.isclose(deltaE, 0.0)) 

33 wfs.kd.comm.sum(r_skvnm) 

34 return r_skvnm 

35 

36 

37def get_momentum_transitions(wfs, savetofile: bool = True) -> ArrayND: 

38 r""" 

39 Finds the momentum matrix elements: 

40 <nk|p|mk> = k \delta_nm - i <nk|\nabla|mk> 

41 

42 Parameters 

43 ---------- 

44 wfs 

45 LCAO WaveFunctions object 

46 savetofile: bool 

47 Determines whether matrix is written to the 

48 file mom_skvnm.npy (default=True) 

49 """ 

50 assert wfs.bd.comm.size == 1 

51 assert wfs.mode == 'lcao' 

52 nbands = wfs.bd.nbands 

53 nspins = wfs.nspins 

54 nk = wfs.kd.nibzkpts 

55 gd = wfs.gd 

56 dtype = wfs.dtype 

57 ksl = wfs.ksl 

58 

59 # print(wfs.kd.comm.size, wfs.gd.comm.size, wfs.bd.comm.size) 

60 

61 mom_skvnm = np.zeros((nspins, nk, 3, nbands, nbands), dtype=complex) 

62 

63 dThetadR_qvMM, _ = wfs.manytci.O_qMM_T_qMM(gd.comm, ksl.Mstart, 

64 ksl.Mstop, False, 

65 derivative=True) 

66 

67 mome_skvnm = np.zeros((nspins, nk, 3, nbands, nbands), dtype=dtype) 

68 momd_skv = np.zeros((nspins, nk, 3), dtype=dtype) 

69 moma_skvnm = np.zeros((nspins, nk, 3, nbands, nbands), dtype=dtype) 

70 

71 B_cv = 2.0 * np.pi * gd.icell_cv 

72 for kpt in wfs.kpt_u: 

73 C_nM = kpt.C_nM 

74 for v in range(3): 

75 dThetadRv_MM = dThetadR_qvMM[kpt.q, v] 

76 nabla_nn = -(C_nM.conj() @ dThetadRv_MM.conj() @ C_nM.T) 

77 gd.comm.sum(nabla_nn) 

78 mome_skvnm[kpt.s, kpt.k, v] = nabla_nn 

79 

80 # augmentation part 

81 moma_vnm = np.zeros((3, nbands, nbands), dtype=dtype) 

82 for a, P_ni in kpt.P_ani.items(): 

83 nabla_iiv = wfs.setups[a].nabla_iiv 

84 moma_vnm += np.einsum('ni,ijv,mj->vnm', 

85 P_ni.conj(), nabla_iiv, P_ni, optimize=True) 

86 gd.comm.sum(moma_vnm) 

87 moma_skvnm[kpt.s, kpt.k] = moma_vnm 

88 

89 # diagonal term 

90 k_v = np.dot(wfs.kd.ibzk_kc[kpt.k], B_cv) 

91 momd_skv[kpt.s, kpt.k, :] = k_v 

92 

93 mom_skvnm = - 1j * (mome_skvnm + moma_skvnm) 

94 # Extract a view of the diagonal elements out 

95 # Einsum with no optimization since that is empirically faster here 

96 mom_diag = np.einsum('skvnn->skvn', mom_skvnm) 

97 # Add the diagonal term. Since mom_diag is a view, this changes mom_skvnm 

98 mom_diag += momd_skv[..., None] 

99 wfs.kd.comm.sum(mom_skvnm) 

100 

101 if world.rank == 0 and savetofile: 

102 np.save('mom_skvnm.npy', mom_skvnm) 

103 return mom_skvnm