diff --git a/tracking/particles.py b/tracking/particles.py index 0839a2f2e55d09b54490317e0b56128c1d174479..fbd684e718a9d9ddfd4d3f4c3dc5b6e823226b87 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.