CODEWITHSUNDEEP
pythonnumpy
jacob
jacob

Reputation: 899

Why does numpy indexing work differently in higher dimensions?

I have a large array and need to change some of its values, but get inconsistent results depending on the number of dimensions. A minimal example looks like this:

a=np.zeros((3,3))
indx=np.array([2,4,6])
a[np.unravel_index(indx, (3,3))] = 1

works as expected. However, the following doesn't:

b=np.zeros((1,3,3))
indx=np.array([2,4,6])
b[0,np.unravel_index(indx, (3,3))] = 1

and gives an array filled with 1.

Questions: Why are these not the same? And how can I mimick the behavior of the first example when I have in fact a higher dimensional array like in the second example?

Upvotes: 4

Views: 150

Answers (3)

shx2
shx2

Reputation: 64368

unravel_index(indx, (3,3)) returns a tuple of length 2, because len((3,3))==2. If you want to use that tuple to index a 3-dim array with your own 1st dim index, you need to prepend your index to the tuple, to make it a tuple of length==3:

b[ ([0,0,0],) + np.unravel_index(indx, (3,3)) ] = 1
b
=> array([[[ 0.,  0.,  1.],
           [ 0.,  1.,  0.],
           [ 1.,  0.,  0.]]])

Now, if the first dim is longer than 1, you can change [0,0,0] to whatever you'd like.

One more way to understand the problem:

Say:

k = np.unravel_index(indx, (3,3))

Then you want:

b[0,k[0],k[1]] = 1

Which is not the same as what you're attempting:

b[0,k] = 1

BTW, a simplified version also works. The first element of the tuple doesn't have to be an array:

j = (0,) + np.unravel_index(indx, (3,3))
j
=> (0, array([0, 1, 2]), array([2, 1, 0]))
b[j] = 1

Upvotes: 1

hpaulj
hpaulj

Reputation: 231665

What does unravel_index produce?

In [126]: indx=np.array([2,4,6])    
In [128]: i=np.unravel_index(indx,(3,3))    
In [129]: i
Out[129]: (array([0, 1, 2]), array([2, 1, 0]))

A tuple of 2 indexing arrays. Apply that to 2d array, and we see it sets the reverse diagonal.

In [130]: a=np.zeros((3,3),int)

In [131]: a[i]=1

In [132]: a
Out[132]: 
array([[0, 0, 1],
       [0, 1, 0],
       [1, 0, 0]])

How could we set the reverse diagonal on the 2nd plane of a 3d array? We need 3 indexes, right?

In [133]: a=np.zeros((2,3,3),int)

In [134]: a[1,[0,1,2],[2,1,0]]=1
#  a[(1,[0,1,2],[2,1,0])]=1  is the same thing

In [135]: a
Out[135]: 
array([[[0, 0, 0],
        [0, 0, 0],
        [0, 0, 0]],

       [[0, 0, 1],
        [0, 1, 0],
        [1, 0, 0]]])

So now the question is - how to go from (array([0, 1, 2]), array([2, 1, 0])) to (1,[0,1,2],[2,1,0])?

For the 1st plane, (1,3,3) works because it produces a 3 element tuple. (2,3,3) does the same thing.

In [140]: np.unravel_index(indx,(1,3,3))
Out[140]: (array([0, 0, 0]), array([0, 1, 2]), array([2, 1, 0]))

How about concatenating 2 tuples?

In [142]: (1,)+np.unravel_index(indx,(3,3))
Out[142]: (1, array([0, 1, 2]), array([2, 1, 0]))

There are number of numpy functions that construct indexes by mixing and matching tuples, lists and arrays.

You could also build the tuple from a list:

In [153]: ind=[1,0,0]
In [154]: ind[1:]=np.unravel_index(indx,(3,3))
In [155]: a[tuple(ind)]=2

Upvotes: 1

eigenchris
eigenchris

Reputation: 5821

The following should work:

b[np.unravel_index(indx, (1,3,3))] = 1

Upvotes: 1

Related Questions