Python: 2D color map with imshow

25,157

Solution 1

Add extent=(-3, 3, 3, -3) to the call to imshow and extent=(-3, 3, -3, 3) (note the unfortunately tricky change in signs!) to the call to contour:

import numpy as np
import matplotlib.pyplot as plt


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

x = np.arange(-3.0, 3.0, 0.1)
y = np.arange(-3.0, 3.0, 0.1)
X, Y = np.meshgrid(x, y)
Z = z_func(X, Y) 


im = plt.imshow(Z, cmap=plt.cm.RdBu, extent=(-3, 3, 3, -3))  
cset = plt.contour(Z, np.arange(-1, 1.5, 0.2), linewidths=2,
                   cmap=plt.cm.Set2,
                   extent=(-3, 3, -3, 3))
plt.clabel(cset, inline=True, fmt='%1.1f', fontsize=10)
plt.colorbar(im)  

plt.title('$z=(1-x^2+y^3) e^{-(x^2+y^2)/2}$')

plt.show()

enter image description here

Solution 2

I would not use imshow for plotting a 2d function, imshow is for showing images.

The Axes labeling comes from the number of "pixels". Your values.

To plot a 2d function, use plt.pcolormesh.

This needs a grid of pixels and you should evaluate your function in the center of that grid. This will result in sharp pixels:

import numpy as np
import matplotlib.pyplot as plt


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

x = np.arange(-3.0, 3.0, 0.1)
y = np.arange(-3.0, 3.0, 0.1)
x_center = 0.5 * (x[:-1] + x[1:])
y_center = 0.5 * (y[:-1] + y[1:])

X, Y = np.meshgrid(x_center, y_center)
Z = z_func(X, Y)

# pcolormesh needs the pixel edges for x and y
# and with default flat shading, Z needs to be evaluated at the pixel center
plot = plt.pcolormesh(x, y, Z, cmap='RdBu', shading='flat')

# contour needs the centers
cset = plt.contour(X, Y, Z, cmap='gray')
plt.clabel(cset, inline=True)

plt.colorbar(plot)
plt.savefig('plot_z_flat.png')


flat shading

You can also use gouraud shading to get a smooth image, but then you have to evaluate your function on the edges like this:


import numpy as np
import matplotlib.pyplot as plt


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

x = np.arange(-3.0, 3.0, 0.1)
y = np.arange(-3.0, 3.0, 0.1)

X, Y = np.meshgrid(x, y)
Z = z_func(X, Y)

# pcolormesh needs the pixel vertices for x and y
# and with gouroud shading, Z has to be evaluated on all vertices
plot = plt.pcolormesh(X, Y, Z, cmap='RdBu', shading='gouraud')

# countour needs the center points
x_center = 0.5 * (x[:-1] + x[1:])
y_center = 0.5 * (y[:-1] + y[1:])
X, Y = np.meshgrid(x, y)
Z = z_func(X, Y)

cset = plt.contour(X, Y, Z, cmap='gray')
plt.clabel(cset, inline=True)

plt.colorbar(plot)
plt.show()

gouraud shading

Solution 3

The extent of the colorbar axis is determined by the min and max value of Z.

>>> Z.min()
-0.96365584108555036

>>> Z.max()
1.4203545446927801

The extent of the x and y axes on the plot are indices in the array Z by default. As Warren pointed out, you can change how the axes are labelled using the extent keyword argument to imshow.

extent = (-3.0,3.0, 3.0,-3.0)

or more generally

extent = (x[0],x[-1], y[-1],y[0] )
imshow( Z,cmap=cm.RdBu, extent=extent )
cset = contour(Z,arange(-1,1.5,0.2),linewidths=2,cmap=cm.Set2,extent=extent)

extent takes the low x coord, then high x, then low y, then high y. The default convention for images is for the origin of the y-axis to start in the upper left corner. This is why the last two entries in extent are "reversed" from what you might expect.

Share:
25,157

Related videos on Youtube

user1830663
Author by

user1830663

Updated on June 05, 2021

Comments

  • user1830663
    user1830663 over 2 years

    I am trying to represent a function of two variable on a two dimensional graph using color. I came across this example here:

    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()
    

    which produces enter image description here

    However, the axis scales and limits do not correspond to the real x and y data(both of which are between -3 and 3). How to make them correspond to the actual data?

    • Warren Weckesser
      Warren Weckesser over 9 years
      Search the matplotlib imshow docs and here at StackOverflow for information about the extent argument of imshow. For example, stackoverflow.com/questions/6999621/… and matplotlib.org/api/pyplot_api.html#matplotlib.pyplot.imshow
    • unutbu
      unutbu over 9 years
      You'll want to add the extent parameter to the call to contour as well.
    • Dimitris
      Dimitris almost 6 years
      A small mistake: the title does not correspond to the function being plotted. The correct version should be title('$z=(1-(x^2+y^3)) e^{-(x^2+y^2)/2}$')
  • Gabriel
    Gabriel over 9 years
    to be consistent with imshow you want extent=(-3,3,3,-3)
  • user1830663
    user1830663 over 9 years
    The contour lines are upside down. "extent" for contour should be (-3,3,-3,3), opposite of imshow.
  • velenos14
    velenos14 over 2 years
    hello, it's a long time since this response, but do you also know how can I add a vertical color bar representing the color->value mapping next to the plot itself? pcolor works for me, but I don't find how to add the colorbar column. Thanks
  • velenos14
    velenos14 over 2 years
    managed to do it via fig, ax = plt.subplots() pos = ax.pcolor(xlocs, ylocs, Efieldvalues.T, cmap="RdBu") fig.colorbar(pos, ax=ax)
  • MaxNoe
    MaxNoe over 2 years
    I updated the answer with the two current best solutions ;)