Dataframe name

rubicon_ml/viz/dataframe_plot.py

@@ -46,17 +46,22 @@ class DataframePlot(VizBase):
 
     def __init__(
         self,
-        dataframe_name,
+        dataframe_name=None,
         experiments=None,
         plotting_func=px.line,
         plotting_func_kwargs={},
         x=None,
         y=None,
     ):
         super().__init__(dash_title="plot dataframes")
-
-        self.dataframe_name = dataframe_name
         self.experiments = experiments
+        if experiments != None:
+            if len(experiments[0].dataframes()) == 0:
+                raise Exception("No dataframe logged to experiment")
+            else:
+                self.dataframe_name = self.experiments[0].dataframes()[0].name
+        else:
+            self.dataframe_name = dataframe_name
         self.plotting_func = plotting_func
         self.plotting_func_kwargs = plotting_func_kwargs
         self.x = x
@@ -92,6 +97,7 @@ def load_experiment_data(self):
         `dataframe_name` must have the same schema.
         """
         self.data_df = None
+        count = 0
 
         for experiment in self.experiments:
             try:
@@ -101,7 +107,7 @@ def load_experiment_data(self):
                     f"Experiment {experiment.id} does not have any dataframes logged to it."
                 )
                 continue
-
+                
             data_df = dataframe.get_data()
             data_df["experiment_id"] = experiment.id
 
@@ -127,6 +133,9 @@ def load_experiment_data(self):
         if self.data_df is None:
             raise RubiconException(f"No dataframe with name {self.dataframe_name} found!")
 
+        if count == 0:
+            raise Exception(f"No dataframe with name {self.dataframe_name} found!")
+
     def register_callbacks(self, link_experiment_table=False):
         outputs = [
             Output("dataframe-plot", "figure"),