Programmatic indexing in numpy

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.

Answer

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.