CODEWITHSUNDEEP
pythonmatplotlibcolorscolormap
AKG
AKG

Reputation: 2459

matplotlib pcolormesh discrete colors

from matplotlib.pyplot import *
import numpy as np

x      = np.linspace( 1., 4., 2 )
y      = np.linspace( 1., 2., 2 )
x,y    = np.meshgrid( x, y )
z      = np.array( [ [ 0.5, 1.5 ], [ 2.5, 3.5 ] ] )
cmap   = matplotlib.colors.ListedColormap( [ 'r', 'y', 'g', 'b' ] )
bounds = [ 0.,  1.,  2.,  3.,  4. ]
ticks  = [ 0.5, 1.5, 2.5, 3.5 ]
norm   = matplotlib.colors.BoundaryNorm( bounds, cmap.N )
fig    = figure( figsize = ( 20, 10 ) )
ax     = fig.add_subplot( 111 )
ax.pcolormesh( x, y, z, cmap = cmap, norm = norm )
ax.axis( [ x.min(), x.max(), y.min(), y.max() ] )
ax.set_xlabel( 'x' )
ax.set_ylabel( 'y' )
m      = cm.ScalarMappable( cmap = cmap )
m.set_array( z )
cbar   = fig.colorbar( m, norm = norm, boundaries = bounds, aspect = 20, ticks = ticks )
cbar.set_ticklabels( [ 'A', 'B', 'C', 'D' ] )
show()

According to the values of z I would like to have a 4-field colormesh. As the first z-entry 0.5 is in the interval [ 0., 1. ], I want it to be red, the second yellow and so on.

The colorbar displays already perfectly what I want to, but the plot itself is wrong.

Maybe I got something wrong with the BoundaryNorm?

Upvotes: 1

Views: 5360

Answers (1)

hitzg
hitzg

Reputation: 12691

This should do it:

from matplotlib.pyplot import *
import numpy as np

x = np.linspace(1.,4.,3)
y = np.linspace(1.,2.,3)
x,y = np.meshgrid(x,y)
z = np.array([[0.5,1.5],[2.5,3.5]])
cmap = matplotlib.colors.ListedColormap(['r', 'y', 'g', 'b'])
bounds = [0.,1.,2.,3.,4.]
ticks = [0.5,1.5,2.5,3.5]
norm = matplotlib.colors.BoundaryNorm(bounds, cmap.N)
fig = figure(figsize = (20,10))
ax = fig.add_subplot(111)
p = ax.pcolormesh(x,y,z, cmap=cmap, norm = norm, vmin=0, vmax=4)
ax.axis([x.min(),x.max(),y.min(),y.max()])
ax.set_xlabel('x')
ax.set_ylabel('y')
cbar = fig.colorbar(p, norm = norm, boundaries = bounds, 
           aspect = 20, ticks = ticks)
cbar.set_ticklabels(['A','B','C','D'])
show()

x and y define the corners of the squares. so they should be 3x3. And you should define the vmin/vmax arguments of pcolormesh.

I also noticed that you dont need to specify the bounds and the norm explicitly if you use vmin/vmax (at least in this simple example), which makes the entire thing quite a bit shorter:

from matplotlib.pyplot import *
import numpy as np

x = np.linspace(1.,4.,3)
y = np.linspace(1.,2.,3)
x,y = np.meshgrid(x,y)
z = np.array([[0.5,1.5],[2.5,3.5]])
cmap = matplotlib.colors.ListedColormap(['r', 'y', 'g', 'b'])
ticks = [0.5,1.5,2.5,3.5]
fig = figure(figsize = (20,10))
ax = fig.add_subplot(111)
p = ax.pcolormesh(x,y,z, cmap=cmap, vmin=0, vmax=4)
ax.axis([x.min(),x.max(),y.min(),y.max()])
ax.set_xlabel('x')
ax.set_ylabel('y')
cbar = fig.colorbar(p, aspect = 20, ticks=ticks)
cbar.set_ticklabels(['A','B','C','D'])
show()

Upvotes: 1

Related Questions