Indexing arrays with other arrays
Since I can't seem to fully internalise how
numpy advanced indexing works, here's the solution to a common indexing task, written out for future reference.
Given an array
a of dimension
[n, m, l] and an index array
idx of dimension
[n, k <= m], I'd like to define a new array
b such that
b[i, j, :] = a[i, idx[i, j], :].
Let's make an example and a naive solution:
import numpy as np a = np.random.random((3, 3, 2)) idx = np.array([[1, 2], [0, 2], [0, 1]]) naive = np.zeros((*idx.shape, 2)) for i in range(idx.shape): for j in range(idx.shape): naive[i, j] = a[i, idx[i, j]]
The indexing-based solution, from this blog post is:
b = a[np.arange(a.shape)[:, None], idx, :] # verify solution np.testing.assert_array_equal(naive, b)
This also works with
torch tensors and seems reasonably fast. If I notice any performance problems, I'll update this post with a (hopefully) more efficient version...