From 7a7d332249079fd8ad9a07f4b4e6e8c8bb4fc126 Mon Sep 17 00:00:00 2001 From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com> Date: Thu, 11 Dec 2025 17:16:30 +0100 Subject: [PATCH 1/3] Modified StageHipace.plot_waterfalls() to use the estimated drive beam trajectory to set the range of the plots. Requires StageHipace.driver_guiding_trajectory() in the branch num_beta_osc_stage_length. --- abel/classes/stage/impl/stage_hipace.py | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py index c62046ff..be89b648 100644 --- a/abel/classes/stage/impl/stage_hipace.py +++ b/abel/classes/stage/impl/stage_hipace.py @@ -717,7 +717,7 @@ def __waterfall_fcn(self, fcns, edges, data_dir, species='beam', remove_halo_nsi Returns - ---------- + ------- waterfalls : list of 2D float ndarrays Each element in ``waterfalls`` corresponds to the output of one function in ``fcns`` applied across all files (i.e., simulation @@ -793,7 +793,7 @@ def plot_waterfalls(self, data_dir, species='beam', remove_halo_nsigma=20, save_ Returns - ---------- + ------- ``None`` ''' @@ -807,14 +807,24 @@ def plot_waterfalls(self, data_dir, species='beam', remove_halo_nsigma=20, save_ num_bins = int(np.sqrt(len(beam0)*2)) nsig = 5 + # Set up the bins for energy and time if species == 'driver': deltaedges = np.linspace(-0.5, 0.5, num_bins) else: deltaedges = np.linspace(-0.05, 0.05, num_bins) tedges = (beam0.z_offset(clean=True) + nsig*beam0.bunch_length(clean=True)*np.linspace(-1, 1, num_bins)) / SI.c - xedges = (nsig*beam0.beam_size_x() + abs(beam0.x_offset()))*np.linspace(-1, 1, num_bins) - yedges = (nsig*beam0.beam_size_y() + abs(beam0.y_offset()))*np.linspace(-1, 1, num_bins) - + + # Set up the bins for x and y based on drive beam trajectory + driver_source = self.get_driver_source() + _, x_trajectory, y_trajectory = self.driver_guiding_trajectory(driver_source.track(), dacc_gradient=0.0e9, num_steps_per_half_osc=100) + xedge_min = np.min(x_trajectory) - nsig*beam0.beam_size_x() + xedge_max = nsig*beam0.beam_size_x() + np.max(x_trajectory) + xedges = np.linspace(xedge_min, xedge_max, num_bins) + yedge_min = np.min(y_trajectory) - nsig*beam0.beam_size_y() + yedge_max = nsig*beam0.beam_size_y() + np.max(y_trajectory) + yedges = np.linspace(yedge_min, yedge_max, num_bins) + + # Caculate the waterfalls waterfalls, locations, bins = self.__waterfall_fcn([Beam.current_profile, Beam.rel_energy_spectrum, Beam.transverse_profile_x, Beam.transverse_profile_y], [tedges, deltaedges, xedges, yedges], data_dir, species=species, remove_halo_nsigma=remove_halo_nsigma, args=[None, None, None, None]) # prepare figure From 5457bacd98d213049d1b0b7f03682e6c6641ec62 Mon Sep 17 00:00:00 2001 From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com> Date: Thu, 12 Feb 2026 15:27:35 +0100 Subject: [PATCH 2/3] Changed self.get_driver_source() to self.get_actual_driver_source() in StageHipace.plot_waterfalls(). --- abel/classes/stage/impl/stage_hipace.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py index be89b648..3228a5ed 100644 --- a/abel/classes/stage/impl/stage_hipace.py +++ b/abel/classes/stage/impl/stage_hipace.py @@ -815,7 +815,7 @@ def plot_waterfalls(self, data_dir, species='beam', remove_halo_nsigma=20, save_ tedges = (beam0.z_offset(clean=True) + nsig*beam0.bunch_length(clean=True)*np.linspace(-1, 1, num_bins)) / SI.c # Set up the bins for x and y based on drive beam trajectory - driver_source = self.get_driver_source() + driver_source = self.get_actual_driver_source() _, x_trajectory, y_trajectory = self.driver_guiding_trajectory(driver_source.track(), dacc_gradient=0.0e9, num_steps_per_half_osc=100) xedge_min = np.min(x_trajectory) - nsig*beam0.beam_size_x() xedge_max = nsig*beam0.beam_size_x() + np.max(x_trajectory) From 5ece8e423aa75c33e10ebdb6f4d59924d9ac9aa9 Mon Sep 17 00:00:00 2001 From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com> Date: Wed, 11 Mar 2026 21:07:58 +0100 Subject: [PATCH 3/3] Correction due to the update in StageHipace.driver_guiding_trajectory(). --- src/abel/classes/stage/impl/stage_hipace.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/abel/classes/stage/impl/stage_hipace.py b/src/abel/classes/stage/impl/stage_hipace.py index 7054cae7..04f52e47 100644 --- a/src/abel/classes/stage/impl/stage_hipace.py +++ b/src/abel/classes/stage/impl/stage_hipace.py @@ -825,8 +825,7 @@ def plot_waterfalls(self, data_dir, species='beam', remove_halo_nsigma=20, save_ tedges = (beam0.z_offset(clean=True) + nsig*beam0.bunch_length(clean=True)*np.linspace(-1, 1, num_bins)) / SI.c # Set up the bins for x and y based on drive beam trajectory - driver_source = self.get_actual_driver_source() - _, x_trajectory, y_trajectory = self.driver_guiding_trajectory(driver_source.track(), dacc_gradient=0.0e9, num_steps_per_half_osc=100) + _, x_trajectory, y_trajectory = self.driver_guiding_trajectory(num_steps=None, dacc_gradient=0.0e9) xedge_min = np.min(x_trajectory) - nsig*beam0.beam_size_x() xedge_max = nsig*beam0.beam_size_x() + np.max(x_trajectory) xedges = np.linspace(xedge_min, xedge_max, num_bins)