Modifying numpy array with `numpy.take`-like function

Question

I know that in numpy, I can use numpy.take to get the subarray at a particular axis value instead of using fancy indexing.

So, if I have a 2D array a=np.array([[1,2],[3,4]]), then I can write a[0, :] or np.take(a, 0, axis=0) to get the same result.

However, the two are not equivalent in the sense that I can use the first one to also set values, while I cannot do the same with the second. In other words, I can do this:

a=np.array([[1,2],[3,4]])
a[0, :] = 10
# I now expect and also get:
# a = [[10, 10],
#      [3,  4]

while I cannot do this:

a=np.array([[1,2],[3,4]])
np.take(a, 0, axis=0) = 10
# returns an error: SyntaxError: cannot assign to function call here. Maybe you meant '==' instead of '='?

So, is there a numpy.take-like function that will allow me to set values, without doing fancy indexing? The motivation here is that I want the axis upon which I am changing values to be settable by the user, so I can't just do a[0, :] because I might need to do a[:, 0]. Now of course I can use an if statement for this, but I am wondering if there is something simpler.

Answer 1

The corollary of numpy.take for setting elements is numpy.put, but unfortunately np.put does not take an axis argument. numpy.put_along_axis exists, but this has the indexing semantics of np.take_along_axis, which is different than what you asked.

I suspect the easiest way to achieve what you have in mind is to use np.take to generate indices that can then be passed to np.put. For example:

>>> a=np.array([[1,2],[3,4]])
>>> i = np.take(np.arange(a.size).reshape(a.shape), 0, axis=0)
>>> np.put(a, i, 10)
>>> print(a)
[[10 10]
 [ 3  4]]

Another possibility would be to combine numpy.apply_along_axis with np.put. For example:

>>> a = np.array([[1,2],[3,4]])
>>> np.apply_along_axis(np.put, arr=a, axis=0, ind=0, v=10)
>>> print(a)
[[10 10]
 [ 3  4]]

Though please be aware that apply_along_axis is implemented via loops rather than vectorized operations, so it may have poor performance for larger arrays.