Mercurial > lcfOS

comparison python/integration_methods.py @ 97:5a965e9664f2

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Movage
author windel
date Mon, 24 Dec 2012 13:32:54 +0100
parents
children
comparison
equal deleted inserted replaced
96:a350055d6119 97:5a965e9664f2
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