From 60b7f241317bb109d80a9b538df4841ab1813f83 Mon Sep 17 00:00:00 2001 From: Angel Ferran Pousa Date: Tue, 30 Apr 2024 17:31:12 +0200 Subject: [PATCH 1/2] Implement `ParticleBunch.copy` --- tests/test_beamline.py | 8 ++++---- tests/test_field_element.py | 2 +- tutorials/00_basic_tutorial.py | 7 +++---- tutorials/01_single_plasma_simulation.py | 7 +++---- wake_t/beamline_elements/tm_elements.py | 5 ++--- wake_t/particles/particle_bunch.py | 25 ++++++++++++++++++++++++ wake_t/tracking/tracker.py | 17 +--------------- 7 files changed, 39 insertions(+), 32 deletions(-) diff --git a/tests/test_beamline.py b/tests/test_beamline.py index 62ca47ac..44d0605b 100644 --- a/tests/test_beamline.py +++ b/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) diff --git a/tests/test_field_element.py b/tests/test_field_element.py index 1cd65f47..b1841dc0 100644 --- a/tests/test_field_element.py +++ b/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) diff --git a/tutorials/00_basic_tutorial.py b/tutorials/00_basic_tutorial.py index 5455887c..5ab9ef4d 100644 --- a/tutorials/00_basic_tutorial.py +++ b/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) diff --git a/tutorials/01_single_plasma_simulation.py b/tutorials/01_single_plasma_simulation.py index 25f256d6..0d9bc75b 100644 --- a/tutorials/01_single_plasma_simulation.py +++ b/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 diff --git a/wake_t/beamline_elements/tm_elements.py b/wake_t/beamline_elements/tm_elements.py index 843d5bb4..72e977b1 100644 --- a/wake_t/beamline_elements/tm_elements.py +++ b/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 diff --git a/wake_t/particles/particle_bunch.py b/wake_t/particles/particle_bunch.py index c212e10e..3a83e57d 100644 --- a/wake_t/particles/particle_bunch.py +++ b/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: diff --git a/wake_t/tracking/tracker.py b/wake_t/tracking/tracker.py index c9b9aaa4..18213c28 100644 --- a/wake_t/tracking/tracker.py +++ b/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: From 210343cabb064de2ca4f0a864e84a596100b2286 Mon Sep 17 00:00:00 2001 From: Angel Ferran Pousa Date: Tue, 30 Apr 2024 17:36:05 +0200 Subject: [PATCH 2/2] Formatting --- wake_t/particles/particle_bunch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wake_t/particles/particle_bunch.py b/wake_t/particles/particle_bunch.py index 3a83e57d..09e30183 100644 --- a/wake_t/particles/particle_bunch.py +++ b/wake_t/particles/particle_bunch.py @@ -299,7 +299,7 @@ def evolve(self, fields, t, dt, pusher='rk4'): 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. """