Numpy vectorizing loop through array indices

Question

i want to make a function that will return an array of booleans. Basically the code will check the equality of elements inside a 4-D array. It will through 2nd and 3rd dimension and compare every array whether they all are equal or not. Here's my current code :

x = np.arange(0,72).reshape(4,2,3,3)

x[:,1,2,:] = 0

def equality(arr):
    return np.array([np.allclose(arr[:,i,j,:],arr[:,i,j,:+1]) for i in range(arr.shape[1]) for j in range(arr.shape[2])])

print(x)
print(equality(x))

Output :

 [[[[ 0  1  2]    [ 3  4  5]    [ 6  7  8]]
 
   [[ 9 10 11]    [12 13 14]    [ 0  0  0]]]
 
 
  [[[18 19 20]    [21 22 23]    [24 25 26]]
 
   [[27 28 29]    [30 31 32]    [ 0  0  0]]]
 
 
  [[[36 37 38]    [39 40 41]    [42 43 44]]
 
   [[45 46 47]    [48 49 50]    [ 0  0  0]]]
 
 
  [[[54 55 56]    [57 58 59]    [60 61 62]]
 
   [[63 64 65]    [66 67 68]    [ 0  0  0]]]] [False False False False
 False  True]

That code works, but is there a way to get rid of that for loop using numpy stuff instead?

Answer 1

Use np.ndarray.reshape and numpy.transpose to get to desired shape. To get an intuition about the order, check this question. Then check if all columns in sub_arrays are equal:

>>> reshaped = np.transpose(x.reshape(4,6,3), axes=(1,0,2))
# 1st dim==4, 2nd*3rd dim==6, 4th dim==3, hence reshape(4,6,3)
>>> reshaped

array([[[ 0,  1,  2],
        [18, 19, 20],
        [36, 37, 38],
        [54, 55, 56]],

       [[ 3,  4,  5],
        [21, 22, 23],
        [39, 40, 41],
        [57, 58, 59]],

       [[ 6,  7,  8],
        [24, 25, 26],
        [42, 43, 44],
        [60, 61, 62]],

       [[ 9, 10, 11],
        [27, 28, 29],
        [45, 46, 47],
        [63, 64, 65]],

       [[12, 13, 14],
        [30, 31, 32],
        [48, 49, 50],
        [66, 67, 68]],

       [[ 0,  0,  0],
        [ 0,  0,  0],
        [ 0,  0,  0],
        [ 0,  0,  0]]])

>>> (reshaped == reshaped[:, :, -1:]).all(1).all(1)
array([False, False, False, False, False,  True])