make learners picklable

adaptive/learner/average_learner.py

@@ -144,3 +144,16 @@ def _get_data(self):
 
     def _set_data(self, data):
         self.data, self.npoints, self.sum_f, self.sum_f_sq = data
+
+    def __getstate__(self):
+        return (
+            self.function,
+            self.atol,
+            self.rtol,
+            self._get_data(),
+        )
+
+    def __setstate__(self, state):
+        function, atol, rtol, data = state
+        self.__init__(function, atol, rtol)
+        self._set_data(data)

adaptive/learner/balancing_learner.py

@@ -440,3 +440,14 @@ def _get_data(self):
     def _set_data(self, data):
         for l, _data in zip(self.learners, data):
             l._set_data(_data)
+
+    def __getstate__(self):
+        return (
+            self.learners,
+            self._cdims_default,
+            self.strategy,
+        )
+
+    def __setstate__(self, state):
+        learners, cdims, strategy = state
+        self.__init__(learners, cdims=cdims, strategy=strategy)

adaptive/learner/data_saver.py

@@ -51,6 +51,18 @@ def _set_data(self, data):
         learner_data, self.extra_data = data
         self.learner._set_data(learner_data)
 
+    def __getstate__(self):
+        return (
+            self.learner,
+            self.arg_picker,
+            self.extra_data,
+        )
+
+    def __setstate__(self, state):
+        learner, arg_picker, extra_data = state
+        self.__init__(learner, arg_picker)
+        self.extra_data = extra_data
+
     @copy_docstring_from(BaseLearner.save)
     def save(self, fname, compress=True):
         # We copy this method because the 'DataSaver' is not a

adaptive/learner/integrator_learner.py

@@ -591,3 +591,16 @@ def _set_data(self, data):
         self.x_mapping = defaultdict(lambda: SortedSet([], key=attrgetter("rdepth")))
         for k, _set in x_mapping.items():
             self.x_mapping[k].update(_set)
+
+    def __getstate__(self):
+        return (
+            self.function,
+            self.bounds,
+            self.tol,
+            self._get_data(),
+        )
+
+    def __setstate__(self, state):
+        function, bounds, tol, data = state
+        self.__init__(function, bounds, tol)
+        self._set_data(data)

adaptive/learner/learner1D.py

@@ -625,6 +625,23 @@ def _set_data(self, data):
         if data:
             self.tell_many(*zip(*data.items()))
 
+    def __getstate__(self):
+        return (
+            self.function,
+            self.bounds,
+            self.loss_per_interval,
+            dict(self.losses),  # SortedDict cannot be pickled
+            dict(self.losses_combined),  # ItemSortedDict cannot be pickled
+            self._get_data(),
+        )
+
+    def __setstate__(self, state):
+        function, bounds, loss_per_interval, losses, losses_combined, data = state
+        self.__init__(function, bounds, loss_per_interval)
+        self._set_data(data)
+        self.losses.update(losses)
+        self.losses_combined.update(losses_combined)
+
 
 def loss_manager(x_scale):
     def sort_key(ival, loss):

adaptive/learner/learner2D.py

@@ -706,3 +706,18 @@ def _set_data(self, data):
         for point in copy(self._stack):
             if point in self.data:
                 self._stack.pop(point)
+
+    def __getstate__(self):
+        return (
+            self.function,
+            self.bounds,
+            self.loss_per_triangle,
+            self._stack,
+            self._get_data(),
+        )
+
+    def __setstate__(self, state):
+        function, bounds, loss_per_triangle, _stack, data = state
+        self.__init__(function, bounds, loss_per_triangle)
+        self._set_data(data)
+        self._stack = _stack

adaptive/learner/sequence_learner.py

@@ -83,16 +83,6 @@ def ask(self, n, tell_pending=True):
 
         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
@@ -128,3 +118,25 @@ def result(self):
     @property
     def npoints(self):
         return len(self.data)
+
+    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 __getstate__(self):
+        return (
+            self._original_function,
+            self.sequence,
+            self._get_data(),
+        )
+
+    def __setstate__(self, state):
+        function, sequence, data = state
+        self.__init__(function, sequence)
+        self._set_data(data)

