Create a stacked 2D histogram using different weights

The weights parameter expects an array of the same length as x and y. np.histogram2d. It will not broadcast a constant value, so even though the mass is the same for each call to np.histogram2d, you still must use something like

weights=np.ones_like(x)*mass

Now, one problem you may run into if you use bin=nbin is that the bin edges, xedges, yedges may change depending on the values of x and y that you pass to np.histogram2d. If you naively add heatmaps together, the final result will accumulate particle density in the wrong places.

So if you want to call np.histogram2d more than once and add partial heatmaps together, you must determine in advance where you want the bin edges.

For example:

import numpy as np
import itertools as IT
import matplotlib.pyplot as plt
N = 50
nbin = 10

xs = [np.array([i,i,i+1,i+1]) for i in range(N)]
ys = [np.array([i,i+1,i,i+1]) for i in range(N)]
masses = np.arange(N)

heatmap = 0
xedges = np.linspace(0, N, nbin)
yedges = np.linspace(0, N, nbin)

for x, y, mass in IT.izip(xs, ys, masses):
    hist, xedges, yedges = np.histogram2d(
        x, y, bins=[xedges, yedges], weights=np.ones_like(x)*mass)
    heatmap += hist

extent = [xedges[0], xedges[-1], yedges[0], yedges[-1]]
heatmap = np.flipud(np.rot90(heatmap))
fig, ax = plt.subplots()
ax.imshow(heatmap, extent=extent, interpolation='nearest')
plt.show()

yields