Merge pull request #112 from AngelFP/fix_laser_envelope_amplitude

Fix laser envelope amplitude

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)
 
 

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'])

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

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,15 +366,14 @@ 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
         self.w_0 = w_0
         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

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

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,

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