Python: 2D color map with imshow

Question

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?

Search the matplotlib `imshow` docs and here at StackOverflow for information about the `extent` argument of `imshow`. For example, http://stackoverflow.com/questions/6999621/how-to-use-extent-in-matplotlib-pyplot-imshow and http://matplotlib.org/api/pyplot_api.html#matplotlib.pyplot.imshow — Warren Weckesser, Jul 22 '14 at 21:05
You'll want to add the `extent` parameter to the call to `contour` as well. — unutbu, Jul 22 '14 at 21:22
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}$')` — Dimitris, Dec 28 '17 at 12:27

MaxNoe · Answer 1 · 2021-06-04T11:24:11.820

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

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

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, Jun 04 '21 at 08:10
managed to do it via fig, ```ax = plt.subplots()``` ```pos = ax.pcolor(xlocs, ylocs, Efieldvalues.T, cmap="RdBu")``` ```fig.colorbar(pos, ax=ax)``` — velenos14, Jun 04 '21 at 08:14

unutbu · Accepted Answer · 2014-07-24T16:14:23.817

3

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

edited Jul 24 '14 at 16:14

answered Jul 22 '14 at 21:24

unutbu

842,883
184
1,785
1,677

to be consistent with `imshow` you want `extent=(-3,3,3,-3)` – Gabriel Jul 22 '14 at 21:26
The contour lines are upside down. "extent" for contour should be (-3,3,-3,3), opposite of imshow. – user1830663 Jul 24 '14 at 15:04

Gabriel · Answer 3 · 2014-07-22T21:27:49.390

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.

Python: 2D color map with imshow

3 Answers3

Linked