Fix simulated learner state restoration

Author: basnijholt

adaptive/learner/balancing_learner.py

@@ -261,8 +261,13 @@ def ask(
             return [], []
 
         if not tell_pending:
-            with restore(*self.learners):
-                return self._ask_and_tell(n)
+            try:
+                with restore(*self.learners):
+                    return self._ask_and_tell(n)
+            finally:
+                self._ask_cache.clear()
+                self._loss.clear()
+                self._pending_loss.clear()
         else:
             return self._ask_and_tell(n)
 

adaptive/learner/learnerND.py

@@ -862,7 +862,7 @@ def _update_range(self, new_output):
             # this is the first point, nothing to do, just set the range
             self._min_value = np.min(new_output)
             self._max_value = np.max(new_output)
-            self._old_scale = self._scale or 1
+            self._old_scale = self._scale
             return False
 
         # if range in one or more directions is doubled, then update all losses
@@ -885,7 +885,10 @@ def _update_range(self, new_output):
 
         self._output_multiplier = scale_multiplier
 
-        scale_factor = self._scale / self._old_scale
+        if self._old_scale == 0:
+            scale_factor = math.inf if self._scale > 0 else 1
+        else:
+            scale_factor = self._scale / self._old_scale
         if scale_factor > self._recompute_losses_factor:
             self._old_scale = self._scale
             self._recompute_all_losses()

adaptive/tests/test_balancing_learner.py

@@ -1,13 +1,24 @@
 from __future__ import annotations
 
+import functools as ft
+import math
+import random
+
+import numpy as np
 import pytest
 
-from adaptive.learner import BalancingLearner, Learner1D
+from adaptive.learner import BalancingLearner, Learner1D, Learner2D
 from adaptive.runner import simple
 
 strategies = ["loss", "loss_improvements", "npoints", "cycle"]
 
 
+def ring_of_fire(xy, d):
+    a = 0.2
+    x, y = xy
+    return x + math.exp(-((x**2 + y**2 - d**2) ** 2) / a**4)
+
+
 def test_balancing_learner_loss_cache():
     learner = Learner1D(lambda x: x, bounds=(-1, 1))
     learner.tell(-1, -1)
@@ -64,3 +75,34 @@ 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_loss_improvements_strategy_with_tell_pending_false_reserves_child_points():
+    random.seed(3104322362)
+    np.random.seed(3104322362 % 2**32)
+
+    learners = [
+        Learner2D(
+            ft.partial(ring_of_fire, d=random.uniform(0.2, 1)),
+            bounds=((-1, 1), (-1, 1)),
+        )
+        for _ in range(4)
+    ]
+    learner = BalancingLearner(learners, strategy="loss_improvements")
+
+    stash = []
+    for n, m in [(1, 1), (4, 4), (2, 0), (4, 4), (8, 6)]:
+        xs, _ = learner.ask(n, tell_pending=False)
+        random.shuffle(xs)
+        for _ in range(m):
+            stash.append(xs.pop())
+
+        for x in xs:
+            learner.tell(x, learner.function(x))
+
+        random.shuffle(stash)
+        for _ in range(m):
+            x = stash.pop()
+            learner.tell(x, learner.function(x))
+
+    assert all(not child.pending_points for child in learners)

adaptive/tests/unit/test_learnernd.py

@@ -153,6 +153,37 @@ def f_vec(x):
     assert np.isclose(result[1], 0.0)
 
 
+def test_learnerND_recomputes_losses_for_small_scale_updates():
+    learner = make_1d_learner()
+    learner._recompute_losses_factor = 1
+
+    for point, value in [((-1,), 0.0), ((0.0,), 0.45), ((1.0,), 0.45)]:
+        learner.tell(point, value)
+
+    simplex = next(
+        simplex
+        for simplex in learner.tri.simplices
+        if {tuple(vertex) for vertex in learner.tri.get_vertices(simplex)}
+        == {(-1.0,), (0.0,)}
+    )
+    cached_before = learner._losses[simplex]
+    assert np.isclose(cached_before, learner._compute_loss(simplex))
+
+    learner.tell((0.5,), 0.67)
+
+    simplex = next(
+        simplex
+        for simplex in learner.tri.simplices
+        if {tuple(vertex) for vertex in learner.tri.get_vertices(simplex)}
+        == {(-1.0,), (0.0,)}
+    )
+    cached_after = learner._losses[simplex]
+
+    assert learner._old_scale == pytest.approx(0.67)
+    assert np.isclose(cached_after, learner._compute_loss(simplex))
+    assert not np.isclose(cached_after, cached_before)
+
+
 def test_learnerND_1d_plot_requires_holoviews(monkeypatch):
     """Test that plotting fails with a clear error without holoviews."""
 

adaptive/utils.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import concurrent.futures as concurrent
+import copy
 import functools
 import gzip
 import inspect
@@ -27,7 +28,7 @@ def named_product(**items: Sequence[Any]):
 
 @contextmanager
 def restore(*learners) -> Iterator[None]:
-    states = [learner.__getstate__() for learner in learners]
+    states = [copy.deepcopy(learner.__getstate__()) for learner in learners]
     try:
         yield
     finally: