diff --git a/mbtrack2/tracking/monitors/plotting.py b/mbtrack2/tracking/monitors/plotting.py index e41904e462e8ccaa4b2ec5c8f7e1b5402c694e86..0f9f287fe0396129d4a2f692ff99df2ca9c8088c 100644 --- a/mbtrack2/tracking/monitors/plotting.py +++ b/mbtrack2/tracking/monitors/plotting.py @@ -21,11 +21,12 @@ def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean", ---------- filenames : str or list of str Names of the HDF5 files to be plotted. - dataset : {"current","emit","mean","std"} + dataset : {"current","emit","mean","std","cs_invariant"} HDF5 file's dataset to be plotted. The default is "mean". dimension : str The dimension of the dataset to plot: for "emit", dimension = {"x","y","s"}, + for "cs_invariant", dimension = {"x","y"}, for "mean" and "std", dimension = {"x","xp","y","yp","tau","delta"}. not used if "current". The default is "tau". @@ -76,6 +77,15 @@ def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean", elif stat_var == "std": y = np.nanstd(data[axis,bunch_index,:],0) y_label = stat_var + " " + label[axis] + elif dataset == "cs_invariant": + dimension_dict = {"x":0, "y":1} + axis = dimension_dict[dimension] + label = ['$J_x$ (m)', '$J_y$ (m)'] + if stat_var == "mean": + y = np.nanmean(data[axis,bunch_index,:],0) + elif stat_var == "std": + y = np.nanstd(data[axis,bunch_index,:],0) + y_label = stat_var + " " + label[axis] elif dataset == "mean" or dataset == "std": dimension_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, "delta":5} @@ -112,6 +122,12 @@ def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean", "$\\epsilon_{s}$ (m.rad)"] y = data[axis,:,idx] y_label = label[axis] + elif dataset == "cs_invariant": + dimension_dict = {"x":0, "y":1} + axis = dimension_dict[dimension] + label = ['$J_x$ (m)', '$J_y$ (m)'] + y = data[axis,:,idx] + y_label = label[axis] elif dataset == "mean" or dataset == "std": dimension_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, "delta":5} @@ -144,11 +160,12 @@ def streak_beamdata(filename, dataset="mean", dimension="tau", cm_lim=None): ---------- filename : str Name of the HDF5 file that contains the data. - dataset : {"current","emit","mean","std"} + dataset : {"current","emit","mean","std","cs_invariant"} HDF5 file's dataset to be plotted. The default is "mean". dimension : str The dimension of the dataset to plot: for "emit", dimension = {"x","y","s"}, + for "cs_invariant", dimension = {"x","y"}, for "mean" and "std", dimension = {"x","xp","y","yp","tau","delta"}. not used if "current". The default is "tau". @@ -184,6 +201,13 @@ def streak_beamdata(filename, dataset="mean", dimension="tau", cm_lim=None): z = np.array(data["emit"][axis,:,:]).T z_label = label[axis] title = z_label + elif dataset == "cs_invariant": + dimension_dict = {"x":0, "y":1} + axis = dimension_dict[dimension] + label = ['$J_x$ (m)', '$J_y$ (m)'] + z = np.array(data["cs_invariant"][axis,:,:]).T + z_label = label[axis] + title = z_label else: dimension_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, "delta":5} @@ -237,6 +261,7 @@ def plot_bunchdata(filenames, bunch_number, dataset, dimension="x", The dimension of the dataset to plot. Use "None" for "current", otherwise use the following : for "emit", dimension = {"x","y","s"}, + for "cs_invariant", dimension = {"x","y"}, for "mean" and "std", dimension = {"x","xp","y","yp","tau","delta"}, for "action", dimension = {"x","y"}. legend : list of str, optional