Saturday, January 14, 2012

How to plot a function of two variables with matplotlib

In this post we will see how to visualize a function of two variables in two ways. First, we will create an intensity image of the function and, second, we will use the 3D plotting capabilities of matplotlib to create a shaded surface plot. So, let's go with the code:
from numpy import exp,arange
from pylab import meshgrid,cm,imshow,contour,clabel,colorbar,axis,title,show

# the function that I'm going to plot
def z_func(x,y):
 return (1-(x**2+y**3))*exp(-(x**2+y**2)/2)
 
x = arange(-3.0,3.0,0.1)
y = arange(-3.0,3.0,0.1)
X,Y = meshgrid(x, y) # grid of point
Z = z_func(X, Y) # evaluation of the function on the grid

im = imshow(Z,cmap=cm.RdBu) # drawing the function
# adding the Contour lines with labels
cset = contour(Z,arange(-1,1.5,0.2),linewidths=2,cmap=cm.Set2)
clabel(cset,inline=True,fmt='%1.1f',fontsize=10)
colorbar(im) # adding the colobar on the right
# latex fashion title
title('$z=(1-x^2+y^3) e^{-(x^2+y^2)/2}$')
show()
The script would have the following output:



And now we are going to use the values stored in X,Y and Z to make a 3D plot using the mplot3d toolkit. Here's the snippet:
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
import matplotlib.pyplot as plt

fig = plt.figure()
ax = fig.gca(projection='3d')
surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, 
                      cmap=cm.RdBu,linewidth=0, antialiased=False)

ax.zaxis.set_major_locator(LinearLocator(10))
ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))

fig.colorbar(surf, shrink=0.5, aspect=5)

plt.show()
And this is the result:

10 comments:

  1. i would like to see a demo using d3py from
    https://github.com/mikedewar/D3py/tree/v2

    ReplyDelete
  2. Which version of matplotlib do you use? I get attribute error

    --> 34 ax.zaxis.set_major_locator(LinearLocator(10))
    AttributeError: 'Axes3DSubplot' object has no attribute 'zaxis'

    ReplyDelete
  3. I made this code using matplotlib 1.1.0

    ReplyDelete
  4. Oh, I have 1.0.1. Thanks

    ReplyDelete
  5. Stack Overflow hints at how to do the 2nd example with matplotlib 0.99:

    http://stackoverflow.com/questions/3810865/need-help-with-matplotlib

    These small changes worked for me:

    http://pastebin.com/PnSMLqDD

    ReplyDelete
  6. Hello there! thank you very much indeed.

    Atte,

    Ignacio Aular

    :)

    ReplyDelete
  7. This comment has been removed by the author.

    ReplyDelete
  8. Thanks! This post was very helpful to me in gaining proficieny with Python plotting. However, there are a number of errors associated with the contour map / intensity image code as presented.

    First, I'm not sure why your contour map axes extend from 0 to 249 (I suppose), rather than from 0 to 59, as set by arange(). But e that as it may, it is really much better to set extents on the contour map's axes so that they can be labeled properly - that is, with the values actually presented to the function z_func(). This can be achieved by replacing the current imshow() with:

    im = imshow(Z,cmap=cm.RdBu, extent=[-3.0,3.0,-3.0,3.0], origin='lower') # drawing the function

    And by adding the x and y arrays to the call to contour() as follows:

    cset = contour(x,y,Z,arange(-1,1.5,0.2),linewidths=2,cmap=cm.Set2)

    Doing so clears up another problem with the contour map as drawn above, and that is that it does not adhere to the Cartesian coordinate system, in that the origin (0,0) should be in the lower left. In the graph above it is in the upper left, resulting in the plot being presented upside down. [Note that it was necessary to set the origin to 'lower' to inform imshow() of the switch.]

    Finally, z_func() differs from the displayed title. In the actual function, due to the parentheses, the negative sign is distributed over x-squared *and* y-cubed, rather than just applying to x-squared, as the title depicts.

    ReplyDelete
  9. With respect to the 3D plot, it is helpful to label its axes and to set its elevation and azimuth angles. It's also useful to reverse the order of the x-coordinates [from the call to arange()] so that the Cartesian orientation is preserved. The code I used to achieve these things appears below:

    def z_func(x,y):
    return (1-(x**2+y**3))*exp(-(x**2+y**2)/2)

    x = arange(3.0,-3.0,-0.1)
    y = arange(-3.0,3.0,0.1)
    X,Y = meshgrid(x, y) # grid of point
    Z = z_func(X, Y) # evaluation of the function on the grid
    fig = plt.figure()
    ax = fig.gca(projection='3d')
    surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=cm.RdBu,linewidth=0, antialiased=False)

    ax.zaxis.set_major_locator(LinearLocator(10))
    ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))

    ax.set_xlabel('x-axis')
    ax.set_ylabel('y-axis')
    ax.set_zlabel('z-axis')
    ax.view_init(elev=25, azim=-120)

    fig.colorbar(surf, shrink=0.5, aspect=5)
    plt.show()

    ReplyDelete
  10. It seems that in the posted contour map image the arange() had a different stepsize, thereby accounting for the different number of elements it displays. Something like:

    x = y = arange(-3.0, 3.0, 0.024)

    ReplyDelete