NumPy: Evaulate chỉ số mảng trong vectorized phân

Question

Tôi muốn vectorize hoạt động này NumPy:

for j in range(yt): 
    for i in range(xt): 
     y[j, i] = x[idx[j, i], j, i] 

nơi idx chứa trục-0 index để một lát x. Có cách nào đơn giản để làm điều này?

Answer 1

Bạn có thể sử dụng:

J, I = np.ogrid[:yt, :xt] 
x[idx, J, I] 

Đây là thử nghiệm:

import numpy as np 

yt, xt = 3, 5 
x = np.random.rand(10, 6, 7) 
y = np.zeros((yt, xt)) 
idx = np.random.randint(0, 10, (yt, xt)) 

for j in range(yt): 
    for i in range(xt): 
     y[j, i] = x[idx[j, i], j, i] 

J, I = np.ogrid[:yt, :xt] 
np.all(x[idx, J, I] == y) 

Answer 2

Dưới đây là một cách tiếp cận sử dụng linear indexing -

zt,yt,xt = x.shape 
out = x.reshape(zt,-1)[idx.ravel(),np.arange(yt*xt)].reshape(-1,xt) 

---

kiểm tra Runtime & xác minh đầu ra

Phần này so sánh phương pháp được đề xuất trong bài đăng này và other orgid based solution về hiệu suất và cũng xác minh kết quả đầu ra.

định nghĩa Chức năng -

def original_app(x,idx): 
    _,yt,xt = x.shape 
    y = np.zeros((yt,xt)) 
    for j in range(yt): 
     for i in range(xt): 
      y[j, i] = x[idx[j, i], j, i] 
    return y 

def ogrid_based(x,idx): 
    _,yt,xt = x.shape 
    J, I = np.ogrid[:yt, :xt] 
    return x[idx, J, I] 

def reshape_based(x,idx):        
    zt,yt,xt = x.shape 
    return x.reshape(zt,-1)[idx.ravel(),np.arange(yt*xt)].reshape(-1,xt) 

cài đặt đầu vào -

In [56]: # Inputs 
    ...: zt,yt,xt = 100,100,100 
    ...: x = np.random.rand(zt,yt,xt) 
    ...: idx = np.random.randint(0,zt,(yt,xt)) 
...: 

Xác nhận kết quả đầu ra -

In [57]: np.allclose(original_app(x,idx),ogrid_based(x,idx)) 
Out[57]: True 

In [58]: np.allclose(original_app(x,idx),reshape_based(x,idx)) 
Out[58]: True 

Thời gian -

In [68]: %timeit original_app(x,idx) 
100 loops, best of 3: 6.97 ms per loop 

In [69]: %timeit ogrid_based(x,idx) 
1000 loops, best of 3: 391 µs per loop 

In [70]: %timeit reshape_based(x,idx) 
1000 loops, best of 3: 230 µs per loop