stevemo
stevemo

Reputation: 1097

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.

Upvotes: 0

Views: 166

Answers (2)

a_guest
a_guest

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

loopy walt
loopy walt

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

Related Questions