Merge branch 'dev' into pr/agolovanov/147

tests/test_active_plasma_lens.py

@@ -29,7 +29,7 @@ def test_active_plasma_lens():
     apl.track(bunch)
     bunch_params = analyze_bunch(bunch)
     gamma_x = bunch_params['gamma_x']
-    assert approx(gamma_x, rel=1e-10) == 92.38646379897074
+    assert approx(gamma_x, rel=1e-10) == 92.38407675999406
 
 
 def test_active_plasma_lens_with_wakefields():
@@ -64,7 +64,7 @@ def test_active_plasma_lens_with_wakefields():
     # Analyze and check results.
     bunch_params = analyze_bunch(bunch)
     gamma_x = bunch_params['gamma_x']
-    assert approx(gamma_x, rel=1e-10) == 77.31995824746237
+    assert approx(gamma_x, rel=1e-10) == 77.32021188373825
 
 
 if __name__ == '__main__':

tests/test_beamline.py

@@ -30,12 +30,12 @@ def test_single_element():
     bl = Beamline([deepcopy(plasma)])
 
     # Track plasma.
-    bunch_1 = deepcopy(bunch)
+    bunch_1 = bunch.copy()
     out_dir_1 = os.path.join(output_folder, 'plasma_diags')
     plasma.track(bunch_1, opmd_diag=True, diag_dir=out_dir_1)
 
     # Track beamline.
-    bunch_2 = deepcopy(bunch)
+    bunch_2 = bunch.copy()
     out_dir_2 = os.path.join(output_folder, 'bl_diags')
     bl.track(bunch_2, opmd_diag=True, diag_dir=out_dir_2)
 
@@ -81,7 +81,7 @@ def test_multiple_element():
     bl = Beamline([deepcopy(d1), deepcopy(plasma), deepcopy(d2)])
 
     # Track elements individually.
-    bunch_1 = deepcopy(bunch)
+    bunch_1 = bunch.copy()
     out_dir_d1 = os.path.join(output_folder, 'd1_diags')
     out_dir_plasma = os.path.join(output_folder, 'plasma_diags')
     out_dir_d2 = os.path.join(output_folder, 'd2_diags')
@@ -90,7 +90,7 @@ def test_multiple_element():
     d2.track(bunch_1, opmd_diag=True, diag_dir=out_dir_d2)
 
     # Track beamline.
-    bunch_2 = deepcopy(bunch)
+    bunch_2 = bunch.copy()
     out_dir_bl = os.path.join(output_folder, 'bl_diags')
     bl.track(bunch_2, opmd_diag=True, diag_dir=out_dir_bl)
 

tests/test_custom_blowout.py

@@ -46,7 +46,7 @@ def test_custom_blowout_wakefield(make_plots=False):
 
     bunch_params = analyze_bunch(bunch)
     rel_ene_sp = bunch_params['rel_ene_spread']
-    assert approx(rel_ene_sp, rel=1e-10) == 0.21192488237458038
+    assert approx(rel_ene_sp, rel=1e-10) == 0.21192494153185745
 
     if make_plots:
         # Analyze bunch evolution.

tests/test_field_element.py

@@ -102,7 +102,7 @@ def test_field_element_error():
     with raises(ValueError) as e_info:
         element.track(bunch, opmd_diag=False)
     # This one should instead work.
-    element.track([bunch, copy.deepcopy(bunch)], opmd_diag=False)
+    element.track([bunch, bunch.copy()], opmd_diag=False)
 
 
 

tests/test_fluid_model.py

@@ -42,7 +42,7 @@ def test_fluid_model(plot=False):
 
     # Check final parameters.
     ene_sp = params_evolution['rel_ene_spread'][-1]
-    assert approx(ene_sp, rel=1e-10) == 0.024179998095119972
+    assert approx(ene_sp, rel=1e-10) == 0.024157374564016194
 
     # Quick plot of results.
     if plot:

tests/test_laser_envelope_model.py

@@ -3,6 +3,7 @@
 import scipy.constants as ct
 import matplotlib.pyplot as plt
 from wake_t.physics_models.laser.laser_pulse import GaussianPulse
+from wake_t.physics_models.laser.utils import unwrap
 
 
 def test_gaussian_laser_in_vacuum(plot=False):
@@ -190,6 +191,15 @@ def test_gaussian_laser_in_vacuum_with_subgrid(plot=False):
         plt.show()
 
 
+def test_unwrap():
+    """Test that the custom function for phase unwrapping implemented for
+    numba agrees with the numpy version.
+    """
+    phase = np.linspace(0, np.pi, num=5)
+    phase[3:] += np.pi
+    np.testing.assert_allclose(np.unwrap(phase), unwrap(phase))
+
+
 def calculate_spot_size(a_env, dr):
     # Project envelope to r
     a_proj = np.sum(np.abs(a_env), axis=0)
@@ -230,3 +240,4 @@ def calculate_a0(a_env):
 if __name__ == "__main__":
     test_gaussian_laser_in_vacuum(plot=True)
     test_gaussian_laser_in_vacuum_with_subgrid(plot=True)
