From 0ac697baae6f0c4c9f06b9ca5e979328a947f49c Mon Sep 17 00:00:00 2001
From: Watanyu Foosang <watanyu.f@gmail.com>
Date: Thu, 4 Jun 2020 18:34:41 +0200
Subject: [PATCH] Add plot_profiledata function
A new function for plotting profile monitor has been added. Two types of plot can be obtained : bunch profile evolution and streak plot.
---
tracking/monitors/plotting.py | 162 ++++++++++++++++++++++++++++++++--
1 file changed, 155 insertions(+), 7 deletions(-)
diff --git a/tracking/monitors/plotting.py b/tracking/monitors/plotting.py
index f28b0df..34486e7 100644
--- a/tracking/monitors/plotting.py
+++ b/tracking/monitors/plotting.py
@@ -9,12 +9,13 @@ tracking.
import numpy as np
import matplotlib.pyplot as plt
+import matplotlib as mpl
import seaborn as sns
import h5py as hp
def plot_beamdata(filename, dataset, option=None, stat_var=None, x_var="time"):
"""
- Plot data recorded from a BeamMonitor.
+ Plot data recorded by BeamMonitor.
Parameters
----------
@@ -112,15 +113,15 @@ def plot_beamdata(filename, dataset, option=None, stat_var=None, x_var="time"):
def plot_bunchdata(filename, bunch_number, dataset, option=None, x_var="time"):
"""
- Plot data recorded from a BunchMonitor.
+ Plot data recorded by BunchMonitor.
Parameters
----------
filename : str
Name of the HDF5 file that contains the data.
bunch_number : int
- The bunch number whose data has been saved in the HDF5 file.
- This has to be identical to 'bunch_number' parameter in 'BunchMonitor' object.
+ Bunch to plot. This has to be identical to 'bunch_number' parameter in
+ 'BunchMonitor' object.
detaset : str {"current","emit","mean","std"}
HDF5 file's dataset to be plotted.
option : str, optional
@@ -177,15 +178,14 @@ def plot_bunchdata(filename, bunch_number, dataset, option=None, x_var="time"):
def plot_phasespacedata(filename, bunch_number, x_var, y_var, turn,
only_alive=True):
"""
- Plot data recorded from a PhaseSpaceMonitor.
+ Plot data recorded by PhaseSpaceMonitor.
Parameters
----------
filename : str
Name of the HDF5 file that contains the data.
bunch_number : int
- Number of the bunch whose data has been saved in the HDF5 file.
- This has to be identical to 'bunch_number' parameter in
+ Bunch to plot. This has to be identical to 'bunch_number' parameter in
'PhaseSpaceMonitor' object.
x_var, y_var : str {"x", "xp", "y", "yp", "tau", "delta"}
If dataset is "particles", the variables to be plotted on the
@@ -240,3 +240,151 @@ 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):
+ """
+ Plot data recorded by ProfileMonitor
+
+ Parameters
+ ----------
+ filename : str
+ Name of the HDF5 file that contains the data.
+ bunch_number : int
+ Bunch to plot. This has to be identical to 'bunch_number' parameter in
+ 'ProfileMonitor' object.
+ stop : int
+ Last turn to plot.
+ start : int, optional
+ First turn to plot. The default is 0.
+ 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
+ equivalent to save_every.
+ profile_plot : bool, optional
+ If Ture, bunch profile plot is plotted.
+ streak_plot : bool, optional
+ If True, strek plot is plotted.
+
+ Returns
+ -------
+ fig : Figure
+ Figure object with the plot on it.
+
+ """
+
+ 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:
+ raise ValueError("stop not found. Choose from {0}"
+ .format(path['time'][:]))
+
+ if start in path['time']:
+ pass
+ else:
+ raise ValueError("start not found. Choose from {0}"
+ .format(path['time'][:]))
+
+ save_every = path['time'][1] - path['time'][0]
+
+ if step is None:
+ step = save_every
+ else:
+ pass
+
+ if step % save_every != 0:
+ raise ValueError("step must be divisible by the recording step "
+ "which is {0}.".format(save_every))
+
+ dimension_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, "delta":5}
+ scale = [1e6, 1e6, 1e6, 1e6, 1e12, 1]
+ label = ["x (um)", "x' ($\\mu$rad)", "y (um)", "y' ($\\mu$rad)",
+ "$\\tau$ (ps)", "$\\delta$"]
+
+ num = int((stop - start)/step)
+ n_bin = len(path[dimension][:,0])
+
+ start_index = np.where(path['time'][:] == start)[0]
+
+ 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.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))
+
+ 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()])
+ 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
+
+ file.close()
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
\ No newline at end of file
--
GitLab