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 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?

the point is that I want to be able to do a[inds, :] = 0, or a[:, inds] = 0 or (in the 3-D case) a[:, :, inds] = 0, etc. programmatically, without writing a bunch of if statements — arghdos
– arghdos, Commented Mar 7, 2017 at 19:50

Community · Accepted Answer · 2017-05-23 11:46:16Z

7

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.

edited May 23, 2017 at 11:46

CommunityBot

11 silver badge

answered Mar 7, 2017 at 19:59

hpaulj

233k14 gold badges260 silver badges392 bronze badges

Sign up to request clarification or add additional context in comments.

3 Comments

user2357112 Over a year ago

You seem to be overlooking the use of the axis argument.

arghdos Over a year ago

Thanks! It'd be nice if put had similar functionality to take with respect to the axis argument, but the construction of the indexing tuple works for me

Paul Panzer Over a year ago

Here is a convenience function for index construction which doesn't need to know the array's ndim: 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)]. @hpaulj feel free to add this to your post if it passes muster.

pyano · Accepted Answer · 2017-07-24 19:50:40Z

0

I have given an example for use of numpy.take in 2 dimensions. Perhaps you can adapt that to your problem

answered Jul 24, 2017 at 19:50

pyano

1,99812 silver badges32 bronze badges

Comments

B. M. · Accepted Answer · 2017-03-07 20:22:16Z

-1

You can juste use indexing in this way :

a[:,[1,2]]=0

answered Mar 7, 2017 at 20:22

B. M.

18.7k2 gold badges40 silver badges56 bronze badges

Collectives™ on Stack Overflow

Numpy assignment like 'numpy.take'

3 Answers 3

PP's suggestion:

earlier thoughts

3 Comments

Comments

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

PP's suggestion:

earlier thoughts

3 Comments

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related