Поиск уникальных 1D-массивов и соответствующих пар 2D-индексов в 3D-массиве (с numpy) ⇐ Python

[Anonymous]

У меня есть трехмерный числовой массив X формы =

Код: Выделить всё

(N,N,6); и двумерную маску логического массива

shape=

Код: Выделить всё

(N,N). Я хотел бы найти все уникальные одномерные подмассивы (длина = 6, вдоль оси = 2) X

, выполняя поиск по первым двум измерениям, которые также замаскированы маской (т. е. имеют значение True в одной и той же позиции в первых двух измерениях).
Это кажется (относительно) простым, но у меня есть дополнительная сложность, заключающаяся в том, что повторяющиеся элементы X, которые должны быть неуникальными, могут быть неуникальными. различаются из-за причин решения с плавающей запятой/числового решения, поэтому мне нужно округлить значения в X перед сравнением. Однако для моих окончательных результатов мне нужны уникальные значения полной точности, поэтому я не могу просто использовать необработанный вывод функции типа np.unique, примененной к округленным значениям, без дополнительной обработки.
С этой целью я решил найти все индексы уникальных массивов в X (они мне также нужны по другим причинам). То есть мне нужны 2D-пары индексов, соответствующие первым двум измерениям X, где расположены одномерные массивы с уникальной+маской полной точности.
Эти спецификации вместе, кажется, не очень хорошо сочетаются с опциями, которые я сейчас могу найти в numpy. Мне потребовалось несколько попыток, чтобы найти что-то, что сработало.
Попытка №1
Моя первоначальная идея была примерно такой:

Код: Выделить всё

X_unique, unique_idx = np.unique(
X[mask].round(decimals=8),
axis=0,
return_index=True
)

Это дает мне правильные результаты, но X_unique, конечно, округляется. Более того, использование двумерного логического массива для индексации первых двух измерений X сжимает их, в результате чего на выходе X_unique выдается форма =

Код: Выделить всё

(,6)

и unique_idx формы=

Код: Выделить всё

(,)

, поэтому я потерял структурную информацию первых двух измерений. В X нет шаблона для уникальных вхождений массива, поэтому это не подойдет.
Попытка № 2
Затем я попытался найти способ замаскировать X, сохранив при этом исходные формы, но понял, что что-то вроде следующего не сработает в любом случае:

Код: Выделить всё

X_unique, unique_idx = np.unique(
hypothetical_X_masked_but_still_3D.round(decimals=8),
axis=,
return_index=True
)

потому что ось kwarg np.unique не может принять кортеж индексов (и я хочу выполнять поиск по осям 0 и 1, сохраняя 2).
Попытка № 3
Далее я попытался сгладить первые два измерения, но вместо того, чтобы опускать элементы, не входящие в маску, я заменил их на Sentinel np.nan, чтобы позже я мог восстановить сплющенный массив в его исходную форму:

Код: Выделить всё

X_masked = np.where(mask[...,None], X, np.nan).reshape(N*N,6)

X_unique, unique_idx = np.unique(
X_masked.round(decimals=8),
axis=0,
equal_nan=True,
return_index=True
)

Очевидно, это не удается, поскольку kwarg равенства_nan np.unique, похоже, не работает для массивов np.nan.
MRE:

Код: Выделить всё

nans = np.full((6), np.nan)
print(nans == nans) # > [False False False False False False]
print(np.all(nans == nans)) # > False

Мне это кажется ошибкой/недосмотром. Я понимаю, что np.nan не сравним сам с собой, поэтому np.nan == np.nan > False имеет для меня смысл. Но, конечно же, весь смысл kwarg равно_nan заключается в том, чтобы специально работать с NaN, и kwarg правильно работает с NaN, не являющимися массивами (якобы; я это не проверял).
В любом случае, ошибка или нет, мне пришлось найти какое-то другое значение для замены дозорного значения. Я не мог использовать нечисловой заполнитель, например строку или ..., потому что мое фактическое приложение требует округления значений X перед проверкой на уникальность (из-за проблем с числовой точностью). Я также не мог использовать мнимую часть комплексных чисел в качестве флага, поскольку мой фактический массив X тоже сложный.
В конце концов я нашел своего рода решение, но оно сложнее, чем мне хотелось бы. Возможно, у кого-то есть лучший способ сделать это.

Подробнее здесь: https://stackoverflow.com/questions/798 ... -with-nump