I'm having real trouble with this. I have a three-dimensional numpy array, and I'd like to re-order it by a two-dimensional indexing array. In reality the arrays will be determined programatically and the three-dimensional array may be two or four-dimensioned, but to keep things simple, here's the desired outcome if both arrays were two-dimensional:
ph = np.array([[1,2,3], [3,2,1]])
ph_idx = np.array([[0,1,2], [2,1,0]])
for sub_dim_n, sub_dim_ph_idx in enumerate(ph_idx):
ph[sub_dim_n] = ph[sub_dim_n][sub_dim_ph_idx]
This makes the ph array into:
array([[1, 2, 3],
[1, 2, 3]])
Which is what I'd like. Could anyone please help if it's the same circumstance, but instead of ph, I have a three-dimensional array (psh), like:
psh = np.array(
[[[1,2,3]],
[[3,2,1]]]
)
Hope that's clear and please ask if it's not. Thanks in advance!
If what you want is to end up with a ph.shape
shaped array, you could simply np.squeeze
ph_ixs
so the shapes match, and use it to index ph
:
print(ph)
[[[1 2 3]]
[[3 2 1]]]
print(ph_idx)
[[0 1 2]
[2 1 0]]
np.take_along_axis(np.squeeze(ph), ph_idx, axis=-1)
array([[1, 2, 3],
[1, 2, 3]])
So, the clues are already in the helpful comments here, but for completeness, it's as simple as using np.take_along_axis and a broadcasted version of the 2d array:
psh = np.array(
[[[1,2,3]],
[[3,2,1]]]
)
ph_idx = np.array(
[[0,1,2],
[2,1,0]]
)
np.take_along_axis(psh, ph_idx[:, None, :], axis=2)
This has the advantage of also working if the 3d array has a dim1 of more than one element:
psh = np.array(
[[[1,2,3],[4,5,6],[7,8,9]],
[[3,2,1],[6,5,4],[9,8,7]]]
)
ph_idx = np.array([[0,1,2], [2,1,0]])
np.take_along_axis(psh, ph_idx[:, None, :], axis=2)
which gives
array([[[1, 2, 3],
[4, 5, 6],
[7, 8, 9]],
[[1, 2, 3],
[4, 5, 6],
[7, 8, 9]]])
The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.