Broadcasting multi-dimensional array indices of the same shape

Question

I have a mask array which represents a 2-dimensional binary image. Let's say it's simply:

mask = np.zeros((9, 9), dtype=np.uint8)
# 0 0 0 | 0 0 0 | 0 0 0
# 0 0 0 | 0 0 0 | 0 0 0
# 0 0 0 | 0 0 0 | 0 0 0
# ------+-------+------
# 0 0 0 | 0 0 0 | 0 0 0
# 0 0 0 | 0 0 0 | 0 0 0
# 0 0 0 | 0 0 0 | 0 0 0
# ------+-------+------
# 0 0 0 | 0 0 0 | 0 0 0
# 0 0 0 | 0 0 0 | 0 0 0
# 0 0 0 | 0 0 0 | 0 0 0

Suppose I want to flip the elements in the middle left ninth:

# 0 0 0 | 0 0 0 | 0 0 0
# 0 0 0 | 0 0 0 | 0 0 0
# 0 0 0 | 0 0 0 | 0 0 0
# ------+-------+------
# 1 1 1 | 0 0 0 | 0 0 0
# 1 1 1 | 0 0 0 | 0 0 0
# 1 1 1 | 0 0 0 | 0 0 0
# ------+-------+------
# 0 0 0 | 0 0 0 | 0 0 0
# 0 0 0 | 0 0 0 | 0 0 0
# 0 0 0 | 0 0 0 | 0 0 0

My incorrect approach was something like this:

x = np.arange(mask.shape[0])
y = np.arange(mask.shape[1])
mask[np.logical_and(y >= 3, y < 6), x < 3] = 1
# 0 0 0 | 0 0 0 | 0 0 0
# 0 0 0 | 0 0 0 | 0 0 0
# 0 0 0 | 0 0 0 | 0 0 0
# ------+-------+------
# 1 0 0 | 0 0 0 | 0 0 0
# 0 1 0 | 0 0 0 | 0 0 0
# 0 0 1 | 0 0 0 | 0 0 0
# ------+-------+------
# 0 0 0 | 0 0 0 | 0 0 0
# 0 0 0 | 0 0 0 | 0 0 0
# 0 0 0 | 0 0 0 | 0 0 0

(This is a simplification of the constraints I'm really dealing with, which would not be easily expressed as something like mask[:3,3:6] = 1 as in this case. Consider the constraints arbitrary, like x % 2 == 0 && y % 3 == 0 if you will.)

Numpy's behavior when the two index arrays are the same shape is to take them pairwise, which ends up only selecting the 3 elements above, rather than 9 I would like.

How would I update the right elements with constraints that apply to different axes? Given that the constraints are independent, can I do this by only evaluating my constraints N+M times, rather than N*M?

Answer 1

You can't broadcast the boolean arrays, but you can construct the equivalent numeric indices with ix_:

In [330]: np.ix_((y>=3)&(y<6), x<3)                                             
Out[330]: 
(array([[3],
        [4],
        [5]]), array([[0, 1, 2]]))

Applying it:

In [331]: arr = np.zeros((9,9),int)                                             
In [332]: arr[_330] = 1                                                         
In [333]: arr                                                                   
Out[333]: 
array([[0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0],
       [1, 1, 1, 0, 0, 0, 0, 0, 0],
       [1, 1, 1, 0, 0, 0, 0, 0, 0],
       [1, 1, 1, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0]])

Attempting to broadcast the booleans directly raises an error (too many indices):

arr[((y>=3)&(y<6))[:,None], x<3]

Answer 2

Per your comment, let's try this fancier example:

mask = np.zeros((90,90), dtype=np.uint8)

# criteria
def f(x,y): return ((x-20)**2 < 50) & ((y-20)**2 < 50)

# ranges
x,y = np.arange(90), np.arange(90)

# meshgrid
xx,yy = np.meshgrid(x,y)

zz = f(xx,yy)

# mask
mask[zz] = 1
plt.imshow(mask, cnap='gray')

Output: