Coverage for dynasor/tools/damped_harmonic_oscillator.py: 100%

1from typing import Optional, Union 

2import numpy as np 

3from numpy.typing import NDArray 

4 

5 

6def _validate_dho_parameters(w0: float, gamma: float): 

7 if not (w0 > 0 and gamma > 0): 

8 raise ValueError('w0 and gamma must be positive.') 

9 

10 

11def acf_position_dho( 

12 t: Union[float, NDArray[float]], 

13 w0: float, 

14 gamma: float, 

15 A: Optional[float] = 1.0, 

16) -> Union[float, NDArray[float]]: 

17 r""" 

18 Calculate the damped harmonic oscillator (DHO) 

19 autocorrelation function for the position. 

20 The definition of this function can be found in the `dynasor documentation 

21 <https://dynasor.materialsmodeling.org/get_started/theory.html#damped-harmonic-oscillator-model>`_. # noqa 

22 

23 Parameters 

24 ---------- 

25 t 

26 Time, usually as an array. 

27 w0 

28 Natural angular frequency of the DHO. 

29 gamma 

30 Damping of DHO. 

31 A 

32 Amplitude of the DHO. 

33 """ 

34 

35 _validate_dho_parameters(w0, gamma) 

36 t = np.abs(t) 

37 

38 if 2 * w0 > gamma: # underdamped 

39 we = np.sqrt(w0**2 - gamma**2 / 4.0) 

40 return A * np.exp(-gamma * t / 2.0) * ( 

41 np.cos(we * t) + 0.5 * gamma / we * np.sin(we * t)) 

42 elif 2 * w0 < gamma: # overdamped 

43 tau = 2 / gamma 

44 tau_S = tau / (1 + np.sqrt(1 - (w0 * tau)**2)) 

45 tau_L = tau / (1 - np.sqrt(1 - (w0 * tau)**2)) 

46 return A / (tau_L - tau_S) * (tau_L * np.exp(-t/tau_L) - tau_S * np.exp(-t/tau_S)) 

47 else: 

48 tau = 2 / gamma 

49 return A * np.exp(-t/tau) * (1 + t / tau) 

50 

51 

52def acf_velocity_dho( 

53 t: Union[float, NDArray[float]], 

54 w0: float, 

55 gamma: float, 

56 A: Optional[float] = 1.0, 

57) -> Union[float, NDArray[float]]: 

58 r""" 

59 Calculate the damped harmonic oscillator (DHO) 

60 autocorrelation function for the velocity. 

61 The definition of this function can be found in the `dynasor documentation 

62 <https://dynasor.materialsmodeling.org/get_started/theory.html#damped-harmonic-oscillator-model>`_. # noqa 

63 

64 Parameters 

65 ---------- 

66 t 

67 Time, usually as an array. 

68 w0 

69 Natural angular frequency of the DHO. 

70 gamma 

71 Damping of DHO. 

72 A 

73 Amplitude of the DHO. 

74 """ 

75 

76 _validate_dho_parameters(w0, gamma) 

77 t = np.abs(t) 

78 

79 if 2 * w0 > gamma: # underdamped 

80 we = np.sqrt(w0**2 - gamma**2 / 4.0) 

81 return A * w0**2 * np.exp(-gamma * t / 2.0) * ( 

82 np.cos(we * t) - 0.5 * gamma / we * np.sin(we * t)) 

83 elif 2 * w0 < gamma: # overdamped 

84 tau = 2 / gamma 

85 tau_S = tau / (1 + np.sqrt(1 - (w0 * tau)**2)) 

86 tau_L = tau / (1 - np.sqrt(1 - (w0 * tau)**2)) 

87 return A / (tau_L - tau_S) * (np.exp(-t/tau_S)/tau_S - np.exp(-t/tau_L)/tau_L) 

88 else: 

89 tau = 2 / gamma 

90 return A * w0**2 * np.exp(-t/tau) * (1 - t / tau) 

91 

92 

93def psd_position_dho( 

94 w: Union[float, NDArray[float]], 

95 w0: float, 

96 gamma: float, 

97 A: Optional[float] = 1.0, 

98) -> Union[float, NDArray[float]]: 

99 r""" 

100 Calculate the power spectral density (PSD) function for the damped harmonic oscillator 

101 (DHO) as the Fourier transform of the position autocorrelation function. 

102 

103 The definition of this function can be found in the `dynasor documentation 

104 <https://dynasor.materialsmodeling.org/get_started/theory.html#damped-harmonic-oscillator-model>`_. # noqa 

105 

106 Parameters 

107 ---------- 

108 w 

109 Angular frequency, usually as an array. 

110 w0 

111 Natural angular frequency of the DHO. 

112 gamma 

113 Damping of DHO. 

114 A 

115 Amplitude of the DHO. 

116 """ 

117 _validate_dho_parameters(w0, gamma) 

118 return 2 * w0**2 * A * gamma / ((w**2 - w0**2)**2 + (w * gamma)**2) 

119 

120 

121def psd_velocity_dho( 

122 w: Union[float, NDArray[float]], 

123 w0: float, 

124 gamma: float, 

125 A: Optional[float] = 1.0, 

126) -> Union[float, NDArray[float]]: 

127 r""" 

128 Calculate the power spectral density (PSD) function for the damped harmonic oscillator 

129 (DHO) as the Fourier transform of the velocity autocorrelation function. 

130 

131 The definition of this function can be found in the `dynasor documentation 

132 <https://dynasor.materialsmodeling.org/get_started/theory.html#damped-harmonic-oscillator-model>`_. # noqa 

133 

134 Parameters 

135 ---------- 

136 w 

137 Angular frequency, usually as an array. 

138 w0 

139 Natural angular frequency of the DHO. 

140 gamma 

141 Damping of DHO. 

142 A 

143 Amplitude of the DHO. 

144 """ 

145 _validate_dho_parameters(w0, gamma) 

146 return w**2 * psd_position_dho(w, w0, gamma, A)