diff --git a/tracking/monitors/plotting.py b/tracking/monitors/plotting.py
index 0c648544ec5c7a4f9749e4ef0c3022cf79bc308b..3d57baf020ff9f625bd33642e4d290d330ee72cc 100644
--- a/tracking/monitors/plotting.py
+++ b/tracking/monitors/plotting.py
@@ -58,16 +58,15 @@ def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean",
for filename in filenames:
file = hp.File(filename, "r")
- data = file["Beam"]
- time = np.array(data["time"])
+ time = np.array(file["Beam"]["time"])
+ data = np.array(file["Beam"][dataset])
if x_var == "time":
x = time
x_label = "Number of turns"
- bunch_index = data["current"][:,0] != 0
-
+ bunch_index = (file["Beam"]["current"][:,0] != 0).nonzero()[0]
if dataset == "current":
- y = np.nansum(data[dataset][bunch_index,:],0)*1e3
+ y = np.nansum(data[bunch_index,:],0)*1e3
y_label = "Total current (mA)"
elif dataset == "emit":
dimension_dict = {"x":0, "y":1, "s":2}
@@ -76,9 +75,9 @@ def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean",
"$\\epsilon_{y}$ (m.rad)",
"$\\epsilon_{s}$ (m.rad)"]
if stat_var == "mean":
- y = np.nanmean(data[dataset][axis,bunch_index,:],0)
+ y = np.nanmean(data[axis,bunch_index,:],0)
elif stat_var == "std":
- y = np.nanstd(data[dataset][axis,bunch_index,:],0)
+ y = np.nanstd(data[axis,bunch_index,:],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,
@@ -88,22 +87,23 @@ def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean",
label = ["x (um)", "x' ($\\mu$rad)", "y (um)",
"y' ($\\mu$rad)", "$\\tau$ (ps)", "$\\delta$"]
if stat_var == "mean":
- y = np.nanmean(data[dataset][axis,bunch_index,:],0)*scale[axis]
+ y = np.nanmean(data[axis,bunch_index,:],0)*scale[axis]
elif stat_var == "std":
- y = np.nanstd(data[dataset][axis,bunch_index,:],0)*scale[axis]
+ y = np.nanstd(data[axis,bunch_index,:],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])
+ h = len(file["Beam"]["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 = data[:,idx]*1e3
y_label = "Bunch current (mA)"
elif dataset == "emit":
dimension_dict = {"x":0, "y":1, "s":2}
@@ -111,7 +111,7 @@ def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean",
label = ["$\\epsilon_{x}$ (m.rad)",
"$\\epsilon_{y}$ (m.rad)",
"$\\epsilon_{s}$ (m.rad)"]
- y = data["emit"][axis,:,idx]
+ y = data[axis,:,idx]
y_label = label[axis]
elif dataset == "mean" or dataset == "std":
dimension_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4,
@@ -120,10 +120,12 @@ def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean",
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 = data[axis,:,idx]*scale[axis]
y_label = dataset + " " + label[axis]
else:
raise ValueError("x_var should be time or index")
+
+ y = np.squeeze(y)
ax.plot(x, y)
ax.set_xlabel(x_label)
@@ -428,19 +430,21 @@ def plot_profiledata(filename, bunch_number, dimension="tau", start=0,
file = hp.File(filename, "r")
path = file['ProfileData_{0}'.format(bunch_number)]
- l_bound = path["{0}_bin".format(dimension)]
+ l_bound = np.array(path["{0}_bin".format(dimension)])
+ data = np.array(path[dimension])
+ time = np.array(path["time"])
if stop is None:
- stop = path['time'][-1]
- elif stop not in path['time']:
+ stop = time[-1]
+ elif stop not in time:
raise ValueError("stop not found. Choose from {0}"
- .format(path['time'][:]))
+ .format(time[:]))
- if start not in path['time']:
+ if start not in time:
raise ValueError("start not found. Choose from {0}"
- .format(path['time'][:]))
+ .format(time[:]))
- save_every = path['time'][1] - path['time'][0]
+ save_every = time[1] - time[0]
if step is None:
step = save_every
@@ -455,9 +459,9 @@ def plot_profiledata(filename, bunch_number, dimension="tau", start=0,
"$\\tau$ (ps)", "$\\delta$"]
num = int((stop - start)/step)
- n_bin = len(path[dimension][:,0])
+ n_bin = len(data[:,0])
- start_index = np.where(path['time'][:] == start)[0][0]
+ start_index = np.where(time[:] == start)[0][0]
x_var = np.zeros((num+1,n_bin))
turn_index_array = np.zeros((num+1,))
@@ -471,16 +475,16 @@ def plot_profiledata(filename, bunch_number, dimension="tau", start=0,
fig, ax = plt.subplots()
for i in range(num+1):
ax.plot(x_var[i]*scale[dimension_dict[dimension]],
- path[dimension][:,turn_index_array[i]],
- label="turn {0}".format(path['time'][turn_index_array[i]]))
+ data[:,turn_index_array[i]],
+ label="turn {0}".format(time[turn_index_array[i]]))
ax.set_xlabel(label[dimension_dict[dimension]])
ax.set_ylabel("number of macro-particles")
ax.legend()
if streak_plot is True:
- turn = np.reshape(path['time'][turn_index_array], (num+1,1))
+ turn = np.reshape(time[turn_index_array], (num+1,1))
y_var = np.ones((num+1,n_bin)) * turn
- z_var = np.transpose(path[dimension][:,turn_index_array])
+ z_var = np.transpose(data[:,turn_index_array])
fig2, ax2 = plt.subplots()
cmap = mpl.cm.cool
c = ax2.imshow(z_var, cmap=cmap, origin='lower' , aspect='auto',
@@ -541,18 +545,19 @@ def plot_wakedata(filename, bunch_number, wake_type="Wlong", start=0,
file = hp.File(filename, "r")
path = file['WakePotentialData_{0}'.format(bunch_number)]
+ time = np.array(path["time"])
if stop is None:
- stop = path['time'][-1]
- elif stop not in path['time']:
+ stop = time[-1]
+ elif stop not in time:
raise ValueError("stop not found. Choose from {0}"
- .format(path['time'][:]))
+ .format(time[:]))
- if start not in path['time']:
+ if start not in time:
raise ValueError("start not found. Choose from {0}"
- .format(path['time'][:]))
+ .format(time[:]))
- save_every = path['time'][1] - path['time'][0]
+ save_every = time[1] -time[0]
if step is None:
step = save_every
@@ -583,33 +588,35 @@ def plot_wakedata(filename, bunch_number, wake_type="Wlong", start=0,
else:
tau_name = "tau_" + wake_type
+ data = np.array(path[wake_type])
+
num = int((stop - start)/step)
- n_bin = len(path[wake_type][:,0])
+ n_bin = len(data[:,0])
- start_index = np.where(path['time'][:] == start)[0][0]
+ start_index = np.where(time[:] == start)[0][0]
x_var = np.zeros((num+1,n_bin))
- turn_index_array = np.zeros((num+1,))
+ turn_index_array = np.zeros((num+1,), dtype=int)
for i in range(num+1):
- turn_index = start_index + i * step / save_every
+ turn_index = int(start_index + i * step / save_every)
turn_index_array[i] = turn_index
# construct an array of bin mids
- x_var[i,:] = path[tau_name][:,turn_index]
-
+ x_var[i,:] = np.array(path[tau_name])[:,turn_index]
+
if profile_plot is True:
fig, ax = plt.subplots()
for i in range(num+1):
ax.plot(x_var[i]*1e12,
- path[wake_type][:,turn_index_array[i]]*scale[dimension_dict[wake_type]],
- label="turn {0}".format(path['time'][turn_index_array[i]]))
+ data[:,turn_index_array[i]]*scale[dimension_dict[wake_type]],
+ label="turn {0}".format(time[turn_index_array[i]]))
ax.set_xlabel("$\\tau$ (ps)")
ax.set_ylabel(label[dimension_dict[wake_type]])
ax.legend()
if streak_plot is True:
- turn = np.reshape(path['time'][turn_index_array], (num+1,1))
+ turn = np.reshape(time[turn_index_array], (num+1,1))
y_var = np.ones((num+1,n_bin)) * turn
- z_var = np.transpose(path[wake_type][:,turn_index_array]*scale[dimension_dict[wake_type]])
+ z_var = np.transpose(data[:,turn_index_array]*scale[dimension_dict[wake_type]])
fig2, ax2 = plt.subplots()
cmap = mpl.cm.cool
c = ax2.imshow(z_var, cmap=cmap, origin='lower' , aspect='auto',
@@ -1062,7 +1069,7 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity",
file = hp.File(filename, "r")
cavity_data = file[cavity_name]
- time = cavity_data["time"]
+ time = np.array(cavity_data["time"])
ph = {"cavity":0, "beam":1}
labels = ["Cavity", "Beam"]
@@ -1093,8 +1100,8 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity",
index = np.array(time) == turn
ph = {"cavity":0, "beam":1}
- data = [cavity_data["cavity_phasor_record"][:,index],
- cavity_data["beam_phasor_record"][:,index]]
+ data = [np.array(cavity_data["cavity_phasor_record"])[:,index],
+ np.array(cavity_data["beam_phasor_record"])[:,index]]
labels = ["Cavity", "Beam"]
h=len(data[0])