From 52c44da7809b4e24554a6b26ec4265d46c1ad71e Mon Sep 17 00:00:00 2001 From: Angel Ferran Pousa Date: Fri, 10 Mar 2023 13:51:27 +0100 Subject: [PATCH 1/3] Use laser polarization only in ponderomotive force --- wake_t/physics_models/laser/laser_pulse.py | 19 +++++++++---------- .../plasma_wakefields/qs_cold_fluid_1x3p.py | 4 ++++ .../qs_rz_baxevanis/wakefield.py | 12 ++++++++---- 3 files changed, 21 insertions(+), 14 deletions(-) diff --git a/wake_t/physics_models/laser/laser_pulse.py b/wake_t/physics_models/laser/laser_pulse.py index 2b1b97a4..3c2f4b1d 100644 --- a/wake_t/physics_models/laser/laser_pulse.py +++ b/wake_t/physics_models/laser/laser_pulse.py @@ -26,14 +26,19 @@ class LaserPulse(): ---------- l_0 : float Laser wavelength in units of m. + polarization : str, optional + Polarization of the laser pulse. Accepted values are 'linear' + (default) or 'circular'. """ def __init__( self, - l_0: float + l_0: float, + polarization: Optional[str] = 'linear' ) -> None: self.l_0 = l_0 + self.polarization = polarization self.a_env = None self.solver_params = None self.n_steps = 0 @@ -361,7 +366,7 @@ def __init__( cep_phase: Optional[float] = 0., polarization: Optional[str] = 'linear' ) -> None: - super().__init__(l_0) + super().__init__(l_0=l_0, polarization=polarization) self.xi_c = xi_c self.a_0 = a_0 self.tau = tau @@ -369,7 +374,6 @@ def __init__( self.z_foc = z_foc self.z_r = np.pi * w_0**2 / l_0 self.cep_phase = cep_phase - self.polarization = polarization def _envelope_function(self, xi, r, z_pos): """ @@ -386,8 +390,6 @@ def _envelope_function(self, xi, r, z_pos): gaussian_profile = np.exp(exp_cep + exp_r + exp_z) # Amplitude avg_amplitude = self.a_0 - if self.polarization == 'linear': - avg_amplitude /= np.sqrt(2) return avg_amplitude / diff_factor * gaussian_profile @@ -441,7 +443,7 @@ def __init__( polarization: Optional[str] = 'linear' ) -> None: # Initialize parent class - super().__init__(l_0) + super().__init__(l_0=l_0, polarization=polarization) # If no focal plane position is given, use xi_c if z_foc is None: @@ -457,7 +459,6 @@ def __init__( self.w0 = w_0 self.cep_phase = cep_phase self.tau = tau - self.polarization = polarization def _envelope_function(self, xi, r, z_pos): """Complex envelope of a Laguerre-Gauss beam.""" @@ -479,8 +480,6 @@ def _envelope_function(self, xi, r, z_pos): * self.laguerre_pm(scaled_radius_squared)) a = self.a0 * profile - if self.polarization == 'linear': - a /= np.sqrt(2) return a @@ -537,7 +536,7 @@ def __init__( polarization: Optional[str] = 'linear' ) -> None: # Initialize parent class. - super().__init__(l_0) + super().__init__(l_0=l_0, polarization=polarization) # Store parameters. self.xi_c = xi_c diff --git a/wake_t/physics_models/plasma_wakefields/qs_cold_fluid_1x3p.py b/wake_t/physics_models/plasma_wakefields/qs_cold_fluid_1x3p.py index 53a09312..388fc773 100644 --- a/wake_t/physics_models/plasma_wakefields/qs_cold_fluid_1x3p.py +++ b/wake_t/physics_models/plasma_wakefields/qs_cold_fluid_1x3p.py @@ -147,6 +147,10 @@ def _calculate_wakefield(self, bunches): # Get laser envelope if self.laser is not None: a_env = np.abs(self.laser.get_envelope()) + # If linearly polarized, divide by sqrt(2) so that the + # ponderomotive force on the plasma particles is correct. + if self.laser.polarization == 'linear': + a_env /= np.sqrt(2) else: a_env = np.zeros((self.n_xi, self.n_r)) diff --git a/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis/wakefield.py b/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis/wakefield.py index e3a8cf1e..c4402805 100644 --- a/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis/wakefield.py +++ b/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis/wakefield.py @@ -165,15 +165,19 @@ def __init__( def _calculate_wakefield(self, bunches): parabolic_coefficient = self.parabolic_coefficient(self.t*ct.c) - # Get laser envelope + # Get square of laser envelope if self.laser is not None: - a_env = np.abs(self.laser.get_envelope()) ** 2 + a_env_2 = np.abs(self.laser.get_envelope()) ** 2 + # If linearly polarized, divide by 2 so that the ponderomotive + # force on the plasma particles is correct. + if self.laser.polarization == 'linear': + a_env_2 /= 2 else: - a_env = np.zeros((self.n_xi, self.n_r)) + a_env_2 = np.zeros((self.n_xi, self.n_r)) # Calculate plasma wakefields calculate_wakefields( - a_env, bunches, self.r_max, self.xi_min, self.xi_max, + a_env_2, bunches, self.r_max, self.xi_min, self.xi_max, self.n_r, self.n_xi, self.ppc, self.n_p, r_max_plasma=self.r_max_plasma, parabolic_coefficient=parabolic_coefficient, From 4a1856b2431b862b85354efdea48772b7493e7cb Mon Sep 17 00:00:00 2001 From: Angel Ferran Pousa Date: Fri, 10 Mar 2023 13:51:40 +0100 Subject: [PATCH 2/3] Fix test --- tests/test_laser_initialization.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_laser_initialization.py b/tests/test_laser_initialization.py index ec107585..3ca6ff52 100644 --- a/tests/test_laser_initialization.py +++ b/tests/test_laser_initialization.py @@ -47,7 +47,7 @@ def test_gaussian_init(): # Check correct a0. z_r = laser_rayleigh_length(w0, l0) a0_analytic = a0 / np.sqrt(1 + (z_foc/z_r)**2) - a0_env = np.max(np.abs(a_env)) * np.sqrt(2) + a0_env = np.max(np.abs(a_env)) assert math.isclose(a0_env, a0_analytic, rel_tol=1e-4) @@ -95,7 +95,7 @@ def test_flattened_gaussian_init(): # Check correct a0. z_r = laser_rayleigh_length(w0, l0) a0_analytic = a0 / np.sqrt(1 + (z_foc/z_r)**2) - a0_env = np.max(np.abs(a_env)) * np.sqrt(2) + a0_env = np.max(np.abs(a_env)) assert math.isclose(a0_env, a0_analytic, rel_tol=1e-4) From e7f0011cc9f3bef2e48cadcc2ddcf882de11762a Mon Sep 17 00:00:00 2001 From: Angel Ferran Pousa Date: Fri, 10 Mar 2023 14:23:11 +0100 Subject: [PATCH 3/3] Allow saving field attributes (like polarization) --- wake_t/diagnostics/openpmd_diag.py | 2 ++ wake_t/fields/rz_wakefield.py | 11 +++++++---- wake_t/utilities/other.py | 7 ++++--- 3 files changed, 13 insertions(+), 7 deletions(-) diff --git a/wake_t/diagnostics/openpmd_diag.py b/wake_t/diagnostics/openpmd_diag.py index cb960e65..51bb9ba0 100644 --- a/wake_t/diagnostics/openpmd_diag.py +++ b/wake_t/diagnostics/openpmd_diag.py @@ -303,6 +303,8 @@ def _write_fields(self, it, wf_data): # properly defined in the openPMD standard. fld.geometry = Geometry.thetaMode # Geometry.cylindrical fld.set_attribute('fieldSmoothing', 'none') + for attr, val in wf_data[field]['attributes'].items(): + fld.set_attribute(attr, val) fld.axis_labels = wf_data[field]['grid']['labels'] fld.grid_spacing = wf_data[field]['grid']['spacing'] global_offset = deepcopy(wf_data[field]['grid']['global_offset']) diff --git a/wake_t/fields/rz_wakefield.py b/wake_t/fields/rz_wakefield.py index 5edbc9dd..338ab1ae 100644 --- a/wake_t/fields/rz_wakefield.py +++ b/wake_t/fields/rz_wakefield.py @@ -173,6 +173,7 @@ def _get_openpmd_diagnostics_data(self, global_time): fld_position = [0.5, 0.] fld_names = ['E', 'B', 'rho'] fld_comps = [['r', 't', 'z'], ['r', 't', 'z'], None] + fld_attrs = [{}, {}, {}] fld_arrays = [ [np.ascontiguousarray(self.e_r.T[2:-2, 2:-2]), np.ascontiguousarray(self.e_t.T[2:-2, 2:-2]), @@ -185,6 +186,7 @@ def _get_openpmd_diagnostics_data(self, global_time): if self.laser is not None: fld_names += ['a_mod', 'a_phase'] fld_comps += [None, None] + fld_attrs += [{'polarization': self.laser.polarization}, {}] fld_arrays += [ [np.ascontiguousarray(np.abs(self.laser.get_envelope().T))], [np.ascontiguousarray(np.angle(self.laser.get_envelope().T))] @@ -193,9 +195,10 @@ def _get_openpmd_diagnostics_data(self, global_time): # Generate dictionary for openPMD diagnostics. diag_data = generate_field_diag_dictionary( - fld_names, fld_comps, fld_arrays, fld_comp_pos, grid_labels, - grid_spacing, grid_global_offset, fld_solver, fld_solver_params, - fld_boundary, fld_boundary_params, part_boundary, - part_boundary_params, current_smoothing, charge_correction) + fld_names, fld_comps, fld_attrs, fld_arrays, fld_comp_pos, + grid_labels, grid_spacing, grid_global_offset, fld_solver, + fld_solver_params, fld_boundary, fld_boundary_params, + part_boundary, part_boundary_params, current_smoothing, + charge_correction) return diag_data diff --git a/wake_t/utilities/other.py b/wake_t/utilities/other.py index 3d9df9a6..c18a17f8 100644 --- a/wake_t/utilities/other.py +++ b/wake_t/utilities/other.py @@ -16,7 +16,7 @@ def print_progress_bar(pre_string, step, total_steps): def generate_field_diag_dictionary( - fld_names, fld_comps, fld_arrays, fld_comp_pos, grid_labels, + fld_names, fld_comps, fld_attrs, fld_arrays, fld_comp_pos, grid_labels, grid_spacing, grid_global_offset, fld_solver, fld_solver_params, fld_boundary, fld_boundary_params, part_boundary, part_boundary_params, current_smoothing, charge_correction): @@ -27,8 +27,8 @@ def generate_field_diag_dictionary( """ diag_data = {} diag_data['fields'] = fld_names - fld_zip = zip(fld_names, fld_comps, fld_arrays, fld_comp_pos) - for fld, comps, arrays, pos in fld_zip: + fld_zip = zip(fld_names, fld_comps, fld_attrs, fld_arrays, fld_comp_pos) + for fld, comps, attrs, arrays, pos in fld_zip: diag_data[fld] = {} if comps is not None: diag_data[fld]['comps'] = {} @@ -43,6 +43,7 @@ def generate_field_diag_dictionary( diag_data[fld]['grid']['spacing'] = grid_spacing diag_data[fld]['grid']['labels'] = grid_labels diag_data[fld]['grid']['global_offset'] = grid_global_offset + diag_data[fld]['attributes'] = attrs diag_data['field_solver'] = fld_solver diag_data['field_solver_params'] = fld_solver_params diag_data['field_boundary'] = fld_boundary