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
Branches
Tags
No related merge requests found
Show 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",
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,11 +120,13 @@ 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)
ax.set_ylabel(y_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])
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment

The place to be for your source code