Howto pass a function to a function in Python?

Question

I am a beginner/intermediate in Python. I have coded a 4th-order Runge-Kutta method (RK4) into Python. It is basically solving a pendulum, but that is not the point here.

I want to improve the RK4 method in the following way: I want to be able to pass the function f directly to the RK4 function, i.e. RK4(y_0, n, h) should become RK4(f,y_0,n,h). This would have the great advantage that I could use RK4 for other f functions that describe other systems, not just this one pendulum.

I have played around with just passing simple functions to RK4, but I am doing something wrong. How do I do this in Python?

import numpy as np

def RK4(y_0, n, h):
    #4th order Runge-Kutta solver, takes as input
    #initial value y_0, the number of steps n and stepsize h
    #returns solution vector y and time vector t
    #right now function f is defined below

    t = np.linspace(0,n*h,n,endpoint = False)   #create time vector t
    y = np.zeros((n,len(y_0))) #create solution vector y
    y[0] = y_0 #assign initial value to first position in y
    for i in range(0,n-1):
        #compute Runge-Kutta weights k_1 till k_4
        k_1 = f(t[i],y[i])
        k_2 = f(t[i] + 0.5*h, y[i] + 0.5*h*k_1)
        k_3 = f(t[i] + 0.5*h, y[i] + 0.5*h*k_2)
        k_4 = f(t[i] + 0.5*h, y[i] + h*k_3)
        #compute next y        
        y[i+1] = y[i] + h / 6. * (k_1 + 2.*k_2 + 2.*k_3 + k_4)
    return t,y

def f(t,vec):
    theta=vec[0]
    omega = vec[1]
    omegaDot = -np.sin(theta) - omega + np.cos(t)
    result = np.array([omega,omegaDot])    
    return result

test = np.array([0,0.5])
t,y = RK4(test,10,0.1)

score 4 · Accepted Answer · answered Mar 21 '13 at 10:52

4

Python functions are objects too. You can pass them around like any other object:

>>> def foo(): print 'Hello world!'
...
>>> foo
<function foo at 0x10c4685f0>
>>> foo()
Hello world!
>>> bar = foo
>>> bar()
Hello world!

Simply pass a function as an extra parameter to your RK4 function and use that as a local variable.

answered Mar 21 '13 at 10:52

Martijn Pieters

1,048,767
296
4,058
3,343

Thanks for the clarification. I'm coming from MatLab (with some C). so Python just doesn't cease to amaze me :-). – seb Mar 21 '13 at 10:57

score 3 · Answer 2 · answered Mar 21 '13 at 10:52

3

You can pass a function to a function in Python just as you might expect:

def call_function(f):
    f()

def my_function():
    print "OK"

call_function(my_function)  # Prints OK

Maybe you should post your failing code?

answered Mar 21 '13 at 10:52

RichieHindle

272,464
47
358
399

David Heffernan · Answer 3 · 2013-03-21T10:59:17.253

2

It's very simple. Change the definition of the RK4 function like so:

def RK4(f, y_0, n, h):

Here, I have added an extra argument, the function.

Then, when you call RK4, pass the function:

t, y = RK4(f, test, 10, 0.1)

And now, of course, you can substitute different functions without having to re-write the integration code.

Functions in Python are just another kind of object. You can pass them around just as you do more prosaic objects.

edited Mar 21 '13 at 10:59

answered Mar 21 '13 at 10:52

David Heffernan

601,492
42
1,072
1,490

ok, I put this in and now it works. I am one bloody newby! anyway, typing this all up and posting it here was a good exercise. – seb Mar 21 '13 at 10:56
Yes, it's very often the case that the discipline of writing a concise example is helpful to your understanding. By the way, this was a very well written question, well done! – David Heffernan Mar 21 '13 at 10:57

Howto pass a function to a function in Python?

3 Answers3

Linked