Test for attrs

Author: basnijholt

adaptive/learner/data_saver.py

@@ -87,7 +87,11 @@ def to_dataframe(
         return df
 
     def load_dataframe(
-        self, df: pandas.DataFrame, extra_data_name: str = "extra_data", **kwargs
+        self,
+        df: pandas.DataFrame,
+        extra_data_name: str = "extra_data",
+        input_names: tuple[str] = (),
+        **kwargs
     ):
         """Load the data from a `pandas.DataFrame` into the learner.
 
@@ -97,11 +101,18 @@ def load_dataframe(
             DataFrame with the data to load.
         extra_data_name : str, optional
             The ``extra_data_name`` used in `to_dataframe`, by default "extra_data".
+        input_names : tuple[str], optional
+            The input names of the child learner. By default the input names are
+            taken from ``df.attrs["inputs"]``, however, metadata is not preserved
+            when saving/loading a DataFrame to/from a file. In that case, the input
+            names can be passed explicitly. For example, for a 2D learner, this would
+            be ``input_names=('x', 'y')``.
         **kwargs : dict
             Keyword arguments passed to each ``child_learner.load_dataframe(**kwargs)``.
         """
         self.learner.load_dataframe(df, **kwargs)
-        for _, x in df[df.attrs["inputs"] + [extra_data_name]].iterrows():
+        keys = df.attrs.get("inputs", list(input_names))
+        for _, x in df[keys + [extra_data_name]].iterrows():
             key = _to_key(x[:-1])
             self.extra_data[key] = x[-1]
 

adaptive/learner/integrator_learner.py

@@ -594,6 +594,8 @@ def to_dataframe(
         if not with_pandas:
             raise ImportError("pandas is not installed.")
         df = pandas.DataFrame(sorted(self.data.items()), columns=[x_name, y_name])
+        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

@@ -440,6 +440,7 @@ def to_dataframe(
         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]
+        df.attrs["output"] = z_name
         if with_default_function_args:
             assign_defaults(self.function, df, function_prefix)
         return df

adaptive/learner/learnerND.py

@@ -445,6 +445,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)
+        df.attrs["output"] = value_name
         if with_default_function_args:
             assign_defaults(self.function, df, function_prefix)
         return df

adaptive/tests/test_learners.py

@@ -729,6 +729,8 @@ def test_to_dataframe(learner_type, f, learner_kwargs):
     # Test empty dataframe
     df = learner.to_dataframe(**kw)
     assert len(df) == 0
+    assert "inputs" in df.attrs
+    assert "output" in df.attrs
 
     # Run the learner
     simple_run(learner, 100)

adaptive/utils.py

@@ -119,6 +119,7 @@ def assign_defaults(
     defaults = _default_parameters(function, function_prefix, start_index)
     for k, v in defaults.items():
         df[k] = len(df) * [v]
+        df[k] = df[k].astype("category")
 
 
 def partial_function_from_dataframe(function, df, function_prefix: str = "function."):