How to plot 3D function as 2D colormap in python?

Question

Are there any python libraries that will let me plot z = f(x,y) where z is represented as the color in a densely rasterized image (as opposed to the color of a bunch of scatterplot points) ? If so, what function do I use?

It looks like some of the contour functions in matplotlib.pyplot come close to what I want, but they draw contour lines and I don't want that.

Solution 1

here's a concrete simple example (works also for functions which can't take matrix arguments for x and y):

# the function to be plotted
def func(x,y):    
    # gives vertical color bars if x is horizontal axis
    return x

import pylab

# define the grid over which the function should be plotted (xx and yy are matrices)
xx, yy = pylab.meshgrid(
    pylab.linspace(-3,3, 101),
    pylab.linspace(-3,3, 111))

# indexing of xx and yy (with the default value for the
# 'indexing' parameter of meshgrid(..) ) is as follows:
#
#   first index  (row index)    is y coordinate index
#   second index (column index) is x coordinate index
#
# as required by pcolor(..)

# fill a matrix with the function values
zz = pylab.zeros(xx.shape)
for i in range(xx.shape[0]):
    for j in range(xx.shape[1]):
        zz[i,j] = func(xx[i,j], yy[i,j])

# plot the calculated function values
pylab.pcolor(xx,yy,zz)

# and a color bar to show the correspondence between function value and color
pylab.colorbar()

pylab.show() 

Solution 2

Take a look at the documentation for pcolor or imshow in matplotlib.

Another good place to start is take a look at the matplotlib gallery and see if there is a plot type that matches what you are looking for and then use the sample code as a jumping off point for your own work:

http://matplotlib.sourceforge.net/gallery.html

Solution 3

To give credit where it's due: this is only a slight variation on Andre Holzner's answer. Please upvote him if you must!

import pylab

def f(x, y):
    return pylab.cos(x) + pylab.sin(y)

xx = pylab.linspace(-5, 5, 100)
yy = pylab.linspace(-5, 5, 100)
zz = pylab.zeros([len(xx), len(yy)])

for i in xrange(len(xx)):
    for j in xrange(len(yy)):
        zz[j, i] = f(xx[i], yy[j])

pylab.pcolor(xx, yy, zz)
pylab.show()

The syntax is perhaps easier to read with the strict minimum of array dimensions and indices. It relies on the following point (quoted from the doc).

If either or both of X and Y are 1-D arrays or column vectors, they will be expanded as needed into the appropriate 2-D arrays, making a rectangular grid.

Solution 4

To expand my comment above, here are some possible ways of computing a function on a grid

boffi@debian:~/Documents/tmp$ cat grid.py 
import numpy as np

def z(x,y):
  return np.sin(np.sqrt(x*x+y*y))

x = np.linspace(-1,1,11)
y = np.linspace(-2,2,21)

# naive

Z0 = np.zeros((len(y), len(x)))
for i, X in enumerate(x):
    for j, Y in enumerate(y):
        Z0[j,i] = z(X,Y)

# trampoline on a double list comprehension,
# it is possibly faster, sure it uses more memory

Z1 = np.array([[z(X,Y) for X in x] for Y in y])

# numpy has meshgrid, 
# meshgrid uses twice memory as the result matrix but
# if used _correctly_ it's FAST

X, Y = np.meshgrid(x, y)

# numpy can avoid you explicit looping,
# but if you are so inclined...

Z2 = np.zeros((len(y), len(x)))
for r in range(len(y)):
    for c in range(len(x)):
        Z2[r, c] = z(X[r, c], Y[r, c])

# numpy has ufuncs, and
# t h i s   i s   t h e   w a y   t o   g o

Z3 = z(X, Y)

# numpy has broadcasting (it's slower than Z = z(X, Y), less memory)

Z4 = z(x, y[:,None])

# note that x is still a _row_ of numbers, indexed by _columns_,
# while y[:,None] is now a _column_ of numbers, indexed by _rows_,
# so that Z4[row,column] <-- z(x[column], y[row])

# a bit of testing

# in previous answers, Z2 (i.e., explicit loops)
# is the preferred method --- here we show that the other four
# possible methods give you exactly the same result

print np.all(Z2==Z0)
print np.all(Z2==Z1)
print np.all(Z2==Z3)
print np.all(Z2==Z4)
boffi@debian:~/Documents/tmp$ python2 grid.py 
True
True
True
True
boffi@debian:~/Documents/tmp$ 

View more solutions

Comments

[marcv81]

The loop code is wrong: you're iterating through the wrong objects. It only works because xx and yy are the same length. Try with other values instead of 101 to see what I mean.

[marcv81]

It does appear to work, thanks for taking the time to revisit your answer. I personally prefer the somewhat simpler approach xx = pylab.linspace(-3,3, 101), same for yy, then iterate with for i in xrange(len(xx)) and for j in xrange(len(yy)), and use zz[i, j] = func(xx[i], yy[j]).

[Andre Holzner]

it's actually rather for i in xrange(len(yy)) instead of xx if you use zz[i,j] (see the added comments in the code). From the help string of pcolor(..): Note that the column index corresponds to the *x*-coordinate, and the row index corresponds to *y*

[marcv81]

I posted a separate answer as I felt it was the best way to have enough space to illustrate my point.

[gboffi]

The double loop, yours and Andre's, is unnecessary and nefarious because you can simply write zz = func(xx, yy) that's possibly orders of magnitude faster.

[marcv81]

I actually tried something like that but could not get it to work! With or without initializing zz with zeros (line before the nested loops) does not make any difference. Can you get it to run?

[gboffi]

See my pseudo answer below

[marcv81]

OK, I was missing meshgrid. Thanks!