Skip to content
Snippets Groups Projects
Commit 96f093c7 authored by Watanyu Foosang's avatar Watanyu Foosang
Browse files

Add a simple plot method to Bunch class

Phase space plotting method was added
parent 869cd286
No related branches found
No related tags found
No related merge requests found
......@@ -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.
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment