use a list for ((learner_index, point), loss_improvement) tuples

Author: basnijholt

adaptive/learner/balancing_learner.py

@@ -113,10 +113,8 @@ def strategy(self, strategy):
                 ' strategy="npoints" is implemented.')
 
     def _ask_and_tell_based_on_loss_improvements(self, n):
-        chosen_points = []
-        chosen_loss_improvements = []
-        npoints = [l.npoints + len(l.pending_points)
-                   for l in self.learners]
+        selected = [] # tuples ((learner_index, point), loss_improvement)
+        npoints = [l.npoints + len(l.pending_points) for l in self.learners]
         for _ in range(n):
             improvements_per_learner = []
             points_per_learner = []
@@ -136,17 +134,15 @@ def _ask_and_tell_based_on_loss_improvements(self, n):
                 zip(points_per_learner, improvements_per_learner),
                 key=itemgetter(1))
             npoints[index] += 1
-            chosen_points.append((index, point))
-            chosen_loss_improvements.append(loss_improvement)
+            selected.append(((index, point), loss_improvement))
             self.tell_pending((index, point))
 
-        return chosen_points, chosen_loss_improvements
+        points, loss_improvements = map(list, zip(*selected))
+        return points, loss_improvements
 
     def _ask_and_tell_based_on_loss(self, n):
-        chosen_points = []
-        chosen_loss_improvements = []
-        npoints = [l.npoints + len(l.pending_points)
-                   for l in self.learners]
+        selected = [] # tuples ((learner_index, point), loss_improvement)
+        npoints = [l.npoints + len(l.pending_points) for l in self.learners]
         for _ in range(n):
             losses = self._losses(real=False)
             priority = zip(losses, (-n for n in npoints))
@@ -158,15 +154,14 @@ def _ask_and_tell_based_on_loss(self, n):
                 self._ask_cache[index] = self.learners[index].ask(n=1)
             points, loss_improvements = self._ask_cache[index]
 
-            chosen_points.append((index, points[0]))
-            chosen_loss_improvements.append(loss_improvements[0])
-        return chosen_points, chosen_loss_improvements
+            selected.append(((index, points[0]), loss_improvements[0]))
+
+        points, loss_improvements = map(list, zip(*selected))
+        return points, loss_improvements
 
     def _ask_and_tell_based_on_npoints(self, n):
-        chosen_points = []
-        chosen_loss_improvements = []
-        npoints = [l.npoints + len(l.pending_points)
-                   for l in self.learners]
+        selected = [] # tuples ((learner_index, point), loss_improvement)
+        npoints = [l.npoints + len(l.pending_points) for l in self.learners]
         n_left = n
         while n_left > 0:
             index = np.argmin(npoints)
@@ -176,9 +171,10 @@ def _ask_and_tell_based_on_npoints(self, n):
             points, loss_improvements = self._ask_cache[index]
             npoints[index] += 1
             n_left -= 1
-            chosen_points.append((index, points[0]))
-            chosen_loss_improvements.append(loss_improvements[0])
-        return chosen_points, chosen_loss_improvements
+            selected.append(((index, points[0]), loss_improvements[0]))
+
+        points, loss_improvements = map(list, zip(*selected))
+        return points, loss_improvements
 
     def ask(self, n, tell_pending=True):
         """Chose points for learners."""