Question

首先，我调查了问题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的数组进行排序，因此它需要尽可能提高内存效率。

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

排序4D numpy数组，但将一根轴绑在一起

1 个答案: