fix point distribution for loss strategy

Author: basnijholt

adaptive/learner/balancing_learner.py

@@ -111,6 +111,7 @@ def strategy(self, strategy):
             raise ValueError(
                 'Only strategy="loss_improvements", strategy="loss", or'
                 ' strategy="npoints" is implemented.')
+        self._points = {}  # reset the cache
 
     def _ask_and_tell_based_on_loss_improvements(self, n):
         chosen_points = []
@@ -125,12 +126,14 @@ def _ask_and_tell_based_on_loss_improvements(self, n):
                     self._points[index] = learner.ask(
                         n=1, tell_pending=False)
                 points, loss_improvements = self._points[index]
-                npoints = npoints_per_learner[index] + learner.npoints
+                npoints = (npoints_per_learner[index]
+                           + learner.npoints
+                           + len(learner.pending_points))
                 priority = (loss_improvements[0], -npoints)
                 improvements_per_learner.append(priority)
                 points_per_learner.append((index, points[0]))
 
-            # Chose the optimal improvement.
+            # Choose the optimal improvement.
             (index, point), (loss_improvement, _) = max(
                 zip(points_per_learner, improvements_per_learner),
                 key=itemgetter(1))
@@ -142,15 +145,23 @@ def _ask_and_tell_based_on_loss_improvements(self, n):
         return chosen_points, chosen_loss_improvements
 
     def _ask_and_tell_based_on_loss(self, n):
-        points = []
-        loss_improvements = []
+        chosen_points = []
+        chosen_loss_improvements = []
+        npoints_per_learner = defaultdict(int)
+
         for _ in range(n):
             losses = self._losses(real=False)
-            max_ind = np.argmax(losses)
-            xs, ls = self.learners[max_ind].ask(1)
-            points.append((max_ind, xs[0]))
-            loss_improvements.append(ls[0])
-        return points, loss_improvements
+            npoints = [-(l.npoints
+                         + npoints_per_learner[i]
+                         + len(l.pending_points))
+                       for i, l in enumerate(self.learners)]
+            priority = zip(losses, npoints)
+            index, (_, npoints) = max(enumerate(priority), key=itemgetter(1))
+            npoints_per_learner[index] += 1
+            points, loss_improvements = self.learners[index].ask(1)
+            chosen_points.append((index, points[0]))
+            chosen_loss_improvements.append(loss_improvements[0])
+        return chosen_points, chosen_loss_improvements
 
     def _ask_and_tell_based_on_npoints(self, n):
         points = []

adaptive/tests/test_balancing_learner.py

@@ -31,7 +31,7 @@ def test_distribute_first_points_over_learners():
         points, _ = learner.ask(100)
         i_learner, xs = zip(*points)
         # assert that are all learners in the suggested points
-        assert len(set(i_learner)) == len(learners)
+        assert len(set(i_learner)) == len(learners), strategy
 
 
 def test_strategies():