Numpy: индекс 2 измерения сразу

Question

Я хочу присвоить новые значения массиву, по позициям, заданным некоторыми индексами. Пример будет более понятным:

import numpy as np

#Dimensions
N = 25
n = 50
d = 100
k = 3
p = 7

A = np.random.uniform(size=(N,n,d,d))

A_new_values = np.random.uniform(size=(N,n,k,p))
indexes_new_values = np.random.choice(np.arange(d),size=k*p,replace=False).reshape((k,p))

print(A[:,:,indexes_new_values].shape) 

Я хотел, чтобы последняя строка выдавала (N,n,k,p), чтобы иметь возможность назначать новые значения как:

A[:,:,indexes_new_values] = A_new_values

Но это приводит к ошибке , Как я могу присвоить A_new_values местам пропппера в A?

Answer 1

Если я правильно понимаю, я думаю, что вы можете делать то, что вы хотите с np.put_along_axis:

import numpy as np

#Dimensions
N = 25
n = 50
d = 100
k = 3
p = 7

np.random.seed(0)
A = np.random.uniform(size=(N, n, d, d))
A_new_values = np.random.uniform(size=(N, n, k, p))
indexes_new_values = np.random.choice(np.arange(d), size=k * p, replace=False).reshape((k, p))
np.put_along_axis(A, indexes_new_values.reshape(1, 1, -1, 1), A_new_values.reshape(N, n, -1, 1), axis=2)
print(np.all(A[10, 20, indexes_new_values[1, 5]] == A_new_values[10, 20, 1, 5]))
# True