How to plot 3D data as 2D grid colormap in Python?

Question

I have 3D data as the columns of a numpy array (that is to say that array[0] = [0,0,0], etc.), e.g.

X     Y     Z

0     0     0
0     1     10
1     0     20
1     1     30

I would like to plot this so that each (X,Y) co-ordinate has a square centered on the co-ordinate, with a colorbar from (e.g.) 0 to 30 showing the Z value.

I would then like to overlay some contour lines, but the first part of the question is most important.

There is help for people who have already-gridded data, but I am not sure of the best matplotlib routine to call for my column data. Also, this is for scientific publication, so needs to be of a good quality and look! Hope someone can help!

Answer 1

You can use griddata from matplotlib.mlab to grid your data properly.

import numpy as np
from matplotlib.mlab import griddata

x = np.array([0,0,1,1])
y = np.array([0,1,0,1])
z = np.array([0,10,20,30])
xi = np.arange(x.min(),x.max()+1)
yi = np.arange(y.min(),y.max()+1)
ar = griddata(x,y,z,xi,yi)

# ar is now
# array([[  0.,  20.],
#        [ 10.,  30.]])

The choice of the mapped xi and yi points is up to you, and they do not have to be integers as griddata can interpolate for you.