heatmaps.py 文件源码

def plot_heatmap(ax, xpoints, ypoints, nbins, title=None, maxcount=None):
    ''' Plot a heatmap of the given data on on the given axes. '''
    # imshow expects y,x for the image, but x,y for the extents,
    # so we have to manage that here...
    bins = np.concatenate( (np.arange(0,1.0,1.0/nbins), [1.0]) )
    heatmap, yedges, xedges = np.histogram2d(ypoints, xpoints, bins=bins)
    extent = [xedges[0],xedges[-1], yedges[0], yedges[-1]]
    # make sure we always show the full extent of the tank and the full extent of the data,
    # whichever is wider.
    ax.set_xlim(min(0, xedges[0]), max(1, xedges[-1]))
    ax.set_ylim(min(0, yedges[0]), max(1, yedges[-1]))
    if title:
        ax.set_title(title)
    if maxcount is not None:
        norm = Normalize(0, maxcount)
    else:
        norm = None
    return ax.imshow(heatmap, extent=extent, cmap=plt.get_cmap('hot'), origin='lower', interpolation='nearest', norm=norm)