CODEWITHSUNDEEP
pythonarraysnumpyheatmap
jokerp
jokerp

Reputation: 157

Create a heat map out of three 1D arrays

I want to create a heatmap out of 3 1dimensional arrays. Something that looks like this:

enter image description here

Up to this point, I was only able to create a scatter plot where the markers have a different color and marker size depending on the third value:

My code:

xf = np.random.rand(1000)
yf = np.random.rand(1000)
zf = 1e5*np.random.rand(1000)    
ms1 = (zf).astype('int')


from matplotlib.colors import LinearSegmentedColormap
# Remove the middle 40% of the RdBu_r colormap
interval = np.hstack([np.linspace(0, 0.4), np.linspace(0.6, 1)])
colors   = plt.cm.RdBu_r(interval)
cmap     = LinearSegmentedColormap.from_list('name', colors)
col      = cmap(np.linspace(0,1,len(ms1)))
#for i in range(len(ms1)):
plt.scatter(xf, yf, c=zf, s=5*ms1/1e4, cmap=cmap,alpha=0.8)#, norm =matplotlib.colors.LogNorm())
ax1 =plt.colorbar(pad=0.01)

is giving me this result:

enter image description here

Any idea how I could make it look like the first figure?

Essentially what I want to do is find the average of the z value for groups of the x and y arrays

Upvotes: 3

Views: 1430

Answers (1)

bb1
bb1

Reputation: 7863

I think the functionality you are looking for is provided by scipy.stats.binned_statistic_2d. You can use it to organize values of xf and yf arrays into 2-dimensional bins, and compute the mean of zf values in each bin:

import numpy as np
from scipy import stats

np.random.seed(0)
xf = np.random.rand(1000)
yf = np.random.rand(1000)
zf = 1e5 * np.random.rand(1000)

means = stats.binned_statistic_2d(xf,
                                  yf,
                                  values=zf,
                                  statistic='mean',
                                  bins=(5, 5))[0]

Then you can use e.g. seaborn to plot a heatmap of the array of mean values:

import matplotlib.pyplot as plt
import seaborn as sns

plt.figure(figsize=(10, 8))
sns.heatmap(means,
            cmap="Reds_r",
            annot=True,
            annot_kws={"fontsize": 16},
            cbar=True,
            linewidth=2,
            square=True)
plt.show()

This gives:

enter image description here

Upvotes: 1

Related Questions