Numpy - Sorting two ndarrays by single axis of first array
Question: I'm trying to sort two large four dimensional arrays in numpy.

I want to sort based on the values axis 2 of the first array, and sort the
second array by the same indices. All other axes should remain in the same
order for both arrays.

The following code does what I want, but relies on looping in python, so it's
slow. The arrays are quite large, so I'd really like to get this working using
compiled numpy operations for performance reasons. Or some other means of
getting this block of code to be compiled (Cython?).

    
    
    import numpy as np
    
    data = np.random.rand(10,6,4,1)
    data2 = np.random.rand(10,6,4,3)
    
    print data[0,0,:,:]
    print data2[0,0,:,:]
    
    for n in range(data.shape[0]):
      for m in range(data.shape[1]):
    
        sort_ids = np.argsort(data[n,m,:,0])
    
        data[n,m,:,:] = data[n,m,sort_ids,:]
        data2[n,m,:,:] = data2[n,m,sort_ids,:]
    
    
    print data[0,0,:,:]
    print data2[0,0,:,:]
    


Answer: Maybe there is a better solution but this should work:

    
    
    sort_ids = np.argsort(data,axis=2)
    
    s1 = data.shape
    s2 = data2.shape
    d1 = data[np.arange(s1[0])[:,None,None,None],np.arange(s1[1])[None,:,None,None],sort_ids,np.arange(s1[3])[None,None,None,:]]
    d2 = data2[np.arange(s2[0])[:,None,None,None],np.arange(s2[1])[None,:,None,None],sort_ids,np.arange(s2[3])[None,None,None,:]]
    

At least the output is identical to your code.

