CODEWITHSUNDEEP
numpyindexing
lrnv
lrnv

Reputation: 1106

Numpy : index 2 dimensions at once

I want to assign new values to an array, on positions given by some indexes. An exemple will be more clear : 

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)

I wanted the last line to yield (N,n,k,p), to be able to assign new values as : 

A[:,:,indexes_new_values] = A_new_values

But this yield an error. How can i assign A_new_values to the propper places in A ?

Upvotes: 1

Views: 66

Answers (1)

javidcf
javidcf

Reputation: 59711

If I understand correctly, I think you can do what you want with 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

Upvotes: 1

Related Questions