diff --git a/tracking/monitors/plotting.py b/tracking/monitors/plotting.py
index 569dee0ecb748fabf8a7643ff3e519ba5682edec..9e601cbf719306429dfb04e6bd4528df003a3d61 100644
--- a/tracking/monitors/plotting.py
+++ b/tracking/monitors/plotting.py
@@ -303,8 +303,6 @@ def plot_bunchdata(filenames, bunch_number, dataset, dimension="x",
x_axis = file[group]["time"][:]
xlabel = "Number of turns"
- if i == 0:
- fig, ax = plt.subplots()
ax.plot(x_axis, y_var)
ax.set_xlabel(xlabel)
ax.set_ylabel(label)
@@ -631,120 +629,177 @@ def plot_wakedata(filename, bunch_number, wake_type="Wlong", start=0,
elif streak_plot is True:
return fig2
-def plot_tunedata(filename, bunch_number, f0, plot_tune=True, plot_fft=False,
- dimension='x', min_tune=0, max_tune=0.5, min_turn=None,
- max_turn=None, streak_plot=True, profile_plot=False):
+def plot_bunchspectrum(filenames, bunch_number, dataset="incoherent", dim="tau",
+ turns=None, fs=None, log_scale=True, legend=None):
"""
- Plot data recorded by TuneMonitor.
-
+ Plot coherent and incoherent spectrum data.
+
Parameters
----------
- filename : str
- Name of the HDF5 file that contains the data.
- bunch_number : int
+ filenames : str or list of str
+ Names of the HDF5 files to be plotted.
+ bunch_number : int or list of int
Bunch to plot. This has to be identical to 'bunch_number' parameter in
- 'BunchMonitor' object.
- f0 : float
- Revolution frequency of the ring used for the tracking in [Hz].
- plot_tune : bool, optional
- If True, tune data usinf NAFF is plotted.
- plot_fft : bool, optional
- If True, FFT data is plotted.
- dimension : {'x', 'y', 's'}
- Option to plot FFT data in horizontal, vertical, or longitudinal plane.
- min_tune, max_tune : int or float, optional
- The minimum and the maximum tune values to plot FFT data.
- min_turn, max_turn : int or float, optional
- The minimum and the maximum number of turns to plot FFT data.
- streak_plot : bool, optional
- If True, the FFT data is plotted as a streak plot.
- profile_plot : bool, optional.
- If True, the FFT data is plotted as line plots.
-
+ 'BunchSpectrumMonitor' object.
+ dataset : {"mean_incoherent", "coherent", "incoherent"}
+ HDF5 file's dataset to be plotted.
+ The default is "incoherent".
+ dim : {"x","y","s"}, optional
+ The dimension of the dataset to plot.
+ The default is "tau".
+ turns : array or None, optional
+ Numbers of the turns to plot.
+ If None, all turns are shown.
+ The default is None.
+ fs : float or None, optional
+ If not None, the frequency axis is noramlised by fs.
+ The default is None.
+ log_scale : bool, optional
+ If True, the spectrum plots are shown in y-log scale.
+ The default is True.
+ legend : list of str, optional
+ Legend to add for each file.
+ The default is None.
+
Return
------
fig : Figure
- Figure objects with the plot on it.
+
+ """
+
+ if isinstance(filenames, str):
+ filenames = [filenames]
+
+ if isinstance(bunch_number, int):
+ ll = []
+ for i in range(len(filenames)):
+ ll.append(bunch_number)
+ bunch_number = ll
+
+ fig, ax = plt.subplots()
+
+ for i, filename in enumerate(filenames):
+ file = hp.File(filename, "r")
+ group = file["BunchSpectrum_{0}".format(bunch_number[i])]
+
+ time = np.array(group["time"])
+ freq = np.array(group["freq"])
+ dim_dict = {"x":0, "y":1, "tau":2}
+
+ if dataset == "mean_incoherent":
+ y_var = group["mean_incoherent"][dim_dict[dim],:]
+ y_err = group["std_incoherent"][dim_dict[dim],:]
+ ax.errorbar(time, y_var, y_err)
+ xlabel = "Turn number"
+ ylabel = "Mean incoherent frequency [Hz]"
+ elif dataset == "incoherent" or dataset == "coherent":
+
+ if turns is None:
+ turn_index = np.where(time == time)[0]
+ else:
+ turn_index = []
+ for turn in turns:
+ idx = np.where(time == turn)[0][0]
+ turn_index.append(idx)
+ turn_index = np.array(turn_index)
+
+ if fs is None:
+ x_var = freq
+ xlabel = "Frequency [Hz]"
+ else:
+ x_var = freq/fs
+ xlabel = r"$f/f_{s}$"
+
+ for idx in turn_index:
+ y_var = group[dataset][dim_dict[dim],:,idx]
+ ax.plot(x_var, y_var)
+
+ if log_scale is True:
+ ax.set_yscale('log')
+
+ ylabel = "FFT amplitude [a.u.]"
+ if dataset == "incoherent":
+ ax.set_title("Incoherent spectrum")
+ elif dataset == "coherent":
+ ax.set_title("Coherent spectrum")
+
+ ax.set_xlabel(xlabel)
+ ax.set_ylabel(ylabel)
+ if legend is not None:
+ plt.legend(legend)
+ file.close()
+
+ return fig
+
+def streak_bunchspectrum(filename, bunch_number, dataset="incoherent",
+ dim="tau", fs=None, log_scale=True):
+ """
+ Plot 3D data recorded by the BunchSpectrumMonitor.
+
+ Parameters
+ ----------
+ filenames : str
+ Name of the HDF5 file to be plotted.
+ bunch_number : int
+ Bunch to plot. This has to be identical to 'bunch_number' parameter in
+ 'BunchSpectrumMonitor' object.
+ dataset : {"coherent", "incoherent"}
+ HDF5 file's dataset to be plotted.
+ The default is "incoherent".
+ dim : {"x","y","s"}, optional
+ The dimension of the dataset to plot.
+ The default is "tau".
+ fs : float or None, optional
+ If not None, the frequency axis is noramlised by fs.
+ log_scale : bool, optional
+ If True, the spectrum plots are shown in y-log scale.
+ The default is True.
+
+ Returns
+ -------
+ fig : Figure
"""
file = hp.File(filename, "r")
+ group = file["BunchSpectrum_{0}".format(bunch_number)]
- group = "TuneData_{0}".format(bunch_number) # Data group of the HDF5 file
- time = file[group]["time"]
+ time = np.array(group["time"])
+ freq = np.array(group["freq"])
+ dim_dict = {"x":0, "y":1, "tau":2}
- fig_to_return = []
+ data = group[dataset][dim_dict[dim], :, :]
- if plot_tune is True:
- tune = file[group]["tune"]
- tune_spread = file[group]["tune_spread"]
-
- fig1, ax1 = plt.subplots()
- ax1.errorbar(x=time, y=tune[0,:], yerr=tune_spread[0,:])
- ax1.errorbar(x=time, y=tune[1,:], yerr=tune_spread[1,:])
- ax1.set_xlabel("Turn number")
- ax1.set_ylabel("Transverse tunes")
- plt.legend(["x","y"])
- fig_to_return.append(fig1)
-
- fig2, ax2 = plt.subplots()
- ax2.errorbar(x=time, y=tune[2,:], yerr=tune_spread[2,:])
- ax2.set_xlabel("Turn number")
- ax2.set_ylabel("Synchrotron tune")
- fig_to_return.append(fig2)
-
- if plot_fft is True:
-
- T0 = 1/f0
- n_freq = file[group]['fft'].shape[1]
- freq = rfftfreq((n_freq-1)*2, T0)
- tune_fft = freq / f0
-
- dimension_dict = {'x':0, 'y':1, 's':2}
- axis = dimension_dict[dimension]
-
- fourier_save = file[group]['fft'][axis]
+ if log_scale is True:
+ option = mpl.colors.LogNorm()
+ else:
+ option = None
+
+ if fs is None:
+ x_var = freq
+ xlabel = "Frequency [Hz]"
+ else:
+ x_var = freq/fs
+ xlabel = r"$f/f_{s}$"
- if max_turn is None:
- max_turn = time[-1]
- if min_turn is None:
- min_turn = time[0]
-
- min_tune_iloc = np.where(tune_fft >= min_tune)[0][0]
- max_tune_iloc = np.where(tune_fft <= max_tune)[0][-1]
- save_every = int(time[1] - time[0])
- min_turn_iloc = min_turn // save_every - 1
- max_turn_iloc = max_turn // save_every
-
-
- if streak_plot is True:
- fig3, ax3 = plt.subplots()
- cmap = plt.get_cmap('Blues')
-
- c = ax3.imshow(np.transpose(np.log(
- fourier_save[min_tune_iloc:max_tune_iloc+1,
- min_turn_iloc:max_turn_iloc+1])),
- cmap=cmap, origin='lower' , aspect='auto',
- extent=[min_tune, max_tune, min_turn, max_turn])
- ax3.set_xlabel('$Q_{}$'.format(dimension))
- ax3.set_ylabel("Turns")
- cbar = fig3.colorbar(c, ax=ax3)
- cbar.set_label("log FFT amplitude")
- fig_to_return.append(fig3)
-
- if profile_plot is True:
- fig4, ax4 = plt.subplots()
- ax4.plot(tune_fft[min_tune_iloc:max_tune_iloc+1],
- fourier_save[min_tune_iloc:max_tune_iloc+1,
- min_turn_iloc:max_turn_iloc+1])
- ax4.set_xlabel('$Q_{}$'.format(dimension))
- ax4.set_ylabel("FFT amplitude")
- ax4.legend(time[min_turn_iloc:max_turn_iloc+1])
- fig_to_return.append(fig4)
-
- file.close()
-
- return tuple(fig_to_return)
+ ylabel = "Turn number"
+
+ fig, ax = plt.subplots()
+ if dataset == "incoherent":
+ ax.set_title("Incoherent spectrum")
+ elif dataset == "coherent":
+ ax.set_title("Coherent spectrum")
+
+ cmap = mpl.cm.inferno # sequential
+ c = ax.imshow(data.T, cmap=cmap, origin='lower' , aspect='auto',
+ extent=[x_var.min(), x_var.max(), time.min(), time.max()],
+ norm=option)
+ cbar = fig.colorbar(c, ax=ax)
+ cbar.ax.set_ylabel("FFT amplitude [a.u.]", rotation=270)
+ ax.set_xlabel(xlabel)
+ ax.set_ylabel(ylabel)
+
+ return fig
def plot_cavitydata(filename, cavity_name, phasor="cavity",
plot_type="bunch", bunch_number=0, turn=None, cm_lim=None):