Skip to content
Snippets Groups Projects
Commit ad0eae65 authored by Alexis GAMELIN's avatar Alexis GAMELIN
Browse files

Repear plotting function for h5py >= 3.1.0 

Convert h5py data into numpy arrays as fancy indexing is broken in h5py >= 3.1.0
parent fd2cd135
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
tracking/monitors/plotting.py
+ 52
− 45
View file @ ad0eae65
...@@ -58,16 +58,15 @@ def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean", ...@@ -58,16 +58,15 @@ def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean",
for filename in filenames: for filename in filenames:
file = hp.File(filename, "r") file = hp.File(filename, "r")
data = file["Beam"] time = np.array(file["Beam"]["time"])
time = np.array(data["time"]) data = np.array(file["Beam"][dataset])
if x_var == "time": if x_var == "time":
x = time x = time
x_label = "Number of turns" x_label = "Number of turns"
bunch_index = data["current"][:,0] != 0 bunch_index = (file["Beam"]["current"][:,0] != 0).nonzero()[0]
if dataset == "current": 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)" y_label = "Total current (mA)"
elif dataset == "emit": elif dataset == "emit":
dimension_dict = {"x":0, "y":1, "s":2} dimension_dict = {"x":0, "y":1, "s":2}
...@@ -76,9 +75,9 @@ def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean", ...@@ -76,9 +75,9 @@ def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean",
"$\\epsilon_{y}$ (m.rad)", "$\\epsilon_{y}$ (m.rad)",
"$\\epsilon_{s}$ (m.rad)"] "$\\epsilon_{s}$ (m.rad)"]
if stat_var == "mean": 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": 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] y_label = stat_var + " " + label[axis]
elif dataset == "mean" or dataset == "std": elif dataset == "mean" or dataset == "std":
dimension_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, 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", ...@@ -88,22 +87,23 @@ def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean",
label = ["x (um)", "x' ($\\mu$rad)", "y (um)", label = ["x (um)", "x' ($\\mu$rad)", "y (um)",
"y' ($\\mu$rad)", "$\\tau$ (ps)", "$\\delta$"] "y' ($\\mu$rad)", "$\\tau$ (ps)", "$\\delta$"]
if stat_var == "mean": 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": 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 "} label_sup = {"mean":"mean of ", "std":"std of "}
y_label = label_sup[stat_var] + dataset + " " + label[axis] y_label = label_sup[stat_var] + dataset + " " + label[axis]
elif x_var == "index": elif x_var == "index":
h = len(data["mean"][0,:,0]) h = len(file["Beam"]["mean"][0,:,0])
x = np.arange(h) x = np.arange(h)
x_label = "Bunch index" x_label = "Bunch index"
if turn is None: if turn is None:
idx = -1 idx = -1
else: else:
idx = np.where(time == int(turn)) idx = np.where(time == int(turn))
if dataset == "current": if dataset == "current":
y = data["current"][:,idx]*1e3 y = data[:,idx]*1e3
y_label = "Bunch current (mA)" y_label = "Bunch current (mA)"
elif dataset == "emit": elif dataset == "emit":
dimension_dict = {"x":0, "y":1, "s":2} dimension_dict = {"x":0, "y":1, "s":2}
...@@ -111,7 +111,7 @@ def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean", ...@@ -111,7 +111,7 @@ def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean",
label = ["$\\epsilon_{x}$ (m.rad)", label = ["$\\epsilon_{x}$ (m.rad)",
"$\\epsilon_{y}$ (m.rad)", "$\\epsilon_{y}$ (m.rad)",
"$\\epsilon_{s}$ (m.rad)"] "$\\epsilon_{s}$ (m.rad)"]
y = data["emit"][axis,:,idx] y = data[axis,:,idx]
y_label = label[axis] y_label = label[axis]
elif dataset == "mean" or dataset == "std": elif dataset == "mean" or dataset == "std":
dimension_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, 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", ...@@ -120,10 +120,12 @@ def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean",
scale = [1e6, 1e6, 1e6, 1e6, 1e12, 1] scale = [1e6, 1e6, 1e6, 1e6, 1e12, 1]
label = ["x (um)", "x' ($\\mu$rad)", "y (um)", label = ["x (um)", "x' ($\\mu$rad)", "y (um)",
"y' ($\\mu$rad)", "$\\tau$ (ps)", "$\\delta$"] "y' ($\\mu$rad)", "$\\tau$ (ps)", "$\\delta$"]
y = data[dataset][axis,:,idx]*scale[axis] y = data[axis,:,idx]*scale[axis]
y_label = dataset + " " + label[axis] y_label = dataset + " " + label[axis]
else: else:
raise ValueError("x_var should be time or index") raise ValueError("x_var should be time or index")
y = np.squeeze(y)
ax.plot(x, y) ax.plot(x, y)
ax.set_xlabel(x_label) ax.set_xlabel(x_label)
...@@ -428,19 +430,21 @@ def plot_profiledata(filename, bunch_number, dimension="tau", start=0, ...@@ -428,19 +430,21 @@ def plot_profiledata(filename, bunch_number, dimension="tau", start=0,
file = hp.File(filename, "r") file = hp.File(filename, "r")
path = file['ProfileData_{0}'.format(bunch_number)] 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: if stop is None:
stop = path['time'][-1] stop = time[-1]
elif stop not in path['time']: elif stop not in time:
raise ValueError("stop not found. Choose from {0}" 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}" 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: if step is None:
step = save_every step = save_every
...@@ -455,9 +459,9 @@ def plot_profiledata(filename, bunch_number, dimension="tau", start=0, ...@@ -455,9 +459,9 @@ def plot_profiledata(filename, bunch_number, dimension="tau", start=0,
"$\\tau$ (ps)", "$\\delta$"] "$\\tau$ (ps)", "$\\delta$"]
num = int((stop - start)/step) 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)) x_var = np.zeros((num+1,n_bin))
turn_index_array = np.zeros((num+1,)) turn_index_array = np.zeros((num+1,))
...@@ -471,16 +475,16 @@ def plot_profiledata(filename, bunch_number, dimension="tau", start=0, ...@@ -471,16 +475,16 @@ def plot_profiledata(filename, bunch_number, dimension="tau", start=0,
fig, ax = plt.subplots() fig, ax = plt.subplots()
for i in range(num+1): for i in range(num+1):
ax.plot(x_var[i]*scale[dimension_dict[dimension]], ax.plot(x_var[i]*scale[dimension_dict[dimension]],
path[dimension][:,turn_index_array[i]], data[:,turn_index_array[i]],
label="turn {0}".format(path['time'][turn_index_array[i]])) label="turn {0}".format(time[turn_index_array[i]]))
ax.set_xlabel(label[dimension_dict[dimension]]) ax.set_xlabel(label[dimension_dict[dimension]])
ax.set_ylabel("number of macro-particles") ax.set_ylabel("number of macro-particles")
ax.legend() ax.legend()
if streak_plot is True: 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 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() fig2, ax2 = plt.subplots()
cmap = mpl.cm.cool cmap = mpl.cm.cool
c = ax2.imshow(z_var, cmap=cmap, origin='lower' , aspect='auto', 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, ...@@ -541,18 +545,19 @@ def plot_wakedata(filename, bunch_number, wake_type="Wlong", start=0,
file = hp.File(filename, "r") file = hp.File(filename, "r")
path = file['WakePotentialData_{0}'.format(bunch_number)] path = file['WakePotentialData_{0}'.format(bunch_number)]
time = np.array(path["time"])
if stop is None: if stop is None:
stop = path['time'][-1] stop = time[-1]
elif stop not in path['time']: elif stop not in time:
raise ValueError("stop not found. Choose from {0}" 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}" 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: if step is None:
step = save_every step = save_every
...@@ -583,33 +588,35 @@ def plot_wakedata(filename, bunch_number, wake_type="Wlong", start=0, ...@@ -583,33 +588,35 @@ def plot_wakedata(filename, bunch_number, wake_type="Wlong", start=0,
else: else:
tau_name = "tau_" + wake_type tau_name = "tau_" + wake_type
data = np.array(path[wake_type])
num = int((stop - start)/step) 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)) 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): 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 turn_index_array[i] = turn_index
# construct an array of bin mids # 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: if profile_plot is True:
fig, ax = plt.subplots() fig, ax = plt.subplots()
for i in range(num+1): for i in range(num+1):
ax.plot(x_var[i]*1e12, ax.plot(x_var[i]*1e12,
path[wake_type][:,turn_index_array[i]]*scale[dimension_dict[wake_type]], data[:,turn_index_array[i]]*scale[dimension_dict[wake_type]],
label="turn {0}".format(path['time'][turn_index_array[i]])) label="turn {0}".format(time[turn_index_array[i]]))
ax.set_xlabel("$\\tau$ (ps)") ax.set_xlabel("$\\tau$ (ps)")
ax.set_ylabel(label[dimension_dict[wake_type]]) ax.set_ylabel(label[dimension_dict[wake_type]])
ax.legend() ax.legend()
if streak_plot is True: 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 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() fig2, ax2 = plt.subplots()
cmap = mpl.cm.cool cmap = mpl.cm.cool
c = ax2.imshow(z_var, cmap=cmap, origin='lower' , aspect='auto', c = ax2.imshow(z_var, cmap=cmap, origin='lower' , aspect='auto',
...@@ -1062,7 +1069,7 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity", ...@@ -1062,7 +1069,7 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity",
file = hp.File(filename, "r") file = hp.File(filename, "r")
cavity_data = file[cavity_name] cavity_data = file[cavity_name]
time = cavity_data["time"] time = np.array(cavity_data["time"])
ph = {"cavity":0, "beam":1} ph = {"cavity":0, "beam":1}
labels = ["Cavity", "Beam"] labels = ["Cavity", "Beam"]
...@@ -1093,8 +1100,8 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity", ...@@ -1093,8 +1100,8 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity",
index = np.array(time) == turn index = np.array(time) == turn
ph = {"cavity":0, "beam":1} ph = {"cavity":0, "beam":1}
data = [cavity_data["cavity_phasor_record"][:,index], data = [np.array(cavity_data["cavity_phasor_record"])[:,index],
cavity_data["beam_phasor_record"][:,index]] np.array(cavity_data["beam_phasor_record"])[:,index]]
labels = ["Cavity", "Beam"] labels = ["Cavity", "Beam"]
h=len(data[0]) h=len(data[0])
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment

The place to be for your source code