add animated slices

plotter/Plotter.py

@@ -116,7 +116,7 @@ class Plotter:
                 self.ndim = 1
             elif args.func == self.plot_contour or args.func == self.ani_plot_contour:
                 self.ndim = 2
-            elif args.func == self.avg_plot or args.func == self.slice:
+            elif args.func == self.avg_plot or args.func == self.slice or args.func == self.ani_slice:
                 self.ndim = 3
 
             if self.var == None and len(self.file_column_names) != 0:
@@ -612,6 +612,114 @@ class Plotter:
             else: plt.close(fig)
             if self.if_save_result: fig.savefig(self.out + self.fig_names[i] + self.fig_end)
 
+    def __ani_slice(self):
+        if self.if_save_result:
+            os.system("mkdir -p " + self.out)
+        os.system("mkdir -p " + self.ani_pngs_dir)
+
+        x_name = self.file_column_names[0]
+        y_name = self.file_column_names[1]
+
+        if self.oname == None:
+            fig_names = self.var
+            fig_end = ".gif"
+        else:
+            fig_names = self.oname
+            fig_end = ""
+
+        X = self.file_data[0].data[x_name]
+        Y = self.file_data[0].data[y_name]
+
+        if self.max_y == None:
+            max_val = {var:self.min_max_var_vals[var][1] for var in self.var}
+        else:
+            max_val = {var:self.max_y[i] for var, i in zip(self.var, list(range(len(self.var))))}
+
+        if self.min_y == None:
+            min_val = {var:self.min_max_var_vals[var][0] for var in self.var}
+        else:
+            min_val = {var:self.min_y[i] for var, i in zip(self.var, list(range(len(self.var))))}
+
+        duration = self.ani_period * len(self.filename)
+
+        if self.plane_type == 'xy':
+            x_name = self.file_column_names[0]
+            y_name = self.file_column_names[1]
+
+            z_name = self.file_column_names[2]
+            arr = self.file_data[0].data[z_name]
+            pos = self.__find_closest(arr, self.slice_position)
+
+        elif self.plane_type == 'yz':
+            x_name = self.file_column_names[1]
+            y_name = self.file_column_names[2]
+
+            z_name = self.file_column_names[0]
+            arr = self.file_data[0].data[z_name]
+            pos = self.__find_closest(arr, self.slice_position)
+
+        elif self.plane_type == 'xz':
+            x_name = self.file_column_names[0]
+            y_name = self.file_column_names[2]
+
+            z_name = self.file_column_names[1]
+            arr = self.file_data[0].data[z_name]
+            pos = self.__find_closest(arr, self.slice_position)
+
+        i = 0
+        for var in self.var:
+            if self.title == None:
+                title = var
+            else:
+                title = self.title
+
+            counter = 0
+            png_names = []
+            levels = np.linspace(min_val[var], max_val[var], 30)
+            for data in self.file_data:
+                fig,ax=plt.subplots(1,1)
+            
+                if self.title == None:
+                    title = self.var
+                else:
+                    title = self.title
+
+                if self.plane_type == 'xy':
+                    Z = data.data[self.var[i]][pos, :, :]
+
+                elif self.plane_type == 'xz':
+                    Z = data.data[self.var[i]][:, pos, :]
+
+                elif self.plane_type == 'yz':
+                    Z = data.data[self.var[i]][:, :, pos]
+
+                if self.ifTranspose:
+                    cp = ax.contourf(Y, X, Z.T)
+                    ax.set_ylabel(x_name)
+                    ax.set_xlabel(y_name)
+                else:
+                    cp = ax.contourf(X, Y, Z)
+                    ax.set_xlabel(x_name)
+                    ax.set_ylabel(y_name)
+                
+                fig.colorbar(cp) # Add a colorbar to a plot
+                ax.set_title(title)
+                figname = var + str(counter)
+                fig.savefig(self.ani_pngs_dir + figname + '.png')
+                name = self.ani_pngs_dir + figname + '.png'
+                png_names.append(name)
+                counter += 1
+                plt.close(fig)
+            
+            images = []
+            for file_name in png_names:
+                images.append(imageio.v2.imread(file_name))
+
+            imageio.mimsave(self.out + fig_names[i] + fig_end, images, duration = 0.25, loop=0)
+            i += 1
+        
+        shutil.rmtree(self.ani_pngs_dir)
+ 
     def dump(self):
         self.__dump()
 
@@ -653,3 +761,7 @@ class Plotter:
     def slice(self):
         self.__get_min_max_bar()
         self.__slice()
+
+    def ani_slice(self):
+        self.__get_min_max_bar()
+        self.__ani_slice()
\ No newline at end of file

plotter/__init__.py

@@ -8,6 +8,7 @@ from .main import plot_contour
 from .main import ani_plot_contour
 from .main import avg_plot
 from .main import slice
+from .main import ani_slice
 from .main import plot_diff
 from .main import get_data
 from .main import multiple_plot

plotter/main.py

@@ -78,6 +78,20 @@ def slice(filename, plane_type = 'xy', slice_position = 0, out = './', oname = [
     Plot.set(args, if_manual_plot=manual_plot, if_save_result=save_result)
     Plot.slice()
 
+def ani_slice(filename, plane_type = 'xy', slice_position = 0, out = './', oname = [None], var=[None], mval=[None], min_y = None, max_y = None, ifTranspose=False):
+    pp = str(PurePath(filename).parent) + '/'
+    name = str(PurePath(filename).name )
+    command = 'find ' + pp + ' -name ' + "'" + name + "'"
+    output_stream = os.popen(command)
+    out_res = output_stream.read().split('\n')[0:-1]
+    output_stream.close()
+    names = natsort.natsorted(out_res,reverse=False)
+    
+    Plot = Plotter()
+    args = Args(names, Plot.ani_slice, var=var, mval=mval, oname=oname, out=out, plane_type=plane_type, slice_position=slice_position, ifTranspose=ifTranspose)
+    Plot.set(args, if_manual_plot=manual_plot, if_save_result=save_result)
+    Plot.ani_slice()
+
 def plot_diff(filename, ndim, out = './', oname = 'fig.png', var=[None], mval=[None], ifTranspose=False):
     Plot = Plotter()
     args = Args(filename, Plot.plot_diff, var=var, mval=mval, oname=oname, out=out, ndim=ndim, ifTranspose=ifTranspose)