Home > OS >  Heatmap plot of array of arrays
Heatmap plot of array of arrays

Time:11-03

I have a numpy array composed of the following arrays (as example):

my_array = [[10,12,0],[10,13,1],[11,12,1],[11,13,0],[12,12,1],[12,13,0]]

Where items are [x,y,value]

I want to plot a heatmap or similar with x values in X, y values i Y, and color depending on value (0 or 1). I want also label color for 0 as 'A points' and color for 1 as 'B points'.

I have tried several heatmaps but can't imagine how to solve this three value aray to be plotted.

Code tried:


xvals = [10,11,12]
yvals = [12,13]

zvals = [[10,12,0],[10,13,1],[11,12,1],[11,13,0],[12,12,1],[12,13,0]]


heatmap, ax = plt.subplots()

im = ax.imshow(zvals,cmap='inferno',extent=[xvals[0],xvals[2],yvals[0],yvals[1]],interpolation='nearest',origin='lower',aspect='auto')
ax.set(xlabel='some x', ylabel='some y')

cbar = heatmap.colorbar(im)
cbar.ax.set_ylabel('stuff')

heatmap.savefig('heatmap.png')

Any help?

Thank you

CodePudding user response:

you could create a 2d list or numpy 2d array and plt that.

import matplotlib.pyplot as plt
xvals = [10,11,12]
yvals = [12,13]
zvals = [[10,12,0],[10,13,1],[11,12,1],[11,13,0],[12,12,1],[12,13,0]]

im = [[0 for j in range(len(xvals))] for i in range(len(yvals))]

for val in zvals:
    im[val[1] - yvals[0]][val[0] - xvals[0]] = val[-1]

plt.imshow(im, cmap='inferno')
plt.xticks(range(len(xvals)), xvals)
plt.yticks(range(len(yvals)), yvals)

and results:

enter image description here

  • Related