diff --git a/tracking/monitors/plotting.py b/tracking/monitors/plotting.py index 63ed6bfc6efe4510a10a0c6b45c5c7e09ea6ff89..f1531e0a26990564e06cc15384a033d50e72cd7f 100644 --- a/tracking/monitors/plotting.py +++ b/tracking/monitors/plotting.py @@ -14,7 +14,7 @@ import seaborn as sns import h5py as hp import random -def plot_beamdata(filename, dataset, option=None, stat_var=None, x_var="time"): +def plot_beamdata(filename, dataset, dimension=None, stat_var=None, x_var="time"): """ Plot data recorded by BeamMonitor. @@ -24,14 +24,14 @@ def plot_beamdata(filename, dataset, option=None, stat_var=None, x_var="time"): 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"} + dimension : str, optional + The dimension of the dataset to plot. Use "None" for "current", + otherwise use the following : + for "emit", dimension = {"x","y","s"}, + for "mean" and "std", dimension = {"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. + Statistical value of the dimension. Unless dataset = "current", stat_var + needs to be specified. x_var : str, optional Variable to be plotted on the horizontal axis. The default is "time". @@ -46,64 +46,43 @@ def plot_beamdata(filename, dataset, option=None, stat_var=None, x_var="time"): path = file["Beam"] if dataset == "current": - total_current = [] - for i in range (len(path["time"])): - total_current.append(np.sum(path["current"][:,i])*1e3) - fig, ax = plt.subplots() - ax.plot(path["time"],total_current) + ax.plot(path["time"], np.nansum(path["current"][:],0)*1e3) ax.set_xlabel("Number of turns") ax.set_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)"] + dimension_dict = {"x":0, "y":1, "s":2} + axis = dimension_dict[dimension] + label = ["$\\epsilon_{x}$ (m.rad)", + "$\\epsilon_{y}$ (m.rad)", + "$\\epsilon_{s}$ (m.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]) - fig, ax = plt.subplots() - ax.plot(path["time"],mean_emit) + ax.plot(path["time"], np.nanmean(path["emit"][axis,:],0)) elif stat_var == "std": - std_emit = [] - for i in range (len(path["time"])): - std_emit.append(np.std(path["emit"][axis,:,i])*scale[axis]) - fig, ax = plt.subplots() - ax.plot(path["time"],std_emit) + ax.plot(path["time"], np.nanstd(path["emit"][axis,:],0)) ax.set_xlabel("Number of turns") ax.set_ylabel(stat_var+" " + label[axis]) elif dataset == "mean" or dataset == "std": - option_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, "delta":5} - axis = option_dict[option] + dimension_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, "delta":5} + axis = dimension_dict[dimension] scale = [1e6, 1e6, 1e6, 1e6, 1e12, 1] label = ["x (um)", "x' ($\\mu$rad)", "y (um)", "y' ($\\mu$rad)", "$\\tau$ (ps)", "$\\delta$"] fig, ax = plt.subplots() - if stat_var == "mean": - mean_list = [] - for i in range (len(path["time"])): - mean_list.append(np.mean(path[dataset][axis,:,i]*scale[axis])) - - ax.plot(path["time"],mean_list) + if stat_var == "mean": + ax.plot(path["time"], np.nanmean(path[dataset][axis,:],0)*scale[axis]) label_sup = {"mean":"", "std":"std of "} # input stat_var - elif stat_var == "std": - std_list = [] - for i in range (len(path["time"])): - std_list.append(np.std(path[dataset][axis,:,i]*scale[axis])) - - ax.plot(path["time"],std_list) + elif stat_var == "std": + ax.plot(path["time"], np.nanstd(path[dataset][axis,:],0)*scale[axis]) label_sup = {"mean":"", "std":"std of "} #input stat_var ax.set_xlabel("Number of turns")