diff --git a/tracking/monitors/plotting.py b/tracking/monitors/plotting.py
index 9ae1449936bf8c12e56b951c543bb48391a3e677..0a9cad7e174ceb1d11d555b3355f5949d3d51ab0 100644
--- a/tracking/monitors/plotting.py
+++ b/tracking/monitors/plotting.py
@@ -12,13 +12,130 @@ import matplotlib.pyplot as plt
import seaborn as sns
import h5py as hp
+def plot_beamdata(filename, dataset, option=None, stat_var=None, x_var="time"):
+ """
+ Plot data recorded from a BeamMonitor.
+
+ Parameters
+ ----------
+ filename : str
+ Name of the HDF5 file that contains the data.
+ dataset : {"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"}
+ stat_var : {"mean", "std"}, optional
+ Statistical value of option. Except when dataset = "current", stat_var
+ needs to be specified. The default is None.
+ x_var : str, optional
+ Variable to be plotted on the horizontal axis. The default is "time".
+
+ """
+
+ file = hp.File(filename, "r")
+ path = file["Beam"]
+
+ if dataset == "current":
+ total_current = []
+ for i in range (len(path["time"])):
+ total_current.append(np.sum(path["current"][:,i])*1e3)
+
+ plt.plot(path["time"],total_current)
+ plt.xlabel("Number of turns")
+ plt.ylabel("total current (mA)")
+
+ elif dataset == "emit":
+ option_dict = {"x":0, "y":1, "s":2} #input option
+ axis = option_dict[option]
+ scale = [1e12, 1e12, 1e15]
+ label = ["$\\epsilon_{x}$ (pm.rad)",
+ "$\\epsilon_{y}$ (pm.rad)",
+ "$\\epsilon_{s}$ (fm.rad)"]
+
+ if stat_var == "mean":
+ mean_emit = []
+ for i in range (len(path["time"])):
+ mean_emit.append(np.mean(path["emit"][axis,:,i])*scale[axis])
+
+ plt.plot(path["time"],mean_emit)
+ label_sup = "avg. "
+
+ elif stat_var == "std":
+ std_emit = []
+ for i in range (len(path["time"])):
+ std_emit.append(np.std(path["emit"][axis,:,i])*scale[axis])
+
+ plt.plot(path["time"],std_emit)
+ label_sup = "std. "
+
+ plt.xlabel("Number of turns")
+ plt.ylabel(label_sup + label[axis])
+
+ elif dataset == "mean":
+ option_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, "delta":5}
+ axis = option_dict[option]
+ scale = [1e6, 1e6, 1e6, 1e6, 1e12, 1]
+ label = ["x (um)", "x' ($\\mu$rad)", "y (um)", "y' ($\\mu$rad)",
+ "$\\tau$ (ps)", "$\\delta$"]
+
+ if stat_var == "mean":
+ mean_mean = []
+ for i in range (len(path["time"])):
+ mean_mean.append(np.mean(path["mean"][axis,:,i]*scale[axis]))
+
+ plt.plot(path["time"],mean_mean)
+ label_sup = "avg. "
+
+ elif stat_var == "std":
+ std_mean = []
+ for i in range (len(path["time"])):
+ std_mean.append(np.std(path["mean"][axis,:,i]*scale[axis]))
+
+ plt.plot(path["time"],std_mean)
+ label_sup = "std. of avg. "
+
+ plt.xlabel("Number of turns")
+ plt.ylabel(label_sup + label[axis])
+
+ elif dataset == "std":
+ option_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, "delta":5}
+ axis = option_dict[option]
+ scale = [1e6, 1e6, 1e6, 1e6, 1e12, 1]
+ label = ["x (um)", "x' ($\\mu$rad)", "y (um)", "y' ($\\mu$rad)",
+ "$\\tau$ (ps)", "$\\delta$"]
+
+ if stat_var == "mean":
+ mean_std = []
+ for i in range (len(path["time"])):
+ mean_std.append(np.mean(path["std"][axis,:,i]*scale[axis]))
+
+ plt.plot(path["time"],mean_std)
+ label_sup = "std. "
+
+ elif stat_var == "std":
+ std_std = []
+ for i in range (len(path["time"])):
+ std_std.append(np.std(path["std"][axis,:,i]*scale[axis]))
+
+ plt.plot(path["time"],std_std)
+ label_sup = "std. of std. "
+
+ plt.xlabel("Number of turns")
+ plt.ylabel(label_sup + label[axis])
+
+ file.close()
+
+
def plot_bunchdata(filename, bunch_number, dataset, option=None, x_var="time"):
"""
Plot data recorded from a BunchMonitor.
Parameters
----------
- filename : str
+ 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.
diff --git a/tracking/optics.py b/tracking/optics.py
index 2e3483657181ebab0073f5d9eb135ac7808d07dd..37bc44eadb88be14086d84c29ea44e0463977a0c 100644
--- a/tracking/optics.py
+++ b/tracking/optics.py
@@ -244,6 +244,58 @@ class Optics:
dispersion = [self.dispX(position), self.disppX(position),
self.dispY(position), self.disppY(position)]
return np.array(dispersion)
+
+ def plot_optics(self, ring, var, option):
+ """
+ Plotting optical variables.
+
+ Parameters
+ ----------
+ ring : Synchrotron object
+ var : {"beta", "alpha", "gamma", "dispersion"}
+ Optical variable to be plotted
+ option : str
+ If var = "beta", "alpha" and "gamma", option = {"x","y"} specifying
+ the axis of interest.
+ If var = "dispersion", option = {"x","px","y","py"} specifying the
+ axis of interest for the dispersion function or its derivative.
+
+ """
+
+ var_dict = {"beta":self.beta, "alpha":self.alpha, "gamma":self.gamma,
+ "dispersion":self.dispersion}
+
+ if var == "dispersion":
+ option_dict = {"x":0, "px":1, "y":2, "py":3}
+
+ label = ["D$_{x}$ (m)", "D'$_{x}$", "D$_{y}$ (m)", "D'$_{y}$"]
+
+ plt.ylabel(label[option_dict[option]])
+
+ elif var=="beta" or var=="alpha" or var=="gamma":
+ option_dict = {"x":0, "y":1}
+ label_dict = {"beta":"$\\beta$", "alpha":"$\\alpha$",
+ "gamma":"$\\gamma$"}
+
+ if option == "x": label_sup = "$_{x}$"
+ elif option == "y": label_sup = "$_{y}$"
+
+ unit = {"beta":" (m)", "alpha":"", "gamma":" (m$^{-1}$)"}
+
+ plt.ylabel(label_dict[var] + label_sup + unit[var])
+
+ else:
+ raise ValueError("Variable name is not found.")
+
+ var_list = []
+ position_list = []
+ for i in range (int(ring.L)):
+ var_list.append(var_dict[var](i)[option_dict[option]])
+ position_list.append(i)
+
+ plt.plot(position_list, var_list)
+ plt.xlabel("position (m)")
+
class PhyisicalModel:
"""