CODEWITHSUNDEEP
pythonarraysnumpymultidimensional-array
Arty
Arty

Reputation: 16747

Is there an opposite of np.take for multi-dimensional case in NumPy?

If I have N-D array a and 1-D array indexes and axis, what is the shortest/easiest way to assign values to sub-array of a indexed with these indexes along given axis?

If N and axis are known/fixed at code writing time then I can easily do

a[:,   :,   :,   indexes,   :,   :] = b
               ^^axis pos^^

but what if N and axis are known only at runtime, e.g. after user inputing some data.

The easiest form that I have figured out is this:

a[(slice(None),) * (axis % a.ndim) + (indexes,)] = b

But maybe there's some built-in function?

There's np.put but it seems to work only on flattened array, doesn't support multi-dimensionality.

There's also np.take, which works for multi-dimensional case, but is the opposite of what I need, it reads data instead of writing, although has exact syntax that I'm looking for in writing. You can do like this:

b = np.take(a, indexes, axis)

So is there an opposite of np.take? Like np.nd_put, that can do:

np.nd_put(a, indexes, axis, b)

Upvotes: 2

Views: 890

Answers (1)

Divakar
Divakar

Reputation: 221614

One way would be to permute axes to bring that axis to the front and simply index -

np.moveaxis(a,axis,0)[indexes] = np.moveaxis(b,axis,0)

Another with np.put_along_axis if you are looking for something built around a bultin -

i = [None]*b.ndim
i[axis] = Ellipsis
np.put_along_axis(a, indexes[tuple(i)], b, axis=axis)

For completeness, here's the slice version -

i = [slice(None)]*b.ndim
i[axis] = indexes
a[tuple(i)] = b

That you can convert to one-liner to get i with list comprehension.

Upvotes: 2

Related Questions