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 acf_position_dho( 

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

8 w0: float, 

9 gamma: float, 

10 A: Optional[float] = 1.0, 

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

12 r""" 

13 Calculate the damped damped harmonic oscillator (DHO) 

14 autocorrelation function for the position. 

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

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

17 

18 Parameters 

19 ---------- 

20 t 

21 Time, usually as an array. 

22 w0 

23 Natural angular frequency of the DHO. 

24 gamma 

25 Damping of DHO. 

26 A 

27 Amplitude of the DHO. 

28 """ 

29 

30 t = np.abs(t) 

31 

32 if 2 * w0 > gamma: # underdamped 

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

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

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

36 elif 2 * w0 < gamma: # overdamped 

37 tau = 2 / gamma 

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

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

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

41 else: 

42 tau = 2 / gamma 

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

44 

45 

46def acf_velocity_dho( 

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

48 w0: float, 

49 gamma: float, 

50 A: Optional[float] = 1.0, 

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

52 r""" 

53 Calculate the damped damped harmonic oscillator (DHO) 

54 autocorrelation function for the velocity. 

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

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

57 

58 Parameters 

59 ---------- 

60 t 

61 Time, usually as an array. 

62 w0 

63 Natural angular frequency of the DHO. 

64 gamma 

65 Damping of DHO. 

66 A 

67 Amplitude of the DHO. 

68 """ 

69 

70 t = np.abs(t) 

71 

72 if 2 * w0 > gamma: # underdamped 

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

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

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

76 elif 2 * w0 < gamma: # overdamped 

77 tau = 2 / gamma 

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

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

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

81 else: 

82 tau = 2 / gamma 

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

84 

85 

86def psd_position_dho( 

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

88 w0: float, 

89 gamma: float, 

90 A: Optional[float] = 1.0, 

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

92 r""" 

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

94 (DHO), i.e., the Fourier transform of the autocorrelation function) for the position. 

95 

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

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

98 

99 Parameters 

100 ---------- 

101 w 

102 Angular frequency, usually as an array. 

103 w0 

104 Natural angular frequency of the DHO. 

105 gamma 

106 Damping of DHO. 

107 A 

108 Amplitude of the DHO. 

109 """ 

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

111 

112 

113def psd_velocity_dho( 

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

115 w0: float, 

116 gamma: float, 

117 A: Optional[float] = 1.0, 

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

119 r""" 

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

121 (DHO), i.e., the Fourier transform of the autocorrelation function) for the position. 

122 

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

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

125 

126 Parameters 

127 ---------- 

128 w 

129 Angular frequency, usually as an array. 

130 w0 

131 Natural angular frequency of the DHO. 

132 gamma 

133 Damping of DHO. 

134 A 

135 Amplitude of the DHO. 

136 """ 

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