Matplotlib: How to make a contour plot?

Question

I have a function that I want to do a contour plot: mixed_model_pdf([[x,y]]).

It takes points of [x,y], instead of X,Y seperately as input. That is, mixed_model_pdf([[-5,-5]]), or mixed_model_pdf([[-5,-5],[-5,-4]])

But to do plt.contourf(X, Y, Z), The Z value should be placed in specific order, how can I do that?

Now, I can get away with this:

X = Y = np.arange(-30, 31, 1)
N, M = len(X), len(Y)
Z = np.zeros((N, M))
for i, (x,y) in enumerate(product(X,Y)):
    Z[np.unravel_index(i, (N,M))] = mixed_model_pdf([[x,y]])
Z = Z.T    
X, Y = np.meshgrid(X, Y)

# Plot 2d prob density
fig, ax = plt.subplots(figsize=(8, 8), dpi=80)
ax.set_aspect('equal')
plt.contourf(X, Y, Z, 10, alpha=.75, cmap='jet')

But this looks really ugly, it transpose the Z, and still need X, Y = np.meshgrid(X, Y).

What would be a better way to do it?

UPDATE:

Now, I can make the code a lot more clear:

X = Y = np.arange(-30, 31, 1)
XX,YY=np.meshgrid(x,y)
coords=np.array((XX.ravel(), YY.ravel())).T  
Z = mixed_model_pdf(coords).reshape(61,61)

fig, ax = plt.subplots(figsize=(8, 8), dpi=80)
ax.set_aspect('equal')
plt.contourf(X, Y, Z, 10, alpha=.75, cmap='jet')

But how can I shorten the .reshape()? It's calculated from Y = np.arange(-30, 31, 1), and I think it may be automated?

Answer 1

Collecting my comments in an answer.

You can call plt.contourf in different ways. The docs make this clear:

X and Y must both be 2-D with the same shape as Z, or they must both be 1-D such that len(X) is the number of columns in Z and len(Y) is the number of rows in Z.

You can get the docs from the interactive Python prompt:

>>> help(plt.contourf)

or, if you use IPython (Notebook), as you should ;):

In [1]: plt.contourf?

You can use:

Z = mixed_model_pdf(coords).reshape(Y.size, Y.size)

to avoid hard coding the 61.