Reputation: 1097
I'd like a way to set values at specific indices and axes of a ndarray
without using fancy indexing.
Say I have an array. I can do:
z = np.zeros((3,4,2,8))
# set z=9 at [1,2] in axis 0, and [0,1] in axis 2
z[[1,2],:,[0,1],:] = 9
But is there a function set_value
that I could instead do something like:
z.set_value(9, axis=(0,2), indices=[[1,2],[0,1]])
I don't think np.put
or np.put_along_axis
meet my need, unless I misunderstand how to use expand_dims
.
Upvotes: 0
Views: 166
Reputation: 36289
You can create such a function yourself:
def set_value(arr, value, *, axis, indices):
index = [slice(None)] * len(arr.shape)
for a, i in zip(axis, indices):
index[a] = i
arr[tuple(index)] = value
Upvotes: 2
Reputation: 968
You can use np.moveaxis
:
z = np.zeros((3,4,2,8))
axes = [0,2]
indices = ([1,2],[0,1])
np.moveaxis(z,axes,range(len(axes)))[indices] = 9
This creates a view of z
with the axes shuffled, such that they can be accessed as axes 0,1. The trick is that the data but not the layout of z
are shared between the original array and this view. By writing to the view z
will be modified as desired.
Upvotes: 3