2

I want to reorder the nth axis of a ndarray of abitrary dimension d according to a list of indices order. If the axis was the last one, the solution from this question (Reordering last dimension of a numpy ndarray) would be enough. In my case, however, the axis is not, in general, the first or the last one, so that Ellipsis alone does not solve the issue.

This is the solution I have come up so far:

axes_list = list(range(d))
axes_list[0], axes_list[i] = axes_list[i], axes_list[0]

ndarr = np.transpose(ndarr, axes=axes_list)[order,...]   # Switches the axis with the first and reorders
ndarr = np.transpose(ndarr, axes=axes_list)              # Switches the axes back

What I don't like about this solution is that I have to manually transpose the ndarray. I was wondering if there is a generalization of the Ellipsis operator such that we can account for a chosen number of axes, such that

ndarr[GenEllipsis(n),order,...]

would skip the n first axes and would reorder the (n+1)th one.

Is such a thing possible?

Lucas Baldo
  • 130
  • 7
  • 4
    Use `n` ":". More generally you could make a `alist = [[slice(None)]*d`; replace one or more of the slices with the order(s). Finally index ing `tupe(alist)`. This approach is quite general. – hpaulj May 20 '21 at 00:54
  • I'm not sure I understand, could you elaborate? Going for ` arr[:,:,:(n times),:,:] ` does not work because I have to specify the number of dimensions when writing the code, and I don't know that in advance (the number of dimensions is arbitrary). – Lucas Baldo May 20 '21 at 01:04

1 Answers1

1

Use the command np.take_along_axis and assign the output to a new variable. See the code below:

arr = np.random.randn(10,3,23,42,3)
ax = 3 #change this to your 'random' axis
order = np.random.permutation(list(range(arr.shape[ax])))
#order needs to have the same number of dims as arr
order = np.expand_dims(order,tuple(i for i in range(len(arr.shape)) if i != ax)) 
shuff_arr = np.take_along_axis(arr,order,ax)

@hpaulj's comment is also a valid answer (and likely better, please give them an upvote!):

arr = np.random.randn(10,3,23,42,3)
ax = 3 #change this to your 'random' axis
order = np.random.permutation(list(range(arr.shape[ax])))
#set the correct dim to 'order'
alist = [slice(None)]*len(arr.shape)
alist[ax] = order
shuff_arr = arr[tuple(alist)]
jhso
  • 3,103
  • 1
  • 5
  • 13
  • The variable i on the last line is not defined. Was it a typo? – Lucas Baldo May 20 '21 at 01:54
  • 1
    *facepalm* - fixed it, was just a typo – jhso May 20 '21 at 01:55
  • Thanks for the input. The first solution works great. I still don't understand it completely, since I had never used the np.take_along_axis function and it seems like it it does a lot of complicated things. The second solution doesn't seem to be working for me, though. – Lucas Baldo May 20 '21 at 02:22
  • 1
    No? Is it an error or are the results incorrect? – jhso May 20 '21 at 02:31
  • 1
    I think there was another typo haha, the second shuff_arr on the last line should be just arr, right? – Lucas Baldo May 20 '21 at 02:46
  • 1
    obviously missed my coffee this morning... good spot though. It's good to work through some errors sometimes... – jhso May 20 '21 at 04:15