John Mee John Mee - 1 month ago 12
Python Question

How to arrange the data for a heatmap pcolormesh?

All the examples I've been turning up use random data. My problem is arranging my actual data to fit the solution.

I'm trying to create a heatmap/scatterplot of x,y,z: x and y are the position, whilst z is the color.
They're in three arrays of equal length.

X = [-0.11, -0.06, -0.07, -0.12, ...]
Y = [0.09, 0.13, 0.17, 0.09, ...]
Z = [0.38, 0.37, 0.44, 0.33, ...]


The pcolormesh documentation doesn't seem to describe what "C" is other than to say it can be a "masked array". Sadly I don't know what that is either (yet).

How do I turn my three arrays into whatever it's looking for? I tried sticking them into a numpy array and passing that, which quieted the error about no "shape", but a three dimensional array doesn't seem to be what it's looking for.

Answer

pcolormesh wants three 2-D arrays.

your X and Y are fine, but need to be run through numpy's meshgrid function, i.e.:

import numpy as np
X = [-0.11, -0.06, -0.07, -0.12, ...]
Y = [0.09, 0.13, 0.17, 0.09, ...]
xx, yy = np.meshgrid(X, Y)

then you just need to get Z to be in a shape that is the same as xx and yy and you'll be all set.

For the scatter plot, X, Y, and, Z are fine:

import matplotlib.pyplot as plt
X = [-0.11, -0.06, -0.07, -0.12, ...]
Y = [0.09, 0.13, 0.17, 0.09, ...]
Z = [0.38, 0.37, 0.44, 0.33, ...]
fig, ax = plt.subplots()
ax.scatter(X, Y, c=Z)
plt.show()
Comments