diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py index 09b2a84c..6e46470c 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,18 +514,41 @@ 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 - + 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=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=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=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]] 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 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) @@ -542,8 +567,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(): diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py index 1309a12f..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): """ @@ -166,10 +165,21 @@ 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' - # ================================================== @property def driver_source(self) -> Source | DriverComplex | None: "Get or set the the driver source or the driver complex of the stage." @@ -192,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." @@ -222,7 +227,6 @@ def upramp(self, upramp : Self | None): _upramp = None - # ================================================== @property def downramp(self) -> Self: "Get or set the downramp." @@ -242,7 +246,6 @@ def downramp(self, downramp : Self | None): _downramp = None - # ================================================== @property def ramp_beta_mag(self) -> Self: """ @@ -267,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 @@ -346,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 @@ -401,7 +402,6 @@ def convert_PlasmaRamp(self, ramp): return trackable_ramp - # ================================================== def _prepare_ramps(self): """ Prepare upramp and downramp parameters before tracking. @@ -453,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." @@ -466,7 +465,6 @@ def is_upramp(self): return False - # ================================================== def is_downramp(self): "Check if this stage is the downramp of its parent" @@ -479,7 +477,6 @@ def is_downramp(self): return False - # ================================================== def has_ramp(self): "Check if there are any ramps attached to this stage." @@ -489,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: @@ -508,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." @@ -523,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: @@ -534,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 @@ -551,7 +544,6 @@ def _getOverallestStage(self) -> Self: ## Tracking methods - # ================================================== @abstractmethod def track(self, beam, savedepth=0, runnable=None, verbose=False): """ @@ -635,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" @@ -676,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" @@ -705,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" @@ -729,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" @@ -755,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: @@ -779,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: @@ -804,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" @@ -831,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" @@ -858,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 @@ -898,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 @@ -906,7 +886,6 @@ def _recalcLengthEnergyGradient(self): self._printVerb() - # ================================================== def _recalcLengthEnergyGradient_helper(self): itrCtr = 0 updateCounter_total = 0 @@ -1232,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." @@ -1240,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) @@ -1256,7 +1233,6 @@ def _printLengthEnergyGradient_internal(stage): ## Various calculations / plots / etc - # ================================================== def get_cost_breakdown(self): """ Get the cost breakdown of the plasma stage. @@ -1267,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 @@ -1302,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. @@ -1322,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 @@ -1340,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() @@ -1386,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() @@ -1593,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') # ================================================== @@ -2047,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. @@ -2089,29 +2012,44 @@ 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') - - num_plots = 1 + int(has_final_step) + 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') + + 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] - # extract initial or final + # extract initial, final and middle if i==0: data_struct = self.initial title = 'Initial step' - elif i==1: + # If middle step exists, extract it + elif i==1 and has_middle_step == True: + data_struct = self.middle + 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: data_struct = self.final title = 'Final step' @@ -2136,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 @@ -2180,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: 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')