Mercurial > lcfOS
comparison python/other/integration_methods.py @ 290:7b38782ed496
File moves
author | Windel Bouwman |
---|---|
date | Sun, 24 Nov 2013 11:24:15 +0100 |
parents | python/integration_methods.py@5a965e9664f2 |
children |
comparison
equal
deleted
inserted
replaced
289:bd2593de3ff8 | 290:7b38782ed496 |
---|---|
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 |