diff --git a/tracking/optics.py b/tracking/optics.py index 37bc44eadb88be14086d84c29ea44e0463977a0c..be326c44431ae43988080dbc4571d4ce922de41b 100644 --- a/tracking/optics.py +++ b/tracking/optics.py @@ -245,20 +245,21 @@ class Optics: self.dispY(position), self.disppY(position)] return np.array(dispersion) - def plot_optics(self, ring, var, option): + def plot_optics(self, var, option, n_points=1000): """ Plotting optical variables. Parameters ---------- - ring : Synchrotron object var : {"beta", "alpha", "gamma", "dispersion"} - Optical variable to be plotted + 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. + n_points : int + Number of points on the plot. The default value is 1000. """ @@ -271,6 +272,7 @@ class Optics: 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} @@ -283,17 +285,15 @@ class Optics: 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) + position = np.linspace(0, self.lattice.circumference, int(n_points)) + var_list = var_dict[var](position)[option_dict[option]] + plt.plot(position,var_list) + plt.xlabel("position (m)") @@ -566,3 +566,4 @@ class PhyisicalModel: axs[1].set(xlabel="Longitudinal position [m]", ylabel="Vertical aperture [mm]") axs[1].legend(["Top","Bottom"]) +