Coverage for strongcoca/response/time_domain.py: 100%

1import logging 

2 

3import numpy as np 

4 

5from .base import BaseResponse 

6from ..types import Array 

7from ..utilities import ClassFormatter 

8from .utilities import fourier_transform 

9from ..units import au_to_eA, eA_to_au, au_to_fs, fs_to_au 

10from .utilities import Broadening 

11 

12 

13logger = logging.getLogger(__name__) 

14 

15 

16class TimeDomainResponse(BaseResponse): 

17 """Objects of this class hold different representations of the response 

18 functions for a time domain system. 

19 

20 Parameters 

21 ---------- 

22 times 

23 Time grid of {t} equally spaced time values 

24 units are fs by default, optionally atomic units (see :attr:`units`). 

25 dipole_moments 

26 Dipole moment tensors at the time values, shape {t}x{3}x{3} 

27 units are eÅ by default, optionally atomic units (see :attr:`units`). 

28 broadening 

29 artificial broadening applied to the response 

30 units 

31 `fseA` to specify times in fs and dipole_moments in eÅ 

32 or `au` to specify inputs in atomic units. 

33 

34 This parameter determines whether conversion should be performed during 

35 initialization and has no effect on instance methods and variables. 

36 name 

37 Name of response. 

38 

39 Examples 

40 -------- 

41 

42 The following snippet illustrates how a time domain response object can be 

43 constructed. Here, artificial dipole moments are used. In practice dipole 

44 moments would usually be imported from, e.g., TDDFT calculations: 

45 

46 >>> import numpy as np 

47 >>> from strongcoca.response import TimeDomainResponse 

48 >>> from strongcoca.response.utilities import LorentzianBroadening 

49 >>> 

50 >>> t = np.linspace(0, 10) 

51 >>> dm = np.ones((t.shape[0], 3, 3)) 

52 >>> broadening = LorentzianBroadening(0.1) 

53 >>> response = TimeDomainResponse(t, dm, broadening, name='Eats') 

54 >>> print(response) 

55 -------------------- TimeDomainResponse -------------------- 

56 name : Eats 

57 times (fs) : [ 0. 0.20408 0.40816 ... 

58 9.59184 9.79592 10. ] 

59 dipole_moments (eÅ) : [[[1. 1. 1.] 

60 [1. 1. 1.] 

61 [1. 1. 1.]] 

62 ... 

63 """ 

64 def __init__(self, 

65 times: np.ndarray, 

66 dipole_moments: np.ndarray, 

67 broadening: Broadening, 

68 units: str = 'fseA', 

69 name: str = 'TimeDomain') -> None: 

70 logger.debug(f'Entering {self.__class__.__name__}.__init__') 

71 super().__init__(broadening=broadening, pbc=False, name=name) 

72 if times.ndim != 1: 

73 raise ValueError('times array must be one dimensional') 

74 if dipole_moments.shape != (times.shape[0], 3, 3): 

75 raise ValueError('dipole_moments array has the wrong shape') 

76 

77 times = times.copy() 

78 dipole_moments = dipole_moments.copy() 

79 

80 if units == 'fseA': 

81 times = times * fs_to_au # TODO Possibly simplify to *= fs_to_au after resolving #57 

82 dipole_moments = dipole_moments * eA_to_au 

83 elif units != 'au': 

84 raise ValueError(f"units has to be 'fseA' or 'au', not '{units}'") 

85 

86 self._time_t = times 

87 dt_t = self._time_t[1:] - self._time_t[:-1] 

88 self._dt = dt_t[0] 

89 if not np.allclose(dt_t, self._dt): 

90 raise ValueError('times array need to be equally spaced') 

91 self._dm_tvv = dipole_moments 

92 

93 def __str__(self) -> str: 

94 fmt = ClassFormatter(self, pad=20) 

95 fmt.append_class_name() 

96 

97 fmt.append_attr('name') 

98 

99 fmt.append_attr('times', unit='fs') 

100 fmt.append_attr('dipole_moments', unit='eÅ') 

101 fmt.append_attr('broadening') 

102 formula = self.atoms.get_chemical_formula() if self.atoms is not None else None 

103 fmt.append_attr('atoms', formula) 

104 

105 return fmt.to_string() 

106 

107 @property 

108 def times(self) -> np.ndarray: 

109 """Times at which dipole moments are provided.""" 

110 return self._time_t * au_to_fs # type: ignore 

111 

112 @property 

113 def dipole_moments(self) -> np.ndarray: 

114 """Dipole moment tensors in the time domain.""" 

115 return self._dm_tvv * au_to_eA # type: ignore 

116 

117 def _get_dynamic_polarizability(self, frequencies: Array) -> np.ndarray: 

118 freq_w = np.asarray(frequencies) 

119 dm_wvv = fourier_transform(self._time_t, self._dm_tvv, freq_w, 

120 broadening=self._broadening) 

121 return dm_wvv 

122 

123 def _get_dynamic_polarizability_imaginary_frequency( 

124 self, frequencies: Array) -> np.ndarray: 

125 freq_w = np.asarray(frequencies) 

126 dm_wvv = fourier_transform(self._time_t, self._dm_tvv, freq_w, 

127 freq_axis='imag') 

128 return dm_wvv