make BalancingLearner picklable

python-adaptive · Apr 9, 2020 · aff10f3 · aff10f3
1 parent ea617a4
commit aff10f3
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/adaptive/learner/balancing_learner.py b/adaptive/learner/balancing_learner.py
@@ -440,3 +440,16 @@ 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,
+            self.strategy,
+            self._get_data(),
+        )
+
+    def __setstate__(self, state):
+        learners, cdims, strategy, data = state
+        self.__init__(learners, cdims=cdims, strategy=strategy)
+        self._set_data(data)