Coverage for gpaw/lcaotddft/laser.py: 60%

1from __future__ import annotations 

2import numpy as np 

3 

4from typing import Type 

5from gpaw.mpi import world 

6from gpaw.tddft.units import as_to_au, eV_to_au 

7 

8 

9known_lasers: dict[str, Type[Laser]] = dict() 

10 

11 

12def create_laser(name, **kwargs): 

13 """ Create Laser from dict """ 

14 

15 if isinstance(name, Laser): 

16 return name 

17 elif isinstance(name, dict): 

18 kwargs.update(name) 

19 return create_laser(**kwargs) 

20 

21 if not known_lasers: 

22 _register_known_lasers() 

23 

24 cls = known_lasers.get(name, None) 

25 if cls is None: 

26 raise ValueError(f'Unknown laser: {name}') 

27 return cls(**kwargs) 

28 

29 

30def register_custom_laser(name: str, 

31 cls: Type[Laser]): 

32 """ Register a custom laser object 

33 

34 This function must be used when restarting TDDFT calculations using 

35 user defined laser classes 

36 

37 Parameters 

38 ---------- 

39 name 

40 Name of laser object. Must be consistent with the name in todict() 

41 cls 

42 Class of the laser object 

43 """ 

44 if not known_lasers: 

45 _register_known_lasers() 

46 

47 known_lasers[name] = cls 

48 

49 

50def _register_known_lasers(): 

51 for name in ['GaussianPulse', 'SincPulse', 'SumLaser']: 

52 known_lasers[name] = globals()[name] 

53 

54 

55class Laser: 

56 def __init__(self): 

57 pass 

58 

59 def strength(self, time): 

60 return 0.0 

61 

62 def derivative(self, time): 

63 return 0.0 

64 

65 def fourier(self, omega): 

66 return 0.0 

67 

68 def write(self, fname, time_t): 

69 """ 

70 Write the values of the pulse to a file. 

71 

72 Parameters 

73 ---------- 

74 fname 

75 filename 

76 time_t 

77 times in attoseconds 

78 """ 

79 if world.rank != 0: 

80 return 

81 time_t = time_t * as_to_au 

82 strength_t = self.strength(time_t) 

83 derivative_t = self.derivative(time_t) 

84 fmt = '%12.6f %20.10e %20.10e' 

85 header = '{:^10} {:^20} {:^20}'.format('time', 'strength', 

86 'derivative') 

87 np.savetxt(fname, np.stack((time_t, strength_t, derivative_t)).T, 

88 fmt=fmt, header=header) 

89 

90 

91class SumLaser(Laser): 

92 def __init__(self, *lasers): 

93 self.laser_i = [] 

94 dict_i = [] 

95 for laser in lasers: 

96 laser = create_laser(laser) 

97 self.laser_i.append(laser) 

98 dict_i.append(laser.todict()) 

99 self.dict = dict(name='SumLaser', 

100 lasers=dict_i) 

101 

102 def strength(self, time): 

103 s = 0.0 

104 for laser in self.laser_i: 

105 s += laser.strength(time) 

106 return s 

107 

108 def fourier(self, omega): 

109 s = 0.0 

110 for laser in self.laser_i: 

111 s += laser.fourier(omega) 

112 return s 

113 

114 def todict(self): 

115 return self.dict 

116 

117 

118class GaussianPulse(Laser): 

119 r""" 

120 Laser pulse with Gaussian envelope: 

121 

122 .. math:: 

123 

124 g(t) = s_0 \sin(\omega_0 (t - t_0)) \exp(-\sigma^2 (t - t_0)^2 / 2) 

125 

126 

127 Parameters 

128 ---------- 

129 strength: float 

130 value of :math:`s_0` in atomic units 

131 time0: float 

132 value of :math:`t_0` in attoseconds 

133 frequency: float 

134 value of :math:`\omega_0` in eV 

135 sigma: float 

136 value of :math:`\sigma` in eV 

137 sincos: 'sin' or 'cos' 

138 use sin or cos function 

139 stoptime: float 

140 pulse is set to zero after this value (in attoseconds) 

141 """ 

142 

143 def __init__(self, strength, time0, frequency, sigma, sincos='sin', 

144 stoptime=np.inf): 

145 self.dict = dict(name='GaussianPulse', 

146 strength=strength, 

147 time0=time0, 

148 frequency=frequency, 

149 sigma=sigma, 

150 sincos=sincos) 

151 self.s0 = strength 

