二维或三维向量场的索引

Question

我有vecfield，这是一个ndarray，在形状为(w, h)或(w, h, d)的（离散的）二维或三维空间中，每个空间位置都有m条目。现在我得到一个ndarray，其中有一个索引列表(i, j)或(i, j, k)，我想读取每个索引对应的m条目。你知道吗

代码（对于2D）可能更有帮助：

vecfield = np.zeros([w, h, m]) 
fill_me_with_data(vecfield)
idx = generate_indices() 
# idx is now for example ndarray with [[0, 0], [9, 8], [15, 6], [9, 1]]
result = vecfield[idx, :]
# result should now be 4 x m (4 because 'idx' has 4 indices)

看起来idx被“线性化”，然后用作索引，所以它不是这样工作的。我怎样才能实现我想要的？

Answer 1

In [258]: vec=np.arange(24).reshape(4,3,2)
In [263]: idx=np.array([[0,0],[1,1],[3,2]])

In [264]: idx.shape
Out[264]: (3, 2)
In [265]: vec[idx,:].shape
Out[265]: (3, 2, 3, 2)

直接用idx索引可以有效地用这个（3,2）数组替换第一个维度（大小为4）。换句话说，它为idx中的所有6个值选择vec[i,:,:]。你知道吗

Divakar建议：

In [266]: vec[idx[:,0], idx[:,1],:]
Out[266]: 
array([[ 0,  1],
       [ 8,  9],
       [22, 23]])

它使用idx的第一列索引第一个维度，第二列索引第二个维度。你知道吗

List zip也可以工作，但需要*：

In [267]: list(zip(*idx))
Out[267]: [(0, 1, 3), (0, 1, 2)]
In [268]: vec[list(zip(*idx))]
Out[268]: 
array([[ 0,  1],
       [ 8,  9],
       [22, 23]])

与第一个答案一样，它用（0,1,3）索引第一个dim。如果idx是一个列表列表而不是一个数组，那么它就可以工作。等效数组为：

In [272]: tuple(idx.T)
Out[272]: (array([0, 1, 3]), array([0, 1, 2]))

总之，要在2维中选取n项，需要2个n元素列表/数组。你知道吗

Answer 2

您应该了解numpy高级索引的工作方式http://docs.scipy.org/doc/numpy/reference/arrays.indexing.html。我想你想要的是：

vecfield[zip(idx)]

zip(idx)返回第一个组件的列表，后跟第二个组件的列表，这是高级索引需要传递的内容。你知道吗