Bilinear binning, supports numpy and CuPy inputs using the np.array(..., like=arr) synax

bilinear_binning.py

def bilinear_binning(points, intensities, subpixel=1, gaussian_blur=False):
    """Bilinear weighting of points onto a grid.
    Extent of grid given by min and max of points in each dimension
    points should be an array of shape (N, 2)
    intensity should be an array of shape (N,)
    subpixel will increase the gridsize by its factor
    gaussian_blur: blur the binned intensity and weighting images before they are divided, avoiding divide-by-zero warnings
    
    TODO: Give a known grid as input
    """

    points = subpixel * points
    floor = np.floor(points)
    ceil = floor + 1
    floored_indices = floor.astype(int)
    minimum_indices = floored_indices.min(0)
    maximum_indices = floored_indices.max(0)

    floored_indices = floored_indices - minimum_indices
    shape = (floored_indices.max(0) - floored_indices.min(0) + 2)
    shape = tuple(s.item() for s in shape)

    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 = np.array([[0, 0], [0, 1], [1, 0], [1, 1]], like=points)
    indices = floored_indices[:, None] + shifts
    indices = (indices * np.array([shape[1], 1], like=points)).sum(-1)
    weights = np.array([w1, w2, w3, w4], like=points).T

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

    intens_image = intens_bins.reshape(shape)
    weight_image = weight_bins.reshape(shape)
    
    if gaussian_blur:
        if gaussian_blur == True:
            gaussiann_blur = 1
        intens_image = gaussian_filter(intens_image, gaussian_blur)
        weight_image = gaussian_filter(weight_image, gaussian_blur)

    
    return intens_image/weight_image