Coverage for gpaw/tddft/spectrum.py: 92%

192 statements  

« prev     ^ index     » next       coverage.py v7.7.1, created at 2025-07-12 00:18 +0000

1import re 

2import numpy as np 

3 

4from gpaw import __version__ as version 

5from gpaw.mpi import world 

6from gpaw.tddft.units import au_to_as, au_to_fs, au_to_eV, rot_au_to_cgs 

7from gpaw.tddft.folding import FoldedFrequencies 

8from gpaw.tddft.folding import Folding 

9 

10 

11def calculate_fourier_transform(x_t, y_ti, foldedfrequencies, velocity=False): 

12 ff = foldedfrequencies 

13 X_w = ff.frequencies 

14 envelope = ff.folding.envelope 

15 

16 # Construct integration weights: 

17 # We use trapezoidal rule except the end point is accounted with 

18 # full weight. This ensures better numerical compatibility 

19 # with the on-the-fly Fourier integrators. 

20 # This shouldn't affect in usual scenarios as the envelope 

21 # should damp the data to zero at the end point in any case. 

22 dx_t1 = x_t[1:] - x_t[:-1] 

23 dx_t = 0.5 * (np.insert(dx_t1, 0, 0.0) + np.append(dx_t1, dx_t1[-1])) 

24 

25 env_t = envelope(x_t) 

26 Ienv = np.sum(dx_t * env_t) 

27 

28 if velocity: 

29 y_ti -= np.sum((dx_t * env_t)[:, None] * y_ti, axis=0) / Ienv 

30 

31 # Integrate 

32 f_wt = np.exp(1.0j * np.outer(X_w, x_t)) 

33 y_it = np.swapaxes(y_ti, 0, 1) 

34 

35 Y_wi = np.tensordot(f_wt, dx_t * env_t * y_it, axes=(1, 1)) 

36 print('Sinc contamination', env_t[-1]) 

37 return Y_wi 

38 

39 

40def read_td_file_data(fname, remove_duplicates=True): 

41 """Read data from time-dependent data file. 

42 

43 Parameters 

44 ---------- 

45 fname 

46 File path 

47 remove_duplicates 

48 If true, remove data from overlapping time values. 

49 The first encountered values are kept. 

50 

51 Returns 

52 ------- 

53 time_t 

54 Array of time values 

55 data_ti 

56 Array of data values 

57 """ 

58 # Read data 

59 data_tj = np.loadtxt(fname) 

60 time_t = data_tj[:, 0] 

61 data_ti = data_tj[:, 1:] 

62 

63 # Remove duplicates due to abruptly stopped and restarted calculation 

64 if remove_duplicates: 

65 flt_t = np.ones_like(time_t, dtype=bool) 

66 maxtime = time_t[0] 

67 for t in range(1, len(time_t)): 

68 # Note about ">=" here: 

69 # The equality is included here in order to 

70 # retain step-like data (for example, the data both just before 

71 # and just after the kick is kept). 

72 if time_t[t] >= maxtime: 

73 maxtime = time_t[t] 

74 else: 

75 flt_t[t] = False 

76 time_t = time_t[flt_t] 

77 data_ti = data_ti[flt_t] 

78 ndup = len(flt_t) - flt_t.sum() 

79 if ndup > 0: 

80 print('Removed %d duplicates' % ndup) 

81 return time_t, data_ti 

82 

83 

84def read_td_file_kicks(fname): 

85 """Read kicks from time-dependent data file. 

86 

87 Parameters 

88 ---------- 

89 fname 

90 File path 

91 

92 Returns 

93 ------- 

94 kick_i 

95 List of kicks. 

96 Each kick is a dictionary with keys 

97 ``strength_v`` and ``time``. 

98 """ 

99 def parse_kick_line(line): 

100 # Kick 

101 regexp = (r"Kick = \[" 

102 r"(?P<k0>[-+0-9\.e\ ]+), " 

103 r"(?P<k1>[-+0-9\.e\ ]+), " 

104 r"(?P<k2>[-+0-9\.e\ ]+)\]") 

105 m = re.search(regexp, line) 

