Coverage for gpaw/elph/raman_calculator.py: 92%

1""" 

2Calculates Raman matrices. 

3 

4i -> j -> m -> n 

5i, n are valence; j, m are conduction, also i=n in the end 

6see https://doi.org/10.1038/s41467-020-16529-6 

7""" 

8import numpy as np 

9 

10from ase.units import invcm, Hartree 

11from ase.utils.filecache import MultiFileJSONCache 

12from gpaw.calculator import GPAW 

13from gpaw.typing import ArrayND 

14 

15# TODO: only take kd, don't need whole calculator 

16 

17 

18class ResonantRamanCalculator: 

19 def __init__( 

20 self, 

21 calc: GPAW, 

22 wph_w: ArrayND, 

23 momname: str = "mom_skvnm.npy", 

24 elphname: str = "gsqklnn.npy", 

25 raman_name: str = "Rlab", 

26 ): 

27 """Resonant Raman Matrix calculator 

28 

29 Parameters 

30 ---------- 

31 calc: GPAW 

32 GPAW calculator object. 

33 w_ph: str, np.ndarray 

34 Zone centre phonon frequencies in eV 

35 momname: str 

36 Name of momentum file 

37 elphname: str 

38 Name of electron-phonon file 

39 raman_name: str 

40 Name of Rlab file cache. Default 'Rlab' 

41 """ 

42 # those won't make sense here 

43 assert calc.wfs.gd.comm.size == 1 

44 assert calc.wfs.bd.comm.size == 1 

45 self.kd = calc.wfs.kd 

46 self.calc = calc 

47 

48 # Phonon frequencies 

49 if isinstance(wph_w, str): 

50 wph_w = np.load(wph_w) 

51 assert max(wph_w) < 1.0 # else not eV units 

52 self.w_ph = wph_w 

53 

54 # Load files 

55 mom_skvnm = np.load(momname, mmap_mode="c") 

56 g_sqklnn = np.load(elphname, mmap_mode="c") # [s,q=0,k,l,n,m] 

57 self.mom_skvnm = mom_skvnm 

58 self.g_sqklnn = g_sqklnn 

59 

60 # Define a few more variables 

61 nspins = g_sqklnn.shape[0] 

62 nk = g_sqklnn.shape[2] 

63 assert wph_w.shape[0] == g_sqklnn.shape[3] 

64 assert mom_skvnm.shape[0] == nspins 

65 assert mom_skvnm.shape[1] == nk 

66 assert mom_skvnm.shape[-1] == g_sqklnn.shape[-1] 

67 

68 # JSON cache 

69 self.raman_cache = MultiFileJSONCache(raman_name) 

70 with self.raman_cache.lock("phonon_frequencies") as handle: 

71 if handle is not None: 

72 handle.save(wph_w) 

73 # Resonant part of Raman tensor does not have a frequency grid 

74 with self.raman_cache.lock("frequency_grid") as handle: 

75 if handle is not None: 

76 handle.save(None) 

77 

78 @classmethod 

79 def resonant_term( 

80 cls, 

81 f_vc: ArrayND, 

82 E_vc: ArrayND, 

83 mom_dnn: ArrayND, 

84 elph_lnn: ArrayND, 

85 nc: int, 

86 nv: int, 

87 w_in: float, 

88 wph_w: ArrayND, 

89 ) -> ArrayND: 

90 """Resonant term of the Raman tensor""" 

91 nmodes = elph_lnn.shape[0] 

92 term_l = np.zeros((nmodes), dtype=complex) 

93 t_ij = f_vc * mom_dnn[0, :nv, nc:] / (w_in - E_vc) 

94 for l in range(nmodes): 

95 t_xx = elph_lnn[l] 

96 t_mn = f_vc.T * mom_dnn[1, nc:, :nv] / (w_in - wph_w[l] - E_vc.T) 

97 term_l[l] += np.einsum("sj,jm,ms", t_ij, t_xx[nc:, nc:], t_mn) 

98 term_l[l] -= np.einsum("is,ni,sn", t_ij, t_xx[:nv, :nv], t_mn) 

