docstrings

pmd_beamphysics/particles.py

@@ -820,17 +820,75 @@ def write(self, h5, name=None):
 
     # Plotting
     # --------
-    # TODO: more general plotting
-    def plot(self, key1='x', key2=None, bins=None, return_figure=False, 
-             tex=True, **kwargs):
+    def plot(self, key1='x', key2=None,
+             bins=None,
+             *,
+             xlim=None,
+             ylim=None,
+             return_figure=False, 
+
+             tex=True,  **kwargs):
         """
         1d or 2d density plot. 
+        
+        If one key is given, this will plot the density of that key.
+        Example:
+            .plot('x')
+        
+        If two keys arg given, this will plot a 2d marginal plot.
+        Example:
+            .plot('x', 'px')
+            
+        
+        Parameters
+        ----------
+        particle_group: ParticleGroup
+            The object to plot
+        
+        key1: str, default = 't'
+            Key to bin on the x-axis
+            
+        key2: str, default = None
+            Key to bin on the y-axis. 
+            
+        bins: int, default = None
+           Number of bins. If None, this will use a heuristic: bins = sqrt(n_particle/4)
+    
+        xlim: tuple, default = None
+            Manual setting of the x-axis limits. 
+            
+        ylim: tuple, default = None
+            Manual setting of the y-axis limits. 
+            
+        tex: bool, defaul = True
+            Use TEX for labels   
+            
+        return_figure: bool, default = False
+            If true, return a matplotlib.figure.Figure object
+            
+        **kwargs
+            Any additional kwargs to send to the the plot in: plt.subplots(**kwargs)
+            
+        
+        Returns
+        -------
+        None or fig: matplotlib.figure.Figure
+            This only returns a figure object if return_figure=T, otherwise returns None
+            
         """
 
         if not key2:
-            fig = density_plot(self, key=key1, bins=bins, tex=tex, **kwargs)
+            fig = density_plot(self, key=key1,
+                               bins=bins,
+                               xlim=xlim,
+                               tex=tex, **kwargs)
         else:
-            fig = marginal_plot(self, key1=key1, key2=key2, bins=bins, tex=tex, **kwargs)
+            fig = marginal_plot(self, key1=key1, key2=key2,
+                                bins=bins,
+                                xlim=xlim,
+                                ylim=ylim,
+                                tex=tex,
+                                **kwargs)
 
         if return_figure:
             return fig

pmd_beamphysics/plot.py

@@ -41,11 +41,37 @@ def slice_plot(particle_group,
                stat_key='sigma_x',
                n_slice=40,
                slice_key='z',
-               tex=True,
                ylim=None,
+               tex=True,
                **kwargs):
     """
     Complete slice plotting routine. Will plot the density of the slice key on the right axis. 
+    
+    Parameters
+    ----------
+    particle_group: ParticleGroup
+        The object to plot
+    
+    stat_key: str, default = 'sigma_x'
+        Key to calculate the statistics
+        
+    n_slice: int, default = 40
+        Number of slices 
+        
+    slice_key: str, default = 'z'
+        Should be 'z' or 't'
+        
+    ylim: tuple, default = None
+        Manual setting of the y-axis limits. 
+        
+    tex: bool, defaul = True
+        Use TEX for labels
+        
+    
+    Returns
+    -------
+    fig: matplotlib.figure.Figure
+
     """
 
     x_key = 'mean_'+slice_key
@@ -105,7 +131,11 @@ def slice_plot(particle_group,
 
 
 
-def density_plot(particle_group, key='x', bins=None, tex=True, **kwargs):
+def density_plot(particle_group, key='x',
+                 bins=None,
+                 *, 
+                 xlim=None,
+                 tex=True, **kwargs):
     """
     1D density plot. Also see: marginal_plot
     
@@ -145,15 +175,63 @@ def density_plot(particle_group, key='x', bins=None, tex=True, **kwargs):
 
     ax.set_xlabel(labelx)  
 
+    # Limits
+    if xlim:
+        xmin = xlim[0]
+        xmax = xlim[1]
+        # Handle None and scaling
+        if xmin is not None:
+            xmin = xmin/f1
+        if xmax is not None:
+            xmax = xmax/f1
+        new_xlim = (xmin, xmax)
+        ax.set_xlim(new_xlim)          
+
     return fig
 
-def marginal_plot(particle_group, key1='t', key2='p', bins=None, tex=True, **kwargs):
+def marginal_plot(particle_group, key1='t', key2='p', 
+                  bins=None,
+                  *,
+                  xlim=None,
+                  ylim=None,
+                  tex=True,
+                  **kwargs):
     """
     Density plot and projections
     
     Example:
     
         marginal_plot(P, 't', 'energy', bins=200)   
+        
+        
+    Parameters
+    ----------
+    particle_group: ParticleGroup
+        The object to plot
+    
+    key1: str, default = 't'
+        Key to bin on the x-axis
+        
+    key2: str, default = 'p'
+        Key to bin on the y-axis        
+        
+    bins: int, default = None
+       Number of bins. If None, this will use a heuristic: bins = sqrt(n_particle/4)
+
+    xlim: tuple, default = None
+        Manual setting of the x-axis limits. 
+        
+    ylim: tuple, default = None
+        Manual setting of the y-axis limits. 
+        
+    tex: bool, defaul = True
+        Use TEX for labels
+        
+    
+    Returns
+    -------
+    fig: matplotlib.figure.Figure        
+        
     
     """
 
@@ -231,6 +309,31 @@ def marginal_plot(particle_group, key1='t', key2='p', bins=None, tex=True, **kwa
     # Set labels on joint
     ax_joint.set_xlabel(labelx)
     ax_joint.set_ylabel(labely)
+
+    # Limits
+    if xlim:
+        xmin = xlim[0]
+        xmax = xlim[1]
+        # Handle None and scaling
+        if xmin is not None:
+            xmin = xmin/f1
+        if xmax is not None:
+            xmax = xmax/f1
+        new_xlim = (xmin, xmax)
+        ax_joint.set_xlim(new_xlim)      
+        ax_marg_x.set_xlim(new_xlim)     
+
+    if ylim:
+        ymin = ylim[0]
+        ymax = ylim[1]
+        # Handle None and scaling
+        if ymin is not None:
+            ymin = ymin/f2
+        if ymax is not None:
+            ymax = ymax/f2
+        new_ylim = (ymin, ymax)
+        ax_joint.set_ylim(new_ylim)      
+        ax_marg_y.set_ylim(new_ylim) 
 
     return fig