Advanced boolean indexing

Question

I wanna select values by mask and changes values by use mask-array.

Code:

import numpy as np

a = np.zeros((2, 2), dtype=(np.uint8, 3))
x = np.arange(4, dtype=int).reshape((2, 2))

mask = np.logical_and(a1 < 3, a1 > 0)

a[mask] = (1, x[mask], 2)

I want result:

a[mask]
>> [[1, 1, 2], [1, 2, 2]]

But i get error: ValueError: setting an array element with a sequence.

If try do things like a[mask] = (1, 2, 2) array will be

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

But i need use values from x. To make it look like

[[[0, 0, 0],
[1, 1, 3]],
[[1, 2, 3],
[0, 0, 0]]]

How i can do it?

Answer 1

It can be done in two steps.

import numpy as np 

a = np.zeros((2, 2), dtype=(np.uint8, 3)) 
x = np.arange(4, dtype=int).reshape((2, 2)) 
a1 = x  # Create an a1 for the mask

mask = np.logical_and(a1 < 3, a1 > 0) 

a[mask] = (1, 0, 2)      # Set the outer columns                                         
a[mask, 1] = x[mask]     # Set the column 1

print( a )                                                              
# [[[0 0 0]
#   [1 1 2]]

#  [[1 2 2]
#  [0 0 0]]]