CODEWITHSUNDEEP
pythonarraysnumpyindexing
Tor
Tor

Reputation: 785

Indexing with two boolean arrays in NumPy

I have some data in a numpy array, and then I want to select a subset of a subset, and update those values.

Let's say the first array, A, has N elements. Then, I use a boolean array, mask1, with N elements as a mask to select the first subset. Then, I use a second boolean array, mask2, with sum(mask1) elements, to select a subset of the first subset. But when I try to update these values, they aren't stored to the original array.

Code example:

N = 10
A = np.arange(N)

# Create mask of size N, with only some True values
mask1 = np.zeros(N, dtype = np.bool)
mask1[:7] = True

# Create mask of size (sum(mask1)), with only some True values
mask2 = np.zeros(np.sum(mask1), dtype = np.bool)
mask2[2:] = True

# Indexing with both masks works as expected:
print(A[mask1][mask2])

This prints

[2 3 4 5 6]

as expected. But when I try to update the values indexed in this way, it doesn't work.

# Trying to update those values does not work:
A[mask1][mask2] = -1
print(A)

This prints the original, unchanged array:

[0 1 2 3 4 5 6 7 8 9]

If I do the indexing with one boolean array only, then the values get updated, so that's a possible workaround. However, I'm implementing an algorithm for a physical process where it's really natural to first create one mask, to subset on one criteria, and then subset that further based on another criteria, so it would be really elegant if I could get this to work.

Upvotes: 3

Views: 822

Answers (1)

Divakar
Divakar

Reputation: 221684

A performant one would be with chained-masking -

mask1[mask1] = mask2
A[mask1] = -1

Not so performant one -

A[np.flatnonzero(mask1)[mask2]] = -1

Upvotes: 2

Related Questions