Index into vector field in 2 or 3 dimensions
<p>I have <code>vecfield</code> which is an <code>ndarray</code> with <code>m</code> entries for each spatial location in a (discrete) 2 or 3 dimensional space with shape <code>(w, h)</code> or <code>(w, h, d)</code>. Now I get an <code>ndarray</code> with a list of indices <code>(i, j)</code> or <code>(i, j, k)</code> and I'd like to read the <code>m</code> entries corresponding to each index.</p>

<p>Code (for 2D) is probably more helpful:</p>

<pre><code>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)
</code></pre>

<p>It seems that idx gets "linearized" and then used as index, so it does not work that way. <strong>How can I achieve what I want?</strong></p>

<pre><code>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)
</code></pre>

<p>Indexing directly with <code>idx</code> effectively replaces that 1st dimension (of size 4) with this (3,2) array.  In other words, it picks <code>vec[i,:,:]</code> for all 6 values in <code>idx</code>.</p>

<p>What Divakar suggested:</p>

<pre><code>In [266]: vec[idx[:,0], idx[:,1],:]
Out[266]: 
array([[ 0,  1],
       [ 8,  9],
       [22, 23]])
</code></pre>

<p>This uses the 1st col of <code>idx</code> to index the 1st dimension, and the 2nd column to index the 2nd dimension.</p>

<p>List zip also works, but needs a <code>*</code>:</p>

<pre><code>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]])
</code></pre>

<p>Like the first answer, it indexes the 1st dim with (0,1,3).  It works if <code>idx</code> is a list of lists rather than an array.  An array equivalent would be:</p>

<pre><code>In [272]: tuple(idx.T)
Out[272]: (array([0, 1, 3]), array([0, 1, 2]))
</code></pre>

<p>In sum, to pick <code>n</code> items in 2 dimensions, you need 2 <code>n</code> element lists/arrays.  </p>

python numpy