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