106 assert m is not None, 'Kick not found' 

107 kick_v = np.array([float(m.group('k%d' % v)) for v in range(3)]) 

108 # Time 

109 regexp = r"Time = (?P<time>[-+0-9\.e\ ]+)" 

110 m = re.search(regexp, line) 

111 if m is None: 

112 print('time not found') 

113 time = 0.0 

114 else: 

115 time = float(m.group('time')) 

116 velocity = 'velocity' in line 

117 return kick_v, time, velocity 

118 

119 # Search kicks 

120 kick_i = [] 

121 with open(fname) as f: 

122 for line in f: 

123 if line.startswith('# Kick'): 

124 kick_v, time, velocity = parse_kick_line(line) 

125 kick_i.append({'strength_v': kick_v, 'time': time, 

126 'velocity': velocity}) 

127 return kick_i 

128 

129 

130def clean_td_data(kick_i, time_t, data_ti): 

131 """Prune time-dependent data. 

132 

133 This function checks that there is only one kick 

134 in the kick list and moves time zero to the kick 

135 time (discarding all preceding data). 

136 

137 Parameters 

138 ---------- 

139 kick_i 

140 List of kicks. 

141 time_t 

142 Array of time values 

143 data_ti 

144 Array of data values 

145 

146 Returns 

147 ------- 

148 kick_i 

149 List of kicks. 

150 Each kick is a dictionary with keys 

151 ``strength_v`` and ``time``. 

152 

153 Raises 

154 ------ 

155 RuntimeError 

156 If kick list contains multiple kicks. 

157 """ 

158 # Check kicks 

159 if len(kick_i) > 1: 

160 raise RuntimeError('Multiple kicks') 

161 kick = kick_i[0] 

162 kick_v = kick['strength_v'] 

163 velocity = kick['velocity'] 

164 kick_time = kick['time'] 

165 

166 # Discard times before kick 

167 flt_t = time_t >= kick_time 

168 time_t = time_t[flt_t] 

169 data_ti = data_ti[flt_t] 

170 

171 # Move time zero to kick time 

172 time_t -= kick_time 

173 assert time_t[0] == 0.0 

174 

175 return kick_v, velocity, time_t, data_ti 

176 

177 

178def read_dipole_moment_file(fname, remove_duplicates=True): 

179 """Read time-dependent dipole moment data file. 

180 

181 Parameters 

182 ---------- 

183 fname 

184 File path 

185 remove_duplicates 

186 If true, remove data from overlapping time values. 

187 The first encountered values are kept. 

188 

189 Returns 

190 ------- 

191 kick_i 

192 List of kicks. 

193 Each kick is a dictionary with keys 

194 ``strength_v`` and ``time``. 

195 time_t 

196 Array of time values 

197 norm_t 

198 Array of norm values 

199 dm_tv 

200 Array of dipole moment values 

201 """ 

202 time_t, data_ti = read_td_file_data(fname, remove_duplicates) 

203 kick_i = read_td_file_kicks(fname) 

204 norm_t = data_ti[:, 0] 

205 dm_tv = data_ti[:, 1:] 

206 return kick_i, time_t, norm_t, dm_tv 

207 

208 

209def calculate_polarizability(kick_v, time_t, dm_tv, 

210 foldedfrequencies, velocity=False): 

211 if not velocity: 

212 dm_tv = dm_tv - dm_tv[0] 

213 

214 alpha_wv = calculate_fourier_transform(time_t, dm_tv, foldedfrequencies, 

215 velocity=velocity) 

216 

217 kick_magnitude = np.sqrt(np.sum(kick_v**2)) 

218 alpha_wv /= kick_magnitude 

219 return alpha_wv 

220 

221 

222def calculate_photoabsorption(kick_v, time_t, dm_tv, 

223 foldedfrequencies, velocity=False): 

224 omega_w = foldedfrequencies.frequencies 

225 alpha_wv = calculate_polarizability(kick_v, time_t, dm_tv, 

226 foldedfrequencies, 

227 velocity=velocity) 

228 if velocity: 

229 abs_wv = 2 / np.pi * alpha_wv.real 

230 else: 

