Add ylim to slice_plot

ChristopherMayes · Jul 31, 2023 · 48426ba · 48426ba
1 parent 8a209eb
commit 48426ba
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 34 deletions.
diff --git a/docs/examples/particle_examples.ipynb b/docs/examples/particle_examples.ipynb
diff --git a/docs/examples/plot_examples.ipynb b/docs/examples/plot_examples.ipynb
diff --git a/pmd_beamphysics/particles.py b/pmd_beamphysics/particles.py
@@ -840,6 +840,7 @@ def slice_plot(self, key='sigma_x',
                    slice_key=None,
                    tex=True,
                    return_figure=False,
+                   ylim=None,
                    **kwargs):
         """
         Slice statistics plot. 
@@ -856,6 +857,7 @@ def slice_plot(self, key='sigma_x',
                          n_slice=n_slice,
                          slice_key=slice_key,
                          tex=tex,
+                         ylim=ylim,
                          **kwargs)
 
         if return_figure:

diff --git a/pmd_beamphysics/plot.py b/pmd_beamphysics/plot.py
@@ -42,6 +42,7 @@ def slice_plot(particle_group,
                n_slice=40,
                slice_key='z',
                tex=True,
+               ylim=None,
                **kwargs):
     """
     Complete slice plotting routine. Will plot the density of the slice key on the right axis. 
@@ -59,7 +60,7 @@ def slice_plot(particle_group,
 
     # Get nice arrays
     x, _, prex = nice_array(slice_dat[x_key])
-    y, _, prey = nice_array(slice_dat[y_key])
+    y, yfactor, prey = nice_array(slice_dat[y_key])
     y2, _, prey2 = nice_array(slice_dat[y2_key])
 
     x_units = f'{prex}{particle_group.units(x_key)}'
@@ -72,8 +73,6 @@ def slice_plot(particle_group,
     y2_units = prey2+y2_units 
 
     # Labels
-
-
     labelx = mathlabel(slice_key, units=x_units, tex=tex)
     labely = mathlabel(y_key, units=y_units, tex=tex)    
     labely2 = mathlabel(y2_key, units=y2_units, tex=tex)        
@@ -83,12 +82,24 @@ def slice_plot(particle_group,
 
     # Main plot
     ax.plot(x, y, color = 'black')
-
-    #ax.set_ylim(0, 1.1*ymax )
 
+    # rhs plot
     ax2 = ax.twinx()
     ax2.set_ylabel(labely2)
     ax2.fill_between(x, 0, y2, color='black', alpha = 0.2)  
+    ax2.set_ylim(0, None)
+
+    # Limits
+    if ylim:
+        ymin = ylim[0]
+        ymax = ylim[1]
+        # Handle None and scaling
+        if ymin is not None:
+            ymin = ymin/yfactor
+        if ymax is not None:
+            ymax = ymax/yfactor
+        new_ylim = (ymin, ymax)
+        ax.set_ylim(new_ylim)     
 
     return fig