Check equality of multiple elements in array

I’m new to Python from Matlab.

I want to create a new variable from a subset of an existing numpy array based on equality to some condition specified by a third numpy array, an ID in this case.

This works fine for one equality.

new_x = old_x[someID == 1]

But if I try to extend it several equalities at once it no longer works:

new_x = old_x[someID == 1:3]

Ideally I want to be able to choose many equalities, like:

new_x = old_x[someID == 1:3,7]

I could loop through each number I want to check but is there a simpler way of doing this?

Answer

You could use np.isin + np.r_:

import numpy as np

# for reproducible results
np.random.seed(42)

# toy data
old_x = np.random.randint(10, size=100)

# create new array by filtering on boolean mask
new_x = old_x[np.isin(old_x, np.r_[1:3,7])]

print(new_x)

Output

[7 2 7 7 7 2 1 7 1 2 2 2 1 1 1 7 2 1 7 1 1 1 7 7 1 7 7 7 7 2 7 2 2 7]

You could substitute np.r_ by something like [1, 2, 7] and use it as below:

new_x = old_x[np.isin(old_x, [1, 2, 7])]

Additionally if the array is 1-dimensional you could use np.in1d:

new_x = old_x[np.in1d(old_x, [1, 2, 7])]
print(new_x) 

Output (from in1d)

[7 2 7 7 7 2 1 7 1 2 2 2 1 1 1 7 2 1 7 1 1 1 7 7 1 7 7 7 7 2 7 2 2 7]