Fill NaN values in an N-dimensional array with their nearest neighbours' values.
The nearest neighbour is determined using the Euclidean distance between array
indices, so there is equal weighting given in all directions (i.e. across all axes).
In the case of tiebreaks, the value from the neighbour with the lowest index is
imputed.
Source code in src/rastr/arr/fill.py
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32 | def fillna_nearest_neighbours(arr: NDArray) -> NDArray:
"""Fill NaN values in an N-dimensional array with their nearest neighbours' values.
The nearest neighbour is determined using the Euclidean distance between array
indices, so there is equal weighting given in all directions (i.e. across all axes).
In the case of tiebreaks, the value from the neighbour with the lowest index is
imputed.
"""
from scipy.interpolate import NearestNDInterpolator
nonnan_mask = np.nonzero(~np.isnan(arr))
nonnan_idxs = np.array(nonnan_mask).transpose()
if nonnan_idxs.size == 0:
# Everything is NaN
return arr
# Interpolate at the array indices
interp = NearestNDInterpolator(nonnan_idxs, arr[nonnan_mask])
filled_arr = interp(*np.indices(arr.shape))
# Preserve the original dtype
return filled_arr.astype(arr.dtype)
|