Bilinear Binning with Numpy

bilinear_bincount_numpy.py

def bilinear_bincount_numpy(points, intensities):
    """Bilinear weighting of points onto a grid.
    Extent of grid given by min and max of points in each dimension
    points should have shape (N, 2)
    intensity should have shape (N,)
    """
    floor = np.floor(points)
    ceil = floor + 1
    floored_indices = np.array(floor, dtype=int)
    low0, low1 = floored_indices.min(0)
    high0, high1 = floored_indices.max(0)
    floored_indices = floored_indices - (low0, low1)
    shape = (high0 - low0 + 2, high1-low1 + 2)

    upper_diff = ceil - points
    lower_diff = points - floor

    w1 = np.prod((upper_diff), axis=1)
    w2 = upper_diff[:,0]*lower_diff[:,1]
    w3 = lower_diff[:,0]*upper_diff[:,1]
    w4 = np.prod((lower_diff), axis=1)

    shifts = np.array([[0, 0], [0, 1], [1, 0], [1, 1]])
    indices = floored_indices[:, None] + shifts
    indices = (indices * (shape[1], 1)).sum(-1)
    weights = np.array([w1, w2, w3, w4]).T

    weight_bins = np.bincount(indices.flatten(), weights=weights.flatten(), minlength = np.prod(shape))
    intens_bins = np.bincount(indices.flatten(), weights=(intensities[:, None]*weights).flatten(), minlength = np.prod(shape))

    weight_image = weight_bins.reshape(shape)
    intens_image = intens_bins.reshape(shape)
    return intens_image, weight_image