annotate python/integration_methods.py @ 176:5fd02aa38b42

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