Axis argument

.travis.yml

@@ -32,7 +32,7 @@ before_install:
 
 # Install packages
 install:
-  - conda install --yes cython numpy scipy h5py
+  - conda install --yes cython numpy scipy h5py pytest
   - pip install matplotlib
   - python setup.py install
   - pip install wget

opmd_viewer/addons/pic/lpa_diagnostics.py

@@ -14,12 +14,7 @@
 import numpy as np
 import scipy.constants as const
 from scipy.optimize import curve_fit
-from opmd_viewer.openpmd_timeseries.plotter import check_matplotlib
-try:
-    import matplotlib.pyplot as plt
-except ImportError:
-    # Any error will be caught later by `check_matplotlib`
-    pass
+from opmd_viewer.openpmd_timeseries.plotter import check_matplotlib_and_axis
 
 
 class LpaDiagnostics( OpenPMDTimeSeries ):
@@ -97,7 +92,8 @@ def get_mean_gamma( self, t=None, iteration=None, species=None,
         return( mean_gamma, std_gamma )
 
     def get_sigma_gamma_slice(self, dz, t=None, iteration=None, species=None,
-                              select=None, plot=False, **kw):
+                              select=None, plot=False, ax=None,
+                              figsize=None, **kw):
         """
         Calculate the standard deviation of gamma for particles in z-slices of
         width dz
@@ -128,6 +124,12 @@ def get_sigma_gamma_slice(self, dz, t=None, iteration=None, species=None,
         plot : bool, optional
            Whether to plot the requested quantity
 
+        ax : matplotlib axis, optional
+            Axis to be used for the plot
+
+        figsize : tuple of two integers, optional
+            Size of the figure for the plot, same as defined in matplotlib
+
         **kw : dict, otional
            Additional options to be passed to matplotlib's `plot` method
 
@@ -158,14 +160,14 @@ def get_sigma_gamma_slice(self, dz, t=None, iteration=None, species=None,
             i += 1
         # Plot the result if needed
         if plot:
-            check_matplotlib()
+            ax = check_matplotlib_and_axis(ax, figsize)
             iteration = self.iterations[ self._current_i ]
             time_fs = 1.e15 * self.t[ self._current_i ]
-            plt.plot(z_pos, spreads, **kw)
-            plt.title("Slice energy spread at %.1f fs   (iteration %d)"
+            ax.plot(z_pos, spreads, **kw)
+            ax.set_title("Slice energy spread at %.1f fs   (iteration %d)"
                 % (time_fs, iteration), fontsize=self.plotter.fontsize)
-            plt.xlabel('$z \;(\mu m)$', fontsize=self.plotter.fontsize)
-            plt.ylabel('$\sigma_\gamma (\Delta_z=%s\mu m)$' % dz,
+            ax.set_xlabel('$z \;(\mu m)$', fontsize=self.plotter.fontsize)
+            ax.set_ylabel('$\sigma_\gamma (\Delta_z=%s\mu m)$' % dz,
                        fontsize=self.plotter.fontsize)
         return(spreads, z_pos)
 
@@ -365,7 +367,7 @@ def get_emittance(self, t=None, iteration=None, species=None,
             return emittance_from_coord(x, y, ux, uy, w)
 
     def get_current( self, t=None, iteration=None, species=None, select=None,
-                     bins=100, plot=False, **kw ):
+                     bins=100, plot=False, ax=None, figsize=None, **kw ):
         """
         Calculate the electric current along the z-axis for selected particles.
 
