Coverage for pySDC/implementations/problem_classes/HarmonicOscillator.py: 71%

1import numpy as np 

2 

3from pySDC.core.Errors import ParameterError 

4from pySDC.core.Problem import ptype 

5from pySDC.implementations.datatype_classes.particles import particles, acceleration 

6 

7 

8# noinspection PyUnusedLocal 

9class harmonic_oscillator(ptype): 

10 r""" 

11 Example implementing the harmonic oscillator with mass :math:`1` 

12 

13 .. math:: 

14 \frac{d^2 x}{dt^2} = -kx - \mu \frac{d x}{dt}, 

15 

16 which is a second-order problem. The unknown function :math:`x` denotes the position of the mass, and the 

17 derivative is the velocity. :math:`\mu` defines the damping and :math:`k` is the spring constant. 

18 

19 Parameters 

20 ---------- 

21 k : float, optional 

22 Spring constant :math:`k`. 

23 mu : float, optional 

24 Damping parameter :math:`\mu`. 

25 u0 : tuple, optional 

26 Initial condition for the position, and the velocity. Should be a tuple, e.g. ``u0=(1, 0)``. 

27 phase : float, optional 

28 Phase of the oscillation. 

29 amp : float, optional 

30 Amplitude of the oscillation. 

31 Source: https://beltoforion.de/en/harmonic_oscillator/ 

32 """ 

33 

34 dtype_u = particles 

35 dtype_f = acceleration 

36 

37 def __init__(self, k=1.0, mu=0.0, u0=(1, 0), phase=0.0, amp=1.0): 

38 """Initialization routine""" 

39 # invoke super init, passing nparts, dtype_u and dtype_f 

40 u0 = np.asarray(u0) 

41 super().__init__((1, None, np.dtype("float64"))) 

42 self._makeAttributeAndRegister('k', 'mu', 'u0', 'phase', 'amp', localVars=locals(), readOnly=True) 

43 

44 def eval_f(self, u, t): 

45 """ 

46 Routine to compute the right-hand side of the problem. 

47 

48 Parameters 

49 ---------- 

50 u : dtype_u 

51 Current values of the particles. 

52 t : float 

53 Current time of the numerical solution is computed (not used here). 

54 

55 Returns 

56 ------- 

57 me : dtype_f 

58 The right-hand side of the problem. 

59 """ 

60 me = self.dtype_f(self.init) 

61 me[:] = -self.k * u.pos - self.mu * u.vel 

62 return me 

63 

64 def u_init(self): 

65 """ 

66 Helper function to compute the initial condition for u. 

67 """ 

68 u0 = self.u0 

69 

70 u = self.dtype_u(self.init) 

71 

72 u.pos[0] = u0[0] 

73 u.vel[0] = u0[1] 

74 

75 return u 

76 

77 def u_exact(self, t): 

78 r""" 

79 Routine to compute the exact trajectory at time :math:`t`. 

80 

81 Parameters 

82 ---------- 

83 t : float 

84 Time of the exact trajectory. 

85 

86 Returns 

87 ------- 

88 me : dtype_u 

89 Exact position and velocity. 

90 """ 

91 me = self.dtype_u(self.init) 

92 delta = self.mu / (2) 

93 omega = np.sqrt(self.k) 

94 

95 U_0 = self.u0 

96 alpha = np.sqrt(np.abs(delta**2 - omega**2)) 

97 print(self.mu) 

98 if delta > omega: 

99 """ 

100 Overdamped case 

101 """ 

102 

103 lam_1 = -delta + alpha 

104 lam_2 = -delta - alpha 

105 L = np.array([[1, 1], [lam_1, lam_2]]) 

106 A, B = np.linalg.solve(L, U_0) 

107 me.pos[:] = A * np.exp(lam_1 * t) + B * np.exp(lam_2 * t) 

108 me.vel[:] = A * lam_1 * np.exp(lam_1 * t) + B * lam_2 * np.exp(lam_2 * t) 

109 

110 elif delta == omega: 

111 """ 

112 Critically damped case 

113 """ 

114 

115 A = U_0[0] 

116 B = U_0[1] + delta * A 

117 me.pos[:] = np.exp(-delta * t) * (A + t * B) 

118 me.vel[:] = -delta * me.pos[:] + np.exp(-delta * t) * B 

119 

120 elif delta < omega: 

121 """ 

122 Underdamped case 

123 """ 

124 

125 lam_1 = -delta + alpha * 1j 

126 lam_2 = -delta - alpha * 1j 

127 

128 M = np.array([[1, 1], [lam_1, lam_2]], dtype=complex) 

129 A, B = np.linalg.solve(M, U_0) 

130 

131 me.pos[:] = np.real(A * np.exp(lam_1 * t) + B * np.exp(lam_2 * t)) 

132 me.vel[:] = np.real(A * lam_1 * np.exp(lam_1 * t) + B * lam_2 * np.exp(lam_2 * t)) 

133 

134 else: 

135 pass 

136 raise ParameterError("Exact solution is not working") 

137 return me 

138 

139 def eval_hamiltonian(self, u): 

140 """ 

141 Routine to compute the Hamiltonian. 

142 

143 Parameters 

144 ---------- 

145 u : dtype_u 

146 Current values of the particles. 

147 

148 Returns 

149 ------- 

150 ham : float 

151 The Hamiltonian. 

152 """ 

153 

154 ham = 0.5 * self.k * u.pos[0] ** 2 + 0.5 * u.vel[0] ** 2 

155 return ham