diff --git a/tracking/monitors/plotting.py b/tracking/monitors/plotting.py index e481411d21147136e80962e134d07f615c1d31b3..dfbe1bb38d658e62f97ac4c892af5ff08048a1fe 100644 --- a/tracking/monitors/plotting.py +++ b/tracking/monitors/plotting.py @@ -100,7 +100,7 @@ def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean", if turn is None: idx = -1 else: - idx = np.where(time == int(turn)) + idx = np.where(time == int(turn))[0] if (idx.size == 0): raise ValueError("Turn is not valid.") @@ -748,7 +748,7 @@ def plot_bunchspectrum(filenames, bunch_number, dataset="incoherent", dim="tau", def streak_bunchspectrum(filename, bunch_number, dataset="incoherent", dim="tau", fs=None, log_scale=True, fmin=None, - fmax=None, turns=None, norm=False): + fmax=None, turns=None, norm=False, ylim=None): """ Plot 3D data recorded by the BunchSpectrumMonitor. @@ -779,6 +779,10 @@ def streak_bunchspectrum(filename, bunch_number, dataset="incoherent", norm : bool, optional If True, normalise the data of each spectrum by its geometric mean. The default is False. + ylim : array, optional + If not None, should be array like in the form [ymin, ymax] where ymin + and ymax are the minimum and maxmimum values used in the y axis. + Returns ------- @@ -795,8 +799,12 @@ def streak_bunchspectrum(filename, bunch_number, dataset="incoherent", if turns is None: turn_index = np.where(time == time)[0] - tmin = time.min() - tmax = time.max() + if ylim is None: + tmin = time.min() + tmax = time.max() + else: + tmin = ylim[0] + tmax = ylim[1] else: tmin = turns.min() tmax = turns.max() @@ -831,8 +839,11 @@ def streak_bunchspectrum(filename, bunch_number, dataset="incoherent", if norm is True: data = data/gmean(data) - - ylabel = "Turn number" + + if ylim is None: + ylabel = "Turn number" + else: + ylabel = "" fig, ax = plt.subplots() if dataset == "incoherent": @@ -937,7 +948,7 @@ def plot_beamspectrum(filenames, dim="tau", turns=None, f0=None, return fig def streak_beamspectrum(filename, dim="tau", f0=None, log_scale=True, fmin=None, - fmax=None, turns=None, norm=False): + fmax=None, turns=None, norm=False, ylim=None): """ Plot 3D data recorded by the BeamSpectrumMonitor. @@ -962,6 +973,9 @@ def streak_beamspectrum(filename, dim="tau", f0=None, log_scale=True, fmin=None, norm : bool, optional If True, normalise the data of each spectrum by its geometric mean. The default is False. + ylim : array, optional + If not None, should be array like in the form [ymin, ymax] where ymin + and ymax are the minimum and maxmimum values used in the y axis. Returns ------- @@ -978,8 +992,12 @@ def streak_beamspectrum(filename, dim="tau", f0=None, log_scale=True, fmin=None, if turns is None: turn_index = np.where(time == time)[0] - tmin = time.min() - tmax = time.max() + if ylim is None: + tmin = time.min() + tmax = time.max() + else: + tmin = ylim[0] + tmax = ylim[1] else: tmin = turns.min() tmax = turns.max() @@ -1015,7 +1033,10 @@ def streak_beamspectrum(filename, dim="tau", f0=None, log_scale=True, fmin=None, if norm is True: data = data/gmean(data) - ylabel = "Turn number" + if ylim is None: + ylabel = "Turn number" + else: + ylabel = "" fig, ax = plt.subplots() ax.set_title("Beam coherent spectrum")