@@ -395,6 +397,12 @@ def get_current( self, t=None, iteration=None, species=None, select=None,
         plot : bool, optional
            Whether to plot the requested quantity
 
+        ax : matplotlib axis, optional
+            Axis to be used for the plot
+
+        figsize : tuple of two integers, optional
+            Size of the figure for the plot, same as defined in matplotlib
+
         **kw : dict, otional
            Additional options to be passed to matplotlib's `plot` method
         Returns
@@ -425,20 +433,21 @@ def get_current( self, t=None, iteration=None, species=None, select=None,
             global_offset=(np.min(z) + len_z / bins / 2,), position=(0,))
         # Plot the result if needed
         if plot:
-            check_matplotlib()
+            ax = check_matplotlib_and_axis(ax, figsize)
             iteration = self.iterations[ self._current_i ]
             time_fs = 1.e15 * self.t[ self._current_i ]
-            plt.plot( info.z, current, **kw)
-            plt.title("Current at %.1f fs   (iteration %d)"
+            ax.plot( info.z, current, **kw)
+            ax.set_title("Current at %.1f fs   (iteration %d)"
                 % (time_fs, iteration ), fontsize=self.plotter.fontsize)
-            plt.xlabel('$z \;(\mu m)$', fontsize=self.plotter.fontsize)
-            plt.ylabel('$I \;(A)$', fontsize=self.plotter.fontsize)
+            ax.set_xlabel('$z \;(\mu m)$', fontsize=self.plotter.fontsize)
+            ax.set_ylabel('$I \;(A)$', fontsize=self.plotter.fontsize)
         # Return the current and bin centers
         return(current, info)
 
     def get_laser_envelope( self, t=None, iteration=None, pol=None, m='all',
                             freq_filter=40, index='center', theta=0,
-                            slicing_dir='y', plot=False, **kw ):
+                            slicing_dir='y', plot=False, ax=None,
+                            figsize=None, **kw ):
         """
         Calculate a laser field by filtering out high frequencies. Can either
         return the envelope slice-wise or a full 2D envelope.
@@ -484,6 +493,12 @@ def get_laser_envelope( self, t=None, iteration=None, pol=None, m='all',
         plot : bool, optional
            Whether to plot the requested quantity
 
+        ax : matplotlib axis, optional
+            Axis to be used for the plot
+
+        figsize : tuple of two integers, optional
+            Size of the figure for the plot, same as defined in matplotlib
+
         **kw : dict, otional
            Additional options to be passed to matplotlib's `plot`(1D) or
            `imshow` (2D) method
@@ -520,22 +535,22 @@ def get_laser_envelope( self, t=None, iteration=None, pol=None, m='all',
 
         # Plot the result if needed
         if plot:
-            check_matplotlib()
+            ax = check_matplotlib_and_axis(ax, figsize)
             iteration = self.iterations[ self._current_i ]
             time_fs = 1.e15 * self.t[ self._current_i ]
             if index != 'all':
-                plt.plot( 1.e6 * info.z, envelope, **kw)
-                plt.ylabel('$E_%s \;(V/m)$' % pol,
-                           fontsize=self.plotter.fontsize)
+                ax.plot( 1.e6 * info.z, envelope, **kw)
+                ax.set_ylabel('$E_%s \;(V/m)$' % pol,
+                              fontsize=self.plotter.fontsize)
             else:
-                plt.imshow( envelope, extent=1.e6 * info.imshow_extent,
-                            aspect='auto', **kw)
-                plt.colorbar()
-                plt.ylabel('$%s \;(\mu m)$' % pol,
+                _ = ax.imshow( envelope, extent=1.e6 * info.imshow_extent,
+                               aspect='auto', **kw)
+                ax.figure.colorbar(mappable=_, ax=ax)
+                ax.set_ylabel('$%s \;(\mu m)$' % pol,
                             fontsize=self.plotter.fontsize)
-            plt.title("Laser envelope at %.1f fs   (iteration %d)"
+            ax.set_title("Laser envelope at %.1f fs   (iteration %d)"
                 % (time_fs, iteration ), fontsize=self.plotter.fontsize)
-            plt.xlabel('$z \;(\mu m)$', fontsize=self.plotter.fontsize)
+            ax.set_xlabel('$z \;(\mu m)$', fontsize=self.plotter.fontsize)
         # Return the result
         return( envelope, info )
 
@@ -655,7 +670,7 @@ def get_main_frequency( self, t=None, iteration=None, pol=None, m='all',
             raise ValueError('Unknown method: {:s}'.format(method))
 
     def get_spectrum( self, t=None, iteration=None, pol=None,
-                      m='all', plot=False, **kw ):
+                      m='all', plot=False, ax=None, figsize=None, **kw ):
         """
         Return the spectrum of the laser
         (Absolute value of the Fourier transform of the fields.)
@@ -682,6 +697,12 @@ def get_spectrum( self, t=None, iteration=None, pol=None,
         plot: bool, optional
            Whether to plot the data
 
+        ax : matplotlib axis, optional
+            Axis to be used for the plot
+
+        figsize : tuple of two integers, optional
+            Size of the figure for the plot, same as defined in matplotlib
+
         **kw : dict, otional
            Additional options to be passed to matplotlib's `plot` method
 
@@ -720,14 +741,14 @@ def get_spectrum( self, t=None, iteration=None, pol=None,
 
         # Plot the field if required
         if plot:
-            check_matplotlib()
+            ax = check_matplotlib_and_axis(ax, figsize)
             iteration = self.iterations[ self._current_i ]
             time_fs = 1.e15 * self.t[ self._current_i ]
-            plt.plot( spect_info.omega, spectrum, **kw )
-            plt.xlabel('$\omega \; (rad.s^{-1})$',
+            ax.plot( spect_info.omega, spectrum, **kw )
+            ax.set_xlabel('$\omega \; (rad.s^{-1})$',
                        fontsize=self.plotter.fontsize )
-            plt.ylabel('Spectrum', fontsize=self.plotter.fontsize )
-            plt.title("Spectrum at %.1f fs   (iteration %d)"
+            ax.set_ylabel('Spectrum', fontsize=self.plotter.fontsize )
+            ax.set_title("Spectrum at %.1f fs   (iteration %d)"
                 % (time_fs, iteration ), fontsize=self.plotter.fontsize)
         return( spectrum, spect_info )
 
@@ -903,7 +924,8 @@ def get_laser_waist( self, t=None, iteration=None, pol=None, theta=0,
             raise ValueError('Unknown method: {:s}'.format(method))
 
     def get_spectrogram( self, t=None, iteration=None, pol=None, theta=0,
-                          slicing_dir='y', plot=False, **kw ):
+                          slicing_dir='y', plot=False, ax=None,
+                          figsize=None, **kw ):
         """
         Calculates the spectrogram of a laserpulse, by the FROG method.
 
@@ -932,6 +954,12 @@ def get_spectrogram( self, t=None, iteration=None, pol=None, theta=0,
         plot: bool, optional
             Whether to plot the spectrogram
 
+        ax : matplotlib axis, optional
+            Axis to be used for the plot
+
+        figsize : tuple of two integers, optional
+            Size of the figure for the plot, same as defined in matplotlib
+
         **kw : dict, otional
            Additional options to be passed to matplotlib's `imshow` method
 
@@ -984,16 +1012,16 @@ def get_spectrogram( self, t=None, iteration=None, pol=None, theta=0,
 
         # Plot the result if needed
         if plot:
-            check_matplotlib()
+            ax = check_matplotlib_and_axis(ax, figsize)
             iteration = self.iterations[ self._current_i ]
             time_fs = 1.e15 * self.t[ self._current_i ]
-            plt.imshow( spectrogram, extent=info.imshow_extent, aspect='auto',
+            ax.imshow( spectrogram, extent=info.imshow_extent, aspect='auto',
                         **kw)
-            plt.title("Spectrogram at %.1f fs   (iteration %d)"
+            ax.set_title("Spectrogram at %.1f fs   (iteration %d)"
                 % (time_fs, iteration ), fontsize=self.plotter.fontsize)
-            plt.xlabel('$t \;(s)$', fontsize=self.plotter.fontsize )
-            plt.ylabel('$\omega \;(rad.s^{-1})$',
-                       fontsize=self.plotter.fontsize )
+            ax.set_xlabel('$t \;(s)$', fontsize=self.plotter.fontsize )
+            ax.set_ylabel('$\omega \;(rad.s^{-1})$',
+                          fontsize=self.plotter.fontsize )
         return( spectrogram, info )
 
 

opmd_viewer/openpmd_timeseries/interactive.py

@@ -82,8 +82,9 @@ def refresh_field(change=None, force=False):
                     do_refresh = True
             # Do the refresh
             if do_refresh:
-                plt.figure(fld_figure_button.value, figsize=figsize)
-                plt.clf()
+                fig_fld = plt.figure(fld_figure_button.value, figsize=figsize)
+                fig_fld.clf()
+                ax_fld = fig_fld.subplots(1, 1)
 
                 # When working in inline mode, in an ipython notebook,
                 # clear the output (prevents the images from stacking
@@ -118,7 +119,7 @@ def refresh_field(change=None, force=False):
                     m=convert_to_int(mode_button.value),
                     slicing=slicing_button.value, theta=theta_button.value,
                     slicing_dir=slicing_dir_button.value,
-                    plot_range=plot_range, **kw_fld )
+                    plot_range=plot_range, ax=ax_fld, **kw_fld )
 
         def refresh_ptcl(change=None, force=False):
             """
@@ -142,8 +143,10 @@ def refresh_ptcl(change=None, force=False):
                     do_refresh = True
             # Do the refresh
             if do_refresh:
-                plt.figure(ptcl_figure_button.value, figsize=figsize)
-                plt.clf()
+                fig_ptcl = plt.figure(ptcl_figure_button.value,
+                                      figsize=figsize)
+                fig_ptcl.clf()
+                ax_ptcl = fig_ptcl.subplots(1, 1)
 
                 # When working in inline mode, in an ipython notebook,
                 # clear the output (prevents the images from stacking
@@ -169,7 +172,8 @@ def refresh_ptcl(change=None, force=False):
                         species=ptcl_species_button.value, plot=True,
                         nbins=ptcl_bins_button.value,
                         plot_range=plot_range,
-                        use_field_mesh=ptcl_use_field_button.value, **kw_ptcl )
+                        use_field_mesh=ptcl_use_field_button.value,
+                        ax=ax_ptcl, **kw_ptcl )
                 else:
                     # 2D histogram
                     self.get_particle( iteration=self.current_iteration,
@@ -179,7 +183,8 @@ def refresh_ptcl(change=None, force=False):
                         species=ptcl_species_button.value, plot=True,
                         nbins=ptcl_bins_button.value,
                         plot_range=plot_range,
-                        use_field_mesh=ptcl_use_field_button.value, **kw_ptcl )
+                        use_field_mesh=ptcl_use_field_button.value,
+                        ax=ax_ptcl, **kw_ptcl )
 
         def refresh_field_type(change):
             """

opmd_viewer/openpmd_timeseries/main.py

@@ -113,7 +113,8 @@ def __init__(self, path_to_dir, check_all_files=True):
     def get_particle(self, var_list=None, species=None, t=None, iteration=None,
             select=None, output=True, plot=False, nbins=150,
             plot_range=[[None, None], [None, None]],
-            use_field_mesh=True, histogram_deposition='cic', **kw):
+            use_field_mesh=True, histogram_deposition='cic',
+            ax=None, figsize=None, **kw):
         """
         Extract a list of particle variables
         from an HDF5 file in the openPMD format.
@@ -193,6 +194,12 @@ def get_particle(self, var_list=None, species=None, t=None, iteration=None,
             particles affects neighboring bins.
             `cic` (which is the default) leads to smoother results than `ngp`.
 
+        ax : matplotlib axis, optional
+            Axis to be used for the plot
+
+        figsize : tuple of two integers, optional
+            Size of the figure for the plot, same as defined in matplotlib
+
         **kw : dict, otional
            Additional options to be passed to matplotlib's
            hist or hist2d.
@@ -331,14 +338,16 @@ def get_particle(self, var_list=None, species=None, t=None, iteration=None,
                 # Do the plotting
                 self.plotter.hist1d(data_list[0], w, var_list[0], species,
                         self._current_i, hist_bins[0], hist_range,
-                        deposition=histogram_deposition, **kw)
+                        deposition=histogram_deposition, ax=ax,
+                        figsize=figsize, **kw)
             # - In the case of two quantities
             elif len(data_list) == 2:
                 # Do the plotting
                 self.plotter.hist2d(data_list[0], data_list[1], w,
                     var_list[0], var_list[1], species,
                     self._current_i, hist_bins, hist_range,
-                    deposition=histogram_deposition, **kw)
+                    deposition=histogram_deposition, ax=ax,
+                    figsize=figsize, **kw)
         # Close the file
         file_handle.close()
 
@@ -349,7 +358,8 @@ def get_particle(self, var_list=None, species=None, t=None, iteration=None,
     def get_field(self, field=None, coord=None, t=None, iteration=None,
                   m='all', theta=0., slicing=0., slicing_dir='y',
                   output=True, plot=False,
-                  plot_range=[[None, None], [None, None]], **kw):
+                  plot_range=[[None, None], [None, None]], ax=None,
+                  figsize=None, **kw):
         """
         Extract a given field from an HDF5 file in the openPMD format.
 
@@ -408,6 +418,12 @@ def get_field(self, field=None, coord=None, t=None, iteration=None,
            along the 1st axis (first list) and 2nd axis (second list)
            Default: plots the full extent of the simulation box
 
+        ax : matplotlib axis, optional
+            Axis to be used for the plot
+
+        figsize : tuple of two integers, optional
+            Size of the figure for the plot, same as defined in matplotlib
+
         **kw : dict, otional
            Additional options to be passed to matplotlib's imshow.
 
@@ -500,11 +516,13 @@ def get_field(self, field=None, coord=None, t=None, iteration=None,
         if plot:
             if geometry == "1dcartesian":
                 self.plotter.show_field_1d(F, info, field_label,
-                self._current_i, plot_range=plot_range, **kw)
+                self._current_i, plot_range=plot_range, ax=ax,
+                figsize=figsize, **kw)
             else:
                 self.plotter.show_field_2d(F, info, slicing_dir, m,
                         field_label, geometry, self._current_i,
-                        plot_range=plot_range, **kw)
+                        plot_range=plot_range, ax=ax,
+                        figsize=figsize, **kw)
 
         # Return the result
         return(F, info)