Maintain 2D structure after logical indexing a numpy array

Question

I am tracking a dynamically changing mask that rolls using some input shift. This mask stores values that determine where I can trust values in another array of the same shape. An example of how the mask changes over each iteration is below. I have a large stack of logical checks that determine how to set the rolled parts of the mask to zero based on whether the x and y values of the shift are equal to 0 or are positive or negative. Here I just hardcoded it all for clarity.

import numpy as np
mask = np.full((8,8), 10)

#Iteration 1
mask = np.roll(mask, (0, 1), axis = (0,1))
mask[:, :1] = 0
#logical indexing happens here
mask += 1
print (mask)

#Iteration 2
mask = np.roll(mask, (1, 0), axis = (0,1))
mask[:1, :] = 0
#logical indexing happens here
mask +=1
print (mask)

#Iteration 3
mask = np.roll(mask, (2, -1), axis = (0,1))
mask[:, -1:] = 0
mask[:2, :] = 0
#logical indexing happens here
mask +=1
print (mask)

After each iteration and before the mask is increased by one, I need to index into and pull the values of a second array where the mask is above some threshold (10 in this case). Since I am rolling and setting values, I always know that the part of the mask that fulfills this condition can be broadcast into a 2d array. A simplified example of what I am doing now is below where arr2 is a flattened array.

import numpy as np

arr1 = np.arange(0, 64, 1).reshape((8,8))
mask = np.full((8,8), 10)
mask[:, 0] = 0

arr2 = arr1[mask >= 10]

How can I keep arr2 as a 2d array where the mask is above the set threshold?

I do not know a priori what the shift will be that is applied to the mask so I have to rely on the values in the mask to determine the shape of the resulting array. My arrays are much larger than this example and the shifts are between -5 and 5 so I know I won't get close to setting the entire array below the threshold. The idea is that after ~10 iterations, some parts of the array become trustworthy again and can be useful information after the logical index.

Answer 1

The answer here is a work around and was obvious now that it has simmered in my mind for a while. Basically, since I know that the resulting area will be square, I can just count across a row and column where each index meets my condition. So continuing my example from before I just add in a couple lines to determine how many values in a row and column meet my condition.

import numpy as np
#Initializing array
arr1 = np.arange(0, 64, 1).reshape((8,8))
#mask array
mask = np.full((8,8), 10)

#Setting some rows and cols to zero to simulate my roll functionality
mask[:, 0] = 0
mask[:2, :] = 0

#Summing across a row and col where condition is met
sizex = np.sum(mask[4, :] >= 10)
sizey = np.sum(mask[:, 4] >= 10)

#Using the mask to index into the original array and reshaping
arr2 = arr1[mask >= 10].reshape((sizey, sizex))