I'm trying to take a list of elements from an 2D numpy array with given list of coordinates and I want to avoid using loop. I saw that np.take works with 1D array but I can't make it work with 2D arrays.
Example:
a = np.array([[1,2,3], [4,5,6]])
print(a)
# [[1 2 3]
# [4 5 6]]
np.take(a, [[1,2]])
# gives [2, 3] but I want just [6]
I want to avoid loop because I think that will be slower (I need speed). But if you can persuade me that a loop is as fast as an existing numpy function solution, then I can go for it.
2 Answers
If I understand it correctly, you have a list of coordinates like this:
coords = [[y0, x0], [y1, x1], ...]
To get the values of array a at these coordinates you need:
a[[y0, y1, ...], [x0, x1, ...]]
So a[coords] will not work. One way to do it is:
Y = [c[0] for c in coords]
X = [c[1] for c in coords]
or
Y = np.transpose(coords)[0]
X = np.transpose(coords)[1]
Then
a[Y, X]
Does fancy indexing do what you want? np.take seems to flatten the array before operating.
import numpy as np
a = np.arange(1, 10).reshape(3,3)
a
# array([[1, 2, 3],
# [4, 5, 6],
# [7, 8, 9]])
rows = [ 1,1,2,0]
cols = [ 0,1,1,2]
# Use the indices to access items in a
a[rows, cols]
# array([4, 5, 8, 3])
a[1,0], a[1,1], a[2,1], a[0,2]
# (4, 5, 8, 3)