231 abs_wv = 2 / np.pi * omega_w[:, np.newaxis] * alpha_wv.imag 

232 

233 kick_magnitude = np.sqrt(np.sum(kick_v**2)) 

234 abs_wv *= kick_v / kick_magnitude 

235 return abs_wv 

236 

237 

238def read_magnetic_moment_file(fname, remove_duplicates=True): 

239 """Read time-dependent magnetic moment data file. 

240 

241 Parameters 

242 ---------- 

243 fname 

244 File path 

245 remove_duplicates 

246 If true, remove data from overlapping time values. 

247 The first encountered values are kept. 

248 

249 Returns 

250 ------- 

251 kick_i 

252 List of kicks. 

253 Each kick is a dictionary with keys 

254 ``strength_v`` and ``time``. 

255 time_t 

256 Array of time values 

257 mm_tv 

258 Array of magnetic moment values 

259 """ 

260 time_t, mm_tv = read_td_file_data(fname, remove_duplicates) 

261 kick_i = read_td_file_kicks(fname) 

262 return kick_i, time_t, mm_tv 

263 

264 

265def calculate_rotatory_strength_components(kick_v, time_t, mm_tv, 

266 foldedfrequencies): 

267 assert np.all(mm_tv[0] == 0.0) 

268 mm_wv = calculate_fourier_transform(time_t, mm_tv, foldedfrequencies) 

269 kick_magnitude = np.sqrt(np.sum(kick_v**2)) 

270 rot_wv = mm_wv.real / (np.pi * kick_magnitude) 

271 return rot_wv 

272 

273 

274def write_spectrum(dipole_moment_file, spectrum_file, 

275 folding, width, e_min, e_max, delta_e, 

276 title, symbol, calculate): 

277 def str_list(v_i, fmt='%g'): 

278 return '[%s]' % ', '.join(map(lambda v: fmt % v, v_i)) 

279 

280 kick_i, time_t, _, dm_tv = read_dipole_moment_file(dipole_moment_file) 

281 kick_v, velocity, time_t, dm_tv = clean_td_data(kick_i, time_t, dm_tv) 

282 dt_t = time_t[1:] - time_t[:-1] 

283 

284 freqs = np.arange(e_min, e_max + 0.5 * delta_e, delta_e) 

285 folding = Folding(folding, width) 

286 ff = FoldedFrequencies(freqs, folding) 

287 omega_w = ff.frequencies 

288 spec_wv = calculate(kick_v, time_t, dm_tv, ff, velocity=velocity) 

289 

290 # Write spectrum file header 

291 with open(spectrum_file, 'w') as f: 

292 def w(s): 

293 f.write('%s\n' % s) 

294 

295 w('# %s spectrum from real-time propagation' % title) 

296 w('# GPAW version: %s' % version) 

297 w('# Total time = %.4f fs, Time steps = %s as' % 

298 (dt_t.sum() * au_to_fs, 

299 str_list(np.unique(np.around(dt_t, 6)) * au_to_as, '%.4f'))) 

300 w('# Kick = %s' % str_list(kick_v)) 

301 w('# %sian folding, Width = %.4f eV = %lf Hartree' 

302 ' <=> FWHM = %lf eV' % 

303 (folding.folding, folding.width * au_to_eV, folding.width, 

304 folding.fwhm * au_to_eV)) 

305 

306 col_i = [] 

307 data_iw = [omega_w * au_to_eV] 

308 for v in range(len(kick_v)): 

309 h = '{}_{}'.format(symbol, 'xyz'[v]) 

310 if spec_wv.dtype == complex: 

311 col_i.append('Re[%s]' % h) 

312 data_iw.append(spec_wv[:, v].real) 

313 col_i.append('Im[%s]' % h) 

314 data_iw.append(spec_wv[:, v].imag) 

315 else: 

316 col_i.append(h) 

317 data_iw.append(spec_wv[:, v]) 

318 

319 w('# %10s %s' % ('om (eV)', ' '.join(['%20s' % s for s in col_i]))) 

320 

321 # Write spectrum file data 

322 with open(spectrum_file, 'ab') as f: 

