Moving Window Calculation Across Multiple Arrays

I have several two-dimensional data arrays loaded into NumPy arrays, all of which have identical dimensions. These shared dimensions are 2606x and 1228y.

I am interested in computing calculations between the first two arrays (a1 & a2) across a moving window, using a window sized 2x by 2y, with the resultant calculation then applied to the third array. Specifically, the workflow would be:

Finding the the maximum & minimum value of a1 at the piece of this moving window
Selecting the corresponding array indices of these values
Extracting the values at these indices of a2
Cast the calculation result to each of the indices in the third array (a3) inside the moving window.

I know that this process involves the following pieces of code to obtain the values I require:

idx1 = np.where(a1 == a1.max())
idx2 = np.where(a1 == a1.min())
val1 = a2[idx1[1], idx1[2]]
val2 = a2[idx2[1], idx2[2]]

What additional code is required to perform this moving window along the identically sized arrays?

Solution

Since your array shape is divisible by your window size, you can use numpy.reshape to split your array up into little windows such that your original array shape of (2606, 1228) becomes (2606/2, 2, 1228/2, 2).

If numpy.argmin accepted sequences of axes, this would be easier, but since it only accepts a single axis (or None but we don't want that), we need to compress the two window axes into a single axes. To do that, we use numpy.moveaxis to make the shape (2606/2, 1228/2, 2, 2) and then numpy.reshape again to flatten the last two axes into (2606/2, 1228/2, 4).

With that headache over with, we can then use numpy.argmin and numpy.argmax on the last axis to compute the indices you're interested in and use advanced indexing to write the corresponding value of a2 to a3. After that, we just have to undo the reshape and moveaxis operations that were done to a3.

import numpy as np

shape = (4, 6)
a1 = np.random.random(shape)
a2 = np.random.random(shape)
a3 = np.zeros(shape)

win_x = 2
win_y = 2

shape_new = (shape[0] // win_x, win_x, shape[1] // win_y, win_y)

a1_r = np.moveaxis(a1.reshape(shape_new), 1, 2).reshape(*shape_new[::2], -1)
a2_r = np.moveaxis(a2.reshape(shape_new), 1, 2).reshape(*shape_new[::2], -1)
a3_r = np.moveaxis(a3.reshape(shape_new), 1, 2).reshape(*shape_new[::2], -1)

index_x, index_y = np.indices(shape_new[::2])
index_min = np.argmin(a1_r, axis=-1)
index_max = np.argmax(a1_r, axis=-1)

a3_r[index_x, index_y, index_min] = a2_r[index_x, index_y, index_min]
a3_r[index_x, index_y, index_max] = a2_r[index_x, index_y, index_max]

a3 = np.moveaxis(a3_r.reshape(*shape_new[::2], win_x, win_y), 2, 1).reshape(shape)

print(a1)
print()
print(a2)
print()
print(a3)

Outputs

[[0.54885307 0.74457945 0.84943538 0.14139329 0.68678556 0.03460323]
 [0.74031057 0.5499962  0.03148748 0.13936734 0.05006111 0.88850868]
 [0.97789608 0.13262023 0.76350358 0.74640822 0.7918286  0.80675845]
 [0.35784598 0.20918229 0.82880072 0.06051794 0.0825886  0.6398353 ]]

[[0.66176657 0.10120202 0.15306892 0.05963046 0.79057051 0.08837686]
 [0.78550049 0.09918834 0.00213652 0.61053454 0.42966757 0.25952916]
 [0.00387273 0.78247644 0.65549303 0.39351233 0.11002493 0.55652453]
 [0.06047582 0.87997514 0.60820023 0.06705212 0.34581512 0.93504438]]

[[0.66176657 0.10120202 0.15306892 0.         0.         0.08837686]
 [0.         0.         0.00213652 0.         0.         0.25952916]
 [0.00387273 0.78247644 0.         0.         0.         0.55652453]
 [0.         0.         0.60820023 0.06705212 0.34581512 0.        ]]