97
|
1 import sys
|
|
2 from random import random
|
|
3 import numpy as np
|
|
4 import pylab
|
|
5 import time
|
|
6 import scipy.integrate
|
|
7
|
|
8 """
|
|
9 This script compares different methods to solve an ODE.
|
|
10
|
|
11 The solved system consists of:
|
|
12 y_dot = f(y, t, p)
|
|
13 y0
|
|
14
|
|
15 where
|
|
16 y: the state vector of the system
|
|
17 y0: the initial state of the system
|
|
18 y_dot: the derivatives of the state vector
|
|
19 f: the function that returns the derivatives.
|
|
20 t: the time instant
|
|
21 p: the parameters of the system
|
|
22 """
|
|
23
|
|
24 def func(y, t, parameters):
|
|
25 """
|
|
26 The example system is a bouncing point mass on a hunt-cossley surface.
|
|
27 State vector:
|
|
28 y[0] = height
|
|
29 y[1] = vertical momentum
|
|
30
|
|
31 y_dot[0] = velocity
|
|
32 y_dot[1] = force
|
|
33 """
|
|
34 # Step 1: extract the states into variables:
|
|
35 h = y[0]
|
|
36 p = y[1]
|
|
37 mass = parameters[0]
|
|
38 Kp_floor = parameters[1]
|
|
39 Kd_floor = parameters[2]
|
|
40
|
|
41 # Step 2: Calculate model:
|
|
42 v = p / mass
|
|
43
|
|
44 # Ground reaction force:
|
|
45 if h > 0:
|
|
46 F = 0.0
|
|
47 else:
|
|
48 F = -h * Kp_floor
|
|
49
|
|
50 # Gravity:
|
|
51 F += -9.81 * mass
|
|
52 y_dot = np.array([v, F], dtype=np.float)
|
|
53 return y_dot
|
|
54
|
|
55 parameters = np.array([3, 50000, 3000], dtype=np.float)
|
|
56 y0 = np.array([1, 0], dtype=np.float) # Start at 1 meter, with 0 momentum
|
|
57
|
|
58 class SCIPY:
|
|
59 def __init__(self, dt):
|
|
60 self.dt = dt
|
|
61 def name(self):
|
|
62 return "scipy.integrate"
|
|
63 def integrate(self, f, y0, parameters, endtime):
|
|
64 timevector = np.arange(0.0, endtime, self.dt)
|
|
65 states = scipy.integrate.odeint(f, y0, timevector, args=(parameters,))
|
|
66 return timevector, states
|
|
67
|
|
68 class RK4:
|
|
69 def __init__(self, dt):
|
|
70 self.dt = dt
|
|
71
|
|
72 def name(self):
|
|
73 return "RK4"
|
|
74 def integrate(self, f, y0, parameters, endtime):
|
|
75 print(y0, parameters)
|
|
76 t = 0.0
|
|
77 states = []
|
|
78 timevector = []
|
|
79 state = y0
|
|
80 states.append(state)
|
|
81 timevector.append(t)
|
|
82 while t < endtime:
|
|
83 #### Runge Kutta integration:
|
|
84 k1 = func(state, t, parameters)
|
|
85 k2 = func(state + 0.5*self.dt*k1, t, parameters)
|
|
86 k3 = func(state + 0.5*self.dt*k1, t, parameters)
|
|
87 k4 = func(state + self.dt*k3, t, parameters)
|
|
88 newstate = state + (self.dt/6.0)*(k1+2*k2+2*k3+k4)
|
|
89
|
|
90 t += self.dt
|
|
91 state = newstate
|
|
92
|
|
93 states.append(state)
|
|
94 timevector.append(t)
|
|
95
|
|
96 return np.array(timevector, dtype=np.float), np.array(states, dtype=np.float)
|
|
97
|
|
98
|
|
99 class BDFfixed:
|
|
100 def __init__(self, order):
|
|
101 self.order = order
|
|
102 def name(self):
|
|
103 return "BDF"
|
|
104 def integrate(self, f, y0, parameters, endtime):
|
|
105 timevector = []
|
|
106 states = []
|
|
107 timevector.append(0.0)
|
|
108 states.append(y0)
|
|
109 t = 0.0
|
|
110 h = 0.01
|
|
111 state = y0
|
|
112 while t < endtime:
|
|
113 if len(states) < 4:
|
|
114 k1 = func(state, t, parameters)
|
|
115 k2 = func(state + 0.5*h*k1, t, parameters)
|
|
116 k3 = func(state + 0.5*h*k1, t, parameters)
|
|
117 k4 = func(state + h*k3, t, parameters)
|
|
118 newstate = state + (h/6.0)*(k1+2*k2+2*k3+k4)
|
|
119 state = newstate
|
|
120 t += h
|
|
121 timevector.append(t)
|
|
122 states.append(state)
|
|
123 else:
|
|
124 t += h
|
|
125 if np.max(state) > 10:
|
|
126 break
|
|
127 hyn = (11.0/6.0)*states[-1] + (-3)*states[-2] + (3/2)*states[-3] + (-1/3)*states[-4]
|
|
128 state = state + hyn
|
|
129 timevector.append(t)
|
|
130 states.append(state)
|
|
131
|
|
132 return np.array(timevector, dtype=np.float), np.array(states, dtype=np.float)
|
|
133
|
|
134 class WODE:
|
|
135 def __init__(self):
|
|
136 pass
|
|
137 def integrate(self, f, y0, parameters, endtime):
|
|
138 timevector = []
|
|
139 states = []
|
|
140 timevector.append(0.0)
|
|
141 states.append(y0)
|
|
142 return np.array(timevector, dtype=np.float), np.array(states, dtype=np.float)
|
|
143
|
|
144 class DOPRI45:
|
|
145 def __init__(self):
|
|
146 pass
|
|
147 def name(self):
|
|
148 return "DOPRI"
|
|
149 def integrate(self, f, y0, parameters, endtime):
|
|
150 timevector = []
|
|
151 states = []
|
|
152 timevector.append(0.0)
|
|
153 states.append(y0)
|
|
154 y = y0
|
|
155 t = 0.0
|
|
156
|
|
157 h = 0.01
|
|
158 iteration = 0
|
|
159 while t < endtime:
|
|
160 k1 = h*f(y, t, parameters)
|
|
161 k2 = h*f(y + (1.0/5.0)*k1, t + (1.0/5.0)*h, parameters)
|
|
162 k3 = h*f(y + (3.0/40.0)*k1 + (9.0/40.0)*k2, t + (3.0/10.0)*h, parameters)
|
|
163 k4 = h*f(y + (44.0/45.0)*k1 - (56.0/15.0)*k2 + (32.0/9.0)*k3, t + (4.0/5.0)*h, parameters)
|
|
164 k5 = h*f(y + (19372.0/6561.0)*k1 - (25360.0/2187.0)*k2 + (64448.0/6561.0)*k3 - (212.0/729.0)*k4, t + (8.0/9.0)*h, parameters)
|
|
165 k6 = h*f(y + (9017.0/3168.0)*k1 - (355.0/33.0)*k2 - (46732.0/5247.0)*k3 + (49.0/176.0)*k4 - (5103.0/18656.0)*k5, t + h, parameters)
|
|
166 k7 = h*f(y + (35.0/384.0)*k1 + (500.0/1113.0)*k3 + (125.0/192.0)*k4 - (2187.0/6784.0)*k5 + (11.0/84.0)*k6, t + h, parameters)
|
|
167
|
|
168
|
|
169 y_next = y + (35.0/384.0)*k1 + (500.0/1113.0)*k3 + (125.0/192.0)*k4 - (2187.0/6784.0)*k5 + (11.0/84.0)*k6
|
|
170 z_next = y + (5179.0/57600.0)*k1 + (7571.0/16695.0)*k3 + (393.0/640.0)*k4 - (92097.0/339200.0)*k5 + (187.0/2100.0)*k6 + (1.0/40.0)*k7
|
|
171 error = np.linalg.norm(y_next - z_next)
|
|
172 #print(error, t)
|
|
173 eps = 0.01
|
|
174 if error < eps:
|
|
175 y = y_next
|
|
176 t = t + h
|
|
177 timevector.append(t)
|
|
178 states.append(y)
|
|
179
|
|
180 s = ((eps*h)/(2*error)) ** (1.0/5.0)
|
|
181 h = s * h
|
|
182 if h < 1e-5:
|
|
183 h = 1e-5
|
|
184 iteration += 1
|
|
185 if iteration > 10000:
|
|
186 break
|
|
187
|
|
188 return np.array(timevector, dtype=np.float), np.array(states, dtype=np.float)
|
|
189
|
|
190 rk4 = RK4(0.01)
|
|
191 sp = SCIPY(0.01)
|
|
192 wode = BDFfixed(4)
|
|
193 dopri = DOPRI45()
|
|
194 methods = [rk4, sp, dopri]
|
|
195
|
|
196 pylab.figure()
|
|
197 pylab.hold(True)
|
|
198 for method in methods:
|
|
199 t, s = method.integrate(func, y0, parameters, 10)
|
|
200 print(t.shape, s.shape)
|
|
201 pylab.plot(t, s[:,0], 'x-', label=method.name())
|
|
202
|
|
203 pylab.legend()
|
|
204 pylab.show()
|
|
205
|