plot mo

Author: bartzbeielstein

pyproject.toml

@@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "spotpython"
-version = "0.28.12"
+version = "0.28.13"
 authors = [
   { name="T. Bartz-Beielstein", email="tbb@bartzundbartz.de" }
 ]

src/spotpython/mo/plot.py

@@ -11,14 +11,22 @@ def plot_mo(
     y_rf: np.ndarray = None,
     pareto_front: bool = False,
     y_best: np.ndarray = None,
+    y_add: np.ndarray = None,
+    y_add2: np.ndarray = None,
+    y_add_color="blue",
+    y_add2_color="green",
     title: str = "",
     y_orig: np.ndarray = None,
     pareto_front_orig: bool = False,
     pareto_label: bool = False,
     y_rf_color="blue",
     y_best_color="red",
-    x_axis_transformation: str = "id",  # New argument for x-axis transformation
-    y_axis_transformation: str = "id",  # New argument for y-axis transformation
+    x_axis_transformation: str = "id",
+    y_axis_transformation: str = "id",
+    y_best_label="Best",
+    y_add_label="Add",
+    y_add2_label="Add2",
+    filename: str = None,
 ) -> None:
     """
     Generates scatter plots for each combination of two targets from a multi-output prediction while highlighting Pareto optimal points.
@@ -30,6 +38,13 @@ def plot_mo(
         pareto (str): Specifies whether to compute Pareto front based on 'min' or 'max' criterion.
         pareto_front (bool): If True, connect Pareto optimal points with a red line for y_rf.
         y_best (np.ndarray, optional): A NumPy array representing the best point to highlight in red. Defaults to None.
+        y_add (np.ndarray, optional): A NumPy array representing the additional points to highlight in blue. Defaults to None.
+        y_add2 (np.ndarray, optional): A NumPy array representing the additional points to highlight in green. Defaults to None.
+        y_add_color (str): The color of the additional points. Defaults to "blue".
+        y_add2_color (str): The color of the additional points. Defaults to "green".
+        y_best_label (str): The label for the best point. Defaults to "Best".
+        y_add_label (str): The label for the additional points. Defaults to "Add".
+        y_add2_label (str): The label for the additional points. Defaults to "Add2".
         title (str): The title of the plot. Defaults to "" (empty string).
         y_orig (np.ndarray, optional): The original target values with shape (n_samples, n_targets). Defaults to None.
         pareto_front_orig (bool): If True, connect Pareto optimal points with a light blue line for y_orig. Defaults to False.
@@ -38,9 +53,11 @@ def plot_mo(
         y_best_color (str): The color of the best point. Defaults to "red".
         x_axis_transformation (str): Transformation for the x-axis. Options are "id" (linear), "log" (logarithmic), and "loglog" (log-log). Defaults to "id".
         y_axis_transformation (str): Transformation for the y-axis. Options are "id" (linear), "log" (logarithmic), and "loglog" (log-log). Defaults to "id".
+        filename (str, optional):
+            If provided, saves the plot to the specified file.  Supports "pdf" and "png" formats. Defaults to None.
 
     Returns:
-        None: Displays the plot.
+        None: Displays or saves the plot.
 
     Examples:
         >>> from spotpython.mo.plot import plot_mo
@@ -50,7 +67,7 @@ def plot_mo(
         >>> pareto = "min"
         >>> y_rf = np.random.rand(100, 2)
         >>> y_orig = np.random.rand(100, 2)
-        >>> plot_mo(target_names, combinations, pareto, y_rf=y_rf, y_orig=y_orig)
+        >>> plot_mo(target_names, combinations, pareto, y_rf=y_rf, y_orig=y_orig, filename="plot.png")
     """
     # Convert y_rf to numpy array if it's a pandas DataFrame
     if isinstance(y_rf, pd.DataFrame):
@@ -72,14 +89,14 @@ def plot_mo(
         if y_orig is not None:
             minimize = pareto == "min"
             pareto_mask_orig = is_pareto_efficient(y_orig[:, [i, j]], minimize)
-            plt.scatter(y_orig[:, i], y_orig[:, j], edgecolor="w", c="gray", s=s, marker="o", alpha=a, label="Original Points")
-            plt.scatter(y_orig[pareto_mask_orig, i], y_orig[pareto_mask_orig, j], edgecolor="k", c="gray", s=pareto_size, marker="o", alpha=a, label="Original Pareto")
+            plt.scatter(y_orig[:, i], y_orig[:, j], edgecolor="w", c="gray", s=s, marker="o", alpha=a, label="Non-Pareto Points")
+            plt.scatter(y_orig[pareto_mask_orig, i], y_orig[pareto_mask_orig, j], edgecolor="k", c="gray", s=pareto_size, marker="o", alpha=a, label="Pareto Points")
             if pareto_label:
                 for idx in np.where(pareto_mask_orig)[0]:
                     plt.text(y_orig[idx, i], y_orig[idx, j], str(idx), color="black", fontsize=8, ha="center", va="center")
             if pareto_front_orig:
                 sorted_indices_orig = np.argsort(y_orig[pareto_mask_orig, i])
-                plt.plot(y_orig[pareto_mask_orig, i][sorted_indices_orig], y_orig[pareto_mask_orig, j][sorted_indices_orig], "k-", alpha=a, label="Original Pareto Front")
+                plt.plot(y_orig[pareto_mask_orig, i][sorted_indices_orig], y_orig[pareto_mask_orig, j][sorted_indices_orig], "k-", alpha=a, label="Pareto Front")
 
         if y_rf is not None:
             minimize = pareto == "min"
@@ -101,7 +118,11 @@ def plot_mo(
                 )
 
         if y_best is not None:
-            plt.scatter(y_best[:, i], y_best[:, j], edgecolor="k", c=y_best_color, s=s, marker="D", alpha=1, label="Best")
+            plt.scatter(y_best[:, i], y_best[:, j], edgecolor="k", c=y_best_color, s=s, marker="D", alpha=1, label=y_best_label)
+        if y_add is not None:
+            plt.scatter(y_add[:, i], y_add[:, j], edgecolor="k", c=y_add_color, s=s, marker="D", alpha=1, label=y_add_label)
+        if y_add2 is not None:
+            plt.scatter(y_add2[:, i], y_add2[:, j], edgecolor="k", c=y_add2_color, s=s, marker="D", alpha=1, label=y_add2_label)
 
         # Apply axis transformations
         if x_axis_transformation == "log":
@@ -117,4 +138,11 @@ def plot_mo(
         plt.grid()
         plt.title(title)
         plt.legend()
-        plt.show()
+        # Save or show the plot
+        if filename:
+            if filename.endswith(".pdf") or filename.endswith(".png"):
+                plt.savefig(filename, format=filename.split(".")[-1])
+            else:
+                raise ValueError("Filename must have a valid suffix: '.pdf' or '.png'.")
+        else:
+            plt.show()