Implement DataSaver.to_dataframe

Author: basnijholt

adaptive/learner/average_learner.py

@@ -116,6 +116,8 @@ def to_dataframe(
         if not with_pandas:
             raise ImportError("pandas is not installed.")
         df = pandas.DataFrame(sorted(self.data.items()), columns=[seed_name, y_name])
+        df.attrs["inputs"] = [seed_name]
+        df.attrs["output"] = y_name
         if with_default_function_args:
             assign_defaults(self.function, df, function_prefix)
         return df

adaptive/learner/average_learner1D.py

@@ -196,6 +196,8 @@ def to_dataframe(
             ]
             columns = [seed_name, x_name, y_name]
         df = pandas.DataFrame(data, columns=columns)
+        df.attrs["inputs"] = [seed_name, x_name]
+        df.attrs["output"] = y_name
         if with_default_function_args:
             assign_defaults(self.function, df, function_prefix)
         return df

adaptive/learner/data_saver.py

@@ -4,6 +4,14 @@
 from adaptive.learner.base_learner import BaseLearner
 from adaptive.utils import copy_docstring_from
 
+try:
+    import pandas
+
+    with_pandas = True
+
+except ModuleNotFoundError:
+    with_pandas = False
+
 
 class DataSaver:
     """Save extra data associated with the values that need to be learned.
@@ -44,6 +52,39 @@ def tell(self, x, result):
     def tell_pending(self, x):
         self.learner.tell_pending(x)
 
+    def to_dataframe(
+        self, extra_data_name: str = "extra_data", **kwargs
+    ) -> pandas.DataFrame:
+        """Return the data as a concatenated `pandas.DataFrame` from child learners.
+
+        Parameters
+        ----------
+        extra_data_name : str, optional
+            The name of the column containing the extra data, by default "extra_data".
+        **kwargs : dict
+            Keyword arguments passed to the ``child_learner.to_dataframe(**kwargs)``.
+
+        Returns
+        -------
+        pandas.DataFrame
+
+        Raises
+        ------
+        ImportError
+            If `pandas` is not installed.
+        """
+        if not with_pandas:
+            raise ImportError("pandas is not installed.")
+        df = self.learner.to_dataframe(**kwargs)
+
+        def to_key(x):
+            return tuple(x.values) if x.values.size > 1 else x.item()
+
+        df[extra_data_name] = [
+            self.extra_data[to_key(x)] for _, x in df[df.attrs["inputs"]].iterrows()
+        ]
+        return df
+
     def _get_data(self):
         return self.learner._get_data(), self.extra_data
 

adaptive/learner/learner1D.py

@@ -369,6 +369,8 @@ def to_dataframe(
         xs, ys = zip(*sorted(self.data.items())) if self.data else ([], [])
         df = pandas.DataFrame(xs, columns=[x_name])
         df[y_name] = ys
+        df.attrs["inputs"] = [x_name]
+        df.attrs["output"] = y_name
         if with_default_function_args:
             assign_defaults(self.function, df, function_prefix)
         return df

adaptive/learner/learner2D.py

@@ -439,6 +439,7 @@ def to_dataframe(
             raise ImportError("pandas is not installed.")
         data = sorted((x, y, z) for (x, y), z in self.data.items())
         df = pandas.DataFrame(data, columns=[x_name, y_name, z_name])
+        df.attrs["inputs"] = [x_name, y_name]
         if with_default_function_args:
             assign_defaults(self.function, df, function_prefix)
         return df

adaptive/learner/learnerND.py

@@ -444,6 +444,7 @@ def to_dataframe(
             )
         data = list((*x, y) for x, y in self.data.items())
         df = pandas.DataFrame(data, columns=[*point_names, value_name])
+        df.attrs["inputs"] = list(point_names)
         if with_default_function_args:
             assign_defaults(self.function, df, function_prefix)
         return df

adaptive/learner/sequence_learner.py

@@ -172,6 +172,8 @@ def to_dataframe(
         df = pandas.DataFrame(indices, columns=[index_name])
         df[x_name] = sequence
         df[y_name] = ys
+        df.attrs["inputs"] = [index_name]
+        df.attrs["output"] = y_name
         if with_default_function_args:
             assign_defaults(self.function, df, function_prefix)
         return df

adaptive/tests/test_learners.py

@@ -8,6 +8,7 @@
 import random
 import shutil
 import tempfile
+import time
 
 import flaky
 import numpy as np
@@ -697,13 +698,25 @@ def test_learner_subdomain(learner_type, f, learner_kwargs):
     raise NotImplementedError()
 
 
+def add_time(f):
+    @ft.wraps(f)
+    def wrapper(*args, **kwargs):
+        t0 = time.time()
+        result = f(*args, **kwargs)
+        return {"result": result, "time": time.time() - t0}
+
+    return wrapper
+
+
 @run_with(
     Learner1D,
     Learner2D,
     LearnerND,
     AverageLearner,
     AverageLearner1D,
     SequenceLearner,
+    IntegratorLearner,
+    with_all_loss_functions=False,
 )
 def test_to_dataframe(learner_type, f, learner_kwargs):
     if learner_type is LearnerND:
@@ -752,3 +765,21 @@ def test_to_dataframe(learner_type, f, learner_kwargs):
     bal_learner2 = BalancingLearner(learners2)
     bal_learner2.load_dataframe(df_bal, **kw)
     assert bal_learner2.npoints == bal_learner.npoints
+
+    if learner_type is SequenceLearner:
+        # We do not test the DataSaver with the SequenceLearner
+        # because the DataSaver is not compatible with the SequenceLearner.
+        return
+
+    # Test with DataSaver
+    learner = learner_type(
+        add_time(generate_random_parametrization(f)), **learner_kwargs
+    )
+    data_saver = DataSaver(learner, operator.itemgetter("result"))
+    df = data_saver.to_dataframe(**kw)  # test if empty dataframe works
+    simple_run(data_saver, 100)
+    df = data_saver.to_dataframe(**kw)
+    if learner_type is AverageLearner1D:
+        assert len(df) == data_saver.nsamples
+    else:
+        assert len(df) == data_saver.npoints