Index numpy array by matrix of two arrays

I have a 2D numpy array like

weights = np.array(
    [
        [1, 2, 3],
        [4, 5, 6],
        [7, 8, 9],
    ]
)

And i have two 1-D numpy arrays that i want to use to index into weights

positions1 = np.array([2, 1, 0])
positions2 = np.array([1, 1, 0])

If i want the results for stepping through the arrays together and indexing into the matrix that way i can do

print(weights[positions1[...], positions2[...]])

And get [8 5 1]

However now i want to index into the matrix with all possible combinations of positions1 and positions2 so that i get a matrix like

[
[weights[2, 1], weights[2, 1], weights[2, 0]],
[weights[1, 1], weights[1, 1], weights[1, 0]],
[weights[0, 1], weights[0, 1], weights[0, 0]],
]

So

[
[weights[pos1[0], pos2[0]], weights[pos1[0], pos2[1]], weights[pos1[0], pos2[2]]],
[weights[pos1[1], pos2[0]], weights[pos1[1], pos2[1]], weights[pos1[1], pos2[2]]],
[weights[pos1[2], pos2[0]], weights[pos1[2], pos2[1]], weights[pos1[2], pos2[2]]],
]

What would be the canonical way to do that in numpy?
I know that this is kinda like an outer product but i dont actually want to multiply the values in my array but just get a tuple of their values to index into the matrix with.

>Solution :

You need change the shape of the first indexer:

weights[positions1[:,None], positions2]

Or for a generalized version, as pointed out by Chrysophylaxs:

weights[np.ix_[positions1, positions2]]

Output:

array([[8, 8, 7],
       [5, 5, 4],
       [2, 2, 1]])

Leave a Reply