From 54ecb29fb076379f2e40995a0e74aeadd44110e2 Mon Sep 17 00:00:00 2001
From: Gamelin Alexis <alexis.gamelin@synchrotron-soleil.fr>
Date: Tue, 31 Aug 2021 11:02:22 +0200
Subject: [PATCH] [Fix] BeamSpectrumMonitor
Fix BeamSpectrumMonitor
Add plotting functions for BeamSpectrumMonitor
---
tracking/monitors/monitors.py | 6 +-
tracking/monitors/plotting.py | 184 +++++++++++++++++++++++++++++++++-
2 files changed, 185 insertions(+), 5 deletions(-)
diff --git a/tracking/monitors/monitors.py b/tracking/monitors/monitors.py
index 45a520c..a543fdf 100644
--- a/tracking/monitors/monitors.py
+++ b/tracking/monitors/monitors.py
@@ -1253,13 +1253,13 @@ class BeamSpectrumMonitor(Monitor):
if (beam.mpi_switch == True):
bunch_num = beam.mpi.bunch_num
bunch = beam[bunch_num]
- self.mean[:, bunch_num, self.buffer_count] = bunch.mean[self.mean_index]
+ self.mean[:, bunch_num, self.save_count] = bunch.mean[self.mean_index]
else:
- self.mean[:, :, self.buffer_count] = beam.bunch_mean[self.mean_index,:]
+ self.mean[:, :, self.save_count] = beam.bunch_mean[self.mean_index,:]
self.save_count += 1
- if self.track_count > 0 and self.track_count % self.save_every == 0:
+ if self.save_count == self.save_every:
self.to_buffer(beam)
self.save_count = 0
diff --git a/tracking/monitors/plotting.py b/tracking/monitors/plotting.py
index d014357..0c64854 100644
--- a/tracking/monitors/plotting.py
+++ b/tracking/monitors/plotting.py
@@ -645,7 +645,7 @@ def plot_bunchspectrum(filenames, bunch_number, dataset="incoherent", dim="tau",
dataset : {"mean_incoherent", "coherent", "incoherent"}
HDF5 file's dataset to be plotted.
The default is "incoherent".
- dim : {"x","y","s"}, optional
+ dim : {"x","y","tau"}, optional
The dimension of the dataset to plot.
The default is "tau".
turns : array or None, optional
@@ -753,7 +753,7 @@ def streak_bunchspectrum(filename, bunch_number, dataset="incoherent",
dataset : {"coherent", "incoherent"}
HDF5 file's dataset to be plotted.
The default is "incoherent".
- dim : {"x","y","s"}, optional
+ dim : {"x","y","tau"}, optional
The dimension of the dataset to plot.
The default is "tau".
fs : float or None, optional
@@ -842,6 +842,186 @@ def streak_bunchspectrum(filename, bunch_number, dataset="incoherent",
return fig
+def plot_beamspectrum(filenames, dim="tau", turns=None, f0=None,
+ log_scale=True, legend=None, norm=False):
+ """
+ Plot coherent beam spectrum data.
+
+ Parameters
+ ----------
+ filenames : str or list of str
+ Names of the HDF5 files to be plotted.
+ dim : {"x","y","tau"}, 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.
+ f0 : float or None, optional
+ If not None, the frequency axis is noramlised by f0.
+ 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.
+ norm : bool, optional
+ If True, normalise the data of each spectrum by its geometric mean.
+ The default is False.
+
+ Return
+ ------
+ fig : Figure
+
+ """
+
+ if isinstance(filenames, str):
+ filenames = [filenames]
+
+ fig, ax = plt.subplots()
+
+ for i, filename in enumerate(filenames):
+ file = hp.File(filename, "r")
+ group = file["BeamSpectrum"]
+
+ dataset = "coherent"
+ time = np.array(group["time"])
+ freq = np.array(group["freq"])
+ dim_dict = {"x":0, "y":1, "tau":2}
+
+ 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 f0 is None:
+ x_var = freq
+ xlabel = "Frequency [Hz]"
+ else:
+ x_var = freq/f0
+ xlabel = r"$f/f_{0}$"
+
+ for idx in turn_index:
+ y_var = group[dataset][dim_dict[dim],:,idx]
+ if norm is True:
+ y_var = y_var/gmean(y_var)
+ ax.plot(x_var, y_var)
+
+ if log_scale is True:
+ ax.set_yscale('log')
+
+ ylabel = "FFT amplitude [a.u.]"
+ ax.set_title("Beam coherent spectrum")
+
+ ax.set_xlabel(xlabel)
+ ax.set_ylabel(ylabel)
+ if legend is not None:
+ plt.legend(legend)
+ file.close()
+
+ return fig
+
+def streak_beamspectrum(filename, dim="tau", f0=None, log_scale=True, fmin=None,
+ fmax=None, turns=None, norm=False):
+ """
+ Plot 3D data recorded by the BeamSpectrumMonitor.
+
+ Parameters
+ ----------
+ filenames : str
+ Name of the HDF5 file to be plotted.
+ dim : {"x","y","tau"}, optional
+ The dimension of the dataset to plot.
+ The default is "tau".
+ f0 : float or None, optional
+ If not None, the frequency axis is noramlised by f0.
+ log_scale : bool, optional
+ If True, the spectrum plots are shown in y-log scale.
+ The default is True.
+ fmin : float, optional
+ If not None, the plot is limitted to values bigger than fmin.
+ fmax : float, optional
+ If not None, the plot is limitted to values smaller than fmax.
+ turns : array, optional
+ If not None, only the turn numbers in the turns array are plotted.
+ norm : bool, optional
+ If True, normalise the data of each spectrum by its geometric mean.
+ The default is False.
+
+ Returns
+ -------
+ fig : Figure
+
+ """
+
+ file = hp.File(filename, "r")
+ group = file["BeamSpectrum"]
+ dataset="coherent"
+ time = np.array(group["time"])
+ freq = np.array(group["freq"])
+ dim_dict = {"x":0, "y":1, "tau":2}
+
+ if turns is None:
+ turn_index = np.where(time == time)[0]
+ tmin = time.min()
+ tmax = time.max()
+ else:
+ tmin = turns.min()
+ tmax = turns.max()
+ turn_index = []
+ for turn in turns:
+ idx = np.where(time == turn)[0][0]
+ turn_index.append(idx)
+ turn_index = np.array(turn_index)
+
+ data = group[dataset][dim_dict[dim], :, turn_index]
+
+ if log_scale is True:
+ option = mpl.colors.LogNorm()
+ else:
+ option = None
+
+ if f0 is None:
+ x_var = freq
+ xlabel = "Frequency [Hz]"
+ else:
+ x_var = freq/f0
+ xlabel = r"$f/f_{0}$"
+
+ if fmin is None:
+ fmin = x_var.min()
+ if fmax is None:
+ fmax = x_var.max()
+
+ ind = (x_var > fmin) & (x_var < fmax)
+ x_var=x_var[ind]
+ data = data[ind,:]
+
+ if norm is True:
+ data = data/gmean(data)
+
+ ylabel = "Turn number"
+
+ fig, ax = plt.subplots()
+ ax.set_title("Beam 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(), tmin, tmax],
+ 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):
"""
--
GitLab