From bb777c1a8e3cd561703c1b9bd610d17aeaef4cf5 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, 24 insertions(+), 2 deletions(-)
diff --git a/tracking/particles.py b/tracking/particles.py
index ee321bb..5de2d6c 100644
--- a/tracking/particles.py
+++ b/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
+
--
GitLab