Binning data into non-grid based bins

Question

I have used voronoi binning: I have the central coordinates for each bin and I want to find the mean of all the pixels contained in each bin. I just can't get my head around how to slice the numpy array to vectorise it.

This is my current code: X and Y are 1D arrays with the x and y coordinates for the centre of each bin; f in the 2d image:

import numpy as np
from scipy.spatial import KDTree

def rebin(f, X, Y):
    s = f.shape
    x_grid = np.arange(s[0])
    y_grid = np.arange(s[1])
    x_grid, y_grid = np.meshgrid(x_grid,y_grid)
    x_grid, y_grid = x_grid.flatten(),  y_grid.flatten()

    tree = KDTree(zip(X,Y))
    _, b = tree.query(zip(x_grid, y_grid))
    out = X*np.nan

    for i in range(max(b)):
        out[i] = np.nanmean(f[x_grid[b==i], y_grid[b==i]])
    return out

The for-loop is currently a huge bottleneck. There is probably a really simple answer - I just can't see it at the moment!

Answer 1

out = X*np.nan
for i in range(max(b)):
    out[i] = np.nanmean(f[x_grid[b==i], y_grid[b==i]])

can be replaced by two calls to np.bincount:

total = np.bincount(b, weights=f[x_grid, y_grid], minlength=len(X))
count = np.bincount(b, minlength=len(X))
out = total/count

or one call to stats.binned_statistic:

out, bin_edges, binnumber = stats.binned_statistic(
    x=b, values=f[x_grid, y_grid], statistic='mean', bins=np.arange(len(X)+1))

---

For example,

import numpy as np
from scipy.spatial import KDTree
import scipy.stats as stats
np.random.seed(2017)

def rebin(f, X, Y):
    s = f.shape
    x_grid = np.arange(s[0])
    y_grid = np.arange(s[1])
    x_grid, y_grid = np.meshgrid(x_grid,y_grid)
    x_grid, y_grid = x_grid.flatten(),  y_grid.flatten()

    tree = KDTree(np.column_stack((X,Y)))
    _, b = tree.query(np.column_stack((x_grid, y_grid)))

    out, bin_edges, binnumber = stats.binned_statistic(
        x=b, values=f[x_grid, y_grid], statistic='mean', bins=np.arange(len(X)+1))
    # total = np.bincount(b, weights=f[x_grid, y_grid], minlength=len(X))
    # count = np.bincount(b, minlength=len(X))
    # out = total/count
    return out

def orig(f, X, Y):
    s = f.shape
    x_grid = np.arange(s[0])
    y_grid = np.arange(s[1])
    x_grid, y_grid = np.meshgrid(x_grid,y_grid)
    x_grid, y_grid = x_grid.flatten(),  y_grid.flatten()

    tree = KDTree(np.column_stack((X,Y)))
    _, b = tree.query(np.column_stack((x_grid, y_grid)))

    out = X*np.nan
    for i in range(len(X)):
        out[i] = np.nanmean(f[x_grid[b==i], y_grid[b==i]])
    return out

N = 100
X, Y = np.random.random((2, N))
f = np.random.random((N, N))

expected = orig(f, X, Y)
result = rebin(f, X, Y)
print(np.allclose(expected, result, equal_nan=True))
# True

Answer 2

I've just found there is a cython feature complete version of KDTree called cKDTree which is much faster.