Skip to content

Improve setting the ranges in StageHipace.plot_waterfalls() #177

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
ben-c-2013 wants to merge 8 commits into
base: main
Choose a base branch
from
Draft

Improve setting the ranges in StageHipace.plot_waterfalls() #177

Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
7a7d332
Modified StageHipace.plot_waterfalls() to use the estimated drive bea…
ben-c-2013 Dec 11, 2025
fb3170d
Merge branch 'main' into StageHipace_waterfallplots_enhacement
ben-c-2013 Dec 19, 2025
d4076f6
Merge branch 'main' into StageHipace_waterfallplots_enhacement
ben-c-2013 Feb 2, 2026
f154e30
Merge branch 'main' into StageHipace_waterfallplots_enhacement
ben-c-2013 Feb 12, 2026
5457bac
Changed self.get_driver_source() to self.get_actual_driver_source() i…
ben-c-2013 Feb 12, 2026
4623a5e
Merge branch 'main' into StageHipace_waterfallplots_enhacement
ben-c-2013 Mar 3, 2026
9e98c8e
Merge branch 'main' into StageHipace_waterfallplots_enhacement
ben-c-2013 Mar 11, 2026
5ece8e4
Correction due to the update in StageHipace.driver_guiding_trajectory().
ben-c-2013 Mar 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 13 additions & 4 deletions src/abel/classes/stage/impl/stage_hipace.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -727,7 +727,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
Expand Down Expand Up @@ -803,7 +803,7 @@ def plot_waterfalls(self, data_dir, species='beam', remove_halo_nsigma=20, save_


Returns
----------
-------
``None``
'''

Expand All @@ -817,14 +817,23 @@ 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
_, 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)
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
Expand Down
Loading