Add SequenceLearner.to_dataframe

Author: basnijholt

adaptive/learner/sequence_learner.py

@@ -4,6 +4,15 @@
 from sortedcontainers import SortedDict, SortedSet
 
 from adaptive.learner.base_learner import BaseLearner
+from adaptive.utils import assign_defaults
+
+try:
+    import pandas
+
+    with_pandas = True
+
+except ModuleNotFoundError:
+    with_pandas = False
 
 
 class _IgnoreFirstArgument:
@@ -120,6 +129,25 @@ def result(self):
     def npoints(self):
         return len(self.data)
 
+    def to_dataframe(
+        self,
+        with_default_function_args: bool = True,
+        function_prefix: str = "function.",
+        index_name: str = "i",
+        x_name: str = "x",
+        y_name: str = "y",
+    ) -> pandas.DataFrame:
+        if not with_pandas:
+            raise ImportError("pandas is not installed.")
+        indices, ys = zip(*self.data.items()) if self.data else ([], [])
+        sequence = [self.sequence[i] for i in indices]
+        df = pandas.DataFrame(indices, columns=[index_name])
+        df[x_name] = sequence
+        df[y_name] = ys
+        if with_default_function_args:
+            assign_defaults(self.function, df, function_prefix)
+        return df
+
     def _get_data(self):
         return self.data
 

adaptive/tests/test_learners.py

@@ -703,11 +703,11 @@ def test_learner_subdomain(learner_type, f, learner_kwargs):
     LearnerND,
     AverageLearner,
     AverageLearner1D,
-    # SequenceLearner,  # TODO: implement this
+    SequenceLearner,
 )
 def test_to_dataframe(learner_type, f, learner_kwargs):
     if learner_type is LearnerND:
-        kw = {"point_names": list("xyz")[: len(learner_kwargs["bounds"])]}
+        kw = {"point_names": tuple("xyz")[: len(learner_kwargs["bounds"])]}
     else:
         kw = {}
 
@@ -734,6 +734,8 @@ def test_to_dataframe(learner_type, f, learner_kwargs):
         learner2.tell_many(df["seed"].values, df["y"])
     elif learner_type is AverageLearner1D:
         learner2.tell_many(df[["seed", "x"]].values, df["y"])
+    elif learner_type is SequenceLearner:
+        learner2.tell_many(df[["i", "x"]].values, df["y"])
     else:
         raise NotImplementedError()