From 325f658a313219ba70e0074a00202ca81af04048 Mon Sep 17 00:00:00 2001 From: Ole Gunnar Finnerud Date: Tue, 7 Oct 2025 20:42:53 +0300 Subject: [PATCH 1/5] Fixed bug when using plasma_density_from_file --- abel/classes/stage/impl/stage_hipace.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py index 09b2a84c..90480426 100644 --- a/abel/classes/stage/impl/stage_hipace.py +++ b/abel/classes/stage/impl/stage_hipace.py @@ -542,8 +542,20 @@ def __extract_initial_and_final_step(self, tmpfolder, beam0, runnable): def _make_ramp_profile(self, tmpfolder): """Prepare the ramps (local to HiPACE).""" - # check that there is not already a plasma density profile set - assert self.plasma_density_from_file is None + # If there is already a density file open and make the plasma profile + if self.plasma_density_from_file: + ss, ns = [], [] + with open(self.plasma_density_from_file, 'r') as f: + for line in f: + s, n = line.strip().split() + ss.append(float(s)) + ns.append(float(n)) + + ss = np.array(ss) + ns = np.array(ns) + self.plasma_profile.ss = ss + self.plasma_profile.ns = ns + return # make the plasma ramp profile if self.has_ramp(): From 8fb69676fc09552da0cd21f7f398683dccaa3db5 Mon Sep 17 00:00:00 2001 From: Ole Gunnar Finnerud Date: Tue, 7 Oct 2025 20:54:04 +0300 Subject: [PATCH 2/5] Fixed bug in stage.py such that it plots the wake for both first and last step --- abel/classes/stage/stage.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py index 1309a12f..cd2f972e 100644 --- a/abel/classes/stage/stage.py +++ b/abel/classes/stage/stage.py @@ -2089,11 +2089,11 @@ def plot_wake(self, aspect='equal', show_beam=True, savefig=None): # make figures has_final_step = self.final is not None \ - and hasattr(self.final, 'plasma.density.extent') \ - and hasattr(self.final, 'plasma.wakefield.onaxis.zs') \ - and hasattr(self.final, 'plasma.wakefield.onaxis.Ezs') \ - and hasattr(self.final, 'plasma.density.rho') \ - and hasattr(self.final, 'beam.density.rho') + and hasattr(self.final.plasma.density, 'extent') \ + and hasattr(self.final.plasma.wakefield.onaxis, 'zs') \ + and hasattr(self.final.plasma.wakefield.onaxis, 'Ezs') \ + and hasattr(self.final.plasma.density, 'rho') \ + and hasattr(self.final.beam.density, 'rho') num_plots = 1 + int(has_final_step) fig, ax = plt.subplots(num_plots,1) From 7c5b4dc0e90f14e4c99f082c59f866a822b00bba Mon Sep 17 00:00:00 2001 From: Ole Gunnar Finnerud Date: Mon, 13 Oct 2025 14:45:20 +0300 Subject: [PATCH 3/5] Added extracting middle step field information as default --- abel/classes/stage/impl/stage_hipace.py | 23 +++++++++++++++++++--- abel/classes/stage/stage.py | 26 +++++++++++++++++++++++-- abel/wrappers/hipace/hipace_wrapper.py | 2 +- 3 files changed, 45 insertions(+), 6 deletions(-) diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py index 90480426..93ca9b44 100644 --- a/abel/classes/stage/impl/stage_hipace.py +++ b/abel/classes/stage/impl/stage_hipace.py @@ -283,7 +283,9 @@ def track(self, beam_incoming, savedepth=0, runnable=None, verbose=False): dz = beta_matched/20 # convert to number of steps (and re-adjust timestep to be divisible) - self.num_steps = np.ceil(self.length_flattop/dz) + self.num_steps = int(np.ceil(self.length_flattop / dz)) + # Adds 1 to make it even if odd + self.num_steps += self.num_steps % 2 # make it even (add 1 if odd) if self.output is not None: remainder = self.num_steps % self.output @@ -351,7 +353,7 @@ def track(self, beam_incoming, savedepth=0, runnable=None, verbose=False): # extract insitu diagnostics and wakefield data self.__extract_evolution(tmpfolder, beam0, runnable) - self.__extract_initial_and_final_step(tmpfolder, beam0, runnable) + self.__extract_initial_middle_and_final_step(tmpfolder, beam0, runnable) # delete temp folder shutil.rmtree(tmpfolder) @@ -495,7 +497,7 @@ def __extract_evolution(self, tmpfolder, beam0, runnable): shutil.move(insitu_path, destination_path) - def __extract_initial_and_final_step(self, tmpfolder, beam0, runnable): + def __extract_initial_middle_and_final_step(self, tmpfolder, beam0, runnable): from openpmd_viewer import OpenPMDTimeSeries @@ -512,6 +514,11 @@ def __extract_initial_and_final_step(self, tmpfolder, beam0, runnable): Ez, metadata = ts.get_field(field='Ez', slice_across=['x'], iteration=max(ts.iterations)) self.final.plasma.wakefield.onaxis.zs = metadata.z self.final.plasma.wakefield.onaxis.Ezs = Ez + + # extract middle on-axis wakefield + Ez, metadata = ts.get_field(field='Ez', slice_across=['x'], iteration=int(max(ts.iterations)/2)) + self.middle.plasma.wakefield.onaxis.zs = metadata.z + self.middle.plasma.wakefield.onaxis.Ezs = Ez # extract initial plasma density rho0_plasma, metadata0_plasma = ts.get_field(field='rho', iteration=min(ts.iterations)) @@ -523,11 +530,21 @@ def __extract_initial_and_final_step(self, tmpfolder, beam0, runnable): self.initial.beam.density.extent = metadata0_beam.imshow_extent[[2,3,0,1]] self.initial.beam.density.rho = -jz0_beam.T/(SI.c*SI.e) + # extract middle beam density + jz0_beam, metadata0_beam = ts.get_field(field='jz_beam', iteration=int(max(ts.iterations)/2)) + self.middle.beam.density.extent = metadata0_beam.imshow_extent[[2,3,0,1]] + self.middle.beam.density.rho = -jz0_beam.T/(SI.c*SI.e) + # extract initial plasma density rho_plasma, metadata_plasma = ts.get_field(field='rho', iteration=max(ts.iterations)) self.final.plasma.density.extent = metadata_plasma.imshow_extent[[2,3,0,1]] self.final.plasma.density.rho = -(rho_plasma.T/SI.e-self.plasma_density) + # extract initial plasma density + rho_plasma, metadata_plasma = ts.get_field(field='rho', iteration=int(max(ts.iterations)/2)) + self.middle.plasma.density.extent = metadata_plasma.imshow_extent[[2,3,0,1]] + self.middle.plasma.density.rho = -(rho_plasma.T/SI.e-self.plasma_density) + # extract final beam density jz_beam, metadata_beam = ts.get_field(field='jz_beam', iteration=max(ts.iterations)) self.final.beam.density.extent = metadata_beam.imshow_extent[[2,3,0,1]] diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py index cd2f972e..265308f8 100644 --- a/abel/classes/stage/stage.py +++ b/abel/classes/stage/stage.py @@ -166,6 +166,18 @@ def __init__(self, nom_accel_gradient, nom_energy_gain, plasma_density, driver_s self.final.plasma.wakefield = SimpleNamespace() self.final.plasma.wakefield.onaxis = SimpleNamespace() + self.middle = SimpleNamespace() + self.middle.driver = SimpleNamespace() + self.middle.driver.instance = SimpleNamespace() + self.middle.beam = SimpleNamespace() + self.middle.beam.instance = SimpleNamespace() + self.middle.beam.current = SimpleNamespace() + self.middle.beam.density = SimpleNamespace() + self.middle.plasma = SimpleNamespace() + self.middle.plasma.density = SimpleNamespace() + self.middle.plasma.wakefield = SimpleNamespace() + self.middle.plasma.wakefield.onaxis = SimpleNamespace() + self.name = 'Plasma stage' @@ -2095,14 +2107,21 @@ def plot_wake(self, aspect='equal', show_beam=True, savefig=None): and hasattr(self.final.plasma.density, 'rho') \ and hasattr(self.final.beam.density, 'rho') - num_plots = 1 + int(has_final_step) + has_middle_step = self.middle is not None \ + and hasattr(self.middle.plasma.density, 'extent') \ + and hasattr(self.middle.plasma.wakefield.onaxis, 'zs') \ + and hasattr(self.middle.plasma.wakefield.onaxis, 'Ezs') \ + and hasattr(self.middle.plasma.density, 'rho') \ + and hasattr(self.middle.beam.density, 'rho') + + num_plots = 1 + int(has_final_step) + int(has_middle_step) fig, ax = plt.subplots(num_plots,1) fig.set_figwidth(CONFIG.plot_width_default*0.7) fig.set_figheight(CONFIG.plot_width_default*0.5*num_plots) # cycle through initial and final step for i in range(num_plots): - if not has_final_step: + if not (has_final_step or has_middle_step): ax1 = ax else: ax1 = ax[i] @@ -2112,6 +2131,9 @@ def plot_wake(self, aspect='equal', show_beam=True, savefig=None): data_struct = self.initial title = 'Initial step' elif i==1: + data_struct = self.middle + title = 'Middle step' + elif i==2: data_struct = self.final title = 'Final step' diff --git a/abel/wrappers/hipace/hipace_wrapper.py b/abel/wrappers/hipace/hipace_wrapper.py index 117d0019..a7544f7d 100644 --- a/abel/wrappers/hipace/hipace_wrapper.py +++ b/abel/wrappers/hipace/hipace_wrapper.py @@ -18,7 +18,7 @@ def hipace_write_inputs(filename_input, filename_beam, filename_driver, plasma_density, num_steps, time_step, box_range_z, box_size_xy, output_period=None, ion_motion=True, ion_species='H', radiation_reaction=False, beam_ionization=True, num_cell_xy=511, num_cell_z=512, driver_only=False, density_table_file=None, no_plasma=False, external_focusing_gradient=0, mesh_refinement=False, do_spin_tracking=False): if output_period is None: - output_period = int(num_steps) + output_period = int(num_steps/2) # locate template file filename_input_template = os.path.join(os.path.dirname(__file__), 'input_template') From 6fb06febca6866e694c063b9a96806133c7b8231 Mon Sep 17 00:00:00 2001 From: Ole Gunnar Finnerud Date: Thu, 4 Dec 2025 13:49:01 +0200 Subject: [PATCH 4/5] Updated pull request --- abel/classes/stage/impl/stage_hipace.py | 34 +++++++++++++------------ abel/classes/stage/stage.py | 5 +++- 2 files changed, 22 insertions(+), 17 deletions(-) diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py index 93ca9b44..7feda9c2 100644 --- a/abel/classes/stage/impl/stage_hipace.py +++ b/abel/classes/stage/impl/stage_hipace.py @@ -515,36 +515,38 @@ def __extract_initial_middle_and_final_step(self, tmpfolder, beam0, runnable): self.final.plasma.wakefield.onaxis.zs = metadata.z self.final.plasma.wakefield.onaxis.Ezs = Ez - # extract middle on-axis wakefield - Ez, metadata = ts.get_field(field='Ez', slice_across=['x'], iteration=int(max(ts.iterations)/2)) - self.middle.plasma.wakefield.onaxis.zs = metadata.z - self.middle.plasma.wakefield.onaxis.Ezs = Ez + # extract middle field information if middle step exists + if int(max(ts.iterations)) % 2 == 0: + # extract middle on-axis wakefield + Ez, metadata = ts.get_field(field='Ez', slice_across=['x'], iteration=int(max(ts.iterations)/2)) + self.middle.plasma.wakefield.onaxis.zs = metadata.z + self.middle.plasma.wakefield.onaxis.Ezs = Ez + + # extract middle beam density + jz0_beam, metadata0_beam = ts.get_field(field='jz_beam', iteration=int(max(ts.iterations)/2)) + self.middle.beam.density.extent = metadata0_beam.imshow_extent[[2,3,0,1]] + self.middle.beam.density.rho = -jz0_beam.T/(SI.c*SI.e) + + # extract middle plasma density + rho_plasma, metadata_plasma = ts.get_field(field='rho', iteration=int(max(ts.iterations)/2)) + self.middle.plasma.density.extent = metadata_plasma.imshow_extent[[2,3,0,1]] + self.middle.plasma.density.rho = -(rho_plasma.T/SI.e-self.plasma_density) # extract initial plasma density rho0_plasma, metadata0_plasma = ts.get_field(field='rho', iteration=min(ts.iterations)) self.initial.plasma.density.extent = metadata0_plasma.imshow_extent[[2,3,0,1]] self.initial.plasma.density.rho = -(rho0_plasma.T/SI.e-self.plasma_density) - # extract final beam density + # extract initial beam density jz0_beam, metadata0_beam = ts.get_field(field='jz_beam', iteration=min(ts.iterations)) self.initial.beam.density.extent = metadata0_beam.imshow_extent[[2,3,0,1]] self.initial.beam.density.rho = -jz0_beam.T/(SI.c*SI.e) - # extract middle beam density - jz0_beam, metadata0_beam = ts.get_field(field='jz_beam', iteration=int(max(ts.iterations)/2)) - self.middle.beam.density.extent = metadata0_beam.imshow_extent[[2,3,0,1]] - self.middle.beam.density.rho = -jz0_beam.T/(SI.c*SI.e) - - # extract initial plasma density + # extract final plasma density rho_plasma, metadata_plasma = ts.get_field(field='rho', iteration=max(ts.iterations)) self.final.plasma.density.extent = metadata_plasma.imshow_extent[[2,3,0,1]] self.final.plasma.density.rho = -(rho_plasma.T/SI.e-self.plasma_density) - # extract initial plasma density - rho_plasma, metadata_plasma = ts.get_field(field='rho', iteration=int(max(ts.iterations)/2)) - self.middle.plasma.density.extent = metadata_plasma.imshow_extent[[2,3,0,1]] - self.middle.plasma.density.rho = -(rho_plasma.T/SI.e-self.plasma_density) - # extract final beam density jz_beam, metadata_beam = ts.get_field(field='jz_beam', iteration=max(ts.iterations)) self.final.beam.density.extent = metadata_beam.imshow_extent[[2,3,0,1]] diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py index 265308f8..542ec8e4 100644 --- a/abel/classes/stage/stage.py +++ b/abel/classes/stage/stage.py @@ -2130,9 +2130,12 @@ def plot_wake(self, aspect='equal', show_beam=True, savefig=None): if i==0: data_struct = self.initial title = 'Initial step' - elif i==1: + elif i==1 & has_middle_step == True: data_struct = self.middle title = 'Middle step' + elif i==1 & has_final_step == True: + data_struct = self.final + title = 'Final step' elif i==2: data_struct = self.final title = 'Final step' From 7bb6ad42aea2d0d9073a3a20a8d4a284e37981f5 Mon Sep 17 00:00:00 2001 From: Ole Gunnar Finnerud Date: Tue, 9 Dec 2025 11:48:05 +0200 Subject: [PATCH 5/5] Changed it on request to always pick the closest to middle and display the position in plot_wake --- abel/classes/stage/impl/stage_hipace.py | 20 ++-- abel/classes/stage/stage.py | 138 +++++------------------- 2 files changed, 39 insertions(+), 119 deletions(-) diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py index 7feda9c2..6e46470c 100644 --- a/abel/classes/stage/impl/stage_hipace.py +++ b/abel/classes/stage/impl/stage_hipace.py @@ -514,24 +514,30 @@ def __extract_initial_middle_and_final_step(self, tmpfolder, beam0, runnable): Ez, metadata = ts.get_field(field='Ez', slice_across=['x'], iteration=max(ts.iterations)) self.final.plasma.wakefield.onaxis.zs = metadata.z self.final.plasma.wakefield.onaxis.Ezs = Ez - - # extract middle field information if middle step exists - if int(max(ts.iterations)) % 2 == 0: + print(ts.iterations) + # extract middle field information if there are more than two steps + if len(ts.iterations) > 2: + # Get step closest to middle + middle_step = ts.iterations[np.argmin(abs(ts.iterations - int(max(ts.iterations)/2)))] + # extract middle on-axis wakefield - Ez, metadata = ts.get_field(field='Ez', slice_across=['x'], iteration=int(max(ts.iterations)/2)) + Ez, metadata = ts.get_field(field='Ez', slice_across=['x'], iteration=middle_step) self.middle.plasma.wakefield.onaxis.zs = metadata.z self.middle.plasma.wakefield.onaxis.Ezs = Ez # extract middle beam density - jz0_beam, metadata0_beam = ts.get_field(field='jz_beam', iteration=int(max(ts.iterations)/2)) + jz0_beam, metadata0_beam = ts.get_field(field='jz_beam', iteration=middle_step) self.middle.beam.density.extent = metadata0_beam.imshow_extent[[2,3,0,1]] self.middle.beam.density.rho = -jz0_beam.T/(SI.c*SI.e) # extract middle plasma density - rho_plasma, metadata_plasma = ts.get_field(field='rho', iteration=int(max(ts.iterations)/2)) + rho_plasma, metadata_plasma = ts.get_field(field='rho', iteration=middle_step) self.middle.plasma.density.extent = metadata_plasma.imshow_extent[[2,3,0,1]] self.middle.plasma.density.rho = -(rho_plasma.T/SI.e-self.plasma_density) - + + # Extract step + self.middle.n_step = middle_step + print(self.middle.n_step) # extract initial plasma density rho0_plasma, metadata0_plasma = ts.get_field(field='rho', iteration=min(ts.iterations)) self.initial.plasma.density.extent = metadata0_plasma.imshow_extent[[2,3,0,1]] diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py index 542ec8e4..b46ab7f2 100644 --- a/abel/classes/stage/stage.py +++ b/abel/classes/stage/stage.py @@ -92,7 +92,6 @@ class Stage(Trackable, CostModeled): """ - # ================================================== @abstractmethod def __init__(self, nom_accel_gradient, nom_energy_gain, plasma_density, driver_source=None, ramp_beta_mag=None): """ @@ -181,7 +180,6 @@ def __init__(self, nom_accel_gradient, nom_energy_gain, plasma_density, driver_s self.name = 'Plasma stage' - # ================================================== @property def driver_source(self) -> Source | DriverComplex | None: "Get or set the the driver source or the driver complex of the stage." @@ -204,17 +202,12 @@ def get_driver_source(self): """ if isinstance(self.driver_source, DriverComplex): - driver_source = self.driver_source.source - elif isinstance(self.driver_source, Source): - driver_source = self.driver_source - - return driver_source - + return self.driver_source.source + else: + return self.driver_source ## Define upramp and downramp, if present - - # ================================================== @property def upramp(self) -> Self | None: "Get or set the upramp." @@ -234,7 +227,6 @@ def upramp(self, upramp : Self | None): _upramp = None - # ================================================== @property def downramp(self) -> Self: "Get or set the downramp." @@ -254,7 +246,6 @@ def downramp(self, downramp : Self | None): _downramp = None - # ================================================== @property def ramp_beta_mag(self) -> Self: """ @@ -279,7 +270,6 @@ def ramp_beta_mag(self, ramp_beta_mag : float | None): _ramp_beta_mag = None - # ================================================== def copy_config2blank_stage(self): """ Return a deep copy of the stage with most physics parameters reset to @@ -358,7 +348,6 @@ def copy_config2blank_stage(self): return stage_copy - # ================================================== def convert_PlasmaRamp(self, ramp): """ Convert a ``PlasmaRamp`` object to a ``Stage`` subclass object of the @@ -413,7 +402,6 @@ def convert_PlasmaRamp(self, ramp): return trackable_ramp - # ================================================== def _prepare_ramps(self): """ Prepare upramp and downramp parameters before tracking. @@ -465,7 +453,6 @@ def _prepare_ramps(self): self.downramp.length_flattop = self._calc_ramp_length(self.downramp) - # ================================================== def is_upramp(self): "Check if this stage is the upramp of its parent." @@ -478,7 +465,6 @@ def is_upramp(self): return False - # ================================================== def is_downramp(self): "Check if this stage is the downramp of its parent" @@ -491,7 +477,6 @@ def is_downramp(self): return False - # ================================================== def has_ramp(self): "Check if there are any ramps attached to this stage." @@ -501,7 +486,6 @@ def has_ramp(self): return False - # ================================================== def _calc_ramp_length(self, ramp : Self) -> float: "Calculate and set the up/down ramp (uniform step ramp) length [m] based on stage nominal energy." if ramp.nom_energy is None: @@ -520,7 +504,6 @@ def _calc_ramp_length(self, ramp : Self) -> float: return ramp_length - # ================================================== @property def parent(self) -> Self | None: "Get or set the parent stage of a ramp." @@ -535,7 +518,6 @@ def parent(self, parent : Self): _parent = None - # ================================================== def _getOtherRamp(self, aRamp : Self) -> Self: "Return the opposite ramp in the same pair." if aRamp == self.upramp: @@ -546,7 +528,6 @@ def _getOtherRamp(self, aRamp : Self) -> Self: raise StageError("Could not find calling ramp?") - # ================================================== def _getOverallestStage(self) -> Self: "Find the top-level parent stage in the hierarchy." bottom_Stage = self @@ -563,7 +544,6 @@ def _getOverallestStage(self) -> Self: ## Tracking methods - # ================================================== @abstractmethod def track(self, beam, savedepth=0, runnable=None, verbose=False): """ @@ -647,7 +627,6 @@ def track(self, beam, savedepth=0, runnable=None, verbose=False): # Methods for setting and getting the variables - # ================================================== @property def length(self) -> float: "Total length of the trackable stage element [m], or None if not set/calculateable" @@ -688,7 +667,6 @@ def get_length(self) -> float: return self.length - # ================================================== @property def nom_energy_gain(self) -> float: "Total nominal energy gain of the stage [eV], or None if not set/calculateable" @@ -717,7 +695,6 @@ def get_nom_energy_gain(self): return self.nom_energy_gain - # ================================================== @property def nom_accel_gradient(self) -> float: "Total nominal accelerating gradient of the stage [eV/m], or None if not set/calculateable" @@ -741,7 +718,6 @@ def nom_accel_gradient(self, nom_accel_gradient : float): _nom_accel_gradient_calc = None - # ================================================== @property def length_flattop(self) -> float: "Length of the plasma flattop [m], or None if not set/calculateable" @@ -767,7 +743,6 @@ def length_flattop(self, length_flattop : float): _length_flattop = None _length_flattop_calc = None - # ================================================== @property def nom_energy_gain_flattop(self) -> float: @@ -791,7 +766,6 @@ def nom_energy_gain_flattop(self,nom_energy_gain_flattop : float): _nom_energy_gain_flattop = None _nom_energy_gain_flattop_calc = None - # ================================================== @property def nom_accel_gradient_flattop(self) -> float: @@ -816,7 +790,6 @@ def nom_accel_gradient_flattop(self, nom_accel_gradient_flattop : float): _nom_accel_gradient_flattop_calc = None - # ================================================== @property def nom_energy(self) -> float: "Nominal energy for the tracking [eV], or ``None`` if not set/calculateable" @@ -843,7 +816,6 @@ def nom_energy(self, nom_energy : float): _nom_energy_calc = None - # ================================================== @property def nom_energy_flattop(self) -> float: "Nominal energy in the flattop for the tracking [eV], or ``None`` if not set/calculateable" @@ -870,15 +842,12 @@ def nom_energy_flattop(self, nom_energy_flattop : float): _nom_energy_flattop_calc = None - # ================================================== ## Recalculation methods - # ================================================== def _resetLengthEnergyGradient(self): "Reset all the calculated values in the current Stage hierarchy" self._getOverallestStage()._resetLengthEnergyGradient_helper() - # ================================================== def _resetLengthEnergyGradient_helper(self): #Climb back up from the bottom and set what we can # directly / to None, from the stored user input @@ -910,7 +879,6 @@ def _resetLengthEnergyGradient_helper(self): self.downramp._resetLengthEnergyGradient_helper() - # ================================================== def _recalcLengthEnergyGradient(self): #Iteratively calculate everything until stability is reached #Note: Before starting calculation, call _resetLengthEnergyGradient() to reset the hierachy @@ -918,7 +886,6 @@ def _recalcLengthEnergyGradient(self): self._printVerb() - # ================================================== def _recalcLengthEnergyGradient_helper(self): itrCtr = 0 updateCounter_total = 0 @@ -1244,7 +1211,6 @@ def _recalcLengthEnergyGradient_helper(self): return updateCounter_total - # ================================================== doVerbosePrint_debug = False def _printVerb(self, *args, **kwargs): "Print() if doVerbosePrint_debug == True, else NOP." @@ -1252,7 +1218,6 @@ def _printVerb(self, *args, **kwargs): print(*args, **kwargs) - # ================================================== def _printLengthEnergyGradient_internal(stage): "For debugging" print("parent/upramp/downramp: ", stage.parent, stage.upramp, stage.downramp) @@ -1268,7 +1233,6 @@ def _printLengthEnergyGradient_internal(stage): ## Various calculations / plots / etc - # ================================================== def get_cost_breakdown(self): """ Get the cost breakdown of the plasma stage. @@ -1279,7 +1243,6 @@ def get_cost_breakdown(self): return (self.name, breakdown) - # ================================================== def matched_beta_function(self, energy_incoming, match_entrance=True): ''' Calculates the matched beta function of the stage. If there is an @@ -1314,7 +1277,6 @@ def matched_beta_function(self, energy_incoming, match_entrance=True): raise ValueError('Downramp ramp_beta_mag not defined.') - # ================================================== def matched_beta_function_flattop(self, energy): ''' Calculates the matched beta function of the flattop stage. @@ -1334,17 +1296,14 @@ def matched_beta_function_flattop(self, energy): return beta_matched(self.plasma_density, energy) - # ================================================== def energy_usage(self): return self.driver_source.energy_usage() - # ================================================== def energy_efficiency(self): return self.efficiency - # ================================================== #@abstractmethod # TODO: calculate the dumped power and use it for the dump cost model. def dumped_power(self): return self.efficiency.dumped_power @@ -1352,37 +1311,17 @@ def dumped_power(self): # ================================================== def calculate_efficiency(self, beam0, driver0, beam, driver): - """ - Calculate the efficiency of energy transfer in the stage. - - This method computes the following efficiency metrics: - - - ``driver_to_wake``: Fraction of driver energy transferred to the wakefield. - - ``wake_to_beam``: Fraction of wakefield energy transferred to the beam. - - ``driver_to_beam``: Overall efficiency from driver to beam (product of the above). - - ``dumped_power``: Remaining driver power dumped after interaction, averaged over - the repetition rate if available. - - Parameters - ---------- - beam0 : ``Beam`` - Input beam before the stage. - - driver0 : ``Beam`` - Input drive beam before the stage. - - beam : ``Beam`` - Output beam after the stage. - - driver : ``Beam`` - Output drive beam after the stage. - - Returns - ------- - ``None`` - Results are stored in ``self.efficiency``. - """ + # abort if no driver or beam + if driver0 is None or beam0 is None: + return + # Added average which assumes no charge loss + driver0_avg = driver0.energy() + driver_avg = driver.energy() + beam0_avg = beam0.energy() + beama_avg = beam.energy() + self.efficiency.wake_to_beam_avg = (beam.charge()*(beama_avg-beam0_avg))/(driver.charge()*(driver0_avg - driver_avg)) + Etot0_beam = beam0.total_energy() Etot_beam = beam.total_energy() Etot0_driver = driver0.total_energy() @@ -1398,38 +1337,10 @@ def calculate_efficiency(self, beam0, driver0, beam, driver): # ================================================== def calculate_beam_current(self, beam0, driver0, beam=None, driver=None): - """ - Calculate and store the beam current profile. - - This method computes the current profile of the initial (input) and - optionally final (output) beams using a temporal binning scheme based on - the RMS bunch lengths of the driver and beam. The results are stored in - ``self.initial.beam.current`` and, if output beams are provided, in - ``self.final.beam.current``. - - Parameters - ---------- - beam0 : ``Beam`` - Input beam before the stage. - driver0 : ``Beam`` - Input drive beam before the stage. - - beam : ``Beam``, optional - Output beam after the stage. - - driver : ``Beam``, optional - Output drive beam after the stage. - - Returns - ------- - ``None`` - Results are stored in ``self.initial.beam.current`` and, if provided, - ``self.final.beam.current`` with attributes: - - ``zs`` : longitudinal positions [m] - - ``Is`` : beam current profile [A] - """ - + if driver0 is None: + driver0 = copy.deepcopy(beam0) + driver0.scale_charge(beam0.charge()*1e-15) dz = 40*np.mean([driver0.bunch_length(clean=True)/np.sqrt(len(driver0)), beam0.bunch_length(clean=True)/np.sqrt(len(beam0))]) num_sigmas = 6 z_min = beam0.z_offset() - num_sigmas * beam0.bunch_length() @@ -1605,7 +1516,7 @@ def undo_beam_coordinate_systems_rotation(self, driver_incoming, driver_outgoing # ================================================== def save_driver_to_file(self, driver, runnable): - driver.save(runnable, beam_name='driver_stage' + str(driver.stage_number+1)) + driver.save(runnable, beam_name = 'driver') # ================================================== @@ -2059,7 +1970,7 @@ def plot_final_wakefield(self): # ================================================== # plot wake - def plot_wake(self, aspect='equal', show_beam=True, savefig=None): + def plot_wake(self, aspect='auto', show_beam=True, savefig=None): """ Plot the wake structure (2D plot) as a new pyplot.figure. @@ -2126,14 +2037,16 @@ def plot_wake(self, aspect='equal', show_beam=True, savefig=None): else: ax1 = ax[i] - # extract initial or final + # extract initial, final and middle if i==0: data_struct = self.initial title = 'Initial step' - elif i==1 & has_middle_step == True: + # If middle step exists, extract it + elif i==1 and has_middle_step == True: data_struct = self.middle - title = 'Middle step' - elif i==1 & has_final_step == True: + title = f"s = {self.evolution.beam.location[self.middle.n_step]:.2f} m" + # If middle step does not exist, simply plot final step + elif i==1 and has_middle_step == False: data_struct = self.final title = 'Final step' elif i==2: @@ -2161,7 +2074,7 @@ def plot_wake(self, aspect='equal', show_beam=True, savefig=None): ax2 = ax1.twinx() ax2.plot(zs0*1e6, Ezs0/1e9, color = 'black') ax2.set_ylabel(r'$E_{z}$' ' [GV/m]') - ax2.set_ylim(bottom=-Ezmax/1e9, top=Ezmax/1e9) + #ax2.set_ylim(bottom=-Ezmax/1e9, top=Ezmax/1e9) axpos = ax1.get_position() pad_fraction = 0.13 # Fraction of the figure width to use as padding between the ax and colorbar cbar_width_fraction = 0.015 # Fraction of the figure width for the colorbar width @@ -2205,6 +2118,7 @@ def plot_wake(self, aspect='equal', show_beam=True, savefig=None): ax1.set_title(title) ax1.grid(False) ax2.grid(False) + #ax1.set_ylim(0, 300) # save the figure if savefig is not None: