I'm trying to manipulate an index and source array such that:

result[ i ][ j ][ k ] = source[ i ][ indices[ i ][ j ][ k ] ]

I know how to do this with for loops but I'm using giant arrays and I'd like to use something more time efficient. I've tried to use numpy's advanced indexing but I don't really understand it.

Example functionality:

source = [[0.0 0.1 0.2 0.3]
 [1.0 1.1 1.2 1.3]
 [2.0 2.1 2.2 2.3]]

indices = [[[3 1 0 1]
 [3 0 0 3]]

 [[0 1 0 2]
 [3 2 1 1]]

 [[1 1 0 1]
 [0 1 2 2]]]

# result[i][j][k] = source[i][indices[i][j][k]]

result = [[[0.3 0.1 0.0 0.1]
 [0.3 0.0 0.0 0.3]]

 [[1.0 1.1 1.0 1.2]
 [1.3 1.2 1.1 1.1]]

 [[2.1 2.1 2.0 2.1]
 [2.0 2.1 2.2 2.2]]]

Solution using Integer Advanced Indexing:

Given:

source = [[0.0, 0.1, 0.2, 0.3],
 [1.0, 1.1, 1.2, 1.3],
 [2.0, 2.1, 2.2, 2.3]]

indices = [[[3, 1, 0, 1],
 [3, 0, 0, 3]],
 [[0, 1, 0, 2],
 [3, 2, 1, 1]],
 [[1, 1, 0, 1],
 [0, 1, 2, 2]]]

Use this:

import numpy as np
nd_source = np.array(source)

source_rows = len(source) # == 3, in above example
source_cols = len(source[0]) # == 4, in above example

row_indices = np.arange(source_rows).reshape(-1,1,1)
result = nd_source [row_indices, indices]

Result:

print (result)
[[[0.3 0.1 0. 0.1]
 [0.3 0. 0. 0.3]]

 [[1. 1.1 1. 1.2]
 [1.3 1.2 1.1 1.1]]

 [[2.1 2.1 2. 2.1]
 [2. 2.1 2.2 2.2]]]

Explanation:

To use Integer Advanced Indexing, the key rules are:

1. We must supply index arrays consisting of integer indices.


2. We must supply as many of these index arrays, as there are dimensions in the source array.


3. The shape of these index arrays must be the same, or, at least all of them must be broadcastable to a single final shape.

How the Integer Advanced Indexing works is:

Given that source array has n dimensions, and that we have therefore supplied n integer index arrays:

1. All of these index arrays, if not in the same uniform shape, will be broadcasted to be in a single uniform shape.


2. To access any element in the source array, we obviously need an n-tuple of indices. Therefore to generate the result array from the source array, we need several n-tuples, one n-tuple for each element-position of the final result array. For each element-position of the result array, the n-tuple of indices will be constructed from the corresponding element-positions in the broadcasted index arrays. (Remember the result array has exactly the same shape as the broadcasted index arrays, as already mentioned above).


3. Thus, by traversing the index arrays in tandem, we get all the n-tuples we need to generate the result array, in the same shape as the broadcasted index arrays.

Applying this explanation to the above example:

1. Our source array is nd_source = np.array(source), which is 2d.



2. Our final result shape is (3,2,4).




3. We therefore need to supply 2 index arrays, and these index arrays must either be in the final result shape of (3,2,4), or broadcastable to the (3,2,4) shape.




4. Our first index array is row_indices = np.arange(source_rows).reshape(-1,1,1). (source_rows is the number of rows in the source, which is 3 in this example) This index array has shape (3,1,1), and actually looks like [[[0]],[[1]],[[2]]]. This is broadcastable to the final result shape of (3,2,4), and the broadcasted array looks like [[[0,0,0,0],[0,0,0,0]],[[1,1,1,1],[1,1,1,1]],[[2,2,2,2],[2,2,2,2]]].




5. Our second index array is indices. Though this is not an array and is only a list of lists, numpy is flexible enough to automatically convert it into the corresponding ndarray, when we pass it as our send index array. Note that this array is already in the final desired result shape of (3,2,4) even without any broadcasting.




6. Traversing these two index arrays in tandem (one a broadcasted array and the other as is), numpy generates all the 2-tuples needed to access our source 2d array nd_source, and generate the final result in the shape (3,2,4).