add 'save' and 'load' to the learners and periodic saving to the Runner

Author: basnijholt

adaptive/learner/average_learner.py

@@ -125,3 +125,9 @@ def plot(self):
         num_bins = int(max(5, sqrt(self.npoints)))
         vals = hv.Points(vals)
         return hv.operation.histogram(vals, num_bins=num_bins, dimension=1)
+
+    def _get_data(self):
+        return (self.data, self.npoints, self.sum_f, self.sum_f_sq)
+
+    def _set_data(self, data):
+        self.data, self.npoints, self.sum_f, self.sum_f_sq = data

adaptive/learner/balancing_learner.py

@@ -3,6 +3,7 @@
 from contextlib import suppress
 from functools import partial
 from operator import itemgetter
+import os.path
 
 import numpy as np
 
@@ -302,3 +303,75 @@ def from_product(cls, f, learner_type, learner_kwargs, combos):
             learner = learner_type(function=partial(f, **combo), **learner_kwargs)
             learners.append(learner)
         return cls(learners, cdims=arguments)
+
+    def save(self, folder, compress=True):
+        """Save the data of the child learners into pickle files
+        in a directory.
+
+        Parameters
+        ----------
+        folder : str
+            Directory in which the learners's data will be saved.
+        compress : bool, default True
+            Compress the data upon saving using 'gzip'. When saving
+            using compression, one must load it with compression too.
+
+        Notes
+        -----
+        The child learners need to have a 'fname' attribute in order to use
+        this method.
+
+        Example
+        -------
+        >>> def combo_fname(val):
+        ...     return '__'.join([f'{k}_{v}.p' for k, v in val.items()])
+        ...
+        ... def f(x, a, b): return a * x**2 + b
+        ...
+        >>> learners = []
+        >>> for combo in adaptive.utils.named_product(a=[1, 2], b=[1]):
+        ...     l = Learner1D(functools.partial(f, combo=combo))
+        ...     l.fname = combo_fname(combo)  # 'a_1__b_1.p', 'a_2__b_1.p' etc.
+        ...     learners.append(l)
+        ... learner = BalancingLearner(learners)
+        ... # Run the learner
+        ... runner = adaptive.Runner(learner)
+        ... # Then save
+        ... learner.save('data_folder')  # use 'load' in the same way
+        """
+        if len(self.learners) != len(set(l.fname for l in self.learners)):
+            raise RuntimeError("The 'learner.fname's are not all unique.")
+
+        for l in self.learners:
+            l.save(os.path.join(folder, l.fname), compress=compress)
+
+    def load(self, folder, compress=True):
+        """Load the data of the child learners from pickle files
+        in a directory.
+
+        Parameters
+        ----------
+        folder : str
+            Directory from which the learners's data will be loaded.
+        compress : bool, default True
+            If the data is compressed when saved, one must load it
+            with compression too.
+
+        Notes
+        -----
+        The child learners need to have a 'fname' attribute in order to use
+        this method.
+
+        Example
+        -------
+        See the example in the 'BalancingLearner.save' doc-string.
+        """
+        for l in self.learners:
+            l.load(os.path.join(folder, l.fname), compress=compress)
+
+    def _get_data(self):
+        return [l._get_data() for l in learner.learners]
+
+    def _set_data(self, data):
+        for l, _data in zip(self.learners, data):
+            l._set_data(_data)

adaptive/learner/base_learner.py

@@ -1,7 +1,10 @@
 # -*- coding: utf-8 -*-
 import abc
+from contextlib import suppress
 from copy import deepcopy
 
+from ..utils import save, load
+
 
 class BaseLearner(metaclass=abc.ABCMeta):
     """Base class for algorithms for learning a function 'f: X → Y'.
@@ -83,8 +86,84 @@ def ask(self, n, tell_pending=True):
         """
         pass
 
+    @abc.abstractmethod
+    def _get_data(self):
+        pass
+
+    @abc.abstractmethod
+    def _set_data(self):
+        pass
+
+    def copy_from(self, other):
+        """Copy over the data from another learner.
+
+        Parameters
+        ----------
+        other : BaseLearner object
+            The learner from which the data is copied.
+        """
+        self._set_data(other._get_data())
+
+    def save(self, fname=None, compress=True):
+        """Save the data of the learner into a pickle file.
+
+        Parameters
+        ----------
+        fname : str, optional
+            The filename of the learner's pickle data file. If None use
+            the 'fname' attribute, like 'learner.fname = "example.p".
+        compress : bool, default True
+            Compress the data upon saving using 'gzip'. When saving
+            using compression, one must load it with compression too.
+
+        Notes
+        -----
+        There are __two ways__ of naming the files:
+        1. Using the 'fname' argument in 'learner.save(fname='example.p')
+        2. Setting the 'fname' attribute, like
+           'learner.fname = "data/example.p"' and then 'learner.save()'.
+        """
+        fname = fname or self.fname
+        data = self._get_data()
+        save(fname, data, compress)
+
+    def load(self, fname=None, compress=True):
+        """Load the data of a learner from a pickle file.
+
+        Parameters
+        ----------
+        fname : str, optional
+            The filename of the saved learner's pickled data file.
+            If None use the 'fname' attribute, like
+            'learner.fname = "example.p".
+        compress : bool, default True
+            If the data is compressed when saved, one must load it
+            with compression too.
+
+        Notes
+        -----
+        See the notes in the 'BaseLearner.save' doc-string.
+        """
+        fname = fname or self.fname
+        with suppress(FileNotFoundError, EOFError):
+            data = load(fname, compress)
+            self._set_data(data)
+
     def __getstate__(self):
         return deepcopy(self.__dict__)
 
     def __setstate__(self, state):
         self.__dict__ = state
+
+    @property
+    def fname(self):
+        # This is a property because then it will be availible in the DataSaver
+        try:
+            return self._fname
+        except AttributeError:
+            raise AttributeError("Set 'learner.fname' or use the 'fname'"
+                " argument when using 'learner.save' or 'learner.load'.")
+
+    @fname.setter
+    def fname(self, fname):
+        self._fname = fname

adaptive/learner/data_saver.py

@@ -2,6 +2,9 @@
 from collections import OrderedDict
 import functools
 
+from .base_learner import BaseLearner
+from ..utils import copy_docstring_from
+
 
 class DataSaver:
     """Save extra data associated with the values that need to be learned.
@@ -40,6 +43,25 @@ def tell(self, x, result):
     def tell_pending(self, x):
         self.learner.tell_pending(x)
 
+    def _get_data(self):
+        return self.learner._get_data(), self.extra_data
+
+    def _set_data(self, data):
+        learner_data, self.extra_data = data
+        self.learner._set_data(learner_data)
+
+    @copy_docstring_from(BaseLearner.save)
+    def save(self, fname=None, compress=True):
+        # We copy this method because the 'DataSaver' is not a
+        # subclass of the 'BaseLearner'.
+        BaseLearner.save(self, fname, compress)
+
+    @copy_docstring_from(BaseLearner.load)
+    def load(self, fname=None, compress=True):
+        # We copy this method because the 'DataSaver' is not a
+        # subclass of the 'BaseLearner'.
+        BaseLearner.load(self, fname, compress)
+
 
 def _ds(learner_type, arg_picker, *args, **kwargs):
     args = args[2:]  # functools.partial passes the first 2 arguments in 'args'!

adaptive/learner/integrator_learner.py

@@ -525,3 +525,30 @@ def plot(self):
         xs, ys = zip(*[(x, y) for ival in ivals
                        for x, y in sorted(ival.done_points.items())])
         return hv.Path((xs, ys))
+
+    def _get_data(self):
+        # Change the defaultdict of SortedSets to a normal dict of sets.
+        x_mapping = {k: set(v) for k, v in self.x_mapping.items()}
+
+        return (self.priority_split,
+                self.done_points,
+                self.pending_points,
+                self._stack,
+                x_mapping,
+                self.ivals,
+                self.first_ival)
+
+    def _set_data(self, data):
+        self.priority_split, self.done_points, self.pending_points, \
+            self._stack, x_mapping, self.ivals, self.first_ival = data
+
+        # Add the pending_points to the _stack such that they are evaluated again
+        for x in self.pending_points:
+            if x not in self._stack:
+                self._stack.append(x)
+
+        # x_mapping is a data structure that can't easily be saved
+        # so we recreate it here
+        self.x_mapping = defaultdict(lambda: SortedSet([], key=attrgetter('rdepth')))
+        for k, _set in x_mapping.items():
+            self.x_mapping[k].update(_set)

adaptive/learner/learner1D.py

@@ -485,3 +485,9 @@ def remove_unfinished(self):
         self.pending_points = set()
         self.losses_combined = deepcopy(self.losses)
         self.neighbors_combined = deepcopy(self.neighbors)
+
+    def _get_data(self):
+        return self.data
+
+    def _set_data(self, data):
+        self.tell_many(*zip(*data.items()))

adaptive/learner/learner2D.py

@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 from collections import OrderedDict
+from copy import copy
 import itertools
 from math import sqrt
 
@@ -522,3 +523,13 @@ def plot(self, n=None, tri_alpha=0):
         no_hover = dict(plot=dict(inspection_policy=None, tools=[]))
 
         return im.opts(style=im_opts) * tris.opts(style=tri_opts, **no_hover)
+
+    def _get_data(self):
+        return self.data
+
+    def _set_data(self, data):
+        self.data = data
+        # Remove points from stack if they already exist
+        for point in copy(self._stack):
+            if point in self.data:
+                self._stack.pop(point)

adaptive/learner/learnerND.py

@@ -572,3 +572,9 @@ def plot_slice(self, cut_mapping, n=None):
             return im.opts(style=dict(cmap='viridis'))
         else:
             raise ValueError("Only 1 or 2-dimensional plots can be generated.")
+
+    def _get_data(self):
+        return self.data
+
+    def _set_data(self, data):
+        self.tell_many(*zip(*data.items()))

adaptive/runner.py

@@ -449,6 +449,7 @@ def goal(_):
                                              self.function)
 
         self.task = self.ioloop.create_task(self._run())
+        self.saving_task = None
         if in_ipynb() and not self.ioloop.is_running():
             warnings.warn("The runner has been scheduled, but the asyncio "
                           "event loop is not running! If you are "
@@ -541,6 +542,31 @@ def elapsed_time(self):
             end_time = time.time()
         return end_time - self.start_time
 
+    def start_periodic_saving(self, save_kwargs, interval):
+        """Periodically save the learner's data.
+
+        Parameters
+        ----------
+        save_kwargs : dict
+            Key-word arguments for 'learner.save(**save_kwargs)'.
+        interval : int
+            Number of seconds between saving the learner.
+
+        Example
+        -------
+        >>> runner = Runner(learner)
+        >>> runner.start_periodic_saving(
+        ...     save_kwargs=dict(fname='data/test.pickle'),
+        ...     interval=600)
+        """
+        async def _saver(save_kwargs=save_kwargs, interval=interval):
+            while self.status() == 'running':
+                self.learner.save(**save_kwargs)
+                await asyncio.sleep(interval)
+            self.learner.save(**save_kwargs)  # one last time
+        self.saving_task = self.ioloop.create_task(_saver())
+        return self.saving_task
+
 
 # Default runner
 Runner = AsyncRunner