Merge pull request #193 from python-adaptive/sequence_learner

Add a SequenceLearner

Author: basnijholt

adaptive/__init__.py

@@ -14,6 +14,7 @@
     Learner2D,
     LearnerND,
     make_datasaver,
+    SequenceLearner,
 )
 from adaptive.notebook_integration import (
     active_plotting_tasks,
@@ -36,6 +37,7 @@
     "Learner2D",
     "LearnerND",
     "make_datasaver",
+    "SequenceLearner",
     "active_plotting_tasks",
     "live_plot",
     "notebook_extension",

adaptive/learner/__init__.py

@@ -10,6 +10,7 @@
 from adaptive.learner.learner1D import Learner1D
 from adaptive.learner.learner2D import Learner2D
 from adaptive.learner.learnerND import LearnerND
+from adaptive.learner.sequence_learner import SequenceLearner
 
 __all__ = [
     "AverageLearner",
@@ -21,6 +22,7 @@
     "Learner1D",
     "Learner2D",
     "LearnerND",
+    "SequenceLearner",
 ]
 
 with suppress(ImportError):

adaptive/learner/sequence_learner.py

@@ -0,0 +1,130 @@
+from copy import copy
+
+from sortedcontainers import SortedSet, SortedDict
+
+from adaptive.learner.base_learner import BaseLearner
+
+
+class _IgnoreFirstArgument:
+    """Remove the first argument from the call signature.
+
+    The SequenceLearner's function receives a tuple ``(index, point)``
+    but the original function only takes ``point``.
+
+    This is the same as `lambda x: function(x[1])`, however, that is not
+    pickable.
+    """
+
+    def __init__(self, function):
+        self.function = function
+
+    def __call__(self, index_point, *args, **kwargs):
+        index, point = index_point
+        return self.function(point, *args, **kwargs)
+
+    def __getstate__(self):
+        return self.function
+
+    def __setstate__(self, function):
+        self.__init__(function)
+
+
+class SequenceLearner(BaseLearner):
+    r"""A learner that will learn a sequence. It simply returns
+    the points in the provided sequence when asked.
+
+    This is useful when your problem cannot be formulated in terms of
+    another adaptive learner, but you still want to use Adaptive's
+    routines to run, save, and plot.
+
+    Parameters
+    ----------
+    function : callable
+        The function to learn. Must take a single element `sequence`.
+    sequence : sequence
+        The sequence to learn.
+
+    Attributes
+    ----------
+    data : dict
+        The data as a mapping from "index of element in sequence" => value.
+
+    Notes
+    -----
+    From primitive tests, the `~adaptive.SequenceLearner` appears to have a
+    similar performance to `ipyparallel`\s ``load_balanced_view().map``. With
+    the added benefit of having results in the local kernel already.
+    """
+
+    def __init__(self, function, sequence):
+        self._original_function = function
+        self.function = _IgnoreFirstArgument(function)
+        self._to_do_indices = SortedSet({i for i, _ in enumerate(sequence)})
+        self._ntotal = len(sequence)
+        self.sequence = copy(sequence)
+        self.data = SortedDict()
+        self.pending_points = set()
+
+    def ask(self, n, tell_pending=True):
+        indices = []
+        points = []
+        loss_improvements = []
+        for index in self._to_do_indices:
+            if len(points) >= n:
+                break
+            point = self.sequence[index]
+            indices.append(index)
+            points.append((index, point))
+            loss_improvements.append(1 / self._ntotal)
+
+        if tell_pending:
+            for i, p in zip(indices, points):
+                self.tell_pending((i, p))
+
+        return points, loss_improvements
+
+    def _get_data(self):
+        return self.data
+
+    def _set_data(self, data):
+        if data:
+            indices, values = zip(*data.items())
+            # the points aren't used by tell, so we can safely pass None
+            points = [(i, None) for i in indices]
+            self.tell_many(points, values)
+
+    def loss(self, real=True):
+        if not (self._to_do_indices or self.pending_points):
+            return 0
+        else:
+            npoints = self.npoints + (0 if real else len(self.pending_points))
+            return (self._ntotal - npoints) / self._ntotal
+
+    def remove_unfinished(self):
+        for i in self.pending_points:
+            self._to_do_indices.add(i)
+        self.pending_points = set()
+
+    def tell(self, point, value):
+        index, point = point
+        self.data[index] = value
+        self.pending_points.discard(index)
+        self._to_do_indices.discard(index)
+
+    def tell_pending(self, point):
+        index, point = point
+        self.pending_points.add(index)
+        self._to_do_indices.discard(index)
+
+    def done(self):
+        return not self._to_do_indices and not self.pending_points
+
+    def result(self):
+        """Get the function values in the same order as ``sequence``."""
+        if not self.done():
+            raise Exception("Learner is not yet complete.")
+        return list(self.data.values())
+
+    @property
+    def npoints(self):
+        return len(self.data)

adaptive/tests/test_learners.py

@@ -24,6 +24,7 @@
     Learner1D,
     Learner2D,
     LearnerND,
+    SequenceLearner,
 )
 from adaptive.runner import simple
 
@@ -116,26 +117,30 @@ def quadratic(x, m: uniform(0, 10), b: uniform(0, 1)):
 
 
 @learn_with(Learner1D, bounds=(-1, 1))
+@learn_with(SequenceLearner, sequence=np.linspace(-1, 1, 201))
 def linear_with_peak(x, d: uniform(-1, 1)):
     a = 0.01
     return x + a ** 2 / (a ** 2 + (x - d) ** 2)
 
 
 @learn_with(LearnerND, bounds=((-1, 1), (-1, 1)))
 @learn_with(Learner2D, bounds=((-1, 1), (-1, 1)))
+@learn_with(SequenceLearner, sequence=np.random.rand(1000, 2))
 def ring_of_fire(xy, d: uniform(0.2, 1)):
     a = 0.2
     x, y = xy
     return x + math.exp(-(x ** 2 + y ** 2 - d ** 2) ** 2 / a ** 4)
 
 
 @learn_with(LearnerND, bounds=((-1, 1), (-1, 1), (-1, 1)))
+@learn_with(SequenceLearner, sequence=np.random.rand(1000, 3))
 def sphere_of_fire(xyz, d: uniform(0.2, 1)):
     a = 0.2
     x, y, z = xyz
     return x + math.exp(-(x ** 2 + y ** 2 + z ** 2 - d ** 2) ** 2 / a ** 4) + z ** 2
 
 
+@learn_with(SequenceLearner, sequence=range(1000))
 @learn_with(AverageLearner, rtol=1)
 def gaussian(n):
     return random.gauss(0, 1)
@@ -247,7 +252,7 @@ def f(x):
     simple(learner, goal=lambda l: l.npoints > 10)
 
 
