Argmax многомерного массива вдоль подмножества измерений в Matlab

Question

Скажем, Y - это 7-мерный массив, и мне нужен эффективный способ максимизировать его по последним 3 измерениям, который будет работать на GPU. В результате мне нужен 4-мерный массив с максимальными значениями Y и три 4-мерных массива с индексами этих значений в трех последних измерениях. я могу сделать

[Y7, X7] = max(Y , [], 7);
[Y6, X6] = max(Y7, [], 6);
[Y5, X5] = max(Y6, [], 5);

Тогда я уже нашел значения (Y5) и индексы по 5-му измерению (X5). Но мне все еще нужны индексы по 6-му и 7-му измерениям.

Answer 1

Вот способ сделать это. Позволять N Обозначим количество измерений, по которым можно максимизировать.

1. Переформуйте Y свернуть последний N Размеры в один.
2. Максимизируйте по свернутым размерам. Это дает argmax как линейный индекс по этим измерениям.
3. Разверните линейный индекс в N подиндексы, по одному для каждого измерения.

Следующий код работает для любого количества измерений (не обязательно 7 а также 3 как в вашем примере). Чтобы достичь этого, он обрабатывает размер Y в общем и использует разделенный запятыми список, полученный из массива ячеек, чтобы получить N выходы из sub2ind,

Y = rand(2,3,2,3,2,3,2); % example 7-dimensional array
N = 3; % last dimensions along which to maximize
D = ndims(Y);
sz = size(Y);
[~, ind] = max(reshape(Y, [sz(1:D-N) prod(sz(D-N+1:end))]), [], D-N+1);
sub = cell(1,N);
[sub{:}] = ind2sub(sz(D-N+1:D), ind);

В качестве проверки, после запуска вышеуказанного кода, посмотрите, например, Y(2,3,1,2,:) (для удобства показан вектор-строка):

>> reshape(Y(2,3,1,2,:), 1, [])
ans =
    0.5621    0.4352    0.3672    0.9011    0.0332    0.5044    0.3416    0.6996    0.0610    0.2638    0.5586    0.3766

Видно, что максимум 0.9011, который происходит на 4th позиция (где "позиция" определяется вдоль N=3 свернутые размеры). По факту,

>> ind(2,3,1,2)
ans =
     4
>> Y(2,3,1,2,ind(2,3,1,2))
ans =
    0.9011

или, с точки зрения N=3 подиндексы,

>> Y(2,3,1,2,sub{1}(2,3,1,2),sub{2}(2,3,1,2),sub{3}(2,3,1,2))
ans =
    0.9011