Sorting a three-dimensional numpy array by a two-dimensional index

Question

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!

What would be an example of ph_idx in your three-dimensional case? — fuglede
– fuglede, Commented Nov 4, 2019 at 9:17
@fuglede - it would be (apologies for the formatting here) np.array([[[1,2,3]],[[1,2,3]]]) — Chris J Harris
– Chris J Harris, Commented Nov 4, 2019 at 9:32
Formatting is as good as it gets. Not [3, 2, 1] this time? Could you also provide an example where ph has shape (2, 2, 3)? — fuglede
– fuglede, Commented Nov 4, 2019 at 9:33
So, I didn't put this into the question because I didn't want to overcomplicate things, but ph might have a dim1 of length greater than 1 (like the 2, 2, 3 shape you suggested). If so I'm trying to apply the 2-d array. In something like numpy pseudocode, this would be ph[ph_idx[:, None, :]] — Chris J Harris
– Chris J Harris, Commented Nov 4, 2019 at 9:36

yatu · Accepted Answer · 2019-11-04 09:45:42Z

2

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]])

edited Nov 4, 2019 at 9:45

answered Nov 4, 2019 at 9:21

yatu

88.6k12 gold badges93 silver badges148 bronze badges

Sign up to request clarification or add additional context in comments.

6 Comments

fuglede Over a year ago

ph is the 3-dimensional array in their example, not ph_idx. That said, some variation of np.take_along_axis is probably what they're after.

kabanus Over a year ago

That's a new function to me! Not available on 1.10 at least.

Chris J Harris Over a year ago

Thanks - I agree - take_along_axis looks like the function to use here.

fuglede Over a year ago

@Chrisper: If all of your examples follow what you've outlined, you won't even need np.squeeze either; plain np.take_along_axis does the job.

fuglede Over a year ago

@yatu: If you use the ph_idx from the comments to OP, you can even drop np.squeeze on ph.

|

Chris J Harris · Accepted Answer · 2019-11-04 11:19:09Z

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]]])

Collectives™ on Stack Overflow

Sorting a three-dimensional numpy array by a two-dimensional index

2 Answers 2

6 Comments

Comments

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

6 Comments

Comments

Related