Contour plotting orbitals in pyquante2 using matplotlib
Question: I'm currently writing line and contour plotting functions for my
[PyQuante](https://github.com/rpmuller/pyquante2) quantum chemistry package
using matplotlib. I have some great functions that evaluate basis sets along a
(npts,3) array of points, e.g.

    
    
    from somewhere import basisset, line
    bfs = basisset(h2) # Generate a basis set
    points = line((0,0,-5),(0,0,5)) # Create a line in 3d space
    bfmesh = bfs.mesh(points)
    for i in range(bfmesh.shape[1]):
        plot(bfmesh[:,i])
    

This is fast because it evaluates all of the basis functions at once, and I
got some great help from stackoverflow
[here](http://stackoverflow.com/questions/17396164/numpythonic-way-to-
make-3d-meshes-for-line-plotting) and
[here](http://stackoverflow.com/questions/17391052/compute-square-distances-
from-numpy-array) to make them extra-nice.

I would now like to update this to do contour plotting as well. The slow way
I've done this in the past is to create two one-d vectors using linspace(),
mesh these into a 2D grid using meshgrid(), and then iterating over all xyz
points and evaluating each one:

    
    
    f = np.empty((50,50),dtype=float)
    xvals = np.linspace(0,10)
    yvals = np.linspace(0,20)
    z = 0
    for x in xvals:
        for y in yvals:
            f = bf(x,y,z)
    X,Y = np.meshgrid(xvals,yvals)
    contourplot(X,Y,f)
    

(this isn't real code -- may have done something dumb)

What I would like to do is to generate the mesh in more or less the same way I
do in the contour plot example, "unravel" it to a (npts,3) list of points,
evaluate the basis functions using my new fast routines, then "re-ravel" it
back to X,Y matrices for plotting with contourplot.

The problem is that I don't have anything that I can simply call .ravel() on:
I either have 1d meshes of xvals and yvals, the 2D versions X,Y, and the
single z value.

Can anyone think of a nice, pythonic way to do this?


Answer: If you can express `f` as a function of `X` and `Y`, you could avoid the
Python `for-loop`s this way:

    
    
    import matplotlib.pyplot as plt
    import numpy as np
    
    
    def bf(x, y):
        return np.sin(np.sqrt(x**2+y**2))
    
    xvals = np.linspace(0,10)
    yvals = np.linspace(0,20)
    X, Y = np.meshgrid(xvals,yvals)
    f = bf(X,Y)
    plt.contour(X,Y,f)
    plt.show()
    

yields

![enter image description here](http://i.stack.imgur.com/kFWHM.png)

