Occupancy normalized hexbin

occupancy_normalized_hexbin.py

def occupancy_normalized_hexbin(x, y, all_x, all_y, ax=None,
                                gridsize=(3, 3), **kwargs):
    '''Bins (x, y) into hexagonal grid and normalizes the
    count by the binned count of (all_x, all_y).
    
    Useful when measuring the frequency of events over time 
    and space when the time spent in each bin is not equal.
    
    Parameters
    ----------
    x : numpy.ndarray
        Grid position at event
    y : numpy.ndarray
        Grid position at event
    all_x : numpy.ndarray
        Grid position over all time. For normalization.
    all_y : numpy.ndarray
        Grid position over all time. For normalization.
    gridsize : array-like, shape (2,), optional
        Size of the grid.
    **kwargs : hexbin keyword arguments, optional
    
    Returns
    -------
    axis_handle : a `~matplotlib.collections.PolyCollection` instance
    
    '''
    if ax is None:
        ax = plt.gca()
    occupancy = ax.hexbin(all_x, all_y, gridsize=gridsize)
    plt.cla()
    occupancy_count = occupancy.get_array()
    hexbin_centers = occupancy.get_offsets()
    event_count = ax.hexbin(x, y, gridsize=gridsize).get_array()
    plt.cla()
    normalized_count = event_count / occupancy_count
    return ax.hexbin(hexbin_centers[:, 0], hexbin_centers[:, 1],
                     C=normalized_count, gridsize=gridsize, **kwargs)