99 return term_l 

100 

101 def calculate(self, w_in, d_i, d_o, gamma_l=0.1, limit_sum=False): 

102 """Calculate resonant Raman matrix 

103 

104 Parameters 

105 ---------- 

106 w_in: float 

107 Laser frequency in eV 

108 d_i: int 

109 Incoming polarization 

110 d_o: int 

111 Outgoing polarization 

112 gamma_l: float 

113 Line broadening in eV, (0.1eV=806rcm) 

114 limit_sum: bool 

115 Limit sum to occupied valence/unoccupied conduction states 

116 Use for debugging and testing. Don't use in production 

117 """ 

118 

119 raman_l = np.zeros((self.w_ph.shape[0]), dtype=complex) 

120 

121 print(f"Calculating Raman spectrum: Laser frequency = {w_in} eV") 

122 

123 # Loop over kpoints - this is parallelised 

124 for kpt in self.calc.wfs.kpt_u: 

125 # print(f"Rank {self.kd.comm.rank}: s={kpt.s}, k={kpt.k}") 

126 

127 f_n = kpt.f_n / kpt.weight 

128 assert np.isclose(max(f_n), 1.0, atol=0.1) 

129 

130 vs = np.arange(0, len(f_n)) 

131 cs = np.arange(0, len(f_n)) 

132 nc = 0 

133 nv = len(f_n) 

134 if limit_sum: # good to test that we got arrays right 

135 vs = np.where(f_n >= 0.1)[0] 

136 cs = np.where(f_n < 0.9)[0] 

137 nv = max(vs) + 1 # VBM+1 index 

138 nc = min(cs) # CBM index 

139 

140 # Precalculate f * (1-f) term 

141 f_vc = np.outer(f_n[vs], 1.0 - f_n[cs]) 

142 

143 # Precalculate E-E term 

144 E_vc = np.zeros((len(vs), len(cs)), dtype=complex) + 1j * gamma_l 

145 for n in range(len(vs)): 

146 E_vc[n] += (kpt.eps_n[cs] - kpt.eps_n[n]) * Hartree 

147 

148 # Obtain appropriate part of mom and g arrays 

149 pols = [d_i, d_o] 

150 mom_dnn = np.ascontiguousarray(self.mom_skvnm[kpt.s, kpt.k, pols]) 

151 assert mom_dnn.shape[0] == 2 

152 g_lnn = np.ascontiguousarray(self.g_sqklnn[kpt.s, 0, kpt.k]) 

153 

154 # Raman contribution of this k-point 

155 raman_l += self.resonant_term( 

156 f_vc, E_vc, mom_dnn, g_lnn, nc, nv, w_in, self.w_ph 

157 ) * kpt.weight 

158 # Collect parallel contributions 

159 self.kd.comm.sum(raman_l) 

160 

161 if self.kd.comm.rank == 0: 

162 print(f"Raman intensities per mode in {'xyz'[d_i]}{'xyz'[d_o]}") 

163 print("--------------------------") 

164 ff = " Phonon {} with energy = {:4.2f} rcm: {:.4f}" 

165 for l in range(self.w_ph.shape[0]): 

166 print(ff.format(l, self.w_ph[l] / invcm, raman_l[l])) 

167 

168 return raman_l 

169 

170 def calculate_raman_tensor(self, w_in, gamma_l=0.1): 

171 """Calculate whole Raman tensor 

172 

173 Parameters 

174 ---------- 

175 w_in: float 

176 Laser frequency in eV 

177 gamma_l: float 

178 Line broadening in eV, (0.1eV=806rcm) 

179 """ 

180 # If exist already, don't recompute 

181 for i in range(3): 

182 for j in range(3): 

183 with self.raman_cache.lock(f"{'xyz'[i]}{'xyz'[j]}") as handle: 

184 if handle is None: 

185 continue 

186 R_l = self.calculate(w_in, i, j, gamma_l) 

187 self.kd.comm.barrier() 

188 handle.save(R_l) 

189 

190 def nm_to_eV(self, laser_wave_length): 

191 return 1.239841e3 / laser_wave_length