323 np.savetxt(f, np.array(data_iw).T, 

324 fmt='%12.6lf' + (' %20.10le' * len(col_i))) 

325 

326 return folding.envelope(time_t[-1]) 

327 

328 

329def photoabsorption_spectrum(dipole_moment_file: str, 

330 spectrum_file: str, 

331 folding: str = 'Gauss', 

332 width: float = 0.2123, 

333 e_min: float = 0.0, 

334 e_max: float = 30.0, 

335 delta_e: float = 0.05): 

336 """Calculates photoabsorption spectrum from the time-dependent 

337 dipole moment. 

338 

339 The spectrum is represented as a dipole strength function 

340 in units of 1/eV. Thus, the resulting spectrum should integrate 

341 to the number of valence electrons in the system. 

342 

343 Parameters 

344 ---------- 

345 dipole_moment_file 

346 Name of the time-dependent dipole moment file from which 

347 the spectrum is calculated 

348 spectrum_file 

349 Name of the spectrum file 

350 folding 

351 Gaussian (``'Gauss'``) or Lorentzian (``'Lorentz'``) folding 

352 width 

353 Width of the Gaussian (sigma) or Lorentzian (Gamma) 

354 Gaussian = 1/(sigma sqrt(2pi)) exp(-(1/2)(omega/sigma)^2) 

355 Lorentzian = (1/pi) (1/2) Gamma / [omega^2 + ((1/2) Gamma)^2] 

356 e_min 

357 Minimum energy shown in the spectrum (eV) 

358 e_max 

359 Maximum energy shown in the spectrum (eV) 

360 delta_e 

361 Energy resolution (eV) 

362 """ 

363 if world.rank == 0: 

364 print('Calculating photoabsorption spectrum from file "%s"' 

365 % dipole_moment_file) 

366 

367 def calculate(*args, **kwargs): 

368 return (calculate_photoabsorption(*args, **kwargs) 

369 / au_to_eV) 

370 sinc = write_spectrum(dipole_moment_file, spectrum_file, 

371 folding, width, e_min, e_max, delta_e, 

372 'Photoabsorption', 'S', calculate) 

373 print('Sinc contamination %.8f' % sinc) 

374 print('Calculated photoabsorption spectrum saved to file "%s"' 

375 % spectrum_file) 

376 

377 

378def polarizability_spectrum(dipole_moment_file, spectrum_file, 

379 folding='Gauss', width=0.2123, 

380 e_min=0.0, e_max=30.0, delta_e=0.05): 

381 """Calculates polarizability spectrum from the time-dependent 

382 dipole moment. 

383 

384 Parameters: 

385 

386 dipole_moment_file: string 

387 Name of the time-dependent dipole moment file from which 

388 the spectrum is calculated 

389 spectrum_file: string 

390 Name of the spectrum file 

391 folding: 'Gauss' or 'Lorentz' 

392 Whether to use Gaussian or Lorentzian folding 

393 width: float 

394 Width of the Gaussian (sigma) or Lorentzian (Gamma) 

395 Gaussian = 1/(sigma sqrt(2pi)) exp(-(1/2)(omega/sigma)^2) 

396 Lorentzian = (1/pi) (1/2) Gamma / [omega^2 + ((1/2) Gamma)^2] 

397 e_min: float 

398 Minimum energy shown in the spectrum (eV) 

399 e_max: float 

400 Maximum energy shown in the spectrum (eV) 

401 delta_e: float 

402 Energy resolution (eV) 

403 """ 

404 if world.rank == 0: 

405 print('Calculating polarizability spectrum from file "%s"' 

406 % dipole_moment_file) 

407 

408 def calculate(*args, **kwargs): 

409 return calculate_polarizability(*args, **kwargs) / au_to_eV**2 

410 sinc = write_spectrum(dipole_moment_file, spectrum_file, 

411 folding, width, e_min, e_max, delta_e, 

412 'Polarizability', 'alpha', calculate) 

413 print('Sinc contamination %.8f' % sinc) 

414 print('Calculated polarizability spectrum saved to file "%s"' 

415 % spectrum_file) 

416 

417 

418def rotatory_strength_spectrum(magnetic_moment_files, spectrum_file, 

419 folding='Gauss', width=0.2123, 

420 e_min=0.0, e_max=30.0, delta_e=0.05): 

421 """Calculates rotatory strength spectrum from the time-dependent 

422 magnetic moment. 

423 

424 Parameters 

425 ---------- 

426 magnetic_moment_files: list of string 

427 Time-dependent magnetic moment files for x, y, and z kicks 

428 spectrum_file: string 

429 Name of the spectrum file 

430 folding: 'Gauss' or 'Lorentz' 

431 Whether to use Gaussian or Lorentzian folding 

432 width: float 

433 Width of the Gaussian (sigma) or Lorentzian (Gamma) 

434 Gaussian = 1/(sigma sqrt(2pi)) exp(-(1/2)(omega/sigma)^2) 

435 Lorentzian = (1/pi) (1/2) Gamma / [omega^2 + ((1/2) Gamma)^2] 

436 e_min: float 

437 Minimum energy shown in the spectrum (eV) 

438 e_max: float 

439 Maximum energy shown in the spectrum (eV) 

440 delta_e: float 

441 Energy resolution (eV) 

442 """ 

443 if world.rank != 0: 

444 return 

445 

446 freqs = np.arange(e_min, e_max + 0.5 * delta_e, delta_e) 

447 folding = Folding(folding, width) 

448 ff = FoldedFrequencies(freqs, folding) 

449 omega_w = ff.frequencies * au_to_eV 

450 rot_w = np.zeros_like(omega_w) 

451 

452 tot_time = np.inf 

453 time_steps = [] 

454 kick_strength = None 

455 for v, fpath in enumerate(magnetic_moment_files): 

456 kick_i, time_t, mm_tv = read_magnetic_moment_file(fpath) 

457 kick_v, velocity, time_t, mm_tv = clean_td_data(kick_i, time_t, mm_tv) 

458 

459 tot_time = min(tot_time, time_t[-1]) 

460 time_steps.append(np.around(time_t[1:] - time_t[:-1], 6)) 

461 

462 # Check kicks 

463 for v0 in range(3): 

464 if v0 == v: 

465 continue 

466 if kick_v[v0] != 0.0: 

467 raise RuntimeError('The magnetic moment files must be ' 

468 'for kicks in x, y, and z directions.') 

469 if kick_strength is None: 

470 kick_strength = np.sqrt(np.sum(kick_v**2)) 

471 if np.sqrt(np.sum(kick_v**2)) != kick_strength: 

472 raise RuntimeError('The magnetic moment files must have ' 

473 'been calculated with the same kick strength.') 

474 

475 rot_wv = calculate_rotatory_strength_components(kick_v, time_t, 

476 mm_tv, ff) 

477 rot_w += rot_wv[:, v] 

478 

479 rot_w *= rot_au_to_cgs * 1e40 / au_to_eV 

480 

481 # Unique non-zero time steps 

482 time_steps = np.unique(time_steps) 

483 time_steps = time_steps[time_steps != 0] 

484 

485 with open(spectrum_file, 'w') as fd: 

486 steps_str = ', '.join(f'{val:.4f}' for val in time_steps * au_to_as) 

487 lines = ['Rotatory strength spectrum from real-time propagations', 

488 f'Total time = {tot_time * au_to_fs:.4f} fs, ' 

489 f'Time steps = [{steps_str}] as', 

490 f'Kick strength = {kick_strength}', 

491 f'{folding.folding}ian folding, ' 

492 f'width = {folding.width * au_to_eV:.4f} eV ' 

493 f'<=> FWHM = {folding.fwhm * au_to_eV:.4f} eV'] 

494 fd.write('# ' + '\n# '.join(lines) + '\n') 

495 fd.write(f'# {"Energy (eV)":>12} {"R (1e-40 cgs / eV)":>20}\n') 

496 

497 data_wi = np.vstack((omega_w, rot_w)).T 

498 np.savetxt(fd, data_wi, 

499 fmt='%14.6lf' + (' %20.10le' * (data_wi.shape[1] - 1)))