Luca
Luca

Reputation: 11016

Python numpy array indexing. How is this working?

I came across this python code (which works) and to me it seems amazing. However, I am unable to figure out what this code is doing. To replicate it, I sort of wrote a test code:

import numpy as np

# Create a random array which represent the 6 unique coeff. 
# of a symmetric 3x3 matrix
x = np.random.rand(10, 10, 6)

So, I have 100 symmetric 3x3 matrices and I am only storing the unique components. Now, I want to generate the full 3x3 matrix and this is where the magic happens.

indices = np.array([[0, 1, 3],
                    [1, 2, 4],
                    [3, 4, 5]])

I see what this is doing. This is how the 0-5 index components should be arranged in the 3x3 matrix to have a symmetric matrix.

mat = x[..., indices]

This line has me lost. So, it is working on the last dimension of the x array but it is not at all clear to me how the rearrangement and reshaping is done but this indeed returns an array of shape (10, 10, 3, 3). I am amazed and confused!

Upvotes: 4

Views: 134

Answers (1)

hpaulj
hpaulj

Reputation: 231690

From the advanced indexing documentation - bi rico's link.

Example

Suppose x.shape is (10,20,30) and ind is a (2,3,4)-shaped indexing intp array, thenresult = x[...,ind,:] has shape (10,2,3,4,30) because the (20,)-shaped subspace has been replaced with a (2,3,4)-shaped broadcasted indexing subspace. If we let i, j, kloop over the (2,3,4)-shaped subspace then result[...,i,j,k,:] =x[...,ind[i,j,k],:]. This example produces the same result as x.take(ind, axis=-2).

Upvotes: 1

Related Questions