2D grid data visualization in Python

I need to visualize some data. It's basic 2D grid, whe开发者_C百科re each cell have float value. I know how to e.g. assign color to value and paint grid in OpenCV. But the point here is that there are so many values so it's nearly impossible to do that. I am looking for some method, where I could use gradient. For example value -5.0 will be represented by blue, 0 - black, and +5.0 as red. Is there any way to do that in Python?

Here is sample data I am talking about

        A       B       C        D
A    -1.045    2.0     3.5    -4.890
B    -5.678    3.2     2.89    5.78

---

Matplotlib has the imshow method for plotting arrays:

import matplotlib as mpl
from matplotlib import pyplot
import numpy as np

# make values from -5 to 5, for this example
zvals = np.random.rand(100,100)*10-5

# make a color map of fixed colors
cmap = mpl.colors.ListedColormap(['blue','black','red'])
bounds=[-6,-2,2,6]
norm = mpl.colors.BoundaryNorm(bounds, cmap.N)

# tell imshow about color map so that only set colors are used
img = pyplot.imshow(zvals,interpolation='nearest',
                    cmap = cmap,norm=norm)

# make a color bar
pyplot.colorbar(img,cmap=cmap,
                norm=norm,boundaries=bounds,ticks=[-5,0,5])

pyplot.show()

This is what it looks like:

The details for the color bar setup were taken from a matplotlib example: colorbar_only.py. It explains that the number of boundaries need to be one larger then then number of colors.

EDIT

You should note, that imshow accepts the origin keyword, which sets the where the first point is assigned. The default is 'upper left', which is why in my posted plot the y axis has 0 in the upper left and 99 (not shown) in the lower left. The alternative is to set origin="lower", so that first point is plotted in the lower left corner.

EDIT 2

If you want a gradient and not a discrete color map, make a color map by linearly interpolating through a series of colors:

fig = pyplot.figure(2)

cmap2 = mpl.colors.LinearSegmentedColormap.from_list('my_colormap',
                                           ['blue','black','red'],
                                           256)

img2 = pyplot.imshow(zvals,interpolation='nearest',
                    cmap = cmap2,
                    origin='lower')

pyplot.colorbar(img2,cmap=cmap2)

fig.savefig("image2.png")

This produces:

EDIT 3

To add a grid, as shown in this example, use the grid method. Setting the grid color to 'white' works well with the colors used by the colormap (ie the default black does not show up well).

pyplot.grid(True,color='white')

Including this before the savefig call produces this plot (made using 11x11 grid for clarity):

There are many options for grid, which are described in the matplotlib documentation. One you might be interested in is linewidth.

---

How about using matplotlib?

from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FixedLocator, FormatStrFormatter
import matplotlib.pyplot as plt
import numpy as np

fig = plt.figure()
ax = Axes3D(fig)

Z = np.array([[-1.045, 2.0, 3.5, -4.890],
              [-5.678, 3.2, 2.89, 5.78]])

X = np.zeros_like(Z)
X[1,:] = 1
Y = np.zeros_like(Z)
Y[:,1] = 1
Y[:,2] = 2
Y[:,3] = 3

surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=cm.jet,
        linewidth=0, antialiased=False)
ax.set_zlim3d(-10.0, 10.0)

ax.w_zaxis.set_major_locator(LinearLocator(10))
ax.w_zaxis.set_major_formatter(FormatStrFormatter('%.03f'))

m = cm.ScalarMappable(cmap=cm.jet)
m.set_array(Z)
fig.colorbar(m)

plt.show()

This shows: