Visualizers: add working default for ax parameter

Author: codingS3b

lib/python/picongpu/plugins/plot_mpl/base_visualizer.py

@@ -29,15 +29,24 @@ def _update_plt_obj(self):
         """
         raise NotImplementedError
 
-    def visualize(self, ax, **kwargs):
+    def _ax_or_gca(self, ax):
+        """
+        Returns the passed ax if it is not None or the current
+        matplotlib axes object otherwise.
+        """
+
+        return ax or plt.gca()
+
+    def visualize(self, ax=None, **kwargs):
         """
         1. Creates the 'plt_obj' if it does not exist
         2. Fills the 'data' parameter by using the reader
         3. Updates the 'plt_obj' with the new data.
         """
-
-        self.data = self.data_reader.get(**kwargs)
+        if ax is None:
+            raise ValueError("A matplotlib axes object needs to be passed!")
 
+        self.data = self.data_reader.get(**kwargs)
         if self.plt_obj is None:
             self._create_plt_obj(ax)
         else:

lib/python/picongpu/plugins/plot_mpl/energy_histogram_visualizer.py

@@ -30,7 +30,7 @@ def _update_plt_obj(self):
         bins, counts = self.data
         self.plt_obj.set_data(bins, counts)
 
-    def visualize(self, ax, **kwargs):
+    def visualize(self, ax=None, **kwargs):
         """
         Creates a semilogy plot on the provided axes object for
         the data of the given iteration using matpotlib.
@@ -52,7 +52,7 @@ def visualize(self, ax, **kwargs):
                 (defined in ``particleFilters.param``)
 
         """
-
+        ax=self._ax_or_gca(ax)
         # this already throws error if no species or iteration in kwargs
         super(Visualizer, self).visualize(ax, **kwargs)
         iteration = kwargs.get('iteration')

lib/python/picongpu/plugins/plot_mpl/phase_space_visualizer.py

@@ -50,7 +50,7 @@ def _update_plt_obj(self):
         dat, meta = self.data
         self.plt_obj.set_data(np.abs(dat).T * meta.dV)
 
-    def visualize(self, ax, **kwargs):
+    def visualize(self, ax=None, **kwargs):
         """
         Creates a phase space plot on the provided axes object for
         the data of the given iteration using matpotlib.
@@ -72,6 +72,7 @@ def visualize(self, ax, **kwargs):
                 phase space selection in order: spatial, momentum component,
                 e.g. 'ypy' or 'ypx'
         """
+        ax = self._ax_or_gca(ax)
         super(Visualizer, self).visualize(ax, **kwargs)
 
         _, meta = self.data

lib/python/picongpu/plugins/plot_mpl/png_visualizer.py

@@ -37,7 +37,7 @@ def _create_plt_obj(self, ax):
     def _update_plt_obj(self):
         self.plt_obj.set_data(self.data)
 
-    def visualize(self, ax, **kwargs):
+    def visualize(self, ax=None, **kwargs):
         """
         Creates a plot on the provided axes object for
         the PNG file of the given iteration using matpotlib.
@@ -65,6 +65,7 @@ def visualize(self, ax, **kwargs):
                 if set to 'None', then return images for all available\
                     iterations
         """
+        ax = self._ax_or_gca(ax)
         super(Visualizer, self).visualize(ax, **kwargs)