Merge pull request #150 from AngelFP/feature/bunch_copy

Implement `ParticleBunch.copy`

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

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.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/tm_elements.py

@@ -1,7 +1,6 @@
 """ Contains the classes of all elements tracked using transfer matrices. """
 from typing import Optional, Union, List
 import time
-from copy import deepcopy
 
 import numpy as np
 import scipy.constants as ct
@@ -140,7 +139,7 @@ def track(
         start_time = time.time()
         output_bunch_list = list()
         if out_initial:
-            output_bunch_list.append(deepcopy(bunch))
+            output_bunch_list.append(bunch.copy())
             if opmd_diag is not None:
                 opmd_diag.write_diagnostics(
                     0., l_step/ct.c, [output_bunch_list[-1]])
@@ -219,7 +218,7 @@ def _update_input_bunch(self, bunch, bunch_mat, output_bunch_list):
             last_bunch = self._create_new_bunch(bunch, new_bunch_mat,
                                                 self.length)
         else:
-            last_bunch = deepcopy(output_bunch_list[-1])
+            last_bunch = output_bunch_list[-1].copy()
         bunch.set_phase_space(last_bunch.x, last_bunch.y, last_bunch.xi,
                               last_bunch.px, last_bunch.py, last_bunch.pz)
         bunch.prop_distance = last_bunch.prop_distance

wake_t/particles/particle_bunch.py

@@ -2,6 +2,8 @@
 This module contains the class defining a particle bunch.
 """
 # TODO: clean methods to set and get bunch matrix
+from __future__ import annotations
+from copy import deepcopy
 from typing import Optional
 
 import numpy as np
@@ -295,6 +297,29 @@ def evolve(self, fields, t, dt, pusher='rk4'):
             )
         self.prop_distance += dt * ct.c
 
+    def copy(self) -> ParticleBunch:
+        """Return a copy of the bunch.
+
+        To improve performance, this copy won't contain copies of auxiliary
+        arrays, only of the particle coordinates and properties.
+        """
+        bunch_copy = ParticleBunch(
+            w=deepcopy(self.w),
+            x=deepcopy(self.x),
+            y=deepcopy(self.y),
+            xi=deepcopy(self.xi),
+            px=deepcopy(self.px),
+            py=deepcopy(self.py),
+            pz=deepcopy(self.pz),
+            prop_distance=deepcopy(self.prop_distance),
+            name=deepcopy(self.name),
+            q_species=deepcopy(self.q_species),
+            m_species=deepcopy(self.m_species)
+        )
+        bunch_copy.x_ref = self.x_ref
+        bunch_copy.theta_ref = self.theta_ref
+        return bunch_copy
+
     def get_field_arrays(self):
         """Get the arrays where the gathered fields will be stored."""
         if not self.__field_arrays_allocated:

wake_t/tracking/tracker.py

@@ -1,6 +1,5 @@
 """ This module contains the Tracker class. """
 from typing import Optional, Callable, List, Literal
-from copy import deepcopy
 
 import numpy as np
 import scipy.constants as ct
@@ -298,21 +297,7 @@ def generate_diagnostics(self) -> None:
         """Generate tracking diagnostics."""
         # Make copy of current bunches and store in output list.
         for i, bunch in enumerate(self.bunches):
-            self.bunch_list[i].append(
-                ParticleBunch(
-                    deepcopy(bunch.w),
-                    deepcopy(bunch.x),
-                    deepcopy(bunch.y),
-                    deepcopy(bunch.xi),
-                    deepcopy(bunch.px),
-                    deepcopy(bunch.py),
-                    deepcopy(bunch.pz),
-                    prop_distance=deepcopy(bunch.prop_distance),
-                    name=deepcopy(bunch.name),
-                    q_species=deepcopy(bunch.q_species),
-                    m_species=deepcopy(bunch.m_species)
-                )
-            )
+            self.bunch_list[i].append(bunch.copy())
 
         # If needed, write also the openPMD diagnostics.
         if self.opmd_diags is not None: