[fix] Sorted out_dir, changed eval_pce_model_3d to plot_metamodel_3d

src/bayesvalidrox/post_processing/post_processing.py

@@ -81,11 +81,6 @@ class PostProcessing:
         # Compute the moments with the PCEModel object
         self.means, self.stds = self.engine.MetaModel.compute_moments()
 
-        # Open a pdf for the plots
-        newpath = (f'Outputs_PostProcessing_{self.name}/')
-        if not os.path.exists(newpath):
-            os.makedirs(newpath)
-
         # Plot the best fit line, set the linewidth (lw), color and
         # transparency (alpha) of the line
         for key in self.engine.out_names:
@@ -871,11 +866,6 @@ class PostProcessing:
             y_val = outputs
         y_pce_val, _ = self.engine.eval_metamodel(samples=samples)
 
-        # Open a pdf for the plots
-        newpath = f'Outputs_PostProcessing_{self.name}/'
-        if not os.path.exists(newpath):
-            os.makedirs(newpath)
-
         # Fit the data(train the model)
         for key in y_pce_val.keys():
 
@@ -964,97 +954,55 @@ class PostProcessing:
             plt.close()
 
     # -------------------------------------------------------------------------
-    def eval_pce_model_3d(self):
-        # This function currently only supports PCE/aPCE
-        PCEModel = self.engine.MetaModel
-        if not hasattr(PCEModel, 'meta_model_type'):
-            raise AttributeError('This evaluation only support PCE-type models!')
-        if PCEModel.meta_model_type.lower() not in ['pce', 'apce']:
-            raise AttributeError('This evaluation only support PCE-type models!')
-
-        self.n_samples = 1000
-        n_samples = self.n_samples
-
-        # Create 3D-Grid
-        # TODO: Make it general
-        x = np.linspace(-5, 10, n_samples)
-        y = np.linspace(0, 15, n_samples)
-
-        X, Y = np.meshgrid(x, y)
-        PCE_Z = np.zeros((self.n_samples, self.n_samples))
-        Model_Z = np.zeros((self.n_samples, self.n_samples))
-
-        for idxMesh in range(self.n_samples):
-            sample_mesh = np.vstack((X[:, idxMesh], Y[:, idxMesh])).T
-
-            univ_p_val = PCEModel.univ_basis_vals(sample_mesh)
-
-            for Outkey, ValuesDict in PCEModel.coeffs_dict.items():
-
-                pce_out_mean = np.zeros((len(sample_mesh), len(ValuesDict)))
-                pce_out_std = np.zeros((len(sample_mesh), len(ValuesDict)))
-                model_outs = np.zeros((len(sample_mesh), len(ValuesDict)))
-
-                for Inkey, InIdxValues in ValuesDict.items():
-                    idx = int(Inkey.split('_')[1]) - 1
-                    basis_deg_ind = PCEModel.basis_dict[Outkey][Inkey]
-                    clf_poly = PCEModel.clf_poly[Outkey][Inkey]
-
-                    PSI_Val = PCEModel.create_psi(basis_deg_ind, univ_p_val)
-
-                    # Perdiction with error bar
-                    y_mean, y_std = clf_poly.predict(PSI_Val, return_std=True)
-
-                    pce_out_mean[:, idx] = y_mean
-                    pce_out_std[:, idx] = y_std
-
-                    # Model evaluation
-                    model_out_dict, _ = self.engine.Model.run_model_parallel(sample_mesh,
-                                                                 key_str='Valid3D')
-                    model_outs[:, idx] = model_out_dict[Outkey].T
-
-                PCE_Z[:, idxMesh] = y_mean
-                Model_Z[:, idxMesh] = model_outs[:, 0]
-
-        # ---------------- 3D plot for PCEModel -----------------------
-        fig_PCE = plt.figure()
-        ax = plt.axes(projection='3d')
-        ax.plot_surface(X, Y, PCE_Z, rstride=1, cstride=1,
-                        cmap='viridis', edgecolor='none')
-        ax.set_title('PCEModel')
-        ax.set_xlabel('$x_1$')
-        ax.set_ylabel('$x_2$')
-        ax.set_zlabel('$f(x_1,x_2)$')
+    def plot_metamodel_3d(self, n_samples = 10):
+        """
+        Visualize the results of a PCE MetaModel as a 3D surface over two input 
+        parameters.
+        
+        Parameters
+        ----------
+        n_samples : int
+            Number of samples that are used to generate the 3D plot.
 
-        plt.grid()
+        Raises
+        ------
+        AttributeError
+            This function is only applicable if the MetaModel input dimension is 2.
 
-        #  Saving the figure
-        newpath = f'Outputs_PostProcessing_{self.name}/'
-        if not os.path.exists(newpath):
-            os.makedirs(newpath)
+        Returns
+        -------
+        None.
 
-        # save the figure to file
-        fig_PCE.savefig(f'./{newpath}/3DPlot_PCEModel.pdf',
-                        bbox_inches='tight')
-        plt.close(fig_PCE)
-
-        # ---------------- 3D plot for Model -----------------------
-        fig_Model = plt.figure()
-        ax = plt.axes(projection='3d')
-        ax.plot_surface(X, Y, PCE_Z, rstride=1, cstride=1,
-                        cmap='viridis', edgecolor='none')
-        ax.set_title('Model')
-        ax.set_xlabel('$x_1$')
-        ax.set_ylabel('$x_2$')
-        ax.set_zlabel('$f(x_1,x_2)$')
-        plt.grid()
+        """
+        if self.engine.ExpDesign.ndim !=2:
+            raise AttributeError('This function is only applicable if the MetaModel input dimension is 2.')
+        samples = self.engine.ExpDesign.generate_samples(n_samples)
+        samples = np.sort(np.sort(samples, axis = 1), axis = 0)
+        mean, stdev = self.engine.eval_metamodel(samples = samples)
+        
+        if self.engine.emulator:
+            title = 'MetaModel'
+        else:
+            title = 'Model'
+        X,Y = np.meshgrid(samples[:,0],samples[:,1])
+        for name in self.engine.out_names:
+            for t in range(mean[name].shape[1]):
+                fig = plt.figure()
+                ax = plt.axes(projection='3d')
+                ax.plot_surface(X, Y, np.atleast_2d(mean[name][:,t]), rstride=1, cstride=1,
+                                cmap='viridis', edgecolor='none')
+                ax.set_title(title)
+                ax.set_xlabel('$x_1$')
+                ax.set_ylabel('$x_2$')
+                ax.set_zlabel('$f(x_1,x_2)$')
+        
+                plt.grid()
+        
+                # save the figure to file
+                fig.savefig(f'./{self.out_dir}/3DPlot_{title}_{name}{t}.pdf',
+                                bbox_inches='tight')
+                plt.close(fig)
         
-        # Save the figure
-        fig_Model.savefig(f'./{newpath}/3DPlot_Model.pdf',
-                          bbox_inches='tight')
-        plt.close(fig_Model)
-
-        return
 
     # -------------------------------------------------------------------------
     def _get_sample(self, n_samples=None):
@@ -1204,10 +1152,6 @@ class PostProcessing:
         if PCEModel.meta_model_type.lower() not in ['pce', 'apce']:
             raise AttributeError('This evaluation only support PCE-type models!')
 
-        newpath = f'Outputs_PostProcessing_{self.name}/'
-        if not os.path.exists(newpath):
-            os.makedirs(newpath)
-
         # List of markers and colors
         color = cycle((['b', 'g', 'r', 'y', 'k']))
         marker = cycle(('x', 'd', '+', 'o', '*'))
@@ -1259,6 +1203,3 @@ class PostProcessing:
 
             # Destroy the current plot
             plt.close()
-
-        # Zip the subdirectories
-        self.engine.Model.zip_subdirs(f'{self.engine.Model.name}valid', f'{self.engine.Model.name}valid_')