Coverage for rhodent/writers/spectrum.py: 100%

1from __future__ import annotations 

2 

3 

4import numpy as np 

5from numpy.typing import NDArray 

6 

7from .. import __version__ 

8from ..perturbation import create_perturbation, PerturbationLike 

9 

10 

11def write_spectrum(out_fname: str, 

12 frequencies: list[float] | NDArray[np.float64], 

13 spectrum: NDArray[np.float64], 

14 frequency_broadening: float, 

15 total_time: float, 

16 timestep: float, 

17 perturbation: PerturbationLike): 

18 """ Write dipole strength function (spectrum) to text file. 

19 

20 Parameters 

21 ---------- 

22 out_fname 

23 File name of the written file. 

24 frequencies 

25 Array of frequencies in units of eV. 

26 spectrum 

27 Spectrum corresponding to the :attr:`frequencies` in units of 1/eV. 

28 frequency_broadening 

29 Gaussian broadening width in units of eV. Default (0) is no broadening. 

30 total_time 

31 Total simulation time in units of as. 

32 timestep 

33 Timestep in units of as. 

34 perturbation 

35 The perturbation that was applied in the :term:`TDDFT` calculation. 

36 """ 

37 frequencies = np.array(frequencies) 

38 osc_wv = np.array(spectrum) 

39 perturbation = create_perturbation(perturbation) 

40 

41 if frequency_broadening == 0: 

42 broadening = 'No broadening' 

43 else: 

44 broadening = f'Gaussian broadening width {frequency_broadening:.2f} eV' 

45 perturbationstr = '\n'.join(['Perturbation during calculation was:'] + 

46 [' ' + line for line in str(perturbation).split('\n')]) 

47 

48 header = ('Photoabsorption spectrum from real-time propagation\n' 

49 f'Calculated using rhodent version: {__version__}\n' 

50 f'Total time = {total_time*1e-3:.4f} fs, Time steps = {timestep:.2f} as\n' 

51 f'{perturbationstr}\n' 

52 f'{broadening}\n' 

53 f'{"om (eV)":>10} {"S_x (1/eV)":>20} {"S_y (1/eV)":>20} {"S_z (1/eV)":>20}' 

54 ) 

55 data_wi = np.zeros((len(frequencies), 1 + 3)) 

56 data_wi[:, 0] = frequencies 

57 data_wi[:, 1:] = osc_wv 

58 fmt = '%12.6f' + (' %20.10le' * 3) 

59 np.savetxt(str(out_fname), data_wi, header=header, fmt=fmt)