implement 'npoints' strategy for the 'BalancingLearner'

Author: basnijholt

adaptive/learner/balancing_learner.py

@@ -36,10 +36,10 @@ class BalancingLearner(BaseLearner):
             >>> cdims = (['A', 'B'], itertools.product([True, False], [0, 1]))
             >>> cdims = (['A', 'B'], [(True, 0), (True, 1),
             ...                       (False, 0), (False, 1)])
-    strategy : 'loss_improvements' (default) or 'loss'
-        The points that the 'BalancingLearner' choses can be either based on
-        the best 'loss_improvements' or the smallest total 'loss' of the
-        child learners.
+    strategy : 'loss_improvements' (default), 'loss', or 'npoints'
+        The points that the 'BalancingLearner' choses can be either based on:
+        the best 'loss_improvements', the smallest total 'loss' of the
+        child learners, or the number of points per learner, using 'npoints'.
 
     Notes
     -----
@@ -73,9 +73,11 @@ def __init__(self, learners, *, cdims=None, strategy='loss_improvements'):
             self._ask_and_tell = self._ask_and_tell_based_on_loss_improvements
         elif strategy == 'loss':
             self._ask_and_tell = self._ask_and_tell_based_on_loss
+        elif strategy == 'npoints':
+            self._ask_and_tell = self._ask_and_tell_based_on_npoints
         else:
-            raise ValueError('Only strategy="loss_improvements" or'
-                             ' strategy="loss" is implemented.')
+            raise ValueError('Only strategy="loss_improvements",'
+                ' strategy="loss", or strategy="npoints" is implemented.')
 
     def _ask_and_tell_based_on_loss_improvements(self, n):
         points = []
@@ -109,6 +111,21 @@ def _ask_and_tell_based_on_loss(self, n):
             loss_improvements.append(ls[0])
         return points, loss_improvements
 
+    def _ask_and_tell_based_on_npoints(self, n):
+        points = []
+        loss_improvements = []
+        npoints = [l.npoints + len(l.pending_points)
+                   for l in self.learners]
+        n_left = n
+        while n_left > 0:
+            i = np.argmin(npoints)
+            xs, ls = self.learners[i].ask(1)
+            npoints[i] += 1
+            n_left -= 1
+            points.append((i, xs[0]))
+            loss_improvements.append(ls[0])
+        return points, loss_improvements
+
     def ask(self, n, tell_pending=True):
         """Chose points for learners."""
         if not tell_pending: