Bilinear binning with CuPy

bilinear_bincount_cupy.py

def bilinear_bincount_cupy(points, intensities, subpixel=1):
    """Bilinear weighting of points onto a grid.
    Extent of grid given by min and max of points in each dimension
    points should be a cupy array of shape (N, 2)
    intensity should be a cupy array of shape (N,)
    """

    points = subpixel * points
    floor = cp.floor(points)
    ceil = floor + 1
    floored_indices = cp.array(floor, dtype=int)
    low0, low1 = floored_indices.min(0)
    high0, high1 = floored_indices.max(0)
    floored_indices = floored_indices - cp.array([low0, low1])
    shape = cp.array([high0 - low0 + 2, high1-low1 + 2])

    upper_diff = ceil - points
    lower_diff = points - floor

    w1 = upper_diff[:, 0] * upper_diff[:, 1]
    w2 = upper_diff[:, 0] * lower_diff[:, 1]
    w3 = lower_diff[:, 0] * upper_diff[:, 1]
    w4 = lower_diff[:, 0] * lower_diff[:, 1]

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

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

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