Implement non centered envelope solver for laser initialization

tests/test_fluid_model.py

@@ -42,7 +42,7 @@ def test_fluid_model(plot=False):
 
     # Check final parameters.
     ene_sp = params_evolution['rel_ene_spread'][-1]
-    assert approx(ene_sp, rel=1e-10) == 0.02386347411720627
+    assert approx(ene_sp, rel=1e-10) == 0.024183646993930535
 
     # Quick plot of results.
     if plot:

tests/test_laser_envelope_model.py

@@ -77,7 +77,7 @@ def test_gaussian_laser_in_vacuum(plot=False):
 
     # Check that solution hasn't changed.
     a_mod = np.abs(a_env)
-    assert_almost_equal(np.sum(a_mod), 7500.380279033873, decimal=10)
+    assert_almost_equal(np.sum(a_mod), 7500.380522059235, decimal=10)
 
     # Make plots.
     if plot:
@@ -175,7 +175,7 @@ def test_gaussian_laser_in_vacuum_with_subgrid(plot=False):
 
     # Check that solution hasn't changed.
     a_mod = np.abs(a_env)
-    assert_almost_equal(np.sum(a_mod), 7500.120030707864, decimal=10)
+    assert_almost_equal(np.sum(a_mod), 7500.120208299948, decimal=10)
 
     # Make plots.
     if plot:

wake_t/physics_models/laser/envelope_solver.py

@@ -10,46 +10,13 @@
 import scipy.constants as ct
 
 from wake_t.utilities.numba import njit_serial
+from .tdma import TDMA
 
 
 @njit_serial(fastmath=True)
