Calculate Distance Between Neighbors Efficiently

I have data geographically scattered without any kind of pattern and I need to create an image where the value of each pixel is an average of the neighbors of that pixel that are l

Solution 1:

This may be one of the cases where KDTrees are not a good solution. This is because you are mapping to a grid, which is a very simple structure meaning there is nothing to gain from the KDTree's sophistication. Nearest grid point and distance can be found by simple arithmetic.

Below is a simple example implementation. I'm using a Gaussian kernel but changing that to IDW if you prefer should be straight-forward.

import numpy as np
from scipy import stats

defrasterize(coords, feature, gu, cutoff, kernel=stats.norm(0, 2.5).pdf):
    # compute overlap (filter size / grid unit)
    ovlp = int(np.ceil(cutoff/gu))

    # compute raster dimensions
    mn, mx = coords.min(axis=0), coords.max(axis=0)
    reso = np.ceil((mx - mn) / gu).astype(int)
    base = (mx + mn - reso * gu) / 2# map coordinates to raster, the residual is the distance
    grid_res = coords - base
    grid_coords = np.rint(grid_res / gu).astype(int)
    grid_res -= gu * grid_coords
    # because of overlap we must add neighboring grid points to the nearest
    gcovlp = np.c_[-ovlp:ovlp+1, np.zeros(2*ovlp+1, dtype=int)]
    grid_coords = (gcovlp[:, None, None, :] + gcovlp[None, :, None, ::-1]
                   + grid_coords).reshape(-1, 2)
    # the corresponding residuals have the same offset with opposite sign
    gdovlp = -gu * (gcovlp+1/2)
    grid_res = (gdovlp[:, None, None, :] + gdovlp[None, :, None, ::-1]
                + grid_res).reshape(-1, 2)
    # discard off fov grid points and points outside the cutoff
    valid, = np.where(((grid_coords>=0) & (grid_coords<=reso)).all(axis=1) & (
        np.einsum('ij,ij->i', grid_res, grid_res) <= cutoff*cutoff))
    grid_res = grid_res[valid]
    feature = feature[valid // (2*ovlp+1)**2]
    # flatten grid so we can use bincount
    grid_flat = np.ravel_multi_index(grid_coords[valid].T, reso+1)
    return np.bincount(
        grid_flat,
        feature * kernel(np.sqrt(np.einsum('ij,ij->i', grid_res, grid_res))),
        (reso + 1).prod()).reshape(reso+1)



gu = 5
cutoff = 10
coords = np.random.randn(10_000, 2) * (100, 20)
coords[:, 1] += 80 * np.sin(coords[:, 0] / 40)
feature = np.random.uniform(0, 1000, (10_000,))

from timeit import timeit

print(timeit("rasterize(coords, feature, gu, cutoff)", globals=globals(), number=100)*10, 'ms')

pic = rasterize(coords, feature, gu, cutoff)

import pylab
pylab.pcolor(pic, cmap=pylab.cm.jet)
pylab.colorbar()
pylab.show()