metric_name arg in plot_performance_curves

Author: sborms

cobra/evaluation/plotting_utils.py

@@ -84,7 +84,8 @@ def plot_performance_curves(model_performance: pd.DataFrame,
                             path: str=None,
                             colors: dict={"train": "#0099bf",
                                           "selection": "#ff9500",
-                                          "validation": "#8064a2"}):
+                                          "validation": "#8064a2"},
+                            metric_name: str=None):
     """Plot performance curves generated by the forward feature selection
     for the train-selection-validation sets.
 
@@ -97,14 +98,21 @@ def plot_performance_curves(model_performance: pd.DataFrame,
         Width and length of the plot.
     path : str, optional
         Path to store the figure.
+    colors : dict, optional
+        Map with colors for train-selection-validation curves.
+    metric_name : str, optional
+        Name to indicate the metric used in model_performance.
+        Defaults to RMSE in case of regression and AUC in case of
+        classification.
     """
 
     model_type = model_performance["model_type"][0]
 
-    if model_type == "classification":
-        metric = "AUC"
-    elif model_type == "regression":
-        metric = "RMSE"
+    if metric_name is None:
+        if model_type == "classification":
+            metric_name = "AUC"
+        elif model_type == "regression":
+            metric_name = "RMSE"
 
     max_metric = np.round(max(max(model_performance['train_performance']),
                               max(model_performance['selection_performance']),
@@ -115,13 +123,13 @@ def plot_performance_curves(model_performance: pd.DataFrame,
         fig, ax = plt.subplots(figsize=dim)
 
         plt.plot(model_performance['train_performance'], marker=".",
-                 markersize=20, linewidth=3, label=metric+" train",
+                 markersize=20, linewidth=3, label="train",
                  color=colors["train"])
         plt.plot(model_performance['selection_performance'], marker=".",
-                 markersize=20, linewidth=3, label=metric+" selection",
+                 markersize=20, linewidth=3, label="selection",
                  color=colors["selection"])
         plt.plot(model_performance['validation_performance'], marker=".",
-                 markersize=20, linewidth=3, label=metric+" validation",
+                 markersize=20, linewidth=3, label="validation",
                  color=colors["validation"])
 
         # Set x- and y-ticks
@@ -141,6 +149,7 @@ def plot_performance_curves(model_performance: pd.DataFrame,
         ax.legend(loc='lower right')
         fig.suptitle('Performance curves forward feature selection',
                      fontsize=20)
+        plt.title("Metric: "+metric_name, fontsize=15, loc="left")
         plt.ylabel('Model performance')
 
         if path is not None: