Merge pull request ComputationalRadiationPhysics#2484 from codingS3b/visualizer

Python MPL Visualization modules

Author: ax3l

docs/source/usage/plugins/phaseSpace.rst

@@ -73,7 +73,7 @@ The easiest way is to load the data in Python:
 
    # load data
    phase_space = PhaseSpace('/home/axel/runs/foil_001')
-   e_ps, e_ps_meta = phase_space.get('e', species_filter='all', ps='ypy', iteration=1000)
+   e_ps, e_ps_meta = phase_space.get('e', species_filter='all', iteration=1000, ps='ypy')
 
    # unit conversion from SI
    mu = 1.e6  # meters to microns

lib/python/picongpu/plugins/energy_histogram.py

@@ -101,7 +101,7 @@ def get_iterations(self, species, species_filter="all"):
                            dtype=np.uint64).as_matrix()[:, 0]
 
     def get(self, species, species_filter="all", iteration=None,
-            include_overflow=False):
+            include_overflow=False, **kwargs):
         """
         Get a histogram.
 

lib/python/picongpu/plugins/phase_space.py

@@ -188,7 +188,8 @@ def get_iterations(self, species, species_filter='all', ps=None):
 
         return iterations
 
-    def get(self, species, species_filter='all', ps=None, iteration=None):
+    def get(self, species, species_filter='all', iteration=None, ps=None,
+            **kwargs):
         """
         Get a phase space histogram.
 

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

