diff --git a/tracking/monitors/plotting.py b/tracking/monitors/plotting.py
index a4c7a50c67b066d41bf56016674aa8be96b54f7c..780dc4875b6fb2074f1132150677d197f016f7a4 100644
--- a/tracking/monitors/plotting.py
+++ b/tracking/monitors/plotting.py
@@ -15,7 +15,8 @@ import h5py as hp
import random
from scipy.fft import rfftfreq
-def plot_beamdata(filename, dataset, dimension=None, stat_var=None, x_var="time"):
+def plot_beamdata(filename, dataset="mean", dimension="tau", stat_var="mean",
+ x_var="time", turn=None, plot_type="normal", cm_lim=None):
"""
Plot data recorded by BeamMonitor.
@@ -24,73 +25,147 @@ def plot_beamdata(filename, dataset, dimension=None, stat_var=None, x_var="time"
filename : str
Name of the HDF5 file that contains the data.
dataset : {"current","emit","mean","std"}
- HDF5 file's dataset to be plotted.
- dimension : str, optional
- The dimension of the dataset to plot. Use "None" for "current",
- otherwise use the following :
+ HDF5 file's dataset to be plotted. The default is "mean".
+ dimension : str
+ The dimension of the dataset to plot:
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 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".
-
+ not used if "current".
+ The default is "tau".
+ stat_var : {"mean", "std"}
+ Statistical value of the dimension.
+ x_var : {"time", "index"}
+ Choice of the horizontal axis for "normal" plot_type.
+ "time" corresponds to turn number.
+ "index" corresponds to bunch index.
+ turn : int or float, optional
+ Choice of the turn to plot for the "normal" plot_type with "index".
+ If None, the last turn is plotted.
+ plot_type : {"normal","streak"}, optional
+ Choice of the type of plot.
+ "normal" is for the standard line 2D plot.
+ "streak" is for the 3D like image.
+ cm_lim : list [vmin, vmax], optional
+ Colormap scale for the "streak" plot.
+
Return
------
fig : Figure
Figure object with the plot on it.
"""
-
file = hp.File(filename, "r")
- path = file["Beam"]
-
- if dataset == "current":
- fig, ax = plt.subplots()
- 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":
- 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":
- fig, ax = plt.subplots()
- ax.plot(path["time"], np.nanmean(path["emit"][axis,:],0))
-
- elif stat_var == "std":
- fig, ax = plt.subplots()
- ax.plot(path["time"], np.nanstd(path["emit"][axis,:],0))
-
- ax.set_xlabel("Number of turns")
- ax.set_ylabel(stat_var+" " + label[axis])
+ data = file["Beam"]
+ time = np.array(data["time"])
+
+ if plot_type == "normal":
+
+ if x_var == "time":
+ x = time
+ x_label = "Number of turns"
+ if dataset == "current":
+ y = np.nansum(data["current"][:],0)*1e3
+ y_label = "Total current (mA)"
+ elif dataset == "emit":
+ 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":
+ y = np.nanmean(data["emit"][axis,:],0)
+ elif stat_var == "std":
+ y = np.nanstd(data["emit"][axis,:],0)
+ y_label = stat_var + " " + label[axis]
+ elif dataset == "mean" or dataset == "std":
+ 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$"]
+ if stat_var == "mean":
+ y = np.nanmean(data[dataset][axis,:],0)*scale[axis]
+ elif stat_var == "std":
+ y = np.nanstd(data[dataset][axis,:],0)*scale[axis]
+ label_sup = {"mean":"mean of ", "std":"std of "}
+ y_label = label_sup[stat_var] + dataset + " " + label[axis]
+
+ elif x_var == "index":
+ h = len(data["mean"][0,:,0])
+ x = np.arange(h)
+ x_label = "Bunch index"
+ if turn is None:
+ idx = -1
+ else:
+ idx = np.where(time == int(turn))
+ if dataset == "current":
+ y = data["current"][:,idx]*1e3
+ y_label = "Bunch current (mA)"
+ elif dataset == "emit":
+ 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)"]
+ y = data["emit"][axis,:,idx]
+ y_label = label[axis]
+ elif dataset == "mean" or dataset == "std":
+ 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$"]
+ y = data[dataset][axis,:,idx]*scale[axis]
+ y_label = dataset + " " + label[axis]
+ else:
+ raise ValueError("x_var should be time or index")
- elif dataset == "mean" or dataset == "std":
- 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()
+ ax.plot(x, y)
+ ax.set_xlabel(x_label)
+ ax.set_ylabel(y_label)
+
+ elif plot_type == "streak":
+ h = len(data["mean"][0,:,0])
+ x = np.arange(h)
+ x_label = "Bunch index"
+ y = time
+ y_label = "Number of turns"
+ if dataset == "current":
+ z = np.array(data["current"]*1e3).T
+ z_label = "Bunch current (mA)"
+ elif dataset == "emit":
+ 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)"]
+ z = np.array(data["emit"][axis,:,:]).T
+ z_label = label[axis]
+ else:
+ 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$"]
+ z = np.array(data[dataset][axis,:,:]).T*scale[axis]
+ z_label = label[axis]
fig, ax = plt.subplots()
- 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":
- 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")
- ax.set_ylabel(label_sup[stat_var] + dataset +" "+ label[axis])
-
- file.close()
- return fig
+ ax.set_xlabel(x_label)
+ ax.set_ylabel(y_label)
+ cmap = mpl.cm.cool
+ c = ax.imshow(z, cmap=cmap, origin='lower' , aspect='auto',
+ extent=[x.min(), x.max(), y.min(), y.max()])
+ if cm_lim is not None:
+ c.set_clim(vmin=cm_lim[0],vmax=cm_lim[1])
+ cbar = fig.colorbar(c, ax=ax)
+ cbar.set_label(z_label)
+
+ return fig
def plot_bunchdata(filename, bunch_number, dataset, dimension="x", x_var="time"):
"""
@@ -609,7 +684,7 @@ def plot_tunedata(filename, bunch_number, f0, plot_tune=True, plot_fft=False,
return tuple(fig_to_return)
def plot_cavitydata(filename, cavity_name, phasor="cavity",
- plot_type="bunch", bunch_number=0, turn=None):
+ plot_type="bunch", bunch_number=0, turn=None, cm_lim=None):
"""
Plot data recorded by CavityMonitor.
@@ -635,6 +710,8 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity",
Bunch number to select. The default is 0.
turn : int, optional
Turn to plot. The default is None.
+ cm_lim : list [vmin, vmax], optional
+ Colormap scale for the "streak" plots.
Returns
-------
@@ -675,8 +752,7 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity",
if plot_type == "turn":
- index = time == turn
-
+ index = np.array(time) == turn
ph = {"cavity":0, "beam":1}
data = [cavity_data["cavity_phasor_record"][:,index],
cavity_data["beam_phasor_record"][:,index]]
@@ -714,6 +790,8 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity",
fig, ax = plt.subplots()
cmap = mpl.cm.cool
c = ax.imshow(data, cmap=cmap, origin='lower' , aspect='auto')
+ if cm_lim is not None:
+ c.set_clim(vmin=cm_lim[0],vmax=cm_lim[1])
ax.set_xlabel("Bunch index")
ax.set_ylabel("Number of turns")
cbar = fig.colorbar(c, ax=ax)