Problem with element-wise operations in numpy

Question

This is my first Question here, let me know if I could've done anything better.

I'm trying to do a element-wise operation between two arrays, but the broadcasting wont work like I want it to.

I have an array of shape (N,4).

square_list = np.array([[1,2,255,255], [255,255,4,4], [255,255,8,8], [255,255,16,16], [255,255,8,4], [255,1,8,8], [1,255,8,8]], dtype='B')

I also have an array of shape (4,).

square = np.array([1, 8, 8, 1], dtype='B')

What I am able to do is compare my square against each element in the square_list and it is being broadcast into shape (N,4) as expected.

Now I want to compare my square in each possible rotation against the square_list . I've written a function which returns an array of shape (4,4), which contains each possible rotation.

square.rotations
    array([[1, 8, 8, 1],
           [1, 1, 8, 8],
           [8, 1, 1, 8],
           [8, 8, 1, 1]], dtype=uint8)

I know how to do this using a loop. I'd prefer however to use an element-wise operator that returns my desired shape.

What I get:

rotations & square_list
ValueError: operands could not be broadcast together with shapes (4,4) (6,4)

What I'd like to get:

rotations & square_list
array([[[1, 0, 8, 1],
        [1, 8, 0, 0],
        [1, 8, 8, 0],
        [1, 8, 0, 0],
        [1, 8, 8, 0],
        [1, 0, 8, 0]],
       [[1, 0, 8, 8],
        [1, 1, 0, 0],
        [1, 1, 8, 8],
        [1, 1, 0, 0],
        [1, 1, 8, 0],
        [1, 1, 8, 8]],
       [[0, 0, 1, 8],
        [8, 1, 0, 0],
        [8, 1, 0, 8],
        [8, 1, 0, 0],
        [8, 1, 0, 0],
        [8, 1, 0, 8],
        [0, 1, 0, 8]],
       [[0, 0, 1, 1],
        [8, 8, 0, 0],
        [8, 8, 0, 0],
        [8, 8, 0, 0],
        [8, 8, 0, 0],
        [8, 0, 0, 0],
        [0, 8, 0, 0]]], dtype=uint8)

This is just to visualize what I want, I don't particularly care about the order of the axis'. A shape of either (4, N, 4) or (N, 4, 4) would be great. I have the feeling that this can be achieved easily by just reshaping one of the input arrays but I couldn't figure it out.

Thanks in advance!

Answer 1

Add an extra dimension to rotations:

square_list & rotations[:,None]

output:

array([[[1, 0, 8, 1],
        [1, 8, 0, 0],
        [1, 8, 8, 0],
        [1, 8, 0, 0],
        [1, 8, 8, 0],
        [1, 0, 8, 0],
        [1, 8, 8, 0]],

       [[1, 0, 8, 8],
        [1, 1, 0, 0],
        [1, 1, 8, 8],
        [1, 1, 0, 0],
        [1, 1, 8, 0],
        [1, 1, 8, 8],
        [1, 1, 8, 8]],

       [[0, 0, 1, 8],
        [8, 1, 0, 0],
        [8, 1, 0, 8],
        [8, 1, 0, 0],
        [8, 1, 0, 0],
        [8, 1, 0, 8],
        [0, 1, 0, 8]],

       [[0, 0, 1, 1],
        [8, 8, 0, 0],
        [8, 8, 0, 0],
        [8, 8, 0, 0],
        [8, 8, 0, 0],
        [8, 0, 0, 0],
        [0, 8, 0, 0]]], dtype=uint8)