Runge Kutta method in python

Question

These are the functions I have written:

def rk4(f, t0, y0, h, N):
    t = t0 + arange(N+1)*h
    y = zeros((N+1, size(y0)))
    y[0] = y0
    for n in range(N):
        xi1 = y[n]
        f1 = f(t[n], xi1)
        xi2 = y[n] + (h/2.)*f1
        f2 = f(t[n+1], xi2)
        xi3 = y[n] + (h/2.)*f2
        f3 = f(t[n+1], xi3)
        xi4 = y[n] + h*f3
        f4 = f(t[n+1], xi4)
        y[n+1] = y[n] + (h/6.)*(f1 + 2*f2 + 2*f3 + f4)
    return y

def lorenzPlot():
    y = rk4(fLorenz, 0, array([0,2,20]), .01, 10000)
    fig = figure()
    ax = Axes3D(fig)
    ax.plot(*y.T)

def fLorenz(t,Y):   
    def Y(x,y,z):
        return array([10*(y-x), x*(28-z)-y, x*y-(8./3.)*z])

By implementing lorenzPlot, it's supposed to graph the numerical solution to fLorenz (the Lorenz system of equations) obtained using rk4 (4th order Runge Kutta method). I'm having a problem with the function fLorenz. When I define it as above and call lorenzPlot, I get an error that says

File "C:/...", line 38, in rk4
    xi2 = y[n] + (h/2.)*f1
TypeError: unsupported operand type(s) for *: 'float' and 'NoneType'

I guessed that this had something to do with not being able to multiply the array correctly.
However, when I changed fLorenz to

def fLorenz(t,x,y,z):   
    return array([10*(y-x), x*(28-z)-y, x*y-(8./3.)*z])

calling lorenzPlot gives me the error stating that fLorenz takes 4 arguments, but only 2 are given.

Additionally, rk4 and lorenzPlot both work correctly for functions composed of a singular equation.

How should I change fLorenz so that it can be used as f in rk4 and lorenzPlot?

Answer 1

Your first fLorenz function defines a sub-function, but doesn't actually return anything. Your second one does, except it's expecting four arguments (t, x, y, z), and you only give it t, Y.

From what I can understand, Y is a three-tuple; you can simply unpack it before you use the values:

def fLorenz(t, Y):
    x, y, z = Y
    return array([10*(y-x), x*(28-z)-y, x*y-(8./3.)*z])