Merge pull request #154 from AngelFP/feature/optional_progress_bar

Make progress bar optional

wake_t/beamline_elements/active_plasma_lens.py

@@ -51,6 +51,20 @@ class ActivePlasmaLens(PlasmaStage):
         A list of values can also be provided. In this case, the list
         should have the same order as the list of bunches given to the
         ``track`` method.
+    push_bunches_before_diags : bool, optional
+        Whether to push the bunches before saving them to the diagnostics.
+        Since the time step of the diagnostics can be different from that
+        of the bunches, it could happen that the bunches appear in the
+        diagnostics as they were at the last push, but not at the actual
+        time of the diagnostics. Setting this parameter to ``True``
+        (default) ensures that an additional push is given to all bunches
+        to evolve them to the diagnostics time before saving.
+        This additional push will always have a time step smaller than
+        the the time step of the bunch, so it has no detrimental impact
+        on the accuracy of the simulation. However, it could make
+        convergence studies more difficult to interpret,
+        since the number of pushes will depend on `n_diags`. Therefore,
+        it is exposed as an option so that it can be disabled if needed.
     n_out : int
         Number of times along the lens in which the particle distribution
         should be returned (A list with all output bunches is returned
@@ -78,6 +92,7 @@ def __init__(
         wakefield_model: Optional[str] = 'quasistatic_2d',
         bunch_pusher: Optional[Literal['boris', 'rk4']] = 'boris',
         dt_bunch: Optional[DtBunchType] = 'auto',
+        push_bunches_before_diags: Optional[bool] = True,
         n_out: Optional[int] = 1,
         name: Optional[str] = 'Active plasma lens',
         **model_params
@@ -100,6 +115,7 @@ def __init__(
             wakefield_model=wakefield_model,
             bunch_pusher=bunch_pusher,
             dt_bunch=dt_bunch,
+            push_bunches_before_diags=push_bunches_before_diags,
             n_out=n_out,
             name=name,
             external_fields=[self.apl_field],

wake_t/beamline_elements/beamline.py

@@ -20,7 +20,8 @@ def track(
         self,
         bunches: Optional[Union[ParticleBunch, List[ParticleBunch]]] = [],
         opmd_diag: Optional[bool] = False,
-        diag_dir: Optional[str] = None
+        diag_dir: Optional[str] = None,
+        show_progress_bar: Optional[bool] = True,
     ) -> Union[List[ParticleBunch], List[List[ParticleBunch]]]:
         """
         Track bunch through beamline.
@@ -40,6 +41,9 @@ def track(
             Directory into which the openPMD output will be written. By default
             this is a 'diags' folder in the current directory. Only needed if
             `opmd_diag=True`.
+        show_progress_bar : bool, optional
+            Whether to show a progress bar of the tracking through each
+            element. By default ``True``.
 
         Returns
         -------
@@ -50,5 +54,11 @@ def track(
         if type(opmd_diag) is not OpenPMDDiagnostics and opmd_diag:
             opmd_diag = OpenPMDDiagnostics(write_dir=diag_dir)
         for element in self.elements:
-            bunch_list.extend(element.track(bunches, opmd_diag=opmd_diag))
+            bunch_list.extend(
+                element.track(
+                    bunches,
+                    opmd_diag=opmd_diag,
+                    show_progress_bar=show_progress_bar,
+                )
+            )
         return bunch_list

wake_t/beamline_elements/field_element.py

@@ -39,6 +39,20 @@ class FieldElement():
         Function used to determine the adaptive time step for bunches in
         which the time step is set to ``'auto'``. The function should take
         solely a ``ParticleBunch`` as argument.
+    push_bunches_before_diags : bool, optional
+        Whether to push the bunches before saving them to the diagnostics.
+        Since the time step of the diagnostics can be different from that
+        of the bunches, it could happen that the bunches appear in the
+        diagnostics as they were at the last push, but not at the actual
+        time of the diagnostics. Setting this parameter to ``True``
+        (default) ensures that an additional push is given to all bunches
+        to evolve them to the diagnostics time before saving.
+        This additional push will always have a time step smaller than
+        the the time step of the bunch, so it has no detrimental impact
+        on the accuracy of the simulation. However, it could make
+        convergence studies more difficult to interpret,
+        since the number of pushes will depend on `n_diags`. Therefore,
+        it is exposed as an option so that it can be disabled if needed.
 
     """
 
@@ -50,7 +64,8 @@ def __init__(
         n_out: Optional[int] = 1,
         name: Optional[str] = 'field element',
         fields: Optional[List[Field]] = [],
-        auto_dt_bunch: Optional[str] = None
+        auto_dt_bunch: Optional[str] = None,
+        push_bunches_before_diags: Optional[bool] = True,
     ) -> None:
         self.length = length
         self.bunch_pusher = bunch_pusher
@@ -59,13 +74,14 @@ def __init__(
         self.name = name
         self.fields = fields
         self.auto_dt_bunch = auto_dt_bunch
+        self.push_bunches_before_diags = push_bunches_before_diags
 
     def track(
         self,
         bunches: Optional[Union[ParticleBunch, List[ParticleBunch]]] = [],
         opmd_diag: Optional[Union[bool, OpenPMDDiagnostics]] = False,
         diag_dir: Optional[str] = None,
-        push_bunches_before_diags: Optional[bool] = True,
+        show_progress_bar: Optional[bool] = True,
     ) -> Union[List[ParticleBunch], List[List[ParticleBunch]]]:
         """
         Track bunch through element.
@@ -84,20 +100,9 @@ def track(
             Directory into which the openPMD output will be written. By default
             this is a 'diags' folder in the current directory. Only needed if
             `opmd_diag=True`.
-        push_bunches_before_diags : bool, optional
-            Whether to push the bunches before saving them to the diagnostics.
-            Since the time step of the diagnostics can be different from that
-            of the bunches, it could happen that the bunches appear in the
-            diagnostics as they were at the last push, but not at the actual
-            time of the diagnostics. Setting this parameter to ``True``
-            (default) ensures that an additional push is given to all bunches
-            to evolve them to the diagnostics time before saving.
-            This additional push will always have a time step smaller than
-            the the time step of the bunch, so it has no detrimental impact
-            on the accuracy of the simulation. However, it could make
-            convergence studies more difficult to interpret,
-            since the number of pushes will depend on `n_diags`. Therefore,
-            it is exposed as an option so that it can be disabled if needed.
+        show_progress_bar : bool, optional
+            Whether to show a progress bar of the tracking. By default
+            ``True``.
 
         Returns
         -------
@@ -134,7 +139,8 @@ def track(
             opmd_diags=opmd_diag,
             bunch_pusher=self.bunch_pusher,
             auto_dt_bunch_f=self.auto_dt_bunch,
-            push_bunches_before_diags=push_bunches_before_diags,
+            push_bunches_before_diags=self.push_bunches_before_diags,
+            show_progress_bar=show_progress_bar,
             section_name=self.name
         )
 

wake_t/beamline_elements/plasma_ramp.py

@@ -96,6 +96,20 @@ class PlasmaRamp(PlasmaStage):
         A list of values can also be provided. In this case, the list
         should have the same order as the list of bunches given to the
         ``track`` method.
+    push_bunches_before_diags : bool, optional
+        Whether to push the bunches before saving them to the diagnostics.
+        Since the time step of the diagnostics can be different from that
+        of the bunches, it could happen that the bunches appear in the
+        diagnostics as they were at the last push, but not at the actual
+        time of the diagnostics. Setting this parameter to ``True``
+        (default) ensures that an additional push is given to all bunches
+        to evolve them to the diagnostics time before saving.
+        This additional push will always have a time step smaller than
+        the the time step of the bunch, so it has no detrimental impact
+        on the accuracy of the simulation. However, it could make
+        convergence studies more difficult to interpret,
+        since the number of pushes will depend on `n_diags`. Therefore,
+        it is exposed as an option so that it can be disabled if needed.
     n_out : int
         Number of times along the stage in which the particle distribution
         should be returned (A list with all output bunches is returned
@@ -126,6 +140,7 @@ def __init__(
         position_down: Optional[float] = None,
         bunch_pusher: Optional[Literal['boris', 'rk4']] = 'boris',
         dt_bunch: Optional[DtBunchType] = 'auto',
+        push_bunches_before_diags: Optional[bool] = True,
         n_out: Optional[int] = 1,
         name: Optional[str] = 'Plasma ramp',
         **model_params
@@ -151,6 +166,7 @@ def __init__(
             wakefield_model=wakefield_model,
             bunch_pusher=bunch_pusher,
             dt_bunch=dt_bunch,
+            push_bunches_before_diags=push_bunches_before_diags,
             n_out=n_out,
             name=name,
             **model_params

wake_t/beamline_elements/plasma_stage.py

@@ -48,6 +48,20 @@ class PlasmaStage(FieldElement):
         stage. A list of values can also be provided. In this case, the list
         should have the same order as the list of bunches given to the
         ``track`` method.
+    push_bunches_before_diags : bool, optional
+        Whether to push the bunches before saving them to the diagnostics.
+        Since the time step of the diagnostics can be different from that
+        of the bunches, it could happen that the bunches appear in the
+        diagnostics as they were at the last push, but not at the actual
+        time of the diagnostics. Setting this parameter to ``True``
+        (default) ensures that an additional push is given to all bunches
+        to evolve them to the diagnostics time before saving.
+        This additional push will always have a time step smaller than
+        the the time step of the bunch, so it has no detrimental impact
+        on the accuracy of the simulation. However, it could make
+        convergence studies more difficult to interpret,
+        since the number of pushes will depend on `n_diags`. Therefore,
+        it is exposed as an option so that it can be disabled if needed.
     n_out : int, optional
         Number of times along the stage in which the particle distribution
         should be returned (A list with all output bunches is returned
@@ -77,6 +91,7 @@ def __init__(
         wakefield_model: Optional[str] = 'simple_blowout',
         bunch_pusher: Optional[Literal['boris', 'rk4']] = 'boris',
         dt_bunch: Optional[DtBunchType] = 'auto',
+        push_bunches_before_diags: Optional[bool] = True,
         n_out: Optional[int] = 1,
         name: Optional[str] = 'Plasma stage',
         external_fields: Optional[List[Field]] = [],
@@ -96,7 +111,8 @@ def __init__(
             n_out=n_out,
             name=name,
             fields=fields,
-            auto_dt_bunch=self._get_optimized_dt
+            auto_dt_bunch=self._get_optimized_dt,
+            push_bunches_before_diags=push_bunches_before_diags,
         )
 
     def _get_density_profile(self, density):

wake_t/beamline_elements/tm_elements.py

@@ -76,7 +76,8 @@ def track(
         backtrack: Optional[bool] = False,
         out_initial: Optional[bool] = False,
         opmd_diag: Optional[Union[bool, OpenPMDDiagnostics]] = False,
-        diag_dir: Optional[str] = None
+        diag_dir: Optional[str] = None,
+        show_progress_bar: Optional[bool] = True,
     ) -> List[ParticleBunch]:
         """
         Track bunch through element.
@@ -99,6 +100,9 @@ def track(
             Directory into which the openPMD output will be written. By default
             this is a 'diags' folder in the current directory. Only needed if
             `opmd_diag=True`.
+        show_progress_bar : bool, optional
+            Whether to show a progress bar of the tracking. By default
+            ``True``.
 
         Returns
         -------
@@ -125,15 +129,16 @@ def track(
             opmd_diag = None
 
         # Print output header
-        print('')
-        print(self.element_name.capitalize())
-        print('-'*len(self.element_name))
-        self._print_element_properties()
-        csr_string = 'on' if self.csr_on else 'off'
-        print('CSR {}.'.format(csr_string))
-        print('')
-        n_steps = len(track_steps)
-        st_0 = 'Tracking in {} step(s)... '.format(n_steps)
+        if show_progress_bar:
+            print('')
+            print(self.element_name.capitalize())
+            print('-'*len(self.element_name))
+            self._print_element_properties()
+            csr_string = 'on' if self.csr_on else 'off'
+            print('CSR {}.'.format(csr_string))
+            print('')
+            n_steps = len(track_steps)
+            st_0 = 'Tracking in {} step(s)... '.format(n_steps)
 
         # Start tracking
         start_time = time.time()
@@ -144,7 +149,8 @@ def track(
                 opmd_diag.write_diagnostics(
                     0., l_step/ct.c, [output_bunch_list[-1]])
         for i in track_steps:
-            print_progress_bar(st_0, i+1, n_steps)
+            if show_progress_bar:
+                print_progress_bar(st_0, i+1, n_steps)
             l_curr = (i+1) * l_step * (1-2*backtrack)
             # Track with transfer matrix
             bunch_mat = track_with_transfer_map(
@@ -173,9 +179,10 @@ def track(
             opmd_diag.increase_z_pos(self.length)
 
         # Finalize
-        tracking_time = time.time() - start_time
-        print('Done ({} s).'.format(tracking_time))
-        print('-'*80)
+        if show_progress_bar:
+            tracking_time = time.time() - start_time
+            print('Done ({} s).'.format(tracking_time))
+            print('-'*80)
         return output_bunch_list
 
     def _get_beam_matrix_for_tracking(self, bunch):

wake_t/tracking/progress_bar.py

@@ -10,7 +10,7 @@
 warnings.filterwarnings('ignore', '.*clamping.*', )
 
 
-def get_progress_bar(description, total_length):
+def get_progress_bar(description, total_length, disable):
     """Get progress bar for the tracker.
 
     Parameters
@@ -19,6 +19,8 @@ def get_progress_bar(description, total_length):
         Description to be appended to start of the progress bar.
     total_length : float
         Total length in metres of the stage to be tracked.
+    disable : bool
+        Whether to disable (not show) the progress bar.
 
     Returns
     -------
@@ -31,6 +33,7 @@ def get_progress_bar(description, total_length):
         total=total_length,
         unit='m',
         bar_format=l_bar + "{bar}" + r_bar,
-        file=sys.stdout
+        file=sys.stdout,
+        disable=disable
     )
     return progress_bar

wake_t/tracking/tracker.py

@@ -75,6 +75,8 @@ class Tracker():
         convergence studies more difficult to interpret,
         since the number of pushes will depend on `n_diags`. Therefore,
         it is exposed as an option so that it can be disabled if needed.
+    show_progress_bar : bool, optional
+        Whether to show a progress bar of the tracking. By default ``True``.
     section_name : str, optional
         Name of the section to be tracked. This will be appended to the
         beginning of the progress bar.
@@ -92,6 +94,7 @@ def __init__(
         auto_dt_bunch_f: Optional[Callable[[ParticleBunch], float]] = None,
         bunch_pusher: Optional[Literal['boris', 'rk4']] = 'boris',
         push_bunches_before_diags: Optional[bool] = True,
+        show_progress_bar: Optional[bool] = True,
         section_name: Optional[str] = 'Simulation'
     ) -> None:
         self.t_final = t_final
@@ -103,6 +106,7 @@ def __init__(
         self.auto_dt_bunch_f = auto_dt_bunch_f
         self.bunch_pusher = bunch_pusher
         self.push_bunches_before_diags = push_bunches_before_diags
+        self.show_progress_bar = show_progress_bar
         self.section_name = section_name
 
         # Get all numerical fields and their time steps.
@@ -149,7 +153,11 @@ def do_tracking(self) -> List[List[ParticleBunch]]:
         set_num_threads(num_threads)
 
         # Initialize progress bar.
-        progress_bar = get_progress_bar(self.section_name, self.t_final*ct.c)
+        progress_bar = get_progress_bar(
+            description=self.section_name,
+            total_length=self.t_final*ct.c,
+            disable=not self.show_progress_bar,
+        )
 
         # Calculate fields at t=0.
         for field in self.num_fields: