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 scat
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
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$
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()
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.