3

Let say I have a NumPy array A of shape (66,5) and B of shape (100, 66, 5).

The elements of A will index the first dimension (axis=0) of B, where the values are from 0 to 99 (i.e. the first dimension of B is 100).

A = 
array([[   1,    0,    0,    1,    0],
       [   0,    2,    0,    2,    4],
       [   1,    7,    0,    5,    5],
       [   2,    1,    0,    1,    7],
       [   0,    7,    0,    1,    4],
       [   0,    0,    3,    6,    0]
       ....                         ]])

For example, A[4,1] will take index 7 of the first dimension of B, index 4 of the second dimension of B and index 1 of the third dimension B.

What I wanted to is to produce array C of shape (66,5) where it contains the elements in B that are selected based on the elements in A.

Improve this question
2
  • A toy example would help us understand better. Is A a boolean mask or index array? and when you say axis=0 do you mean the first element or first dimension? (axis refers to dimension and I cannot see how a 2-D index on all elements of axis=0 of 3-D array will result in a 2-D array. Commented Aug 26, 2020 at 8:37
  • @Ehsan It is hard to give a toy example. But, I have tried my best to explain in the edit. Commented Aug 26, 2020 at 9:45

3 Answers 3

Reset to default
6

You can use np.take_along_axis to do that:

import numpy as np
np.random.seed(0)
a = np.random.randint(100, size=(66, 5))
b = np.random.random(size=(100, 66, 5))
c = np.take_along_axis(b, a[np.newaxis], axis=0)[0]
# Test some element
print(c[25, 3] == b[a[25, 3], 25, 3])
# True
Improve this answer
Sign up to request clarification or add additional context in comments.

2 Comments

What is [0] at the end of c = np.take_along_axis(b, a[np.newaxis], axis=0)[0]?
@Aqee np.take_along_axis requires both arguments to have the same number of dimensions, although they can be broadcasted. So, I add an initial singleton dimension to a with [np.newaxis], and then the result is (1, 66, 5), so I squeeze out that dimension from the result with [0]. It would have been equivalent to do np.expand_dims(a, 0) and np.squeeze(np.take_along_axis(...), 0).
0

If I understand correctly, you are looking for advances indexing of first dimension of B. You can use np.indices to create the indices required for the other two dimensions of B and use advanced indexing:

idx = np.indices(A.shape)
C = B[A,idx[0],idx[1]]

Example:

B = np.random.rand(10,20,30)
A = np.array([[   1,    0,    0,    1,    0],
       [   0,    2,    0,    2,    4],
       [   1,    7,    0,    5,    5],
       [   2,    1,    0,    1,    7],
       [   0,    7,    0,    1,    4],
       [   0,    0,    3,    6,    0]])

print(C[4,1]==B[7,4,1])
#True
Improve this answer

Comments

-1

Use the following (using functions of NumPy library):

print(A)
# array([[2, 0],
#        [1, 1],
#        [2, 0]])

print(B)
# array([[[ 5,  7],
#         [ 0,  0],
#         [ 0,  0]],

#        [[ 1,  8],
#         [ 1,  9],
#         [10,  1]],

#        [[12, 22],
#         [ 2,  2],
#         [ 2,  2]]])

temp = A.reshape(-1) + np.cumsum(np.ones([A.reshape(-1).shape[0]])*B.shape[0], dtype = 'int') - 3
C = B.swapaxes(0, 1).swapaxes(2, 1).reshape(-1)[temp].reshape(A.shape)

print(C)
# array([[12,  7],
#        [ 1,  9],
#        [ 2,  0]])
Improve this answer

3 Comments

I would suggest adding explanation on what your code does.
Hopefully, you can explain what you suggested
It does what you asked for. I have added an example, do check.

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.