Update domain_space.py

domain_space.py

@@ -193,12 +193,12 @@ def interpolate_all(params_list, df):
     return df_filtered, filter_cols
 
 
-def make_domain_plot(df_synth_pred, explored_domain_space, x_list):
+def make_domain_plot(df_synth_pred, explored_domain_space, x_list, target="y_pred"):
     """
     Create a plot with the uncertainty space and the training space
     """
-    uncertainty_min = df_synth_pred["y_pred"].min()
-    uncertainty_max = df_synth_pred["y_pred"].max()
+    uncertainty_min = df_synth_pred[target].min()
+    uncertainty_max = df_synth_pred[target].max()
 
     # df_synth_pred2, filter_cols = filter_dataframe(x_list[:6], df_synth_pred)
     df_synth_pred2, filter_cols = interpolate_all(x_list[:6], df_synth_pred)
@@ -210,9 +210,9 @@ def make_domain_plot(df_synth_pred, explored_domain_space, x_list):
         x=x_col,
         y=y_col,
         z=z_col,
-        color="y_pred",
+        color= target,
         range_color=[uncertainty_min, uncertainty_max],
-        hover_data={"y_pred": ":.3f"},
+        hover_data={target: ":.3f"},
     )
 
     # Filter domain space
@@ -256,7 +256,7 @@ def make_domain_plot(df_synth_pred, explored_domain_space, x_list):
     return fig
 
 
-def create_plot(df_synth_pred, explored_space_dict):
+def create_plot(df_synth_pred, explored_space_dict, target):
     """
     Wrapper to create the function to generate the plotly plots
     """
@@ -264,7 +264,7 @@ def create_plot(df_synth_pred, explored_space_dict):
 
     def plot_figure(x):
         x_params = x[:6]
-        fig = make_domain_plot(df_synth_pred, explored_space_dict, x_params)
+        fig = make_domain_plot(df_synth_pred, explored_space_dict, x_params, target)
         if len(x) == 6:
             return fig