diff --git a/tracking/monitors/plotting.py b/tracking/monitors/plotting.py index 34486e7986f8e28fb725dc7b0b4b404428eae2b5..d77049ac0a894460e5c07bc7832b5996ff08121d 100644 --- a/tracking/monitors/plotting.py +++ b/tracking/monitors/plotting.py @@ -241,8 +241,8 @@ def plot_phasespacedata(filename, bunch_number, x_var, y_var, turn, file.close() return fig -def plot_profiledata(filename, bunch_number, stop, start=0, step=None, - profile_plot=True, streak_plot=True): +def plot_profiledata(filename, bunch_number, dimension="tau", start=0, + stop=None, step=None, profile_plot=True, streak_plot=True): """ Plot data recorded by ProfileMonitor @@ -253,10 +253,12 @@ def plot_profiledata(filename, bunch_number, stop, start=0, step=None, bunch_number : int Bunch to plot. This has to be identical to 'bunch_number' parameter in 'ProfileMonitor' object. - stop : int - Last turn to plot. + dimension : str, optional + Dimension to plot. The default is "tau" start : int, optional First turn to plot. The default is 0. + stop : int, optional + Last turn to plot. If None, the last turn of the record is selected. step : int, optional Plotting step. This has to be divisible by 'save_every' parameter in 'ProfileMonitor' object, i.e. step % save_every == 0. If None, step is @@ -275,18 +277,15 @@ def plot_profiledata(filename, bunch_number, stop, start=0, step=None, file = hp.File(filename, "r") path = file['ProfileData_{0}'.format(bunch_number)] - dimension = list(path)[0] l_bound = path["{0}_bin".format(dimension)] - if stop in path['time']: - pass - else: + if stop is None: + stop = path['time'][-1] + elif stop not in path['time']: raise ValueError("stop not found. Choose from {0}" .format(path['time'][:])) - if start in path['time']: - pass - else: + if start not in path['time']: raise ValueError("start not found. Choose from {0}" .format(path['time'][:])) @@ -294,8 +293,6 @@ def plot_profiledata(filename, bunch_number, stop, start=0, step=None, if step is None: step = save_every - else: - pass if step % save_every != 0: raise ValueError("step must be divisible by the recording step " @@ -309,66 +306,49 @@ def plot_profiledata(filename, bunch_number, stop, start=0, step=None, num = int((stop - start)/step) n_bin = len(path[dimension][:,0]) - start_index = np.where(path['time'][:] == start)[0] + start_index = np.where(path['time'][:] == start)[0][0] + x_var = np.zeros((num+1,n_bin)) + turn_index_array = np.zeros((num+1,)) + for i in range(num+1): + turn_index = start_index + i * step / save_every + turn_index_array[i] = turn_index + # construct an array of bin mids + x_var[i,:] = 0.5*(l_bound[:-1,turn_index]+l_bound[1:,turn_index]) + if profile_plot is True: fig, ax = plt.subplots() - turn = start for i in range(num+1): - index = start_index + i * step / save_every - x_var = np.zeros((n_bin,)) - for j in range (n_bin): - # constructing an array of midpoints - x_var[j] = l_bound[j,i] + 0.5*(l_bound[j+1,i]-l_bound[j,i]) - ax.plot(x_var*scale[dimension_dict[dimension]], - path[dimension][:,index], label="turn {0}".format(turn)) - turn = turn + step - + 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]])) ax.set_xlabel(label[dimension_dict[dimension]]) ax.set_ylabel("number of macro-particles") ax.legend() - - return fig - - else: - pass - + if streak_plot is True: - X = np.zeros((num+1,n_bin)) - Y = np.zeros((num+1,n_bin)) - Z = np.zeros((num+1,n_bin)) - + turn = np.reshape(path['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]) fig2, ax2 = plt.subplots() - turn = start - for i in range(num+1): - index = start_index + i * step / save_every - x_var = np.zeros((n_bin,)) - for j in range (n_bin): - # constructing an array of midpoints - x_var[j] = l_bound[j,i] + 0.5*(l_bound[j+1,i]-l_bound[j,i]) - - X[i,:] = x_var - Y[i,:] = np.ones((1,n_bin)) * turn - Z[i,:] = np.reshape(path[dimension][:,index], (1,n_bin)) - - turn = turn + step - cmap = mpl.cm.cool - c = ax2.imshow(Z, cmap=cmap, origin='lower' , aspect='auto', - extent=[X.min()*1e12,X.max()*1e12,Y.min(),Y.max()]) + c = ax2.imshow(z_var, cmap=cmap, origin='lower' , aspect='auto', + extent=[x_var.min()*scale[dimension_dict[dimension]], + x_var.max()*scale[dimension_dict[dimension]], + y_var.min(),y_var.max()]) ax2.set_xlabel(label[dimension_dict[dimension]]) ax2.set_ylabel("Number of turns") - cbar = fig2.colorbar(c, ax=ax2) - cbar.set_label("Number of macro-particles") - - return fig2 - - else: - pass + cbar.set_label("Number of macro-particles") file.close() - + if profile_plot is True and streak_plot is True: + return fig, fig2 + elif profile_plot is True: + return fig + elif streak_plot is True: + return fig2 +