Merge pull request #41 from AngelFP/bug/fix-default-dz_fields

Make `dz_fields` equal to simulation box length by default
AngelFP · Oct 11, 2021 · 032ffbf · 032ffbf
2 parents 7775e3e + a50c41a
commit 032ffbf
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 18 deletions.
diff --git a/wake_t/beamline_elements/plasma_stage.py b/wake_t/beamline_elements/plasma_stage.py
@@ -31,17 +31,14 @@ def __init__(self, length, density, wakefield_model='simple_blowout',
         """
         Initialize plasma stage.
 
-        Parameters:
-        -----------
+        Parameters
+        ----------
         length : float
             Length of the plasma stage in m.
 
-        n_p : float
+        density : float
             Plasma density in units of m^{-3}.
 
-        laser : LaserPulse
-            Laser driver of the plasma stage.
-
         wakefield_model : str
             Wakefield model to be used. Possible values are 'blowout',
             'custom_blowout', 'focusing_blowout', 'cold_fluid_1d' and
@@ -58,7 +55,7 @@ def __init__(self, length, density, wakefield_model='simple_blowout',
             below:
 
         Model 'focusing_blowout'
-        ----------------------
+        ------------------------
         No additional parameters required.
 
         Model 'simple_blowout'
@@ -129,9 +126,9 @@ def __init__(self, length, density, wakefield_model='simple_blowout',
             (most expensive option). If specified, the wakefields are only
             updated in steps determined by dz_fields. For example, if
             dz_fields=10e-6, the plasma wakefields are only updated every time
-            the simulation window advances by 10 micron. If dz_fields=None, the
-            wakefields are only computed once (at the start of the plasma) and
-            never updated throughout the simulation.
+            the simulation window advances by 10 micron. By default, if not
+            specified, the value of `dz_fields` will be `xi_max-xi_min`, i.e.,
+            the length the simulation box.
 
         p_shape : str
             Particle shape to be used for the beam charge deposition. Possible
@@ -175,9 +172,9 @@ def __init__(self, length, density, wakefield_model='simple_blowout',
             (most expensive option). If specified, the wakefields are only
             updated in steps determined by dz_fields. For example, if
             dz_fields=10e-6, the plasma wakefields are only updated every time
-            the simulation window advances by 10 micron. If dz_fields=None, the
-            wakefields are only computed once (at the start of the plasma) and
-            never updated throughout the simulation.
+            the simulation window advances by 10 micron. By default, if not
+            specified, the value of `dz_fields` will be `xi_max-xi_min`, i.e.,
+            the length the simulation box.
 
         r_max_plasma : float
             Maximum radial extension of the plasma column. If `None`, the
@@ -194,7 +191,7 @@ def __init__(self, length, density, wakefield_model='simple_blowout',
 
         p_shape : str
             Particle shape to be used for the beam charge deposition. Possible
-            values are 'linear' or 'cubic'.
+            values are 'linear' or 'cubic' (default).
 
         max_gamma : float
             Plasma particles whose `gamma` exceeds `max_gamma` are considered

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
@@ -12,7 +12,7 @@
 class NonLinearColdFluidWakefield(Wakefield):
     def __init__(self, density_function, laser=None, laser_evolution=False,
                  r_max=None, xi_min=None, xi_max=None, n_r=100,
-                 n_xi=100, dz_fields=0, beam_wakefields=False,
+                 n_xi=100, dz_fields=None, beam_wakefields=False,
                  p_shape='linear'):
         super().__init__()
         self.openpmd_diag_supported = True
@@ -24,7 +24,7 @@ def __init__(self, density_function, laser=None, laser_evolution=False,
         self.xi_max = xi_max
         self.n_r = n_r
         self.n_xi = n_xi
-        self.dz_fields = np.inf if dz_fields is None else dz_fields
+        self.dz_fields = xi_max - xi_min if dz_fields is None else dz_fields
         self.beam_wakefields = beam_wakefields
         self.p_shape = p_shape
         # Last time at which the fields where requested.

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
@@ -12,7 +12,7 @@ class Quasistatic2DWakefield(Wakefield):
 
     def __init__(self, density_function, laser=None, laser_evolution=False,
                  r_max=None, xi_min=None, xi_max=None, n_r=100,
-                 n_xi=100, ppc=2, dz_fields=0, r_max_plasma=None,
+                 n_xi=100, ppc=2, dz_fields=None, r_max_plasma=None,
                  parabolic_coefficient=0., p_shape='cubic', max_gamma=10):
         super().__init__()
         self.openpmd_diag_supported = True
@@ -25,7 +25,7 @@ def __init__(self, density_function, laser=None, laser_evolution=False,
         self.n_r = n_r
         self.n_xi = n_xi
         self.ppc = ppc
-        self.dz_fields = np.inf if dz_fields is None else dz_fields
+        self.dz_fields = xi_max - xi_min if dz_fields is None else dz_fields
         self.r_max_plasma = r_max_plasma
         self.parabolic_coefficient = self._get_parabolic_coefficient_fn(
             parabolic_coefficient)