Numpy assignment like 'numpy.take'

Question

Is it possible to assign to a numpy array along the lines of how the take functionality works?

E.g. if I have a an array a, a list of indices inds, and a desired axis, I can use take as follows:

import numpy as np
a = np.arange(12).reshape((3, -1))
inds = np.array([1, 2])
print(np.take(a, inds, axis=1))

[[ 1  2]
 [ 5  6]
 [ 9 10]]

This is extremely useful when the indices / axis needed may change at runtime.

However, numpy does not let you do this:

np.take(a, inds, axis=1) = 0
print(a)

It looks like there is some limited (1-D) support for this via numpy.put, but I was wondering if there was a cleaner way to do this?

Answer 1

In [222]: a = np.arange(12).reshape((3, -1))
     ...: inds = np.array([1, 2])
     ...: 
In [223]: np.take(a, inds, axis=1)
Out[223]: 
array([[ 1,  2],
       [ 5,  6],
       [ 9, 10]])
In [225]: a[:,inds]
Out[225]: 
array([[ 1,  2],
       [ 5,  6],
       [ 9, 10]])

construct an indexing tuple

In [226]: idx=[slice(None)]*a.ndim
In [227]: axis=1
In [228]: idx[axis]=inds
In [229]: a[tuple(idx)]
Out[229]: 
array([[ 1,  2],
       [ 5,  6],
       [ 9, 10]])
In [230]: a[tuple(idx)] = 0
In [231]: a
Out[231]: 
array([[ 0,  0,  0,  3],
       [ 4,  0,  0,  7],
       [ 8,  0,  0, 11]])

Or for a[inds,:]:

In [232]: idx=[slice(None)]*a.ndim
In [233]: idx[0]=inds
In [234]: a[tuple(idx)]
Out[234]: 
array([[ 4,  0,  0,  7],
       [ 8,  0,  0, 11]])
In [235]: a[tuple(idx)]=1
In [236]: a
Out[236]: 
array([[0, 0, 0, 3],
       [1, 1, 1, 1],
       [1, 1, 1, 1]])

PP's suggestion:

def put_at(inds, axis=-1, slc=(slice(None),)): 
    return (axis<0)*(Ellipsis,) + axis*slc + (inds,) + (-1-axis)*slc 

To be used as in a[put_at(ind_list,axis=axis)]

I've seen both styles on numpy functions. This looks like one used for extend_dims, mine was used in apply_along/over_axis.

earlier thoughts

In a recent take question I/we figured out that it was equivalent to arr.flat[ind] for some some raveled index. I'll have to look that up.

There is an np.put that is equivalent to assignment to the flat:

Signature: np.put(a, ind, v, mode='raise')
Docstring:
Replaces specified elements of an array with given values.

The indexing works on the flattened target array. `put` is roughly
equivalent to:

    a.flat[ind] = v

Its docs also mention place and putmask (and copyto).

numpy multidimensional indexing and the function 'take'

I commented take (without axis) is equivalent to:

lut.flat[np.ravel_multi_index(arr.T, lut.shape)].T

with ravel:

In [257]: a = np.arange(12).reshape((3, -1))
In [258]: IJ=np.ix_(np.arange(a.shape[0]), inds)
In [259]: np.ravel_multi_index(IJ, a.shape)
Out[259]: 
array([[ 1,  2],
       [ 5,  6],
       [ 9, 10]], dtype=int32)
In [260]: np.take(a,np.ravel_multi_index(IJ, a.shape))
Out[260]: 
array([[ 1,  2],
       [ 5,  6],
       [ 9, 10]])
In [261]: a.flat[np.ravel_multi_index(IJ, a.shape)] = 100
In [262]: a
Out[262]: 
array([[  0, 100, 100,   3],
       [  4, 100, 100,   7],
       [  8, 100, 100,  11]])

and to use put:

In [264]: np.put(a, np.ravel_multi_index(IJ, a.shape), np.arange(1,7))
In [265]: a
Out[265]: 
array([[ 0,  1,  2,  3],
       [ 4,  3,  4,  7],
       [ 8,  5,  6, 11]])

Use of ravel is unecessary in this case but might useful in others.