adaptive/tests/test_pickling.py

@@ -0,0 +1,117 @@
+import pickle
+
+import pytest
+
+from adaptive.learner import (
+    AverageLearner,
+    BalancingLearner,
+    DataSaver,
+    IntegratorLearner,
+    Learner1D,
+    Learner2D,
+    SequenceLearner,
+)
+from adaptive.runner import simple
+
+try:
+    import cloudpickle
+
+    with_cloudpickle = True
+except ModuleNotFoundError:
+    with_cloudpickle = False
+
+try:
+    import dill
+
+    with_dill = True
+except ModuleNotFoundError:
+    with_dill = False
+
+
+def goal_1(learner):
+    return learner.npoints == 10
+
+
+def goal_2(learner):
+    return learner.npoints == 20
+
+
+def pickleable_f(x):
+    return hash(str(x)) / 2 ** 63
+
+
+nonpickleable_f = lambda x: hash(str(x)) / 2 ** 63  # noqa: E731
+
+
+def identity_function(x):
+    return x
+
+
+def datasaver(f, learner_type, learner_kwargs):
+    return DataSaver(
+        learner=learner_type(f, **learner_kwargs), arg_picker=identity_function
+    )
+
+
+def balancing_learner(f, learner_type, learner_kwargs):
+    learner_1 = learner_type(f, **learner_kwargs)
+    learner_2 = learner_type(f, **learner_kwargs)
+    return BalancingLearner([learner_1, learner_2])
+
+
+learners_pairs = [
+    (Learner1D, dict(bounds=(-1, 1))),
+    (Learner2D, dict(bounds=[(-1, 1), (-1, 1)])),
+    (SequenceLearner, dict(sequence=list(range(100)))),
+    (IntegratorLearner, dict(bounds=(0, 1), tol=1e-3)),
+    (AverageLearner, dict(atol=0.1)),
+    (datasaver, dict(learner_type=Learner1D, learner_kwargs=dict(bounds=(-1, 1)))),
+    (
+        balancing_learner,
+        dict(learner_type=Learner1D, learner_kwargs=dict(bounds=(-1, 1))),
+    ),
+]
+
+serializers = [(pickle, pickleable_f)]
+if with_cloudpickle:
+    serializers.append((cloudpickle, nonpickleable_f))
+if with_dill:
+    serializers.append((dill, nonpickleable_f))
+
+
+learners = [
+    (learner_type, learner_kwargs, serializer, f)
+    for serializer, f in serializers
+    for learner_type, learner_kwargs in learners_pairs
+]
+
+
+@pytest.mark.parametrize(
+    "learner_type, learner_kwargs, serializer, f", learners,
+)
+def test_serialization_for(learner_type, learner_kwargs, serializer, f):
+    """Test serializing a learner using different serializers."""
+
+    learner = learner_type(f, **learner_kwargs)
+
+    simple(learner, goal_1)
+    learner_bytes = serializer.dumps(learner)
+    loss = learner.loss()
+    asked = learner.ask(10)
+    data = learner.data
+
+    del f
+    del learner
+
+    learner_loaded = serializer.loads(learner_bytes)
+    assert learner_loaded.npoints == 10
+    assert loss == learner_loaded.loss()
+    assert data == learner_loaded.data
+
+    assert asked == learner_loaded.ask(10)
+
+    # load again to undo the ask
+    learner_loaded = serializer.loads(learner_bytes)
+
+    simple(learner_loaded, goal_2)
+    assert learner_loaded.npoints == 20

setup.py

@@ -51,8 +51,10 @@ def get_version_and_cmdclass(package_name):
         "pre_commit",
     ],
     "other": [
-        "ipyparallel>=6.2.5",  # because of https://github.com/ipython/ipyparallel/issues/404
+        "cloudpickle",
+        "dill",
         "distributed",
+        "ipyparallel>=6.2.5",  # because of https://github.com/ipython/ipyparallel/issues/404
         "loky",
         "scikit-optimize",
         "wexpect" if os.name == "nt" else "pexpect",