Repository: python-adaptive/adaptive

Issue:
BalancingLearner puts all points in the first child-learner when asking for points with no data present
```python
def h(x, offset=0):
    a = 0.01
    return x + a**2 / (a**2 + (x - offset)**2)

learners = [adaptive.Learner1D(partial(h, offset=random.uniform(-1, 1)),
            bounds=(-1, 1)) for i in range(10)]

bal_learner = adaptive.BalancingLearner(learners)

bal_learner.ask(10)
```
```
([(0, -1),
  (0, 1),
  (0, 0.0),
  (0, -0.5),
  (0, 0.5),
  (0, -0.75),
  (0, -0.25),
  (0, 0.25),
  (0, 0.75),
  (0, -0.875)],
 [inf, inf, inf, inf, inf, inf, inf, inf, inf, inf])
```

Instead, it should balance the points over all learners. This happens because the loss improvements are `inf` for all points.


A temporary solution would be to force all learners to have the same amount of points by using:
```
bal_learner = adaptive.BalancingLearner(learners, strategy='npoints')
```

Discussion:
akhmerov: This fix seems to be only partial: consider for example the case when one learner is much cheaper than all the others. Then as soon as one point arrives, the balancing learner will request all points from another random learner in the queue.

If I understand correctly, an empty learner has both an infinite loss and an infinite loss improvement, however some non-empty learners also have the same loss and loss improvement. In that case the problem would arise as well.

Therefore it seems to suggest that a better fix would be to sort on `(-loss, n_points)` (or loss improvement, depending on the balancing learner setting).
basnijholt: @akhmerov, good suggestion! 

I've done this in 6e172ef85228e5f4469488fe07d6ae0773b4b498. This is what you meant right?

I'll also do it in `_ask_and_tell_based_on_loss`.

Patch:
diff --git a/adaptive/learner/balancing_learner.py b/adaptive/learner/balancing_learner.py
index b49a25f..f77da2a 100644
--- a/adaptive/learner/balancing_learner.py
+++ b/adaptive/learner/balancing_learner.py
@@ -77,7 +77,7 @@ class BalancingLearner(BaseLearner):
         # pickle the whole learner.
         self.function = partial(dispatch, [l.function for l in self.learners])
 
-        self._points = {}
+        self._ask_cache = {}
         self._loss = {}
         self._pending_loss = {}
         self._cdims_default = cdims
@@ -113,54 +113,74 @@ class BalancingLearner(BaseLearner):
                 ' strategy="npoints" is implemented.')
 
     def _ask_and_tell_based_on_loss_improvements(self, n):
-        points = []
-        loss_improvements = []
+        selected = []  # tuples ((learner_index, point), loss_improvement)
+        total_points = [l.npoints + len(l.pending_points) for l in self.learners]
         for _ in range(n):
-            improvements_per_learner = []
-            pairs = []
+            to_select = []
             for index, learner in enumerate(self.learners):
-                if index not in self._points:
-                    self._points[index] = learner.ask(
+                # Take the points from the cache
+                if index not in self._ask_cache:
+                    self._ask_cache[index] = learner.ask(
                         n=1, tell_pending=False)
-                point, loss_improvement = self._points[index]
-                improvements_per_learner.append(loss_improvement[0])
-                pairs.append((index, point[0]))
-            x, l = max(zip(pairs, improvements_per_learner),
-                       key=itemgetter(1))
-            points.append(x)
-            loss_improvements.append(l)
-            self.tell_pending(x)
-
+                points, loss_improvements = self._ask_cache[index]
+                to_select.append(
+                    ((index, points[0]),
+                     (loss_improvements[0], -total_points[index]))
+                )
+
+            # Choose the optimal improvement.
+            (index, point), (loss_improvement, _) = max(
+                to_select, key=itemgetter(1))
+            total_points[index] += 1
+            selected.append(((index, point), loss_improvement))
+            self.tell_pending((index, point))
+
+        points, loss_improvements = map(list, zip(*selected))
         return points, loss_improvements
 
     def _ask_and_tell_based_on_loss(self, n):
-        points = []
-        loss_improvements = []
+        selected = []   # tuples ((learner_index, point), loss_improvement)
+        total_points = [l.npoints + len(l.pending_points) for l in self.learners]
         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])
