diff python/integration_methods.py @ 97:5a965e9664f2

Movage
author windel
date Mon, 24 Dec 2012 13:32:54 +0100
parents
children
line wrap: on
line diff
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/python/integration_methods.py	Mon Dec 24 13:32:54 2012 +0100
@@ -0,0 +1,205 @@
+import sys
+from random import random
+import numpy as np
+import pylab
+import time
+import scipy.integrate
+
+"""
+  This script compares different methods to solve an ODE.
+
+  The solved system consists of:
+  y_dot = f(y, t, p)
+  y0
+
+  where
+   y: the state vector of the system
+   y0: the initial state of the system
+   y_dot: the derivatives of the state vector
+   f: the function that returns the derivatives.
+   t: the time instant
+   p: the parameters of the system
+"""
+
+def func(y, t, parameters):
+   """
+    The example system is a bouncing point mass on a hunt-cossley surface.
+    State vector:
+    y[0] = height
+    y[1] = vertical momentum
+
+    y_dot[0] = velocity
+    y_dot[1] = force
+   """
+   # Step 1: extract the states into variables:
+   h = y[0]
+   p = y[1]
+   mass = parameters[0]
+   Kp_floor = parameters[1]
+   Kd_floor = parameters[2]
+
+   # Step 2: Calculate model:
+   v = p / mass
+
+   # Ground reaction force:
+   if h > 0:
+      F = 0.0
+   else:
+      F = -h * Kp_floor
+
+   # Gravity:
+   F += -9.81 * mass 
+   y_dot = np.array([v, F], dtype=np.float)
+   return y_dot
+
+parameters = np.array([3, 50000, 3000], dtype=np.float)
+y0 = np.array([1, 0], dtype=np.float) # Start at 1 meter, with 0 momentum
+
+class SCIPY:
+   def __init__(self, dt):
+      self.dt = dt
+   def name(self):
+      return "scipy.integrate"
+   def integrate(self, f, y0, parameters, endtime):
+      timevector = np.arange(0.0, endtime, self.dt)
+      states = scipy.integrate.odeint(f, y0, timevector, args=(parameters,))
+      return timevector, states
+
+class RK4:
+   def __init__(self, dt):
+      self.dt = dt
+
+   def name(self):
+      return "RK4"
+   def integrate(self, f, y0, parameters, endtime):
+      print(y0, parameters)
+      t = 0.0
+      states = []
+      timevector = []
+      state = y0
+      states.append(state)
+      timevector.append(t)
+      while t < endtime:
+         #### Runge Kutta integration:
+         k1 = func(state,             t, parameters)
+         k2 = func(state + 0.5*self.dt*k1, t, parameters)
+         k3 = func(state + 0.5*self.dt*k1, t, parameters)
+         k4 = func(state + self.dt*k3,     t, parameters)
+         newstate = state + (self.dt/6.0)*(k1+2*k2+2*k3+k4)
+
+         t += self.dt
+         state = newstate
+
+         states.append(state)
+         timevector.append(t)
+
+      return np.array(timevector, dtype=np.float), np.array(states, dtype=np.float)
+
+
+class BDFfixed:
+   def __init__(self, order):
+      self.order = order
+   def name(self):
+      return "BDF"
+   def integrate(self, f, y0, parameters, endtime):
+      timevector = []
+      states = []
+      timevector.append(0.0)
+      states.append(y0)
+      t = 0.0
+      h = 0.01
+      state = y0
+      while t < endtime:
+         if len(states) < 4:
+            k1 = func(state,             t, parameters)
+            k2 = func(state + 0.5*h*k1, t, parameters)
+            k3 = func(state + 0.5*h*k1, t, parameters)
+            k4 = func(state + h*k3,     t, parameters)
+            newstate = state + (h/6.0)*(k1+2*k2+2*k3+k4)
+            state = newstate
+            t += h
+            timevector.append(t)
+            states.append(state)
+         else:
+            t += h
+            if np.max(state) > 10:
+               break
+            hyn = (11.0/6.0)*states[-1] + (-3)*states[-2] + (3/2)*states[-3] + (-1/3)*states[-4]
+            state = state + hyn
+            timevector.append(t)
+            states.append(state)
+
+      return np.array(timevector, dtype=np.float), np.array(states, dtype=np.float)
+
+class WODE:
+   def __init__(self):
+      pass
+   def integrate(self, f, y0, parameters, endtime):
+      timevector = []
+      states = []
+      timevector.append(0.0)
+      states.append(y0)
+      return np.array(timevector, dtype=np.float), np.array(states, dtype=np.float)
+
+class DOPRI45:
+   def __init__(self):
+      pass
+   def name(self):
+      return "DOPRI"
+   def integrate(self, f, y0, parameters, endtime):
+      timevector = []
+      states = []
+      timevector.append(0.0)
+      states.append(y0)
+      y = y0
+      t = 0.0
+
+      h = 0.01
+      iteration = 0
+      while t < endtime:
+         k1 = h*f(y, t, parameters)
+         k2 = h*f(y + (1.0/5.0)*k1, t + (1.0/5.0)*h, parameters)
+         k3 = h*f(y + (3.0/40.0)*k1 + (9.0/40.0)*k2, t + (3.0/10.0)*h, parameters)
+         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)
+         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)
+         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)
+         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)
+
+
+         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
+         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
+         error = np.linalg.norm(y_next - z_next)
+         #print(error, t)
+         eps = 0.01
+         if error < eps:
+            y = y_next
+            t = t + h
+            timevector.append(t)
+            states.append(y)
+
+         s = ((eps*h)/(2*error)) ** (1.0/5.0)
+         h = s * h
+         if h < 1e-5:
+            h = 1e-5
+         iteration += 1
+         if iteration > 10000:
+            break
+
+      return np.array(timevector, dtype=np.float), np.array(states, dtype=np.float)
+
+rk4 = RK4(0.01)
+sp = SCIPY(0.01)
+wode = BDFfixed(4)
+dopri = DOPRI45()
+methods = [rk4, sp, dopri]
+
+pylab.figure()
+pylab.hold(True)
+for method in methods:
+   t, s = method.integrate(func, y0, parameters, 10)
+   print(t.shape, s.shape)
+   pylab.plot(t, s[:,0], 'x-', label=method.name())
+
+pylab.legend()
+pylab.show()
+