python - Plotting a heat map from three lists: X, Y, Intensity

Question

I don't get how to create a heatmap (or contour plot) when I have x, y, intensity. I have a file which looks like this:

0,1,6
0,2,10
....

So far:

with open('eye_.txt', 'r') as f:
        for line in f:
                for word in line.split():
                        l = word.strip().split(',')
                        x.append(l[0])
                        y.append(l[1])
                        z.append(l[2])

Tried using pcolormesh but it wants a shape object and I'm unsure how to convert these lists into a NumPy array.

I tried:

i,j = np.meshgrid(x,y)
arr = np.array(z)
plt.pcolormesh(i,j,arr)
plt.show()

It tells me that:

IndexError: too many indices

Can someone stop me from bashing my head against a keyboard, please?

Answer 1

OK, there's a few steps to this.

First, a much simpler way to read your data file is with numpy.genfromtxt. You can set the delimiter to be a comma with the delimiter argument.

Next, we want to make a 2D mesh of x and y, so we need to just store the unique values from those to arrays to feed to numpy.meshgrid.

Finally, we can use the length of those two arrays to reshape our z array.

(NOTE: This method assumes you have a regular grid, with an x, y and z for every point on the grid).

For example:

import matplotlib.pyplot as plt
import numpy as np

data = np.genfromtxt('eye_.txt',delimiter=',')

x=data[:,0]
y=data[:,1]
z=data[:,2]

## Equivalently, we could do that all in one line with:
# x,y,z = np.genfromtxt('eye_.txt', delimiter=',', usecols=(0,1,2))

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

Z=z.reshape(len(y),len(x))

plt.pcolormesh(X,Y,Z)

plt.show()