Coverage for gpaw/tddft/folding.py: 82%

1import numpy as np 

2 

3from gpaw.tddft.units import au_to_eV 

4from gpaw.tddft.units import eV_to_au 

5 

6 

7def frequencies(frequencies, foldings, widths, units='eV'): 

8 f_w = [] 

9 for folding in np.array([foldings]).ravel(): 

10 for width in np.array([widths]).ravel(): 

11 folding = Folding(folding, width, units) 

12 f_w.append(FoldedFrequencies(frequencies, folding, units)) 

13 return f_w 

14 

15 

16def convert_to_au(val, units='eV'): 

17 if units == 'eV': 

18 return val * eV_to_au 

19 elif units == 'au': 

20 return val 

21 raise RuntimeError('Unknown units: %s' % units) 

22 

23 

24class FoldedFrequencies: 

25 def __init__(self, frequencies, folding, units='eV'): 

26 freqs = np.array([frequencies], dtype=float).ravel() 

27 self.frequencies = convert_to_au(freqs, units=units) 

28 if isinstance(folding, dict): 

29 self.folding = Folding(**folding) 

30 else: 

31 self.folding = folding 

32 

33 def todict(self): 

34 d = dict(units='au') 

35 for arg in ['frequencies', 'folding']: 

36 val = getattr(self, arg) 

37 if hasattr(val, 'todict'): 

38 val = val.todict() 

39 d[arg] = val 

40 return d 

41 

42 def __repr__(self): 

43 s = 'With %s: ' % self.folding 

44 s += ', '.join(['%.5f' % (f * au_to_eV) for f in self.frequencies]) 

45 s += ' eV' 

46 return s 

47 

48 

49class Frequency: 

50 def __init__(self, freq, folding, units='eV'): 

51 ff = FoldedFrequencies(freq, folding, units) 

52 self.folding = ff.folding 

53 self.freq = ff.frequencies[0] 

54 

55 

56class Folding: 

57 def __init__(self, folding, width, units='eV'): 

58 if width is None: 

59 folding = None 

60 

61 self.folding = folding 

62 if self.folding is None: 

63 self.width = None 

64 else: 

65 self.width = convert_to_au(float(width), units=units) 

66 

67 if self.folding not in [None, 'Gauss', 'Lorentz']: 

68 raise RuntimeError('Unknown folding: %s' % self.folding) 

69 

70 if self.folding is None: 

71 self.fwhm = 0.0 

72 elif self.folding == 'Gauss': 

73 self.fwhm = self.width * (2. * np.sqrt(2. * np.log(2.0))) 

74 elif self.folding == 'Lorentz': 

75 self.fwhm = 2. * self.width 

76 

77 def envelope(self, time): 

78 if self.folding is None: 

79 return 0 * time + 1 

80 elif self.folding == 'Gauss': 

81 return np.exp(- 0.5 * self.width**2 * time**2) 

82 elif self.folding == 'Lorentz': 

83 return np.exp(- self.width * time) 

84 

85 def todict(self): 

86 d = dict(units='au') 

87 for arg in ['folding', 'width']: 

88 d[arg] = getattr(self, arg) 

89 return d 

90 

91 def __repr__(self): 

92 if self.folding is None: 

93 return 'No folding' 

94 return f'{self.folding}({self.width * au_to_eV:.5f} eV)'