+            index, _ = max(
+                enumerate(zip(losses, (-n for n in total_points))),
+                key=itemgetter(1)
+            )
+            total_points[index] += 1
+
+            # Take the points from the cache
+            if index not in self._ask_cache:
+                self._ask_cache[index] = self.learners[index].ask(n=1)
+            points, loss_improvements = self._ask_cache[index]
+
+            selected.append(((index, points[0]), loss_improvements[0]))
+            self.tell_pending((index, points[0]))
+
+        points, loss_improvements = map(list, zip(*selected))
         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])
+        selected = []  # tuples ((learner_index, point), loss_improvement)
+        total_points = [l.npoints + len(l.pending_points) for l in self.learners]
+        for _ in range(n):
+            index = np.argmin(total_points)
+            # Take the points from the cache
+            if index not in self._ask_cache:
+                self._ask_cache[index] = self.learners[index].ask(n=1)
+            points, loss_improvements = self._ask_cache[index]
+            total_points[index] += 1
+            selected.append(((index, points[0]), loss_improvements[0]))
+            self.tell_pending((index, points[0]))
+
+        points, loss_improvements = map(list, zip(*selected))
         return points, loss_improvements
 
     def ask(self, n, tell_pending=True):
         """Chose points for learners."""
+        if n == 0:
+            return [], []
+
         if not tell_pending:
             with restore(*self.learners):
                 return self._ask_and_tell(n)
@@ -169,14 +189,14 @@ class BalancingLearner(BaseLearner):
 
     def tell(self, x, y):
         index, x = x
-        self._points.pop(index, None)
+        self._ask_cache.pop(index, None)
         self._loss.pop(index, None)
         self._pending_loss.pop(index, None)
         self.learners[index].tell(x, y)
 
     def tell_pending(self, x):
         index, x = x
-        self._points.pop(index, None)
+        self._ask_cache.pop(index, None)
         self._loss.pop(index, None)
         self.learners[index].tell_pending(x)
 


Test patch:
diff --git a/adaptive/tests/test_balancing_learner.py b/adaptive/tests/test_balancing_learner.py
index 2b5569d..2354f63 100644
--- a/adaptive/tests/test_balancing_learner.py
+++ b/adaptive/tests/test_balancing_learner.py
@@ -1,6 +1,9 @@
 # -*- coding: utf-8 -*-
 
+import pytest
+
 from adaptive.learner import Learner1D, BalancingLearner
+from adaptive.runner import simple
 
 
 def test_balancing_learner_loss_cache():
@@ -21,3 +24,37 @@ def test_balancing_learner_loss_cache():
     bl = BalancingLearner([learner])
     assert bl.loss(real=False) == pending_loss
     assert bl.loss(real=True) == real_loss
+
+
+@pytest.mark.parametrize('strategy', ['loss', 'loss_improvements', 'npoints'])
+def test_distribute_first_points_over_learners(strategy):
+    for initial_points in [0, 3]:
+        learners = [Learner1D(lambda x: x, bounds=(-1, 1)) for i in range(10)]
+        learner = BalancingLearner(learners, strategy=strategy)
+
+        points = learner.ask(initial_points)[0]
+        learner.tell_many(points, points)
+
+        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)
+
+
+@pytest.mark.parametrize('strategy', ['loss', 'loss_improvements', 'npoints'])
+def test_ask_0(strategy):
+    learners = [Learner1D(lambda x: x, bounds=(-1, 1)) for i in range(10)]
+    learner = BalancingLearner(learners, strategy=strategy)
+    points, _ = learner.ask(0)
+    assert len(points) == 0
+
+
+@pytest.mark.parametrize('strategy, goal', [
+    ('loss', lambda l: l.loss() < 0.1),
+    ('loss_improvements', lambda l: l.loss() < 0.1),
+    ('npoints', lambda bl: all(l.npoints > 10 for l in bl.learners)),
+])
+def test_strategies(strategy, goal):
+    learners = [Learner1D(lambda x: x, bounds=(-1, 1)) for i in range(10)]
+    learner = BalancingLearner(learners, strategy=strategy)
+    simple(learner, goal=goal)