do not do the entire loop but only until "n"

python-adaptive · Apr 22, 2020 · 473dd42 · 473dd42
1 parent 61525c0
commit 473dd42
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/adaptive/runner.py b/adaptive/runner.py
@@ -170,11 +170,10 @@ def do_log(self):
         return self.log is not None
 
     def _ask(self, n):
-        pids = [
-            pid
-            for pid in self._to_retry.keys()
-            if pid not in self._pending_points.values()
-        ][:n]
+        pending_ids = self._pending_points.values()
+        pids_gen = (pid for pid in self._to_retry.keys() if pid not in pending_ids)
+        pids = list(itertools.islice(pids_gen, n))
+
         loss_improvements = len(pids) * [float("inf")]
 
         if len(pids) < n: