Сначала я изучил вопрос NumPy: сортировка трехмерного массива с сохранением 2-го измерения, присвоенного первому , но принятый ответ не совсем подходит для моей проблемы, так как мне нужен полный диапазон возможных значений в 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] ... это, вероятно, простоверно?!
В конце концов, я хочу иметь возможностьСортируйте массив размером 50 x 11000 x 11000 x 2 в uint16 , чтобы он был максимально эффективным с точки зрения памяти.