Add BalancingLearner.load_dataframe

Author: basnijholt

adaptive/learner/balancing_learner.py

@@ -389,11 +389,14 @@ def from_product(cls, f, learner_type, learner_kwargs, combos):
             learners.append(learner)
         return cls(learners, cdims=arguments)
 
-    def to_dataframe(self, **kwargs):
+    def to_dataframe(self, index_name: str = "learner_index", **kwargs):
         """Return the data as a concatenated `pandas.DataFrame` from child learners.
 
         Parameters
         ----------
+        index_name : str, optional
+            The name of the index column indicating the learner index,
+            by default "learner_index".
         **kwargs : dict
             Keyword arguments passed to each ``child_learner.to_dataframe(**kwargs)``.
 
@@ -408,10 +411,33 @@ def to_dataframe(self, **kwargs):
         """
         if not with_pandas:
             raise ImportError("pandas is not installed.")
-        dfs = [learner.to_dataframe(**kwargs) for learner in self.learners]
+        dfs = []
+        for i, learner in enumerate(self.learners):
+            df = learner.to_dataframe(**kwargs)
+            cols = list(df.columns)
+            df[index_name] = i
+            df = df[[index_name] + cols]
+            dfs.append(df)
         df = pandas.concat(dfs, axis=0, ignore_index=True)
         return df
 
+    def load_dataframe(
+        self, df: pandas.DataFrame, index_name: str = "learner_index", **kwargs
+    ):
+        """Load the data from a `pandas.DataFrame` into the child learners.
+
+        Parameters
+        ----------
+        df : pandas.DataFrame
+            DataFrame with the data to load.
+        index_name : str, optional
+            The ``index_name`` used in `to_dataframe`, by default "learner_index".
+        **kwargs : dict
+            Keyword arguments passed to each ``child_learner.load_dataframe(**kwargs)``.
+        """
+        for i, gr in df.groupby(index_name):
+            self.learners[i].load_dataframe(gr, **kwargs)
+
     def save(self, fname, compress=True):
         """Save the data of the child learners into pickle files
         in a directory.

adaptive/tests/test_learners.py

@@ -736,12 +736,19 @@ def test_to_dataframe(learner_type, f, learner_kwargs):
         learner_type(generate_random_parametrization(f), **learner_kwargs)
         for _ in range(2)
     ]
-    learner = BalancingLearner(learners)
-    simple_run(learner, 100)
-    df = learner.to_dataframe(**kw)
-    assert isinstance(df, pandas.DataFrame)
+    bal_learner = BalancingLearner(learners)
+    simple_run(bal_learner, 100)
+    df_bal = bal_learner.to_dataframe(**kw)
+    assert isinstance(df_bal, pandas.DataFrame)
 
     if learner_type is not AverageLearner1D:
-        assert len(df) == learner.npoints
+        assert len(df_bal) == bal_learner.npoints
 
-    # TODO: Test this for a learner in a DataSaver
+    # Test loading from a DataFrame into the BalancingLearner
+    learners2 = [
+        learner_type(generate_random_parametrization(f), **learner_kwargs)
+        for _ in range(2)
+    ]
+    bal_learner2 = BalancingLearner(learners2)
+    bal_learner2.load_dataframe(df_bal, **kw)
+    assert bal_learner2.npoints == bal_learner.npoints