Indexing a 4D NumPy Array with two 2D arrays-CodePudding

I have a 4D target NumPy array which I want to fill with values from a 2D source array, using two additional 2D arrays which specify the position in the second and third axis of the target array where the value from the source array should be placed. The code below with some sample values can do this using a for-loop approach:

import numpy as np

T, M, N, D = 5, 4, 3, 2  # Example dimensions
target = np.zeros((T, M, N, D))  # Target array
source = np.arange(T*D).reshape(T, D)  # Example source array with arbitrary numbers
index_dim_1 = np.array([  # Array of shape T, D with values between 0 and M-1
    [0, 3],
    [1, 2],
    [2, 1],
    [3, 0],
    [0, 3]
])
index_dim_2 = np.array([  # Array of shape T, D with values between 0 and N-1
    [0, 2],
    [1, 1],
    [2, 0],
    [0, 2],
    [1, 1]
])

for t in range(T):
    for d in range(D):
        m = index_dim_1[t, d]
        n = index_dim_2[t, d]
        target[t, m, n, d] = source[t, d]

How could I vectorise this approach? I know if the target array was 2D, and the source and index arrays were both 1D, I could use this approach, but I'm not sure how to generalise it to my case.

CodePudding user response：

for t in range(T):
    for d in range(D):
        n = index_dim_1[t, d]
        m = index_dim_2[t, d]
        target[t, n, m, d] = source[t, d]

Since you provide code, but no example, I'll skip that step myself, and 'eyeball' an answer - without testing.

 target[np.arange(T)[:,None], index_dim_1, index_dim_2, np.arange(D)] = source

The idea is to create index arrays for t and d that broadcast with the (T,D) shaped arrays.