Defining functions in python

| categories: python | tags:

Compare what's here to the Matlab implementation.

We often need to make functions in our codes to do things.

def f(x):
    "return the inverse square of x"
    return 1.0 / x**2

print f(3)
print f([4,5])
... ... >>> 0.111111111111
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 3, in f
TypeError: unsupported operand type(s) for ** or pow(): 'list' and 'int'

Note that functions are not automatically vectorized. That is why we see the error above. There are a few ways to achieve that. One is to “cast” the input variables to objects that support vectorized operations, such as numpy.array objects.

import numpy as np

def f(x):
    "return the inverse square of x"
    x = np.array(x)
    return 1.0 / x**2

print f(3)
print f([4,5])
>>> ... ... ... ... >>> 0.111111111111
[ 0.0625  0.04  ]

It is possible to have more than one variable.

import numpy as np

def func(x, y):
    "return product of x and y"
    return x * y

print func(2, 3)
print func(np.array([2, 3]), np.array([3, 4]))
6
[ 6 12]

You can define “lambda” functions, which are also known as inline or anonymous functions. The syntax is lambda var:f(var). I think these are hard to read and discourage their use. Here is a typical usage where you have to define a simple function that is passed to another function, e.g. scipy.integrate.quad to perform an integral.

from scipy.integrate import quad
print quad(lambda x:x**3, 0 ,2)
(4.0, 4.440892098500626e-14)

It is possible to nest functions inside of functions like this.

def wrapper(x):
    a = 4
    def func(x, a):
        return a * x

    return func(x, a)

print wrapper(4)
16

An alternative approach is to “wrap” a function, say to fix a parameter. You might do this so you can integrate the wrapped function, which depends on only a single variable, whereas the original function depends on two variables.

def func(x, a):
        return a * x
 
def wrapper(x):
    a = 4
    return func(x, a)

print wrapper(4)
16

Last example, defining a function for an ode

from scipy.integrate import odeint
import numpy as np
import matplotlib.pyplot as plt

k = 2.2
def myode(t,y):
    "ode defining exponential growth"
    return k * t

y0 = 3
tspan = np.linspace(0,1)
y =  odeint(myode, y0, tspan)

plt.plot(tspan, y)
plt.xlabel('Time')
plt.ylabel('y')
plt.savefig('images/funcs-ode.png')

Copyright (C) 2013 by John Kitchin. See the License for information about copying.

org-mode source

Discuss on Twitter