def testTopKIndices2DSorted(self):
scores = np.array([[0.3, 0.1, 0.4, 0.2], [0.1, 0.2, 0.3, 0.6]])
got = metric_util.top_k_indices(2, scores, sort=True)
# Indices are in ([row_index,...], [col_index, ...]) format.
self.assertAllClose(got,
(np.array([0, 0, 1, 1]), np.array([2, 0, 3, 2])))
self.assertAllClose(scores[got], np.array([0.4, 0.3, 0.6, 0.3]))
def testTopKIndicesWithBinaryClassification(self):
scores = np.array([0.2, 0.8])
got = metric_util.top_k_indices(1, scores)
self.assertAllClose(got, np.array([1]))
self.assertAllClose(scores[got], np.array([0.8]))
def testTopKIndicesSorted(self):
scores = np.array([0.1, 0.3, 0.4, 0.2])
got = metric_util.top_k_indices(2, scores, sort=True)
self.assertAllClose(got, np.array([2, 1]))
self.assertAllClose(scores[got], np.array([0.4, 0.3]))
def testTopKIndices2D(self):
scores = np.array([[0.4, 0.1, 0.2, 0.3], [0.1, 0.2, 0.1, 0.6]])
got = metric_util.top_k_indices(2, scores)
scores[got] = -1.0
self.assertAllClose(
scores, np.array([[-1.0, 0.1, 0.2, -1.0], [0.1, -1.0, 0.1, -1.0]]))
def testTopKIndices(self):
scores = np.array([0.4, 0.1, 0.2, 0.3])
got = metric_util.top_k_indices(2, scores)
# Indices could be in any order, test by overwritting the original scores
scores[got] = -1.0
self.assertAllClose(scores, np.array([-1.0, 0.1, 0.2, -1.0]))
