diff --git a/src/abel/classes/stage/impl/stage_wake_t.py b/src/abel/classes/stage/impl/stage_wake_t.py index 191625c1..2020c4f3 100644 --- a/src/abel/classes/stage/impl/stage_wake_t.py +++ b/src/abel/classes/stage/impl/stage_wake_t.py @@ -22,7 +22,7 @@ class StageWakeT(Stage): simulations using Wake-T. It prepares input files, launches Wake-T runs, extracts beam and plasma diagnostics, and post-processes simulation data. - Inherits all attributes from ``Stage``. + Inherits all attributes from :class:`Stage `. Attributes @@ -31,30 +31,45 @@ class StageWakeT(Stage): The source of the drive beam. The beam axis is always aligned to its propagation direction. Defaults to ``None``. - num_cell_xy : int - Number of transverse grid cells in Wake-T. + num_cell_xy : int, optional + Number of transverse grid cells in Wake-T. Defaults to 256. - keep_data : bool - Flag for whether to keep raw Wake-T output after simulation. + output_period : int, optional + Number of times along the stage in which the Wake-T diagnostics data + is written to file. If ``None``, ``output_period`` is set to + ``round(stage length/step size/8)``. Defaults to ``None``. - ion_motion : bool - Flag to include ion motion in the plasma. + keep_data : bool, optional + Flag for whether to keep raw Wake-T output after simulation. Defaults to + ``False``. - run_path : str - Path to store plots and outputs. + ion_motion : bool, optional + Flag to include ion motion in the plasma. Defaults to ``False``. + + run_path : str, optional + Path to store plots and outputs. Defaults to ``None``. use_single_beam : bool Enable tracking only the main beam. Defaults to ``False``. stage_number : int Keeps track of which stage it is in the beamline. + + + Notes + ----- + - The box sizes are set based on beam extent and estimated blowout radius. + - The step size is set to matched beta function/8 """ + + # TODO: allow the user to set the box size and step size. - def __init__(self, nom_accel_gradient=None, nom_energy_gain=None, plasma_density=None, driver_source=None, ramp_beta_mag=None, num_cell_xy=256, keep_data=False, ion_motion=False, run_path=None, use_single_beam=False): + def __init__(self, nom_accel_gradient=None, nom_energy_gain=None, plasma_density=None, driver_source=None, ramp_beta_mag=None, num_cell_xy=256, output_period=None, keep_data=False, ion_motion=False, run_path=None, use_single_beam=False): super().__init__(nom_accel_gradient=nom_accel_gradient, nom_energy_gain=nom_energy_gain, plasma_density=plasma_density, driver_source=driver_source, ramp_beta_mag=ramp_beta_mag) self.num_cell_xy = num_cell_xy + self.output_period = output_period self.keep_data = keep_data # physics flags @@ -82,6 +97,11 @@ def track(self, beam0, savedepth=0, runnable=None, verbose=False): if not os.path.exists(tmpfolder): os.mkdir(tmpfolder) + # make diagnostics folder + diag_dir = os.path.join(tmpfolder, 'hdf5') + if not os.path.exists(diag_dir): + os.mkdir(diag_dir) + # Set ramp lengths, nominal energies, nominal energy gains # and flattop nominal energy if not already done self._prepare_ramps() @@ -169,13 +189,14 @@ def track(self, beam0, savedepth=0, runnable=None, verbose=False): else: # Exclude ion motion effects wakefield_model='quasistatic_2d' - + if self.output_period is None: + self.output_period = round(self.length/dz/8) + # ========== Set up a Wake-T plasma stage ========== - n_out = round(self.length/dz/8) plasma = wake_t.PlasmaStage(length=self.length, density=plasma_profile, wakefield_model=wakefield_model, r_max=box_size_r, r_max_plasma=box_size_r, xi_min=box_min_z, xi_max=box_max_z, - n_out=n_out, n_r=int(self.num_cell_xy), n_xi=int(num_cell_z), dz_fields=dz, ppc=1) + n_out=self.output_period, n_r=int(self.num_cell_xy), n_xi=int(num_cell_z), dz_fields=dz, ppc=1) # ========== Do tracking ========== @@ -190,13 +211,13 @@ def track(self, beam0, savedepth=0, runnable=None, verbose=False): self.__extract_evolution(bunches) else: self.__extract_evolution([bunches]) - self.__extract_initial_and_final_step(tmpfolder) + self.__extract_initial_and_final_step(diag_dir) # delete or move data - #if self.keep_data: - # shot_path = runnable.shot_path() # TODO: this does not work yet - # destination_path = runnable.shot_path() + 'stage_' + str(bunches[1].stage_number) + '/insitu' - # shutil.move(tmpfolder, destination_path) + if self.keep_data: + destination_path = runnable.shot_path() + 'stage_' + str(self.stage_number) + #destination_path = runnable.shot_path() + 'stage_' + str(self.stage_number) + '/insitu' + shutil.move(diag_dir, destination_path) # remove temporary directory shutil.rmtree(tmpfolder) @@ -400,12 +421,11 @@ def __extract_evolution(self, bunches): # ================================================== - def __extract_initial_and_final_step(self, tmpfolder): + def __extract_initial_and_final_step(self, source_path): from openpmd_viewer import OpenPMDTimeSeries # prepare to read simulation data - source_path = tmpfolder + 'hdf5/' ts = OpenPMDTimeSeries(source_path) # extract initial on-axis wakefield @@ -547,7 +567,7 @@ def energy_usage(self): # ================================================== # Apply waterfall function to all beam dump files - def __waterfall_fcn(self, fcns, edges, data_dir, species='beam', clean=False, remove_halo_nsigma=20, args=None): + def __waterfall_fcn(self, fcns, edges, data_dir, species='beam', remove_halo_nsigma=None, args=None): """ Applies waterfall function to all Wake-T HDF5 output files in ``data_dir``. @@ -567,18 +587,17 @@ def __waterfall_fcn(self, fcns, edges, data_dir, species='beam', clean=False, re Path to the directory containing all Wake-T HDF5 output files. species : str, optional - Specifies the name of the beam to be extracted. - - clean : bool, optional - Determines whether the extracted beams from the Wake-T HDF5 output - files should be cleaned before further processing. + Specifies the name of the beam to be extracted. Defaults to + ``'beam'``. remove_halo_nsigma : float, optional - Defines a threshold for identifying and removing "halo" particles - based on their deviation from the core of the particle beam. + If not ``None``, defines a threshold for identifying and excluding + outlier particles based on their deviation from the core of the + particle beam. Defaults to ``None``. args : float list, optional - Allows passing additional arguments to the functions in ``fcns``. + Allows passing additional arguments to the functions in ``fcns``. + Defaults to ``None``. Returns @@ -597,7 +616,7 @@ def __waterfall_fcn(self, fcns, edges, data_dir, species='beam', clean=False, re ``waterfalls``. """ - from abel.apis.wake_t.wake_t_api import wake_t_hdf5_load + from abel.wrappers.wake_t.wake_t_wrapper import wake_t_hdf5_load # find number of beam outputs to plot files = sorted(os.listdir(data_dir)) @@ -620,7 +639,7 @@ def __waterfall_fcn(self, fcns, edges, data_dir, species='beam', clean=False, re file_path = data_dir + files[index] beam = wake_t_hdf5_load(file_path=file_path, species=species) - if clean: + if remove_halo_nsigma is not None and remove_halo_nsigma > 0.0: beam.remove_halo_particles(nsigma=remove_halo_nsigma) # find beam location @@ -637,7 +656,7 @@ def __waterfall_fcn(self, fcns, edges, data_dir, species='beam', clean=False, re # ================================================== - def extract_waterfalls(self, data_dir, species='beam', clean=False, remove_halo_nsigma=20, nsig=5, args=None): + def extract_waterfalls(self, data_dir, species='beam', remove_halo_nsigma=None, nsig=5, args=[None, None, None, None]): ''' Extracts data for waterfall plots for current profile, relative energy spectrum, horizontal transverse profile and vertical transverse profile. @@ -648,22 +667,21 @@ def extract_waterfalls(self, data_dir, species='beam', clean=False, remove_halo_ Path to the directory containing all Wake-T HDF5 output files. species : str, optional - Specifies the name of the beam to be extracted. - - clean : bool, optional - Determines whether the extracted beams from the Wake-T HDF5 output - files should be cleaned before further processing. + Specifies the name of the beam to be extracted. Defaults to + ``'beam'``. remove_halo_nsigma : float, optional - Defines a threshold for identifying and removing "halo" particles - based on their deviation from the core of the particle beam. + If not ``None``, defines a threshold for identifying and excluding + outlier particles based on their deviation from the core of the + particle beam. Defaults to ``None``. nsig : float, optional Helps define the range of the histograms. E.g. the range in x is defined by the beam x offset +- ``nsig`` * x beam size. args : float list, optional - Allows passing additional arguments to the functions in ``fcns``. + Allows passing additional arguments to the functions in ``fcns``. + Defaults to ``[None, None, None, None]``. Returns @@ -682,7 +700,7 @@ def extract_waterfalls(self, data_dir, species='beam', clean=False, remove_halo_ ``waterfalls``. ''' - from abel.apis.wake_t.wake_t_api import wake_t_hdf5_load + from abel.wrappers.wake_t.wake_t_wrapper import wake_t_hdf5_load files = sorted(os.listdir(data_dir)) file_path = data_dir + files[0] @@ -697,34 +715,34 @@ def extract_waterfalls(self, data_dir, species='beam', clean=False, remove_halo_ 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) - 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, clean=clean, remove_halo_nsigma=remove_halo_nsigma, args=[None, None, None, None]) + 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=args) return waterfalls, locations, bins # ================================================== - def plot_waterfalls(self, waterfalls, locations, bins, save_fig=False): + def plot_waterfalls(self, data_dir, species='beam', remove_halo_nsigma=20, save_path=None): # TODO move this to Stage ''' - Makes waterfall plots for current profile, relative energy spectrum, + Create waterfall plots for current profile, relative energy spectrum, horizontal transverse profile and vertical transverse profile. Parameters ---------- - 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 - outputs). The dimension of element i is determined by the length of - edges and the number of simulation outputs. - - locations : [m] 1D float ndarray - Stores the location for each slice of the ``waterfalls``. - - bins : list of 1D float ndarrays - Each element contains the bins used for the slices/histograms in - ``waterfalls``. + data_dir : str + Path to the directory containing all Wake-T HDF5 output files. + + species : str, optional + Specifies the name of the beam to be extracted. Defaults to + ``'beam'``. - save_fig : bool, optional - Flag for saving the output figure. + remove_halo_nsigma : float, optional + If not ``None``, defines a threshold for identifying and excluding + outlier particles based on their deviation from the core of the + particle beam. Defaults to 20. + + save_path : str, optional + If not ``None``, saves the output figure to the specified file path. + Defaults to ``None``. Returns @@ -734,6 +752,8 @@ def plot_waterfalls(self, waterfalls, locations, bins, save_fig=False): from matplotlib import pyplot as plt + waterfalls, locations, bins = self.extract_waterfalls(data_dir, species=species, remove_halo_nsigma=remove_halo_nsigma, nsig=5, args=[None, None, None, None]) + # prepare figure fig, axs = plt.subplots(4,1) fig.set_figwidth(8) @@ -772,11 +792,7 @@ def plot_waterfalls(self, waterfalls, locations, bins, save_fig=False): axs[3].set_xlabel('Location along the stage [m]') plt.show() - if save_fig: - plot_path = self.run_path + 'plots' + os.sep - if not os.path.exists(plot_path): - os.makedirs(plot_path) - filename = plot_path + 'waterfalls' + '.png' - fig.savefig(filename, format='png', dpi=600, bbox_inches='tight', transparent=False) - + if save_path is not None: + fig.savefig(save_path, format='png', dpi=600, bbox_inches='tight', transparent=False) + \ No newline at end of file diff --git a/src/abel/wrappers/wake_t/wake_t_wrapper.py b/src/abel/wrappers/wake_t/wake_t_wrapper.py index f99e2a27..c75f5ac9 100644 --- a/src/abel/wrappers/wake_t/wake_t_wrapper.py +++ b/src/abel/wrappers/wake_t/wake_t_wrapper.py @@ -369,7 +369,7 @@ def run_single_step_wake_t(plasma_density, drive_beam, beam): wakeT_xy_res = 0.1*beam.bunch_length() wakeT_max_box_r = 4/k_p(plasma_density) wakeT_num_cell_xy = int(wakeT_max_box_r/wakeT_xy_res) - plasma_stage = plasma_stage_setup(plasma_density, drive_beam, beam, stage_length=None, dz_fields=None, num_cell_xy=wakeT_num_cell_xy) + plasma_stage = plasma_stage_setup(plasma_density, drive_beam, beam, stage_length=None, dz_fields=None, num_cell_xy=wakeT_num_cell_xy, n_out=1) # Make temp folder if not os.path.exists(CONFIG.temp_path):