Преобразование прогнозов в метки с порогом

Question

Предполагая, что у меня есть следующий массив NumPy

import numpy as np

arr = np.array([[0.2, 0.8], [0.99, 0.01], [0.08, 0.92]])
arr
Out[57]: 
array([[0.2 , 0.8 ],
       [0.99, 0.01],
       [0.08, 0.92]])

Если бы я хотел преобразовать этот вывод в «классы» (или индекс наибольшего значения в каждой строке), я бы просто использовал:

arr.argmax(axis=1)
Out[58]: array([1, 0, 1], dtype=int64)

Проблема в том, что я хочу ограничить определенную границу. Для примера давайте используем 0,9 . Таким образом, каждая строка, которая не соответствует пороговому ограничению, вернет метку -1.

Выход для приведенного выше примера будет тогда [-1, 0, 1] (поскольку ни 0,8, ни 0,2 не превышают 0,9).

Какой самый питонический способ сделать это? надеюсь (но не обязательно), используя numpy.

Answer 1

Вы можете использовать np.where:

m = arr > 0.9
np.where(m.any(axis=1), m.argmax(axis=1), -1)
array([-1,  0,  1])

Подробности

(arr > 0.9) возвращает ndarray с той же формой, указывающей, где выполняется условие:

array([[False, False],
       [ True, False],
       [False,  True]])

m.argmax(axis=1) возвращает, где m равно True:

array([0, 0, 1])

np.where вернет m.argmax(axis=1) для тех строк, которые удовлетворяют m.any(axis=1), таким образом, где хотя бы один элемент больше, чем пороговое значение. Здесь m.any(axis=1) дает:

array([False,  True,  True])

В противном случае np.where вернет -1