-@run_with(Learner1D, Learner2D, LearnerND)
+@run_with(Learner1D, Learner2D, LearnerND, SequenceLearner)
 def test_adding_existing_data_is_idempotent(learner_type, f, learner_kwargs):
     """Adding already existing data is an idempotent operation.
 
@@ -264,7 +269,7 @@ def test_adding_existing_data_is_idempotent(learner_type, f, learner_kwargs):
     N = random.randint(10, 30)
     control.ask(N)
     xs, _ = learner.ask(N)
-    points = [(x, f(x)) for x in xs]
+    points = [(x, learner.function(x)) for x in xs]
 
     for p in points:
         control.tell(*p)
@@ -277,13 +282,24 @@ def test_adding_existing_data_is_idempotent(learner_type, f, learner_kwargs):
     M = random.randint(10, 30)
     pls = zip(*learner.ask(M))
     cpls = zip(*control.ask(M))
-    # Point ordering is not defined, so compare as sets
-    assert set(pls) == set(cpls)
+    if learner_type is SequenceLearner:
+        # The SequenceLearner's points might not be hasable
+        points, values = zip(*pls)
+        indices, points = zip(*points)
+
+        cpoints, cvalues = zip(*cpls)
+        cindices, cpoints = zip(*cpoints)
+        assert (np.array(points) == np.array(cpoints)).all()
+        assert values == cvalues
+        assert indices == cindices
+    else:
+        # Point ordering is not defined, so compare as sets
+        assert set(pls) == set(cpls)
 
 
 # XXX: This *should* pass (https://github.com/python-adaptive/adaptive/issues/55)
 #      but we xfail it now, as Learner2D will be deprecated anyway
-@run_with(Learner1D, xfail(Learner2D), LearnerND, AverageLearner)
+@run_with(Learner1D, xfail(Learner2D), LearnerND, AverageLearner, SequenceLearner)
 def test_adding_non_chosen_data(learner_type, f, learner_kwargs):
     """Adding data for a point that was not returned by 'ask'."""
     # XXX: learner, control and bounds are not defined
@@ -300,17 +316,29 @@ def test_adding_non_chosen_data(learner_type, f, learner_kwargs):
     N = random.randint(10, 30)
     xs, _ = control.ask(N)
 
-    ys = [f(x) for x in xs]
+    ys = [learner.function(x) for x in xs]
     for x, y in zip(xs, ys):
         control.tell(x, y)
         learner.tell(x, y)
 
     M = random.randint(10, 30)
     pls = zip(*learner.ask(M))
     cpls = zip(*control.ask(M))
-    # Point ordering within a single call to 'ask'
-    # is not guaranteed to be the same by the API.
-    assert set(pls) == set(cpls)
+
+    if learner_type is SequenceLearner:
+        # The SequenceLearner's points might not be hasable
+        points, values = zip(*pls)
+        indices, points = zip(*points)
+
+        cpoints, cvalues = zip(*cpls)
+        cindices, cpoints = zip(*cpoints)
+        assert (np.array(points) == np.array(cpoints)).all()
+        assert values == cvalues
+        assert indices == cindices
+    else:
+        # Point ordering within a single call to 'ask'
+        # is not guaranteed to be the same by the API.
+        assert set(pls) == set(cpls)
 
 
 @run_with(Learner1D, xfail(Learner2D), xfail(LearnerND), AverageLearner)
@@ -334,7 +362,7 @@ def test_point_adding_order_is_irrelevant(learner_type, f, learner_kwargs):
     N = random.randint(10, 30)
     control.ask(N)
     xs, _ = learner.ask(N)
-    points = [(x, f(x)) for x in xs]
+    points = [(x, learner.function(x)) for x in xs]
 
     for p in points:
         control.tell(*p)
@@ -366,7 +394,7 @@ def test_expected_loss_improvement_is_less_than_total_loss(
     xs, loss_improvements = learner.ask(N)
 
     for x in xs:
-        learner.tell(x, f(x))
+        learner.tell(x, learner.function(x))
 
     M = random.randint(50, 100)
     _, loss_improvements = learner.ask(M)
@@ -429,7 +457,12 @@ def test_learner_performance_is_invariant_under_scaling(
 
 
 @run_with(
-    Learner1D, Learner2D, LearnerND, AverageLearner, with_all_loss_functions=False
+    Learner1D,
+    Learner2D,
+    LearnerND,
+    AverageLearner,
+    SequenceLearner,
+    with_all_loss_functions=False,
 )
 def test_balancing_learner(learner_type, f, learner_kwargs):
     """Test if the BalancingLearner works with the different types of learners."""
@@ -474,6 +507,7 @@ def test_balancing_learner(learner_type, f, learner_kwargs):
     AverageLearner,
     maybe_skip(SKOptLearner),
     IntegratorLearner,
+    SequenceLearner,
     with_all_loss_functions=False,
 )
 def test_saving(learner_type, f, learner_kwargs):
@@ -504,6 +538,7 @@ def test_saving(learner_type, f, learner_kwargs):
     AverageLearner,
     maybe_skip(SKOptLearner),
     IntegratorLearner,
+    SequenceLearner,
     with_all_loss_functions=False,
 )
 def test_saving_of_balancing_learner(learner_type, f, learner_kwargs):

docs/source/reference/adaptive.learner.sequence_learner.rst

@@ -0,0 +1,7 @@
+adaptive.SequenceLearner
+========================
+
+.. autoclass:: adaptive.SequenceLearner
+    :members:
+    :undoc-members:
+    :show-inheritance:

docs/source/reference/adaptive.rst

@@ -14,6 +14,7 @@ Learners
     adaptive.learner.learner1D
     adaptive.learner.learner2D
     adaptive.learner.learnerND
+    adaptive.learner.sequence_learner
     adaptive.learner.skopt_learner
 
 Runners