From 96f093c7a401b28e89200ce0417481742b838142 Mon Sep 17 00:00:00 2001
From: Watanyu Foosang <watanyu.f@gmail.com>
Date: Wed, 1 Apr 2020 12:44:59 +0200
Subject: [PATCH] Add a simple plot method to Bunch class

Phase space plotting method was added
---
 tracking/particles.py | 26 +++++++++++++++++++++++++-
 1 file changed, 25 insertions(+), 1 deletion(-)

diff --git a/tracking/particles.py b/tracking/particles.py
index 0839a2f..fbd684e 100644
--- a/tracking/particles.py
+++ b/tracking/particles.py
@@ -4,12 +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
-from tracking.parallel import Mpi
+import matplotlib.pyplot as plt
+from stage.mbtrack2.tracking.parallel import Mpi
 from scipy.constants import c, m_e, m_p, e
+import seaborn as sns
 
 class Particle:
     """ Define a particle object
@@ -270,6 +273,27 @@ class Bunch:
         self.particles["tau"] = values[:,4]
         self.particles["delta"] = values[:,5]
         
+    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. 
-- 
GitLab