+    test_unwrap()

tests/test_multibunch.py

@@ -55,8 +55,8 @@ def test_multibunch_plasma_simulation(plot=False):
     # Assert final parameters are correct.
     final_energy_driver = driver_params['avg_ene'][-1]
     final_energy_witness = witness_params['avg_ene'][-1]
-    assert approx(final_energy_driver, rel=1e-10) == 1700.3927190416732
-    assert approx(final_energy_witness, rel=1e-10) == 636.330857261392
+    assert approx(final_energy_driver, rel=1e-10) == 1700.3843657635728
+    assert approx(final_energy_witness, rel=1e-10) == 636.3260426124102
 
     if plot:
         z = driver_params['prop_dist'] * 1e2

tests/test_ramps.py

@@ -33,7 +33,7 @@ def test_downramp():
     downramp.track(bunch)
     bunch_params = analyze_bunch(bunch)
     beta_x = bunch_params['beta_x']
-    assert beta_x == 0.009750309290619276
+    assert beta_x == 0.009750308724018872
 
 
 def test_upramp():
@@ -64,7 +64,7 @@ def test_upramp():
     downramp.track(bunch)
     bunch_params = analyze_bunch(bunch)
     beta_x = bunch_params['beta_x']
-    assert beta_x == 0.0007631600676104024
+    assert beta_x == 0.000763155045965493
 
 
 if __name__ == '__main__':

tests/test_simple_blowout.py

@@ -45,7 +45,7 @@ def test_simple_blowout_wakefield(make_plots=False):
 
     bunch_params = analyze_bunch(bunch)
     rel_ene_sp = bunch_params['rel_ene_spread']
-    assert approx(rel_ene_sp, rel=1e-10) == 0.3637648484576557
+    assert approx(rel_ene_sp, rel=1e-10) == 0.3637651769109033
 
     if make_plots:
         # Analyze bunch evolution.

tutorials/00_basic_tutorial.py

@@ -63,7 +63,6 @@
 # average energy of :math:`100 \ \mathrm{MeV}` with a :math:`1 \ \%` spread
 # and a total charge of :math:`30 \ \mathrm{pC}`.
 
-from copy import deepcopy
 from wake_t.utilities.bunch_generation import get_gaussian_bunch_from_size
 
 # Beam parameters.
@@ -82,7 +81,7 @@
     q_bunch, n_part, name='elec_bunch')
 
 # Store bunch copy (will be needed later).
-bunch_bkp = deepcopy(bunch)
+bunch_bkp = bunch.copy()
 
 # Show phase space.
 bunch.show()
@@ -148,7 +147,7 @@
 # :math:`0.4 \ cm`, :math:`0.6 \ cm`, :math:`0.8 \ cm` and :math:`1.0 \ cm`.
 
 # Get again the original distribution.
-bunch = deepcopy(bunch_bkp)
+bunch = bunch_bkp.copy()
 
 # Create a 1 cm drift with 5 outputs (one every 0.2 cm).
 drift = Drift(length=1e-2, n_out=5)
@@ -195,7 +194,7 @@
 
 
 # Get again the original distribution.
-bunch = deepcopy(bunch_bkp)
+bunch = bunch_bkp.copy()
 
 # Create a 1 cm drift with 5 outputs (one every 0.2 cm).
 drift = Drift(length=1e-2, n_out=5)

tutorials/01_single_plasma_simulation.py

@@ -27,7 +27,6 @@
 # As a first step, let's generate a gaussian electron beam and keep a copy
 # of it for later use:
 
-from copy import deepcopy
 from wake_t.utilities.bunch_generation import get_gaussian_bunch_from_size
 
 # Beam parameters.
@@ -46,7 +45,7 @@
     q_bunch, n_part, name='elec_bunch')
 
 # Store bunch copy (will be needed later).
-bunch_bkp = deepcopy(bunch)
+bunch_bkp = bunch.copy()
 
 # Show phase space.
 bunch.show()
@@ -105,7 +104,7 @@
 from wake_t import GaussianPulse
 
 # Get again the original distribution.
-bunch = deepcopy(bunch_bkp)
+bunch = bunch_bkp.copy()
 
 # Laser parameters.
 laser_xi_c = 60e-6  # m (laser centroid in simulation box)
@@ -211,7 +210,7 @@ def density_profile(z):
 import scipy.constants as ct
 
 # Get again the original distribution.
-bunch = deepcopy(bunch_bkp)
+bunch = bunch_bkp.copy()
 
 # Calculate transverse parabolic profile.
 r_e = ct.e**2 / (4. * np.pi * ct.epsilon_0 * ct.m_e * ct.c**2)  # elec. radius

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,12 +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
+        diag_dir: Optional[str] = None,
+        show_progress_bar: Optional[bool] = True,
     ) -> Union[List[ParticleBunch], List[List[ParticleBunch]]]:
         """
         Track bunch through element.
@@ -83,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
         -------
@@ -119,6 +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=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