modAL.utils.selection.multi_argmax bugfix, now it works well with objects having shape (n, 1)

Author: cosmic-cortex

modAL/utils/selection.py

@@ -25,7 +25,7 @@ def multi_argmax(values, n_instances=1):
     """
     assert n_instances <= len(values), 'n_instances must be less or equal than the size of utility'
 
-    max_idx = np.argpartition(-values, n_instances-1)[:n_instances]
+    max_idx = np.argpartition(-values, n_instances-1, axis=0)[:n_instances]
     return max_idx
 
 

tests/core_tests.py

@@ -366,7 +366,7 @@ def test_multi_argmax(self):
             for n_instances in range(1, n_pool):
                 utility = np.zeros(n_pool)
                 max_idx = np.random.choice(range(n_pool), size=n_instances, replace=False)
-                utility[max_idx] = 1.0
+                utility[max_idx] = 1e-10 + np.random.rand(n_instances, )
                 np.testing.assert_equal(
                     np.sort(modAL.utils.selection.multi_argmax(utility, n_instances)),
                     np.sort(max_idx)