Python: 2D color map with imshow
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()
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')
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()
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.
Related videos on Youtube
user1830663
Updated on June 05, 2021Comments
-
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
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 over 9 yearsSearch the matplotlib
imshow
docs and here at StackOverflow for information about theextent
argument ofimshow
. For example, stackoverflow.com/questions/6999621/… and matplotlib.org/api/pyplot_api.html#matplotlib.pyplot.imshow -
unutbu over 9 yearsYou'll want to add the
extent
parameter to the call tocontour
as well. -
Dimitris almost 6 yearsA 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 over 9 yearsto be consistent with
imshow
you wantextent=(-3,3,3,-3)
-
user1830663 over 9 yearsThe contour lines are upside down. "extent" for contour should be (-3,3,-3,3), opposite of imshow.
-
velenos14 over 2 yearshello, 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 over 2 yearsmanaged to do it via fig,
ax = plt.subplots()
pos = ax.pcolor(xlocs, ylocs, Efieldvalues.T, cmap="RdBu")
fig.colorbar(pos, ax=ax)
-
MaxNoe over 2 yearsI updated the answer with the two current best solutions ;)