diff --git a/tracking/particles.py b/tracking/particles.py
index 9c2fdf82cef43332b1e8e4cc47c1d627c771b6d2..06aa16159584e270d4b7e8bab0e03319b79c7f26 100644
--- a/tracking/particles.py
+++ b/tracking/particles.py
@@ -336,17 +336,21 @@ class Bunch:
Can be modified to "sc" for a scatter plot.
"""
+ label_dict = {"x":"x (mm)", "xp":"x' (mrad)", "y":"y (mm)", "yp":"y' (mrad)",
+ "tau":"$\\tau$ (ps)", "delta":"$\\delta$"}
+ scale = {"x": 1e3, "xp":1e3, "y":1e3, "yp":1e3, "tau":1e12, "delta":1}
+
if plot_type == "sc":
- plt.scatter(self.particles["tau"]*1e12,
- self.particles["delta"])
- plt.xlabel("$\\tau$ (ps)")
- plt.ylabel("$\\delta$")
+ plt.scatter(self.particles[x_var]*scale[x_var],
+ self.particles[y_var]*scale[y_var])
+ plt.xlabel(label_dict[x_var])
+ plt.ylabel(label_dict[y_var])
else:
- sns.jointplot(self.particles["tau"]*1e12,
- self.particles["delta"],kind="kde")
- plt.xlabel("$\\tau$ (ps)")
- plt.ylabel("$\\delta$")
+ sns.jointplot(self.particles[x_var]*scale[x_var],
+ self.particles[y_var]*scale[y_var],kind="kde")
+ plt.xlabel(label_dict[x_var])
+ plt.ylabel(label_dict[y_var])
class Beam:
"""
@@ -604,163 +608,84 @@ class Beam:
self.mpi_switch = False
self.mpi = None
- def long_phasespace(self,bunch_number,x_var="tau",y_var="delta",
- plot_type="j"):
- """
- Plot longitudinal phase space.
-
- Parameters
- ----------
- bunch_number : int
- Specify a bunch among those in beam object to be displayed.
- The value must not exceed the total length of filling_pattern object.
- x_var : str
- name from Bunch object to plot on horizontal axis.
- y_var : str
- name from Bunch object to plot on ver. axis.
- plot_type : str {"j" , "sc"}
- type of the plot. The defualt value is "j" for a joint plot.
- Can be modified to "sc" for a scatter plot.
- """
-
- if plot_type == "sc":
- plt.scatter(self[bunch_number-1]["tau"]*1e12,
- self[bunch_number-1]["delta"])
- plt.xlabel("$\\tau$ (ps)")
- plt.ylabel("$\\delta$")
-
- else:
- sns.jointplot(self[bunch_number-1]["tau"]*1e12,
- self[bunch_number-1]["delta"],kind="kde")
- plt.xlabel("$\\tau$ (ps)")
- plt.ylabel("$\\delta$")
-
- def plot_bunchdata(self, filename, detaset, x_var="time"):
+ def plot_bunchnumber(self, var, option=None):
"""
- Plot the evolution of the variables from the Beam object.
-
+ Plot varviables with respect to bunch number.
+
Parameters
----------
- filename : str
- Name of the HDF5 file that contains the data.
- detaset : str {"current","emit","mean","std"}
- HDF5 file's dataset to be plotted.
- x_var : str, optional
- The variable to be plotted on horizontal axis. The default is "time".
+ var : str {"bunch_currebt", "bunch_charge", "bunch_particle", "bunch_mean", "bunch_std", "bunch_emit"}
+ Variable to be plotted.
+ option : str, optional
+ If var is "bunch_mean", "bunch_std", or "bunch_emit, option needs tobe specified.
+ For "bunch_mean" and "bunch_std", option = {"x","xp","y","yp","tau","delta"}.
+ For "bunch_emit", option = {"x","y","s"}.
+ The default is None.
"""
- file = hp.File(filename, "r")
-
- group = "BunchData_0" # Data group of the HDF5 file
-
- if detaset == "current":
- y_var = file[group][detaset][:]*1e3
- label = "current (mA)"
+ var_dict = {"bunch_current":self.bunch_current,
+ "bunch_charge":self.bunch_charge,
+ "bunch_particle":self.bunch_particle,
+ "bunch_mean":self.bunch_mean,
+ "bunch_std":self.bunch_std,
+ "bunch_emit":self.bunch_emit}
+
+ if var == "bunch_mean" or var == "bunch_std":
+ value_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, "delta":5}
+ scale = [1e6, 1e6, 1e6, 1e6, 1e12, 1e6]
+ label_mean = ["x (um)", "x' ($\\mu$rad)", "y (um)", "y' ($\\mu$rad)",
+ "$\\tau$ (ps)", "$\\delta(\\times 10^{-6})$"]
+ label_std = ["std x (um)", "std x' ($\\mu$rad)", "std y (um)",
+ "std y' ($\\mu$rad)", "std $\\tau$ (ps)",
+ "std $\\delta(\\times 10^{-6})$"]
+
+ y_axis = var_dict[var][value_dict[option]]
- elif detaset == "emit":
- axis = int(input("""Specify the axis by entering the corresponding number :
- horizontal axis (x) -> 0
- vertical axis (y) -> 1
- longitudinal axis (s) -> 2
- : """))
-
- y_var = file[group][detaset][axis]*1e9
+ # Convert NaN in y_axis array into zero
+ where_is_nan = np.isnan(y_axis)
+ y_axis[where_is_nan] = 0
- if axis == 0: label = "hor. emittance (nm.rad)"
- elif axis == 1: label = "ver. emittance (nm.rad)"
- elif axis == 2: label = "long. emittance (nm.rad)"
+ plt.plot(np.arange(len(self.filling_pattern)),
+ y_axis*scale[value_dict[option]])
+ plt.xlabel('bunch number')
+ if var == "bunch_mean":
+ plt.ylabel(label_mean[value_dict[option]])
+ else:
+ plt.ylabel(label_std[value_dict[option]])
+ elif var == "bunch_emit":
+ value_dict = {"x":0, "y":1, "s":2}
+ scale = [1e9, 1e9, 1e15]
- elif detaset == "mean" or "std":
- param = int(input("""Specify the variable by entering the corresponding number: \
- x -> 0 \
- xp -> 1 \
- y -> 2 \
- yp -> 3 \
- tau -> 4 \
- delta -> 5 \
- :"""))
- if param == 0:
- y_var = file[group][detaset][param]*1e6
- label = "x (um)"
- elif param == 1:
- y_var = file[group][detaset][param]*1e6
- label = "x' ($\\mu$rad)"
- elif param == 2:
- y_var = file[group][detaset][param]*1e6
- label = "y (um)"
- elif param == 3:
- y_var = file[group][detaset][param]*1e6
- label = "y' ($\\mu$rad)"
- elif param == 4:
- y_var = file[group][detaset][param]*1e12
- label = "$\\tau$ (ps)"
- else :
- y_var = file[group][detaset][param]
- label = "$\\delta$"
-
- plt.plot(file[group]["time"][:],y_var)
- plt.xlabel("number of turns")
- plt.ylabel(label)
-
-
- def plot_phasespacedata(self, filename, total_size, save_every, dataset):
- """
- Plot data from PhaseSpaceData_0 group of the HDF5 file.
-
- Parameters
- ----------
- filename : str
- Name of the HDF5 file that contains the data.
- total_size : int
- Total size of the save regarding to 'PhaseSpaceMonitor' object.
- save_every : int
- The frequency of the save regarding to 'PhaseSpaceMonitor' object.
- dataset : str {'alive', 'partcicles'}
- HDF5 file's dataset to be plotted.
-
- """
-
- file = hp.File(filename, "r")
-
- data_points = int(total_size/save_every)
- timelist = file["PhaseSpaceData_0"]["time"][0:data_points]
-
- if dataset == "alive":
- alive_at_a_time = []
- for i in range (data_points):
- alive_at_a_time.append(np.sum(file["PhaseSpaceData_0"][dataset][:,i]))
-
- plt.plot(timelist,alive_at_a_time)
- plt.xlabel("number of turns")
- plt.ylabel("number of alive particles")
+ y_axis = var_dict[var][value_dict[option]]
- elif dataset == "particles":
- print("""Specify the parameter to be plotted by entering the corresponding number:
- x -> 0
- xp -> 1
- y -> 2
- yp -> 3
- tau -> 4
- delta -> 5 """)
-
- x_var = int(input("Horizontal axis: "))
- y_var = int(input("Vertical axis: "))
+ # Convert NaN in y_axis array into zero
+ where_is_nan = np.isnan(y_axis)
+ y_axis[where_is_nan] = 0
- print("Specify the time at turn =", timelist)
- turn = int(input(": "))
- index_in_timelist = np.where(timelist==turn)
-
- sns.jointplot(file["PhaseSpaceData_0"][dataset][:,x_var,index_in_timelist],
- file["PhaseSpaceData_0"][dataset][:,y_var,index_in_timelist],
- kind="kde")
+ plt.plot(np.arange(len(self.filling_pattern)),
+ y_axis*scale[value_dict[option]])
- name_list = ["x (m)","xp (rad)","y (m)","yp (rad)","$\\tau$ (s)","$\\delta$"]
- plt.xlabel(name_list[x_var])
- plt.ylabel(name_list[y_var])
-
-
+ if option == "x": label_y = "hor. emittance (nm.rad)"
+ elif option == "y": label_y = "ver. emittance (nm.rad)"
+ elif option == "s": label_y = "long. emittance (fm.rad)"
+
+ plt.xlabel('bunch number')
+ plt.ylabel(label_y)
+
+ elif var == "bunch_current" or var == "bunch_charge" or var =="bunch_particle":
+ plt.plot(np.arange(len(self.filling_pattern)), var_dict[var])
+ plt.xlabel('bunch number')
+
+ if var == "bunch_current": label_y = "bunch current (A)"
+ elif var == "bunch_charge": label_y = "bunch chagre (C)"
+ else: label_y = "number of particles"
+
+ plt.ylabel(label_y)
+
+ elif var == "current" or var=="charge" or var=="particle_number":
+ print("'{0}'is a total value and cannot be plotted".format(var))
diff --git a/tracking/plotting.py b/tracking/plotting.py
new file mode 100644
index 0000000000000000000000000000000000000000..a194acceff243a1e39c1001e2ec78f74eb6656bc
--- /dev/null
+++ b/tracking/plotting.py
@@ -0,0 +1,179 @@
+# -*- coding: utf-8 -*-
+"""
+Module for plotting Bunch and Beam objects.
+
+@author: Watanyu Foosang
+@Date: 10/04/2020
+"""
+
+import numpy as np
+import matplotlib.pyplot as plt
+import seaborn as sns
+import h5py as hp
+
+
+def plot_bunchdata(filename, bunch_number, dataset, option=None, x_var="time"):
+ """
+ Plot the evolution of the variables from the Beam object.
+
+ Parameters
+ ----------
+ filename : str
+ Name of the HDF5 file that contains the data.
+ bunch_number : int
+ The bunch number whose data has been saved in the HDF5 file.
+ This has to be identical to 'bunch_number' parameter in 'BunchMonitor' object.
+ detaset : str {"current","emit","mean","std"}
+ HDF5 file's dataset to be plotted.
+ option : str, optional
+ If dataset is "emit", "mean", or "std", the variable name to be plotted
+ needs to be specified :
+ for "emit", option = {"x","y","s"}
+ for "mean" and "std", option = {"x","xp","y","yp","tau","delta"}
+ x_var : str, optional
+ Variable to be plotted on the horizontal axis. The default is "time".
+
+ """
+
+ file = hp.File(filename, "r")
+
+ group = "BunchData_{0}".format(bunch_number) # Data group of the HDF5 file
+
+ if dataset == "current":
+ y_var = file[group][dataset][:]*1e3
+ label = "current (mA)"
+
+ elif dataset == "emit":
+ # Specifying the axis
+ # horizontal axis (x) -> 0
+ # vertical axis (y) -> 1
+ # longitudinal axis (s) -> 2
+
+ emit_axis = {"x":0, "y":1, "s":2}
+
+ y_var = file[group][dataset][emit_axis[option]]*1e9
+
+ if option == "x": label = "hor. emittance (nm.rad)"
+ elif option == "y": label = "ver. emittance (nm.rad)"
+ elif option == "s": label = "long. emittance (nm.rad)"
+
+
+ elif dataset == "mean" or "std":
+ # Specify the variable
+ # x -> 0
+ # xp -> 1
+ # y -> 2
+ # yp -> 3
+ # tau -> 4
+ # delta -> 5
+
+ var_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, "delta":5}
+ scale = [1e6, 1e6, 1e6, 1e6, 1e12, 1e6]
+ label_list = ["x (um)", "x' ($\\mu$rad)", "y (um)", "y' ($\\mu$rad)",
+ "$\\tau$ (ps)", "$\\delta(\\times 10^{-6})$"]
+
+ y_var = file[group][dataset][var_dict[option]]*scale[var_dict[option]]
+ label = label_list[var_dict[option]]
+
+
+ plt.plot(file[group]["time"][:],y_var)
+ plt.xlabel("number of turns")
+ plt.ylabel(label)
+
+ file.close()
+
+def plot_phasespacedata(filename, bunch_number, total_size, save_every, dataset,
+ x_var=None, y_var=None, turn=None, only_alive=False):
+ """
+ Plot data from PhaseSpaceData_0 group of the HDF5 file.
+
+ Parameters
+ ----------
+ filename : str
+ Name of the HDF5 file that contains the data.
+ bunch_number : int
+ Number of the bunch whose data has been saved in the HDF5 file.
+ This has to be identical to 'bunch_number' parameter in 'PhaseSpaceMonitor' object.
+ total_size : int
+ Total size of the save regarding to 'PhaseSpaceMonitor' object.
+ save_every : int
+ Frequency of the save regarding to 'PhaseSpaceMonitor' object.
+ dataset : str {'alive', 'partcicles'}
+ HDF5 file's dataset to be plotted.
+ x_var, y_var : str {"x", "xp", "y", "yp", "tau", "delta"}, optional
+ If dataset is "particles", the variables to be plotted on the horizontal
+ and the vertical axes need to be specified.
+ turn : int, optional
+ Turn at which the data will be extracted.
+ only_alive : bool
+ When only_alive is True, only alive particles are plotted and dead particles will be discarded.
+
+ """
+
+ file = hp.File(filename, "r")
+
+ group = "PhaseSpaceData_{0}".format(bunch_number)
+
+ data_points = int(total_size/save_every)
+ timelist = file[group]["time"][0:data_points] # example output: [0,10,20,30,40]
+
+ if dataset == "alive":
+ alive_at_a_time = []
+ for i in range (data_points):
+ alive_at_a_time.append(np.sum(file[group][dataset][:,i]))
+
+ plt.plot(timelist,alive_at_a_time)
+ plt.xlabel("number of turns")
+ plt.ylabel("number of alive particles")
+
+ elif dataset == "particles":
+ # Specify the parameter
+ # x -> 0
+ # xp -> 1
+ # y -> 2
+ # yp -> 3
+ # tau -> 4
+ # delta -> 5
+
+ var_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, "delta":5}
+ scale = [1e3,1e3,1e3,1e3,1e12,1e3]
+ label = ["x (mm)","x' (mrad)","y (mm)","y' (mrad)","$\\tau$ (ps)","$\\delta \\times 10^{-3}$"]
+
+ # find the index of "turn" in timelist
+ index_in_timelist = np.where(timelist==turn)
+
+ path = file[group][dataset]
+
+ if only_alive == False:
+ # format : sns.jointplot(x_axis, yaxis, kind)
+ x_axis = path[:,var_dict[x_var],index_in_timelist]
+ y_axis = path[:,var_dict[y_var],index_in_timelist]
+
+
+ elif only_alive == True:
+ x_alive = []
+ y_alive = []
+ for i in path[:,var_dict[x_var],index_in_timelist]:
+ if bool(i) == True:
+ x_alive.append(i)
+ else: pass
+ for i in path[:,var_dict[y_var],index_in_timelist]:
+ if bool(i) == True:
+ y_alive.append(i)
+ else: pass
+ x_axis = np.array(x_alive)
+ y_axis = np.array(y_alive)
+
+ sns.jointplot(x_axis*scale[var_dict[x_var]], y_axis*scale[var_dict[y_var]],
+ kind="kde")
+
+ plt.xlabel(label[var_dict[x_var]])
+ plt.ylabel(label[var_dict[y_var]])
+
+
+ file.close()
+
+
+
+
+