Added capability to StageWakeT to save simulation outputs along the stage and some bugfixes

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 <abel.classes.stage.stage.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)
+

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):