CODEWITHSUNDEEP
arraysnumpybooleanvectorizationtensor
phntm
phntm

Reputation: 541

Create a boolean array comparing each matrix in a Numpy tensor to a different integer in a Numpy Array

Say I have a Numpy tensor X that is 3*3*3 (actual dimensions would vary). I want to test each matrix in the tensor against a different value in a set of integers.

For example if

X=np.array([1,2,3]*9).reshape(3,3,3)
test=np.array([1,2,3])

The desired output would be: 

 [[[ True, False, False],
    [True, False, False],
    [True, False, False]],

   [[False,  True, False],
    [False,  True, False],
    [False,  True, False]],

   [[False, False,  True],
    [False, False,  True],
    [False, False,  True]]])

However I can't seem to get this result. X==test returns:

array([[[ True,  True,  True],
        [ True,  True,  True],
        [ True,  True,  True]],

       [[ True,  True,  True],
        [ True,  True,  True],
        [ True,  True,  True]],

       [[ True,  True,  True],
        [ True,  True,  True],
        [ True,  True,  True]]])

If 

test=[[1],[2],[3]]

I get:

array([[[ True, False, False],
        [False,  True, False],
        [False, False,  True]],

       [[ True, False, False],
        [False,  True, False],
        [False, False,  True]],

       [[ True, False, False],
        [False,  True, False],
        [False, False,  True]]])

The same result holds true for np.equal. Is there any direct way to do this without using any loops? It seems like there would be a way given that with indexing

X[[0,1,2],[0,2,1]]

would yield

np.array([X[0][0],X[1][2],X[2][1]])

rather than 

X[:,[0,2,1]]

Upvotes: 1

Views: 93

Answers (1)

cs95
cs95

Reputation: 402483

This is a simple equality comparison, but the tricky part is figuring out how to broadcast the operation. You can do this as, 

X == test[:, None, None]

array([[[ True, False, False],
        [ True, False, False],
        [ True, False, False]],

       [[False,  True, False],
        [False,  True, False],
        [False,  True, False]],

       [[False, False,  True],
        [False, False,  True],
        [False, False,  True]]])

Where,

test[:, None, None]

array([[[1]],

       [[2]],

       [[3]]])

The idea is to make the dimensions of X and test match, that way we can broadcast the equality comparison so first item of test is compared with the first sub-matrix of X, second item compared with the second sub-matrix, and so on.

Upvotes: 1

Related Questions