From 850f73a029086bebef515d46cebd0dafd5085a17 Mon Sep 17 00:00:00 2001
From: Gamelin Alexis <alexis.gamelin@synchrotron-soleil.fr>
Date: Fri, 13 Aug 2021 11:22:06 +0200
Subject: [PATCH] Fix plotting imports and add options to streak_bunchspectrum
---
tracking/monitors/__init__.py | 5 +++--
tracking/monitors/plotting.py | 35 +++++++++++++++++++++++++++++++----
2 files changed, 34 insertions(+), 6 deletions(-)
diff --git a/tracking/monitors/__init__.py b/tracking/monitors/__init__.py
index 69a1b8b..126dbbc 100644
--- a/tracking/monitors/__init__.py
+++ b/tracking/monitors/__init__.py
@@ -18,8 +18,9 @@ from mbtrack2.tracking.monitors.plotting import (plot_bunchdata,
plot_profiledata,
plot_beamdata,
plot_wakedata,
- plot_tunedata,
plot_cavitydata,
- streak_beamdata)
+ streak_beamdata,
+ plot_bunchspectrum,
+ streak_bunchspectrum)
from mbtrack2.tracking.monitors.tools import merge_files
\ No newline at end of file
diff --git a/tracking/monitors/plotting.py b/tracking/monitors/plotting.py
index 9e601cb..ebbfe8b 100644
--- a/tracking/monitors/plotting.py
+++ b/tracking/monitors/plotting.py
@@ -13,7 +13,6 @@ import matplotlib as mpl
import seaborn as sns
import h5py as hp
import random
-from scipy.fft import rfftfreq
def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean",
x_var="time", turn=None, legend=None):
@@ -732,7 +731,7 @@ def plot_bunchspectrum(filenames, bunch_number, dataset="incoherent", dim="tau",
return fig
def streak_bunchspectrum(filename, bunch_number, dataset="incoherent",
- dim="tau", fs=None, log_scale=True):
+ dim="tau", fs=None, log_scale=True, fmin=None, fmax=None, turns=None):
"""
Plot 3D data recorded by the BunchSpectrumMonitor.
@@ -754,6 +753,12 @@ def streak_bunchspectrum(filename, bunch_number, dataset="incoherent",
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.
Returns
-------
@@ -768,7 +773,20 @@ def streak_bunchspectrum(filename, bunch_number, dataset="incoherent",
freq = np.array(group["freq"])
dim_dict = {"x":0, "y":1, "tau":2}
- data = group[dataset][dim_dict[dim], :, :]
+ 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()
@@ -782,6 +800,15 @@ def streak_bunchspectrum(filename, bunch_number, dataset="incoherent",
x_var = freq/fs
xlabel = r"$f/f_{s}$"
+ 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,:]
+
ylabel = "Turn number"
fig, ax = plt.subplots()
@@ -792,7 +819,7 @@ def streak_bunchspectrum(filename, bunch_number, dataset="incoherent",
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()],
+ 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)
--
GitLab