Add a simple plot method to Bunch class

Phase space plotting method was added

tracking/particles.py

@@ -4,13 +4,15 @@ Module where particles, bunches and beams are described as objects.
 
 @author: Alexis Gamelin
 @date: 11/03/2020
+
 """
 
 import numpy as np
 import pandas as pd
+import matplotlib.pyplot as plt
 from tracking.parallel import Mpi
+import seaborn as sns
 from scipy.constants import c, m_e, m_p, e
-import matplotlib.pyplot as plt
 
 class Particle:
     """
@@ -319,6 +321,26 @@ class Bunch:
         ax = fig.gca()
         ax.plot(bins.mid, profile)
         
+    def getplot(self,x,y,p_type):
+        """
+        Plot the parameters from particles object.
+        
+        Parameters
+        ----------
+        x: str, name from Bunch object to plot on hor. axis.
+        y: str, name from Bunch object to plot on ver. axis.
+        p_type: str, "sc" for a scatter plot or "j" for a joint plot.
+        """
+   
+        x = self.particles[x]
+        y = self.particles[y]
+        
+        if p_type == "sc":
+            plt.scatter(x,y)
+        
+        elif p_type == "j":
+            sns.jointplot(x,y,kind="kde")
+        
 class Beam:
     """
     Define a Beam object composed of several Bunch objects. 
@@ -578,4 +600,4 @@ class Beam:
         
         
         
-    
+    
\ No newline at end of file