Numpy double-slice assignment with integer indexing followed by boolean indexing

user_15335086 0
Viewed 71

I already know that Numpy "double-slice" with fancy indexing creates copies instead of views, and the solution seems to be to convert them to one single slice (e.g. This question). However, I am facing this particular problem where i need to deal with an integer indexing followed by boolean indexing and I am at a loss what to do. The problem (simplified) is as follows:

a = np.random.randn(2, 3, 4, 4)
idx_x = np.array([[1, 2], [1, 2], [1, 2]])
idx_y = np.array([[0, 0], [1, 1], [2, 2]])
print(a[..., idx_y, idx_x].shape) # (2, 3, 3, 2)
mask = (np.random.randn(2, 3, 3, 2) > 0)
a[..., idx_y, idx_x][mask] = 1 # assignment doesn't work

How can I make the assignment work?

slice numpy python
3 Answers

Not sure, but an idea is to do the broadcasting manually and adding the mask respectively just like Tim suggests. idx_x and idx_y both have the same shape (3,2) which will be broadcasted to the shape (6,6) from the cartesian product (3*2)^2.

x = np.broadcast_to(idx_x.ravel(), (6,6))
y = np.broadcast_to(idx_y.ravel(), (6,6))

# this should be the same as
x,y = np.meshgrid(idx_x, idx_y)

Now reshape the mask to the broadcasted indices and use it to select

mask = mask.reshape(6,6)
a[..., x[mask], y[mask]] = 1

The assignment now works, but I am not sure if this is the exact assignment you wanted.

user_149232270

Ok apparently I am making things complicated. No need to combine the indexing. The following code solves the problem elegantly:

b = a[..., idx_y, idx_x]
b[mask] = 1
a[..., idx_y, idx_x] = b
print(a[..., idx_y, idx_x][mask]) # all 1s
user_153350860

EDIT: Use @Kevin's solution which actually gets the dimensions correct!

I haven't tried it specifically on your sample code but I had a similar issue before. I think I solved it by applying the mask to the indices instead, something like:

a[..., idx_y[mask], idx_x[mask]] = 1

-that way, numpy can assign the values to the a array correctly.

EDIT2: Post some test code as comments remove formatting.

a = np.arange(27).reshape([3, 3, 3])
ind_x = np.array([[0, 0], [1, 2]])
ind_y = np.array([[1, 2], [1, 1]])
x = np.broadcast_to(ind_x.ravel(), (4, 4))
y = np.broadcast_to(ind_y.ravel(), (4, 4)).T
# x1, y2 = np.meshgrid(ind_x, ind_y)  # above should be the same as this
mask = a[:, ind_y, ind_x] % 2 == 0  # what should this reshape to?
# a[..., x[mask], y[mask]] = 1  # Then you can mask away (may also need to reshape a or the masked x or y)
user_79858550
Related