Search code examples
pythonplotimshow

Change x- and y-numbering in imshow

I would like to plot a function of two variables in python. Similar to this article, we can obtain an output like

enter image description here

using this code:

from numpy import exp,arange
from pylab import meshgrid,cm,imshow,contour,clabel,colorbar,axis,title,show
from matplotlib import pyplot

# 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)
z = [[0] * y.__len__() for i in range(x.__len__())]
for i in range(0, x.__len__()):
    for j in range(0, y.__len__()):
        z[j][i] = z_func(x[i], y[j])
   
im = imshow(z,cmap=cm.RdBu, extent = [-3, 3, -3, 3], interpolation = "none", origin='lower') # 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()

As you can see, the x- and y-labels go from 0 to 59 (which is the count of elements in x and y). How can I correct these values such that they range from -3 to 3?

A minor sub-question: Why do I need to "transpose" in z[j][i] = z_func(x[i], y[j])? Does Python treat the first dimension as "column" and the second as "row"?

Solution

  • You're trying to plot both the z-function and the countour plots. You need to add the "extent" parameter to matplotlib.pyplot.countour plot too.

    cset = contour(z, arange(-1,1.5,0.2), 
               extent = [-3, 3, -3, 3], 
               linewidths = 2, 
               cmap = cm.Set2)