-def TDMA(a, b, c, d, p):
-    """TriDiagonal Matrix Algorithm: solve a linear system Ax=b,
-    where A is a tridiagonal matrix. Source:
-    https://stackoverflow.com/questions/8733015/tridiagonal-matrix-algorithm-
-    tdma-aka-thomas-algorithm-using-python-with-nump
-
-    Parameters
-    ----------
-    a : array
-        Lower diagonal of A. Dimension: nr-1.
-    b : array
-        Main diagonal of A. Dimension: nr.
-    c : array
-        Upper diagonal of A. Dimension: nr-1.
-    d : array
-        Solution vector. Dimension: nr.
-
-    """
-    n = len(d)
-    w = np.empty(n - 1, dtype=np.complex128)
-    g = np.empty(n, dtype=np.complex128)
-
-    w[0] = c[0] / b[0]  # MAKE SURE THAT b[0]!=0
-    g[0] = d[0] / b[0]
-
-    for i in range(1, n - 1):
-        w[i] = c[i] / (b[i] - a[i - 1] * w[i - 1])
-    for i in range(1, n):
-        g[i] = (d[i] - a[i - 1] * g[i - 1]) / (b[i] - a[i - 1] * w[i - 1])
-    p[n - 1] = g[n - 1]
-    for i in range(n - 1, 0, -1):
-        p[i - 1] = g[i - 1] - w[i - 1] * p[i]
-
-
-@njit_serial(fastmath=True)
-def evolve_envelope(a, a_old, chi, k0, kp, zmin, zmax, nz, rmax, nr, dt, nt,
-                    start_outside_plasma=False, use_phase=True):
+def evolve_envelope(
+        a, a_old, chi, k0, kp, zmin, zmax, nz, rmax, nr, dt, nt,
+        use_phase=True):
     """
     Solve the 2D envelope equation
     (\nabla_tr^2+2i*k0/kp*d/dt+2*d^2/(dzdt)-d^2/dt^2)â = chi*â
@@ -80,10 +47,6 @@ def evolve_envelope(a, a_old, chi, k0, kp, zmin, zmax, nz, rmax, nr, dt, nt,
         Tau step size.
     nt : int
         Number of tau steps.
-    start_outside_plasma : bool
-        If `True`, it indicates that the laser is outside of the plasma at
-        `t=-dt`. This will then force the plasma susceptibility to be zero
-        at that time.
     use_phase : bool
         Determines whether to take into account the terms related to the
         longitudinal derivative of the complex phase.
@@ -130,12 +93,6 @@ def evolve_envelope(a, a_old, chi, k0, kp, zmin, zmax, nz, rmax, nr, dt, nt,
         if use_phase:
             phases = np.angle(a[:, 0])
 
-        # If laser starts outside plasma, make chi^{n-1} = 0.
-        if start_outside_plasma and n == 0:
-            chi_nm1 = 0. * chi
-        else:
-            chi_nm1 = chi
-
         # Loop over z.
         for j in range(nz - 1, -1, -1):
 
@@ -161,15 +118,17 @@ def evolve_envelope(a, a_old, chi, k0, kp, zmin, zmax, nz, rmax, nr, dt, nt,
             for k in range(nr):
                 rhs[k] = (
                     - 2 * inv_dt ** 2 * a[j, k]
-                    - ((C_minus - chi_nm1[j, k] * 0.5 - 1j * inv_dt * D_jkn)
+                    - ((C_minus - chi[j, k] * 0.5 - 1j * inv_dt * D_jkn)
                        * a_old[j, k])
                     - (2 * np.exp(-1j * d_theta1) * inv_dzdt
                        * (a_new_jp1[k] - a_old[j + 1, k]))
                     + (0.5 * np.exp(-1j * (d_theta2 + d_theta1)) * inv_dzdt
                        * (a_new_jp2[k] - a_old[j + 2, k]))
-                    - L_plus_over_2[k] * a_old[j, k + 1] * (k + 1 < nr)
-                    - L_minus_over_2[k] * a_old[j, k - 1] * (k > 0)
                 )
+                if k > 0:
+                    rhs[k] -= L_minus_over_2[k] * a_old[j, k - 1]
+                if k + 1 < nr:
+                    rhs[k] -= L_plus_over_2[k] * a_old[j, k + 1]
 
             # Calculate diagonals.
             d_main = C_plus - chi[j] * 0.5 + 1j * inv_dt * D_jkn

wake_t/physics_models/laser/envelope_solver_non_centered.py

@@ -0,0 +1,153 @@
+"""
+This module contains a non-centered (in time) version of the envelope solver
+described in 'An accurate and efficient laser-envelope solver for the modeling
+of laser-plasma accelerators', written by C Benedetti et al in 2018.
+
+Authors: Wilbert den Hertog, Ángel Ferran Pousa, Carlo Benedetti
+"""
+
+import numpy as np
+import scipy.constants as ct
+
+from wake_t.utilities.numba import njit_serial
+from .tdma import TDMA
+
+
+@njit_serial(fastmath=True)
+def evolve_envelope_non_centered(
+        a, a_old, chi, k0, kp, zmin, zmax, nz, rmax, nr, dt, nt,
+        use_phase=True):
+    """
+    Solve the 2D envelope equation using a non-centered time discretization.
+    (\nabla_tr^2+2i*k0/kp*d/dt+2*d^2/(dzdt)-d^2/dt^2)â = chi*â
+
+    Parameters
+    ----------
+    a0 : array
+        Initial value for â at tau=0. Dimension: nz*nr.
+    aold : array
+        Initial value for â at tau=-1. Dimension: nz*nr.
+    chi : array
+        Arrays of the values of the susceptibility. Dimension nz*nr.
+    k0 : float
+        Laser wave number.
+    kp : float
+        Plasma skin depth.
+    zmin : float
+        Minimum value for zeta.
+    zmax : float
+        Maximum value for zeta.
+    nz : int
+        Number of grid points in zeta-direction.
+    rmax : float
+        Maximum value for rho (minimum value is always 0).
+    nr : int
+        Amount of points in the rho direction.
+    dt : float
+        Tau step size.
+    nt : int
+        Number of tau steps.
+    use_phase : bool
+        Determines whether to take into account the terms related to the
+        longitudinal derivative of the complex phase.
+
+    """
+    # Preallocate arrays. a_old and a include 2 ghost cells in the z direction.
+    rhs = np.empty(nr, dtype=np.complex128)
+
+    # Calculate step sizes.
+    dz = (zmax - zmin) * kp / (nz - 1)
+    dr = rmax * kp / nr
+    dt = dt * ct.c * kp
+
+    # Precalculate common fractions.
+    inv_dt = 1 / dt
+    inv_dr = 1 / dr
+    inv_dz = 1 / dz
+    inv_dzdt = inv_dt * inv_dz
+    k0_over_kp = k0 / kp
+
+    # Initialize phase difference
+    d_theta1 = 0.
+    d_theta2 = 0.
+
+    # Calculate C^+ and C^- [Eq. (8)].
+    C_minus = (-2. * inv_dr ** 2. * 0.5 - 2j * k0_over_kp * inv_dt
+               + 3 * inv_dzdt)
+    C_plus = (-2. * inv_dr ** 2. * 0.5 + 2j * k0_over_kp * inv_dt
+              - 3 * inv_dzdt)
+
+    # Calculate L^+ and L^-. Change wrt Benedetti - 2018: in Wake-T we use
+    # cell-centered nodes in the radial direction.
+    L_base = 1. / (2. * (np.arange(nr) + 0.5))
+    L_minus_over_2 = (1. - L_base) * inv_dr ** 2. * 0.5
+    L_plus_over_2 = (1. + L_base) * inv_dr ** 2. * 0.5
+
+    # Loop over time iterations.
+    for n in range(nt):
+        # a_new_jp1 is equivalent to a_new[j+1] and a_new_jp2 to a_new[j+2].
+        a_new_jp1 = np.zeros(nr, dtype=np.complex128)
+        a_new_jp2 = np.zeros(nr, dtype=np.complex128)
+
+        # Getting the phase of the envelope on axis.
+        if use_phase:
+            phases = np.angle(a[:, 0])
+
+        # Loop over z.
+        for j in range(nz - 1, -1, -1):
+
+            # Calculate phase differences between adjacent points.
+            if use_phase:
+                d_theta1 = phases[j + 1] - phases[j]
+                d_theta2 = phases[j + 2] - phases[j + 1]
+
+                # Prevent phase jumps bigger than 1.5*pi.
+                if d_theta1 < -1.5 * np.pi:
+                    d_theta1 += 2 * np.pi
+                if d_theta2 < -1.5 * np.pi:
+                    d_theta2 += 2 * np.pi
+                if d_theta1 > 1.5 * np.pi:
+                    d_theta1 -= 2 * np.pi
+                if d_theta2 > 1.5 * np.pi:
+                    d_theta2 -= 2 * np.pi
+
+            # Calculate D factor [Eq. (6)].
+            D_jkn = (1.5 * d_theta1 - 0.5 * d_theta2) * inv_dz
+
+            # Calculate right-hand side of Eq (7).
+            for k in range(nr):
+                rhs[k] = (
+                    - (C_minus - chi[j, k] * 0.5 - 2j * inv_dt * D_jkn)
+                    * a[j, k]
+                    - (4 * np.exp(-1j * d_theta1) * inv_dzdt
+                       * (a_new_jp1[k] - a[j + 1, k]))
+                    + (1 * np.exp(-1j * (d_theta2 + d_theta1)) * inv_dzdt
+                       * (a_new_jp2[k] - a[j + 2, k]))
+                )
+                if k > 0:
+                    rhs[k] -= L_minus_over_2[k] * a[j, k - 1]
+                if k + 1 < nr:
+                    rhs[k] -= L_plus_over_2[k] * a[j, k + 1]
+
+            # Calculate diagonals.
+            d_main = C_plus - chi[j] * 0.5 + 2j * inv_dt * D_jkn
+            d_upper = L_plus_over_2[:nr - 1]
+            d_lower = L_minus_over_2[1:nr]
+
+            # Update a_old and a at j+2 with the current values of a a_new.
+            a_old[j + 2] = a[j + 2]
+            a[j + 2] = a_new_jp2
+
+            # Shift a_new[j+1] to a_new[j+2] in preparation for the next
+            # iteration.
+            a_new_jp2[:] = a_new_jp1
+
+            # Compute a_new at the current j using the TDMA method and store
+            # result in a_new_jp1 to use it in the next iteration.
+            TDMA(d_lower, d_main, d_upper, rhs, a_new_jp1)
+
+        # When the left of the computational domain is reached, paste the last
+        # few values in the a_old and a arrays.
+        a_old[0:2] = a[0:2]
+        a[0] = a_new_jp1
+        a[1] = a_new_jp2

wake_t/physics_models/laser/laser_pulse.py

@@ -13,6 +13,7 @@
 from scipy.special import genlaguerre, binom
 
 from .envelope_solver import evolve_envelope
+from .envelope_solver_non_centered import evolve_envelope_non_centered
 from wake_t.fields.interpolation import interpolate_rz_field
 
 
@@ -31,7 +32,6 @@ def __init__(self, l_0):
         self.l_0 = l_0
         self.a_env = None
         self.solver_params = None
-        self.init_outside_plasma = False
         self.n_steps = 0
 
     def __add__(self, pulse_2):
@@ -116,17 +116,14 @@ def initialize_envelope(self):
             r_max = self.solver_params['rmax']
             nz = self.solver_params['nz']
             nr = self.solver_params['nr']
-            dt = self.solver_params['dt']
             dr = r_max / nr
             z = np.linspace(z_min, z_max, nz)
             r = np.linspace(dr/2, r_max-dr/2, nr)
             ZZ, RR = np.meshgrid(z, r, indexing='ij')
             self._a_env_old = np.zeros((nz + 2, nr), dtype=np.complex128)
             self._a_env = np.zeros((nz + 2, nr), dtype=np.complex128)
             self._a_env[0:-2] = self.envelope_function(ZZ, RR, 0.)
-            self._a_env_old[0:-2] = self.envelope_function(ZZ, RR, -dt*ct.c)
             self.__update_output_envelope()
-            self.init_outside_plasma = True
 
     def get_envelope(self):
         """Get the current laser envelope array."""
@@ -146,17 +143,19 @@ def evolve(self, chi, n_p):
         k_0 = 2*np.pi / self.l_0
         k_p = np.sqrt(ct.e**2 * n_p / (ct.m_e*ct.epsilon_0)) / ct.c
 
-        # Determine if laser starts evolution outside plasma.
-        start_outside_plasma = (self.n_steps == 0 and self.init_outside_plasma)
-
         # If needed, interpolate chi to subgrid.
         if self.use_subgrid:
             chi = self.__interpolate_chi_to_subgrid(chi)
 
         # Compute evolution.
-        evolve_envelope(
-            self._a_env, self._a_env_old, chi, k_0, k_p,
-            **self.solver_params, start_outside_plasma=start_outside_plasma)
+        if self.n_steps == 0:
+            evolve_envelope_non_centered(
+                self._a_env, self._a_env_old, chi, k_0, k_p,
+                **self.solver_params)
+        else:
+            evolve_envelope(
+                self._a_env, self._a_env_old, chi, k_0, k_p,
+                **self.solver_params)
 
         # Update arrays and step count.
         self.__update_output_envelope()

wake_t/physics_models/laser/tdma.py

@@ -0,0 +1,45 @@
+"""
+This module contais the tridiagonal matrix algorithm (TDMA) used by the laser
+envelope solver.
+
+Authors: Wilbert den Hertog, Ángel Ferran Pousa
+"""
+
+import numpy as np
+
+from wake_t.utilities.numba import njit_serial
+
+
+@njit_serial(fastmath=True)
+def TDMA(a, b, c, d, p):
+    """TriDiagonal Matrix Algorithm: solve a linear system Ax=b,
+    where A is a tridiagonal matrix. Source:
+    https://stackoverflow.com/questions/8733015/tridiagonal-matrix-algorithm-
+    tdma-aka-thomas-algorithm-using-python-with-nump
+
+    Parameters
+    ----------
+    a : array
+        Lower diagonal of A. Dimension: nr-1.
+    b : array
+        Main diagonal of A. Dimension: nr.
+    c : array
+        Upper diagonal of A. Dimension: nr-1.
+    d : array
+        Solution vector. Dimension: nr.
+
+    """
+    n = len(d)
+    w = np.empty(n - 1, dtype=np.complex128)
+    g = np.empty(n, dtype=np.complex128)
+
+    w[0] = c[0] / b[0]  # MAKE SURE THAT b[0]!=0
+    g[0] = d[0] / b[0]
+
+    for i in range(1, n - 1):
+        w[i] = c[i] / (b[i] - a[i - 1] * w[i - 1])
+    for i in range(1, n):
+        g[i] = (d[i] - a[i - 1] * g[i - 1]) / (b[i] - a[i - 1] * w[i - 1])
+    p[n - 1] = g[n - 1]
+    for i in range(n - 1, 0, -1):
+        p[i - 1] = g[i - 1] - w[i - 1] * p[i]