时间戳

首页 > 解决方案 > 对 4D numpy 数组进行排序，但将一个轴连接在一起

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

问题描述

首先，我研究了问题NumPy:sorting 3D array but keep 2nd dimension assignment 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 的数组进行排序，因此它需要尽可能地提高内存效率。

标签： pythonarrayssortingnumpy

解决方案

使用新的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)

推荐阅读