对 4D numpy 数组进行排序，但将一个轴连接在一起答案

【问题标题】：Sorting 4D numpy array but keeping one axis tied together对 4D numpy 数组进行排序，但将一个轴连接在一起
【发布时间】：2018-10-23 15:46:05
【问题描述】：

首先，我查看了问题NumPy: sorting 3D array but keeping 2nd dimension assigned to first，但接受的答案不太适合我的问题，因为我需要 uint16 中可能的全部值并且不想去到 int32 以避免使用过多的内存。

我的问题是我有一堆 3D 数组（它们是每个具有两个波段的图像），我想沿着堆栈的轴（按第一个波段的值）进行排序，但是通过将两个波段保留在一起每张图片的...我希望这有点清楚。

生成类似于我所拥有的数组的代码：

import numpy as np 
# Here a stack of three 3x2 images containing two bands each
arr = np.zeros((3,3,2,2), 'uint16')

np.random.seed(5)
arr[0,:,:,0] = np.random.randint(10, 90, 6).reshape(3,2)
arr[0,:,:,1] = 51
arr[1,:,:,0] = np.random.randint(10, 90, 6).reshape(3,2)
arr[1,:,:,1] = 52
arr[2,:,:,0] = np.random.randint(10, 90, 6).reshape(3,2)
arr[2,:,:,1] = 50
arr[np.where(arr >= 85)] = 99 #just to have couple identical values like my dataset has

>>> arr
array([[[[99, 51],
         [71, 51]],

        [[26, 51],
         [83, 51]],

        [[18, 51],
         [72, 51]]],


       [[[37, 52],
         [40, 52]],

        [[17, 52],
         [99, 52]],

        [[25, 52],
         [63, 52]]],


       [[[37, 50],
         [54, 50]],

        [[99, 50],
         [99, 50]],

        [[75, 50],
         [57, 50]]]], dtype=uint16)

因为我希望对堆栈进行排序，所以我使用了arr_sorted = np.sort(arr, axis=0)，但这打破了每个栅格的两个波段之间的链接：

>>> arr[0,2,1,:]
array([72, 51], dtype=uint16)

>>> arr_sorted[2,2,1,:]
array([72, 52], dtype=uint16) #value 72 is no longer tied to 51 but to 52

我可以使用idx = np.argsort(arr[:,:,:,0], axis=0) 来获得我想要的排序索引，但我没有找到如何使用idx 来对arr[:,:,:,0] 和arr[:,:,:,1] 进行相同的排序...这可能很容易吧？！

最终，我希望能够在 uint16 中对一个 50 x 11000 x 11000 x 2 的数组进行排序，因此它需要尽可能节省内存。

【问题讨论】：

可以构造一组用idx 广播的索引来做这种排序，但是有点乱。最近的numpy 版本添加了一个np.take_along_axis 函数来处理细节。尝试使用idx[...,None] 以获得正确的维数。

标签： python arrays sorting numpy

【解决方案1】：

使用新的take_along_axis：

In [351]: arr = np.random.randint(0,10,(3,3,2,2))
In [352]: idx = np.argsort(arr[...,0], axis=0)
In [353]: idx.shape
Out[353]: (3, 3, 2)
In [354]: arr1 = np.take_along_axis(arr, idx[...,None], axis=0)

【讨论】：