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

Question

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]]

cs95 · Accepted Answer · 2020-06-01 05:45:50Z

1

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.

answered Jun 1, 2020 at 5:45

cs95

406k106 gold badges744 silver badges797 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

phntm Over a year ago

Thanks so much!! I've been staring at this for days. So None is essentially equivalent to np.newaxis and each bracket corresponds to an axis correct? Any recommendations on additional sources to turn to if stuck on things like this in the future? because if you hadn't posted I'd be starting at the screen in perpetuity.

cs95 Over a year ago

@phntm I learned by asking questions on here, same as you :) Eventually started answering as well because you learn faster that way. Some people recommend various tutorials and videos online, but the documentation to me is the source of truth. Often good documentation has tutorials which take you through the basics, and numpy is no exception. On the specific topic of indexing, try reading the indexing docs.

Collectives™ on Stack Overflow

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

1 Answer 1

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related