NumPy Nearest Neighbor Line Fitting Across Moving Window

Question

I have two two-dimensional arrays loaded into NumPy, both of which are 80i x 80j in size. I'm looking to do a moving window polyfit calculation across these arrays, I've nailed down how to conduct the polyfit but am stuck on the specific moving window approach I'm looking to accomplish. I'm aiming for:

1) At each index of Array 1 (a1), code isolates all the values of the index & its closest 8 neighbors into a separate 1D array, and repeats over the same window at Array 2 (a2)

2) With these two new arrays isolated, perform linear regression line-fitting using NumPy's polyfit in the following approach: model = np.polyfit(a1slice, a2slice, 1)

3) Cast the resulting regression coefficient and intercept (example output doing this manually: array([-0.02114911, 10.02127152])) to the same index of two other arrays, where model[0] would be placed into the first new array and model[1] into the second new array at this index.

4) The code then moves sequentially to the next index and performs steps 1-3 again, or a1(i+1, j+0, etc.)

I've provided a graphical example of what I'm trying achieve for two random index selections across Array 1 and the calculation across the index's eight nearest neighbors, should this make the desired result easier to understand:

Answer 1

I've written a function to get the window, but there is a troublesome edge case when the index is on the edge and one distance away from the corner. You'll probably want to modify the function to get exactly what you're looking for.

def get_neighbors(arr, origin, num_neighbors = 8):

    coords = np.array([[i,j] for (i,j),value in np.ndenumerate(arr)]).reshape(arr.shape + (2,))
    distances = np.linalg.norm(coords - origin, axis = -1)
    neighbor_limit = np.sort(distances.ravel())[num_neighbors]

    window = np.where(distances <= neighbor_limit)
    exclude_window = np.where(distances > neighbor_limit)
    
    return window, exclude_window, distances


test = np.zeros((5, 5))

plt.close('all')

cases = [[0,0], [0, 1], [0, 2]]

for case in cases:
    
    window, exclude, distances = get_neighbors(test, case)
    distances[exclude] = 0
    
    plt.figure()
    plt.imshow(distances)

Image Outputs: