Add model_out to PostProcessing._plot_validation_multi

src/bayesvalidrox/post_processing/post_processing.py

@@ -220,16 +220,14 @@ class PostProcessing:
         else:
             n_samples = samples.shape[0]
 
-        if model_out_dict is not None:
-            self.model_out_dict = model_out_dict
-        else:
-            self.model_out_dict, _ = self.engine.Model.run_model_parallel(
+        if model_out_dict is None:
+            model_out_dict, _ = self.engine.Model.run_model_parallel(
                 samples, key_str="valid"
             )
 
         out_mean, out_std = self.engine.eval_metamodel(samples)
 
-        self._plot_validation_multi(out_mean=out_mean, out_std=out_std)
+        self._plot_validation_multi(out_mean, out_std, model_out_dict)
 
         # TODO: should this be kept here?
         # Zip the subdirectories
@@ -986,48 +984,36 @@ class PostProcessing:
                 )
                 plt.close(fig)
     # -------------------------------------------------------------------------
-    def _plot_validation_multi(self, out_mean, out_std):
+    def _plot_validation_multi(self, out_mean, out_std, model_out):
         """
         Plots outputs for visual comparison of metamodel outputs with that of
         the (full) multioutput original model
 
         Parameters
         ----------
-        x_values : list or array
-            List of x values.
-        x_axis : str, optional
-            Label of the x axis. The default is "x [m]".
-
-        Raises
-        ------
-        AttributeError: This evaluation only support PCE-type models!
+        out_mean : dict
+            MetaModel mean outputs.
+        out_std : dict
+            MetaModel stdev outputs.
+        model_out : dict
+            Model outputs.
 
         Returns
         -------
         None.
 
         """
-        # This function currently only supports PCE/aPCE
-        if not hasattr(self.engine.MetaModel, "meta_model_type"):
-            raise AttributeError("This evaluation only support PCE-type models!")
-        if self.engine.MetaModel.meta_model_type.lower() not in ["pce", "apce", "gpe"]:
-            raise AttributeError("This evaluation only support PCE-type models!")
-
         # List of markers and colors
         color = cycle((["b", "g", "r", "y", "k"]))
         marker = cycle(("x", "d", "+", "o", "*"))
-        
-        if self.engine.MetaModel.meta_model_type.lower() == "gpe":
-            lower = "GPE"
-        elif self.engine.MetaModel.meta_model_type.lower() == "pce" or self.engine.MetaModel.meta_model_type.lower() == "apce":
-            lower = "PCE"
+        metamod_name = self.engine.MetaModel.meta_model_type.lower() 
         
         # Plot the model vs PCE model
         fig = plt.figure()
         for _, key in enumerate(self.engine.out_names):
             y_val = out_mean[key]
             y_val_std = out_std[key]
-            y_val = self.model_out_dict[key]
+            y_val = model_out[key]
             
             for idx in range(y_val.shape[0]):
                 plt.plot(
@@ -1037,26 +1023,14 @@ class PostProcessing:
                     marker=next(marker),
                     label="$Y_{%s}^M$" % (idx + 1),
                 )
-                if self.engine.MetaModel.meta_model_type.lower() == "pce" \
-                    or self.engine.MetaModel.meta_model_type.lower() == "apce":
-                    plt.plot(
-                        self.x_values,
-                        y_val[idx],
-                        color=next(color),
-                        marker=next(marker),
-                        linestyle="--",
-                        label="$Y_{{{}}}^{{{}}}$".format(idx + 1, lower)
+                plt.plot(
+                    self.x_values,
+                    y_val[idx],
+                    color=next(color),
+                    marker=next(marker),
+                    linestyle="--",
+                    label="$Y_{{{}}}^{{{}}}$".format(idx + 1, metamod_name)
                 )
-                
-                elif self.engine.MetaModel.meta_model_type.lower() == "gpe":
-                    plt.plot(
-                        self.x_values,
-                        y_val[idx],
-                        color=next(color),
-                        marker=next(marker),
-                        linestyle="--",
-                        label="$Y_{{{}}}^{{{}}}$".format(idx + 1, lower)
-                    )
                 plt.fill_between(
                     self.x_values,
                     y_val[idx] - 1.96 * y_val_std[idx],
@@ -1080,6 +1054,7 @@ class PostProcessing:
             plt.grid()
             key = key.replace(" ", "_")
             fig.savefig(
-                f"./{self.out_dir}/Model_vs_{lower}Model_{key}.{self.out_format}", bbox_inches="tight"
+                f"./{self.out_dir}/Model_vs_{metamod_name}Model_{key}.{self.out_format}",
+                bbox_inches="tight"
             )
             plt.close()

tests/test_PostProcessing.py

@@ -417,11 +417,16 @@ def test_plot_validation_multi(pce_engine_3d_plot):
     """
     engine = pce_engine_3d_plot
     post = PostProcessing(engine)
-    y_val = {'Z': [1, 2, 3, 4, 5]}
-    y_val_std = {'Z': [0.1, 0.2, 0.3, 0.4, 0.5]}
-    post._plot_validation_multi(y_val, y_val_std)
+    y_val = {'Z': np.array([[1], [2], [3], [4], [5]]),
+             'Y': np.array([[1], [2], [3], [4], [5]])}
+    y_val_std = {'Z': np.array([[0.1], [0.2], [0.3], [0.4], [0.5]]),
+                'Y': np.array([[0.1], [0.2], [0.3], [0.4], [0.5]])}
+    model_out = {'Z': np.array([[1.5],[2],[3.5],[4],[4.5]]),
+                 'Y': np.array([[1.5],[2],[3.5],[4],[4.5]])}
+    post._plot_validation_multi(y_val, y_val_std, model_out)
     # Check if the plot was created and saved
-    assert os.path.exists(f"./{engine.out_dir}/Model_vs_Model_key1.{engine.out_format}")
+    assert os.path.exists(f"./{post.out_dir}/Model_vs_pceModel_Y.{post.out_format}")
+    assert os.path.exists(f"./{post.out_dir}/Model_vs_pceModel_Z.{post.out_format}")
 
 def test_plot_validation_multi_with_multiple_keys(pce_engine_3d_plot):
     """