@@ -0,0 +1,96 @@
+"""
+This file is part of the PIConGPU.
+
+Copyright 2017-2018 PIConGPU contributors
+Authors: Sebastian Starke
+License: GPLv3+
+"""
+
+import matplotlib.pyplot as plt
+
+
+class Visualizer(object):
+    """
+    Abstract base class for matplotlib visualizers that implements
+    the visualization logic.
+    Classes that derive from this class need to write their own implementations
+    for the following functions in order to work:
+        _create_data_reader(self, run_directory)
+        _create_plt_obj(self, ax)
+        _update_plt_obj(self)
+    """
+
+    def __init__(self, run_directory):
+        """
+        Initialize the reader and data as member parameters.
+
+        Parameters
+        ----------
+        run_directory : string
+            path to the run directory of PIConGPU
+            (the path before ``simOutput/``)
+        """
+        if run_directory is None:
+            raise ValueError('The run_directory parameter can not be None!')
+
+        self.plt_obj = None
+        self.data_reader = self._create_data_reader(run_directory)
+        self.data = None
+
+    def _create_data_reader(self, run_directory):
+        """
+        Needs to return an instance of a picongpu data reader
+        (as defined in the ../plugin directory) which implements
+        a 'get()' method.
+        """
+        raise NotImplementedError
+
+    def _create_plt_obj(self, ax):
+        """
+        Sets 'self.plt_obj' to an instance of a matplotlib.artist.Artist
+        object (or derived classes) which can be updated
+        by feeding new data into it.
+        Only called on the first call for visualization.
+
+        Parameters
+        ----------
+        ax: matplotlib.axes.Axes instance
+        """
+        raise NotImplementedError
+
+    def _update_plt_obj(self):
+        """
+        Take the 'self.data' member, interpret it and feed it into the
+        'self.plt_obj'.
+        """
+        raise NotImplementedError
+
+    def _ax_or_gca(self, ax):
+        """
+        Returns the passed ax if it is not None or the current
+        matplotlib axes object otherwise.
+
+        Parameters
+        ----------
+        ax: None or matplotlib.axes.Axes instance
+        Returns
+        -------
+        A matplotlib.axes.Axes instance
+        """
+
+        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.
+        """
+        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:
+            self._update_plt_obj()

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

@@ -0,0 +1,146 @@
+"""
+This file is part of the PIConGPU.
+
+Copyright 2017-2018 PIConGPU contributors
+Authors: Sebastian Starke
+License: GPLv3+
+"""
+
+from picongpu.plugins.energy_histogram import EnergyHistogram
+from picongpu.plugins.plot_mpl.base_visualizer import Visualizer as\
+    BaseVisualizer, plt
+
+
+class Visualizer(BaseVisualizer):
+    """
+    Class for creation of histogram plots on a logscaled y-axis.
+    """
+
+    def __init__(self, run_directory):
+        """
+        Parameters
+        ----------
+        run_directory : string
+            path to the run directory of PIConGPU
+            (the path before ``simOutput/``)
+        """
+        super(Visualizer, self).__init__(run_directory)
+
+    def _create_data_reader(self, run_directory):
+        """
+        Implementation of base class function.
+        """
+        return EnergyHistogram(run_directory)
+
+    def _create_plt_obj(self, ax):
+        """
+        Implementation of base class function.
+        Turns 'self.plt_obj' into a matplotlib.pyplot.plot object.
+        """
+        bins, counts = self.data
+        self.plt_obj = ax.semilogy(bins, counts, nonposy='clip')[0]
+
+    def _update_plt_obj(self):
+        """
+        Implementation of base class function.
+        """
+        bins, counts = self.data
+        self.plt_obj.set_data(bins, counts)
+
+    def visualize(self, ax=None, **kwargs):
+        """
+        Creates a semilogy plot on the provided axes object for
+        the data of the given iteration using matpotlib.
+
+        Parameters
+        ----------
+        iteration: int
+            the iteration number for which data will be plotted.
+        ax: matplotlib axes object
+            the part of the figure where this plot will be shown.
+        kwargs: dictionary with further keyword arguments, valid are:
+            species: string
+                short name of the particle species, e.g. 'e' for electrons
+                (defined in ``speciesDefinition.param``)
+            iteration: int
+                number of the iteration
+            species_filter: string
+                name of the particle species filter, default is 'all'
+                (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')
+        species = kwargs.get('species')
+        species_filter = kwargs.get('species_filter', 'all')
+
+        if iteration is None or species is None:
+            raise ValueError("Iteration and species have to be provided as\
+            keyword arguments!")
+
+        ax.set_xlabel('Energy [keV]')
+        ax.set_ylabel('Count')
+        ax.set_title('Energy Histogram for species ' +
+                     species + ', filter = ' + species_filter +
+                     ', iteration ' + str(iteration))
+
+
+if __name__ == '__main__':
+
+    def main():
+        import sys
+        import getopt
+
+        def usage():
+            print("usage:")
+            print(
+                "python", sys.argv[0],
+                "-p <path to run_directory> -i <iteration>"
+                " -s <particle species> -f <species_filter>")
+
+        path = None
+        iteration = None
+        species = None
+        filtr = None
+
+        try:
+            opts, args = getopt.getopt(sys.argv[1:], "hp:i:s:f:", [
+                "help", "path", "iteration", "species", "filter"])
+        except getopt.GetoptError as err:
+            print(err)
+            usage()
+            sys.exit(2)
+
+        for opt, arg in opts:
+            if opt in ["-h", "--help"]:
+                usage()
+                sys.exit()
+            elif opt in ["-p", "--path"]:
+                path = arg
+            elif opt in ["-i", "--iteration"]:
+                iteration = int(arg)
+            elif opt in ["-s", "--species"]:
+                species = arg
+            elif opt in ["-f", "--filter"]:
+                filtr = arg
+
+        # check that we got all args that we need
+        if path is None or iteration is None:
+            print("Path to 'run' directory and iteration have to be provided!")
+            usage()
+            sys.exit(2)
+        if species is None:
+            species = 'e'
+            print("Particle species was not given, will use", species)
+        if filtr is None:
+            filtr = 'all'
+            print("Species filter was not given, will use", filtr)
+
+        fig, ax = plt.subplots(1, 1)
+        Visualizer(path).visualize(ax, iteration=iteration, species=species,
+                                   species_filter=filtr)
+        plt.show()
+
+    main()