Mercurial > lcfOS
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() +