152 self.t0 = time0 * as_to_au 

153 self.omega0 = frequency * eV_to_au 

154 self.sigma = sigma * eV_to_au 

155 self.stoptime = stoptime * as_to_au 

156 assert sincos in ['sin', 'cos'] 

157 self.sincos = sincos 

158 

159 def strength(self, t): 

160 """ 

161 Return the value of the pulse :math:`g(t)`. 

162 

163 Parameters 

164 ---------- 

165 t 

166 time in atomic units 

167 

168 Returns 

169 ------- 

170 The value of the pulse. 

171 """ 

172 s = self.s0 * np.exp(-0.5 * self.sigma**2 * (t - self.t0)**2) 

173 if self.sincos == 'sin': 

174 s *= np.sin(self.omega0 * (t - self.t0)) 

175 else: 

176 s *= np.cos(self.omega0 * (t - self.t0)) 

177 flt = t < self.stoptime 

178 

179 return s * flt 

180 

181 def derivative(self, t): 

182 """ 

183 Return the derivative of the pulse :math:`g'(t)`. 

184 

185 Parameters 

186 ---------- 

187 t 

188 time in atomic units 

189 

190 Returns 

191 ------- 

192 The derivative of the pulse. 

193 """ 

194 dt = t - self.t0 

195 s = self.s0 * np.exp(-0.5 * self.sigma**2 * dt**2) 

196 if self.sincos == 'sin': 

197 s *= (-self.sigma**2 * dt * np.sin(self.omega0 * dt) + 

198 self.omega0 * np.cos(self.omega0 * dt)) 

199 else: 

200 s *= (-self.sigma**2 * dt * np.cos(self.omega0 * dt) + 

201 -self.omega0 * np.sin(self.omega0 * dt)) 

202 return s 

203 

204 def fourier(self, omega): 

205 r""" 

206 Return Fourier transform of the pulse :math:`g(\omega)`. 

207 

208 Parameters 

209 ---------- 

210 omega 

211 frequency in atomic units 

212 

213 Returns 

214 ------- 

215 Fourier transform of the pulse. 

216 """ 

217 s = (self.s0 * np.sqrt(np.pi / 2) / self.sigma * 

218 np.exp(-0.5 * (omega - self.omega0)**2 / self.sigma**2) * 

219 np.exp(1.0j * self.t0 * omega)) 

220 if self.sincos == 'sin': 

221 s *= 1.0j 

222 return s 

223 

224 def todict(self): 

225 return self.dict 

226 

227 

228class SincPulse(Laser): 

229 r""" 

230 Laser pulse with sinc envelope: 

231 

232 .. math:: 

233 

234 g(t) = s_0 \frac{\sin(\pi \omega_{cut} (t - t_0))} 

235 {\pi \omega_{cut} (t - t_0)} 

236 

237 

238 Parameters 

239 ---------- 

240 strength: float 

241 value of :math:`s_0` in atomic units 

242 time0: float 

243 value of :math:`t_0` in attoseconds, or in units of 

244 :math:`\omega_{cut} / 2` if relative_t0 is True 

245 cutoff_freq: float 

246 Cutoff frequency: value of :math:`\omega_{cut}` in eV 

247 relative_t0: 

248 Specify time0 in units relative to the cutoff frequency 

249 """ 

250 

251 def __init__(self, 

252 strength: float, 

253 time0: float, 

254 cutoff_freq: float, 

255 relative_t0: bool): 

256 self.dict = dict(name='SincPulse', 

257 strength=strength, 

258 time0=time0, 

259 cutoff_freq=cutoff_freq, 

260 relative_t0=relative_t0) 

261 self.s0 = strength 

262 self.omega_cut = (cutoff_freq / np.pi) * eV_to_au 

263 if relative_t0: 

264 self.t0 = 2 * time0 / self.omega_cut 

265 else: 

266 self.t0 = time0 * as_to_au 

267 

268 def strength(self, t): 

269 """ 

270 Return the value of the pulse :math:`g(t)`. 

271 

272 Parameters 

273 ---------- 

274 t 

275 time in atomic units 

276 

277 Returns 

278 ------- 

279 The value of the pulse. 

280 """ 

281 s = self.s0 * np.sinc(self.omega_cut * (t - self.t0)) 

282 

283 return s 

284 

285 def derivative(self, t): 

286 raise NotImplementedError 

287 

288 def fourier(self, omega): 

289 raise NotImplementedError 

290 

291 def todict(self): 

292 return self.dict