I have a numpy array and corresponding row and column indices:
mat = np.array([[0, 1, 2],
[3, 4, 5],
[6, 7, 8]])
row_idx = np.array([0, 0, 0,
0, 0, 0,
1, 1, 1])
col_idx = np.array([0, 1, 2,
0, 1, 2,
0, 1, 2])
I would like to unravel the matrix by the groups specified in row_idx and col_idx:
result = np.array([0, 3, 1, 4, 2, 5, 6, 7, 8])
I can do this with numpy_groupies.aggregate(), but can have weird return values:
import numpy as np
import numpy_groupies as npg
mat = np.array([[0, 1, 2],
[3, 4, 5],
[6, 7, 8]])
row_idx = np.array([0, 0, 0,
0, 0, 0,
1, 1, 1])
col_idx = np.array([0, 1, 2,
0, 1, 2,
0, 1, 2])
result = npg.aggregate(group_idx=np.vstack([row_idx, col_idx]), a=mat.reshape(-1), func='array')
Produces:
[[array([0, 3]) array([1, 4]) array([2, 5])]
[array([6]) array([7]) array([8])]]
I have tried aggregate() from the numpy_groupies package, but it is both slow and returns an array combining np.arrays and can also include ints rather than arrays in certain cirumstances (which makes further manipulation difficult and slow).
Why ravel doesn't work
The problem with ravel is that it does not generalise to where the matrix has areas that are grouped by rows and other areas by columns based on the row_idx and col_idx. There can be a combination of both row and column grouping.