Suppose we want to evaluate a function of two variables on a regularly spaced grid and visualize the results.
There are tools that make this easy to do in the numpy
and matplotlib
modules.
import numpy as np
import matplotlib.pyplot as plt
Let's define some arbitary function of two (scalar) variables.
def arbitrary_function(x, y):
if x < 0:
return -1.0
else:
return np.sin(x) * np.exp(y)
If given arrays of input values rather than scalars, this function doesn't automatically operate element-wise like MATLAB. (If it did, the first element of the output array would be -1.)
xs = [-1, 1, 2, 3, 4.]
ys = [1, 1, 2, 3, 4.]
arbitrary_function(xs, ys)
array([ -2.28735529, 2.28735529, 6.7188497 , 2.83447113, -41.32001618])
np.vectorize
takes as input a function and returns a new function that does have this automatic element-wise behavior.
vectorized_function = np.vectorize(arbitrary_function)
vectorized_function(xs, ys)
array([ -1. , 2.28735529, 6.7188497 , 2.83447113, -41.32001618])
np.linspace
produces a regulary spaced 1-D "grid".
xs = np.linspace(-1, 1, num=5)
ys = np.linspace(-1, 1, num=5)
print xs
print ys
[-1. -0.5 0. 0.5 1. ] [-1. -0.5 0. 0.5 1. ]
np.meshgrid
makes a convenient representation of an n-D grid out of n 1-D inputs.
xx, yy = np.meshgrid(xs, ys)
print xx
print yy
[[-1. -0.5 0. 0.5 1. ] [-1. -0.5 0. 0.5 1. ] [-1. -0.5 0. 0.5 1. ] [-1. -0.5 0. 0.5 1. ] [-1. -0.5 0. 0.5 1. ]] [[-1. -1. -1. -1. -1. ] [-0.5 -0.5 -0.5 -0.5 -0.5] [ 0. 0. 0. 0. 0. ] [ 0.5 0.5 0.5 0.5 0.5] [ 1. 1. 1. 1. 1. ]]
plt.imshow
can be used to visualize functions of two variables, using color as the third "dimension".
plt.imshow(vectorized_function(xx, yy),
interpolation='nearest',
extent=[-1, 1, -1, 1],
origin='lower',
)
<matplotlib.image.AxesImage at 0x9ffadcc>
Lets make the grid finer for a better looking picture.
fine_xs = np.linspace(-1, 1, num=200)
fine_ys = np.linspace(-1, 1, num=200)
fine_xx, fine_yy = np.meshgrid(fine_xs, fine_ys)
plt.imshow(vectorized_function(fine_xx, fine_yy),
interpolation='nearest',
extent=[-1, 1, -1, 1],
origin='lower',
)
<matplotlib.image.AxesImage at 0xa3642cc>
We can mix scalar and array arguments to a vectorized function.
def arbitrary_function_2(x, y, a):
if x < 0:
return -1.0
else:
return np.sin(a*x) * np.exp(y)
vectorized_function_2 = np.vectorize(arbitrary_function_2)
plt.imshow(vectorized_function_2(fine_xx, fine_yy, 3.0),
extent=[-1, 1, -1, 1],
origin='lower',
)
<matplotlib.image.AxesImage at 0xa54248c>
plt.imshow(vectorized_function_2(fine_xx, fine_yy, 6.0),
extent=[-1, 1, -1, 1],
origin='lower',
)
<matplotlib.image.AxesImage at 0xa762c8c>