Implement "sequential" strategy for BalancingLearner

Author: basnijholt

adaptive/learner/balancing_learner.py

@@ -36,7 +36,9 @@ def dispatch(child_functions: list[Callable], arg: Any) -> Any:
     return child_functions[index](x)
 
 
-STRATEGY_TYPE: TypeAlias = Literal["loss_improvements", "loss", "npoints", "cycle"]
+STRATEGY_TYPE: TypeAlias = Literal[
+    "loss_improvements", "loss", "npoints", "cycle", "sequential"
+]
 
 CDIMS_TYPE: TypeAlias = Union[
     Sequence[dict[str, Any]],
@@ -77,13 +79,21 @@ class BalancingLearner(BaseLearner):
     function : callable
         A function that calls the functions of the underlying learners.
         Its signature is ``function(learner_index, point)``.
-    strategy : 'loss_improvements' (default), 'loss', 'npoints', or 'cycle'.
+    strategy : 'loss_improvements' (default), 'loss', 'npoints', 'cycle', or 'sequential'
         The points that the `BalancingLearner` choses can be either based on:
-        the best 'loss_improvements', the smallest total 'loss' of the
-        child learners, the number of points per learner, using 'npoints',
-        or by cycling through the learners one by one using 'cycle'.
-        One can dynamically change the strategy while the simulation is
-        running by changing the ``learner.strategy`` attribute.
+
+        - 'loss_improvements': This strategy selects the points with the best
+          improvement in loss.
+        - 'loss': This strategy selects the points with the smallest total loss
+          from the child learners.
+        - 'npoints': This strategy selects points based on the number of points
+          per learner.
+        - 'cycle': This strategy cycles through all learners one by one.
+        - 'sequential': This strategy goes through learners in a sequential
+          order. Only works with learners that have a `done` method.
+
+        You can change the strategy dynamically while the simulation is
+        running by modifying the `learner.strategy` attribute.
 
     Notes
     -----
@@ -159,13 +169,19 @@ def nsamples(self):
 
     @property
     def strategy(self) -> STRATEGY_TYPE:
-        """Can be either 'loss_improvements' (default), 'loss', 'npoints', or
-        'cycle'. The points that the `BalancingLearner` choses can be either
-        based on: the best 'loss_improvements', the smallest total 'loss' of
-        the child learners, the number of points per learner, using 'npoints',
-        or by going through all learners one by one using 'cycle'.
-        One can dynamically change the strategy while the simulation is
-        running by changing the ``learner.strategy`` attribute."""
+        """The `BalancingLearner` can choose points based on different strategies.
+
+        The strategies are:
+
+        - 'loss_improvements': This strategy selects the points with the best improvement in loss.
+        - 'loss': This strategy selects the points with the smallest total loss from the child learners.
+        - 'npoints': This strategy selects points based on the number of points per learner.
+        - 'cycle': This strategy cycles through all learners one by one.
+        - 'sequential': This strategy goes through learners in a sequential order.
+
+        You can change the strategy dynamically while the simulation is
+        running by modifying the `learner.strategy` attribute.
+        """
         return self._strategy
 
     @strategy.setter
@@ -180,6 +196,9 @@ def strategy(self, strategy: STRATEGY_TYPE) -> None:
         elif strategy == "cycle":
             self._ask_and_tell = self._ask_and_tell_based_on_cycle
             self._cycle = itertools.cycle(range(len(self.learners)))
+        elif strategy == "sequential":
+            self._ask_and_tell = self._ask_and_tell_based_on_sequential
+            ...
         else:
             raise ValueError(
                 'Only strategy="loss_improvements", strategy="loss",'
@@ -255,7 +274,8 @@ def _ask_and_tell_based_on_npoints(
     def _ask_and_tell_based_on_cycle(
         self, n: int
     ) -> tuple[list[tuple[Int, Any]], list[float]]:
-        points, loss_improvements = [], []
+        points: list[tuple[Int, Any]] = []
+        loss_improvements: list[float] = []
         for _ in range(n):
             index = next(self._cycle)
             point, loss_improvement = self.learners[index].ask(n=1)
@@ -265,6 +285,33 @@ def _ask_and_tell_based_on_cycle(
 
         return points, loss_improvements
 
+    def _ask_and_tell_based_on_sequential(
+        self, n: int
+    ) -> tuple[list[tuple[Int, Any]], list[float]]:
+        points: list[tuple[Int, Any]] = []
+        loss_improvements: list[float] = []
+        learner_index = 0
+
+        while len(points) < n:
+            learner = self.learners[learner_index]
+            if learner.done():  # type: ignore[attr-defined]
+                if learner_index == len(self.learners) - 1:
+                    break
+                learner_index += 1
+                continue
+
+            point, loss_improvement = learner.ask(n=1)
+            if not point:  # if learner is exhausted, we don't get points
+                if learner_index == len(self.learners) - 1:
+                    break
+                learner_index += 1
+                continue
+            points.append((learner_index, point[0]))
+            loss_improvements.append(loss_improvement[0])
+            self.tell_pending((learner_index, point[0]))
+
+        return points, loss_improvements
+
     def ask(
         self, n: int, tell_pending: bool = True
     ) -> tuple[list[tuple[Int, Any]], list[float]]:

adaptive/tests/test_balancing_learner.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from adaptive.learner import BalancingLearner, Learner1D
+from adaptive.learner import BalancingLearner, Learner1D, SequenceLearner
 from adaptive.runner import simple
 
 strategies = ["loss", "loss_improvements", "npoints", "cycle"]
@@ -64,3 +64,11 @@ def test_strategies(strategy, goal_type, goal):
     learners = [Learner1D(lambda x: x, bounds=(-1, 1)) for i in range(10)]
     learner = BalancingLearner(learners, strategy=strategy)
     simple(learner, **{goal_type: goal})
+
+
+def test_sequential_strategy() -> None:
+    learners = [SequenceLearner(lambda x: x, sequence=[0, 1, 2, 3]) for i in range(10)]
+    learner = BalancingLearner(learners, strategy="sequential")  # type: ignore[arg-type]
+    simple(learner, goal=lambda lrn: sum(x.npoints for x in lrn.learners) >= 4 * 5)
+    assert all(lrn.done() for lrn in learners[:5])
+    assert all(not lrn.done() for lrn in learners[5:])