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

Options added to plot_bunchdata and plot_phasespacedata

- Option to plot bunch current on x-axis has been included in plot_bunchdata.
- Option to select the plot size for plot_phasespacedata has been added.
parent 1c750a8d
No related branches found
No related tags found
No related merge requests found
......@@ -12,6 +12,7 @@ import matplotlib.pyplot as plt
import matplotlib as mpl
import seaborn as sns
import h5py as hp
import random
def plot_beamdata(filename, dataset, option=None, stat_var=None, x_var="time"):
"""
......@@ -122,14 +123,14 @@ def plot_bunchdata(filename, bunch_number, dataset, option=None, x_var="time"):
bunch_number : int
Bunch to plot. This has to be identical to 'bunch_number' parameter in
'BunchMonitor' object.
detaset : str {"current","emit","mean","std"}
detaset : {"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
x_var : {"time", "current"}, optional
Variable to be plotted on the horizontal axis. The default is "time".
Return
......@@ -160,23 +161,36 @@ def plot_bunchdata(filename, bunch_number, dataset, option=None, x_var="time"):
elif dataset == "mean" or "std":
option_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, "delta":5}
scale = [1e6, 1e6, 1e6, 1e6, 1e12, 1]
label_list = ["x (um)", "x' ($\\mu$rad)", "y (um)", "y' ($\\mu$rad)",
"$\\tau$ (ps)", "$\\delta$"]
y_var = file[group][dataset][option_dict[option]]\
*scale[option_dict[option]]
label = label_list[option_dict[option]]
axis_index = option_dict[option]
y_var = file[group][dataset][axis_index]*scale[axis_index]
if dataset == "mean":
label_list = ["x ($\\mu$m)", "x' ($\\mu$rad)", "y ($\\mu$m)",
"y' ($\\mu$rad)", "$\\tau$ (ps)", "$\\delta$"]
else:
label_list = ["$\\sigma_x$ ($\\mu$m)", "$\\sigma_{x'}$ ($\\mu$rad)",
"$\\sigma_y$ ($\\mu$m)", "$\\sigma_{y'}$ ($\\mu$rad)",
"$\\sigma_{\\tau}$ (ps)", "$\\sigma_{\\delta}$"]
label = label_list[axis_index]
if x_var == "current":
x_axis = file[group]["current"][:] * 1e3
xlabel = "current (mA)"
elif x_var == "time":
x_axis = file[group]["time"][:]
xlabel = "number of turns"
fig, ax = plt.subplots()
ax.plot(file[group]["time"][:],y_var)
ax.set_xlabel("number of turns")
ax.plot(x_axis,y_var)
ax.set_xlabel(xlabel)
ax.set_ylabel(label)
file.close()
return fig
def plot_phasespacedata(filename, bunch_number, x_var, y_var, turn,
only_alive=True):
only_alive=True, plot_size=1):
"""
Plot data recorded by PhaseSpaceMonitor.
......@@ -195,6 +209,10 @@ def plot_phasespacedata(filename, bunch_number, x_var, y_var, turn,
only_alive : bool, optional
When only_alive is True, only alive particles are plotted and dead
particles will be discarded.
plot_size : [0,1], optional
Number of macro-particles to plot relative to the total number
of macro-particles recorded. This option helps reduce processing time
when the data is big.
Return
------
......@@ -220,17 +238,24 @@ def plot_phasespacedata(filename, bunch_number, x_var, y_var, turn,
format(turn, file[group]["time"][:]))
path = file[group][dataset]
mp_number = path[:,0,0].size
if only_alive is False:
# format : sns.jointplot(x_axis, yaxis, kind)
x_axis = path[:,option_dict[x_var],turn_index[0][0]]
y_axis = path[:,option_dict[y_var],turn_index[0][0]]
if only_alive is True:
index = np.where(file[group]["alive"][:,turn_index])[0]
else:
index = np.arange(mp_number)
elif only_alive is True:
alive_index = np.where(file[group]["alive"][:,turn_index])
if plot_size == 1:
samples = index
elif plot_size < 1:
samples_meta = random.sample(list(index), int(plot_size*mp_number))
samples = sorted(samples_meta)
else:
raise ValueError("plot_size must be in range [0,1].")
x_axis = path[alive_index[0],option_dict[x_var],turn_index[0][0]]
y_axis = path[alive_index[0],option_dict[y_var],turn_index[0][0]]
# format : sns.jointplot(x_axis, yaxis, kind)
x_axis = path[samples,option_dict[x_var],turn_index[0][0]]
y_axis = path[samples,option_dict[y_var],turn_index[0][0]]
fig = sns.jointplot(x_axis*scale[option_dict[x_var]],
y_axis*scale[option_dict[y_var]], kind="kde")
......@@ -349,22 +374,3 @@ def plot_profiledata(filename, bunch_number, dimension="tau", start=0,
elif streak_plot is True:
return fig2
\ No newline at end of file
\ No newline at end of file
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment