Return to first-order pred-corr scheme

Author: kargapolov

lcode2dPy/plasma3d/fields.py

@@ -83,9 +83,11 @@ def dx_dy(arr, grid_step_size):
     return dx / (grid_step_size * 2), dy / (grid_step_size * 2)
 
 
-def calculate_Ex_Ey_Bx_By(grid_step_size, xi_step_size, trick, variant_A,
-                          const: Const_Arrays, fields_avg: Fields,
-                          beam_ro, currents: Currents, currents_prev: Currents):
+def calculate_Ex_Ey_Bx_By(
+    grid_step_size, xi_step_size, trick, variant_A, const: Const_Arrays,
+    fields_avg: Fields, ro_beam_full, ro_beam_prev, currents_full: Currents,
+    currents_prev: Currents
+):
     """
     Calculate transverse fields as iDST-DCT(mixed_matrix * DST-DCT(RHS.T)).T,
     with and without transposition depending on the field component.
@@ -95,14 +97,16 @@ def calculate_Ex_Ey_Bx_By(grid_step_size, xi_step_size, trick, variant_A,
            minus the coarse plasma particle cloud width).
     """
     jx_prev, jy_prev = currents_prev.jx, currents_prev.jy
-    jx,      jy      = currents.jx,      currents.jy
+    jx,      jy      = currents_full.jx, currents_full.jy
 
-    ro = currents.ro if not variant_A else (currents.ro + currents_prev.ro) / 2
-    jz = currents.jz if not variant_A else (currents.jz + currents_prev.jz) / 2
+    ro = (currents_full.ro + ro_beam_full +
+               currents_prev.ro + ro_beam_prev) / 2
+    jz = (currents_full.jz + ro_beam_full +
+               currents_prev.jz + ro_beam_prev) / 2
 
     # 1. Calculate gradients and RHS.
-    dro_dx, dro_dy = dx_dy(ro + beam_ro, grid_step_size)
-    djz_dx, djz_dy = dx_dy(jz + beam_ro, grid_step_size)
+    dro_dx, dro_dy = dx_dy(ro, grid_step_size)
+    djz_dx, djz_dy = dx_dy(jz, grid_step_size)
     djx_dxi = (jx_prev - jx) / xi_step_size  # - ?
     djy_dxi = (jy_prev - jy) / xi_step_size  # - ?
 
@@ -191,29 +195,29 @@ def __init__(self, config: Config):
         self.trick = config.getfloat('field-solver-subtraction-trick')
         self.variant_A = config.getbool('field-solver-variant-A')
 
-    def compute_fields(self, flds: Fields, flds_prev: Fields,
-                       const: Const_Arrays, rho_beam,
-                       curr_prev: Currents, curr: Currents):
+    def compute_fields(
+        self, fields: Fields, flds_prev: Fields, const: Const_Arrays,
+        rho_beam_full, rho_beam_prev, currents_prev: Currents,
+        currents_full: Currents
+    ):
         # Looks terrible! TODO: rewrite this function entirely
-        new_flds = Fields((flds.Ex).shape[0])
+        new_flds = Fields((fields.Ex).shape[0])
 
-        (new_flds.Ex,
-         new_flds.Ey,
-         new_flds.Bx,
-         new_flds.By) = calculate_Ex_Ey_Bx_By(self.grid_step_size,
-                                              self.xi_step_size,
-                                              self.trick, self.variant_A, const,
-                                              flds, rho_beam, curr, curr_prev)
+        new_flds.Ex, new_flds.Ey, new_flds.Bx, new_flds.By =\
+            calculate_Ex_Ey_Bx_By(
+                self.grid_step_size, self.xi_step_size, self.trick,
+                self.variant_A, const, fields, rho_beam_full, rho_beam_prev,
+                currents_full, currents_prev
+            )
 
         if self.variant_A:
             new_flds.Ex = 2 * new_flds.Ex - flds_prev.Ex
             new_flds.Ey = 2 * new_flds.Ey - flds_prev.Ey
             new_flds.Bx = 2 * new_flds.Bx - flds_prev.Bx
             new_flds.By = 2 * new_flds.By - flds_prev.By
 
-        new_flds.Ez = calculate_Ez(self.grid_step_size, const, curr)
-        new_flds.Bz = calculate_Bz(self.grid_step_size, const, curr)
-
-        new_flds.Phi = calculate_Phi(const, curr)
+        new_flds.Ez = calculate_Ez(self.grid_step_size, const, currents_full)
+        new_flds.Bz = calculate_Bz(self.grid_step_size, const, currents_full)
+        new_flds.Phi = calculate_Phi(const, currents_full)
 
         return new_flds, new_flds.average(flds_prev)

lcode2dPy/plasma3d/solver.py

@@ -11,31 +11,50 @@ def __init__(self, config: Config):
         self.PMover = ParticleMover(config)
         self.CComputer = RhoJComputer(config)
 
-    # Perfoms one full step along xi
-    def step_dxi(self, particles_prev: Particles, fields_prev: Fields,
-                 currents_prev: Currents, const_arrays: Const_Arrays, rho_beam):
-        particles = self.PMover.move_particles_wo_fields(particles_prev)
-
-
-        particles = self.PMover.move_particles(fields_prev,
-                                               particles_prev, particles)
-        currents = self.CComputer.compute_rhoj(particles, const_arrays)
-
-        _, fields_avg = self.FComputer.compute_fields(fields_prev, fields_prev,
-                                               const_arrays, rho_beam,
-                                               currents_prev, currents)
-
-
-        particles = self.PMover.move_particles(fields_avg,
-                                               particles_prev, particles)
-        currents = self.CComputer.compute_rhoj(particles, const_arrays)
-
-        fields, fields_avg = self.FComputer.compute_fields(fields_avg, fields_prev,
-                                               const_arrays, rho_beam,
-                                               currents_prev, currents)
-
-        particles = self.PMover.move_particles(fields_avg,
-                                               particles_prev, particles)
-        currents = self.CComputer.compute_rhoj(particles, const_arrays)
-
-        return particles, fields, currents
+    # Perfoms one full step along xi.
+    # To understand the numerical scheme, read values as following:
+    # *_prev = * on the previous xi step, an index number = k
+    # *_half = * on the halfstep, an index number = k + 1/2
+    # *_full = * on the next xi step (fullstep), an index number = k + 1
+    # *_prevprev = * on the xi step with an index number k - 1
+    def step_dxi(
+        self, particles_prev: Particles, fields_prev: Fields,
+        currents_prev: Currents, const_arrays: Const_Arrays,
+        rho_beam_full, rho_beam_prev
+    ):
+        particles_full = self.PMover.move_particles_wo_fields(particles_prev)
+
+
+        particles_full = self.PMover.move_particles(
+            fields_prev, particles_prev, particles_full
+        )
+        currents_full = self.CComputer.compute_rhoj(
+            particles_full, const_arrays
+        )
+
+        _, fields_half = self.FComputer.compute_fields(
+            fields_prev, fields_prev, const_arrays, rho_beam_full, rho_beam_prev,
+            currents_prev, currents_full
+        )
+
+
+        particles_full = self.PMover.move_particles(
+            fields_half, particles_prev, particles_full
+        )
+        currents_full = self.CComputer.compute_rhoj(
+            particles_full, const_arrays
+        )
+
+        fields_full, fields_half = self.FComputer.compute_fields(
+            fields_half, fields_prev, const_arrays, rho_beam_full, rho_beam_prev,
+            currents_prev, currents_full
+        )
+
+        particles_full = self.PMover.move_particles(
+            fields_half, particles_prev, particles_full
+        )
+        currents_full = self.CComputer.compute_rhoj(
+            particles_full, const_arrays
+        )
+
+        return particles_full, fields_full, currents_full

lcode2dPy/plasma3d_gpu/fields.py

@@ -111,8 +111,8 @@ def dx_dy(arr, grid_step_size):
     return dx / (grid_step_size * 2), dy / (grid_step_size * 2)
 
 
-def calculate_Ex_Ey_Bx_By_predictor(
-    grid_step_size, xi_step_size, const, fields_prev, ro_beam_full, ro_beam_prev,
+def calculate_Ex_Ey_Bx_By(
+    grid_step_size, xi_step_size, const, fields, ro_beam_full, ro_beam_prev,
     currents_full, currents_prev
 ):
     """
@@ -138,10 +138,10 @@ def calculate_Ex_Ey_Bx_By_predictor(
     djy_dxi = (jy_prev - jy_full) / xi_step_size  # - ?
 
     # We are solving a Helmholtz equation
-    Ex_rhs = -(dro_dx - djx_dxi - fields_prev.Ex)  # -?
-    Ey_rhs = -(dro_dy - djy_dxi - fields_prev.Ey)
-    Bx_rhs = +(djz_dy - djy_dxi + fields_prev.Bx)
-    By_rhs = -(djz_dx - djx_dxi - fields_prev.By)
+    Ex_rhs = -(dro_dx - djx_dxi - fields.Ex)  # -?
+    Ey_rhs = -(dro_dy - djy_dxi - fields.Ey)
+    Bx_rhs = +(djz_dy - djy_dxi + fields.Bx)
+    By_rhs = -(djz_dx - djx_dxi - fields.By)
 
     # Boundary conditions application (for future reference, ours are zero):
     # rhs[:, 0] -= bound_bottom[:] * (2 / grid_step_size)
@@ -155,68 +155,14 @@ def calculate_Ex_Ey_Bx_By_predictor(
     Ey_f *= mix_mat
 
     # 3. Apply our mixed DCT-DST transform again.
-    Ey_half = mix2d(Ey_f)
+    Ey = mix2d(Ey_f)
 
     # Likewise for other fields:
-    Bx_half = mix2d(mix_mat * mix2d(Bx_rhs[1:-1, :])[1:-1, :])
-    By_half = mix2d(mix_mat * mix2d(By_rhs.T[1:-1, :])[1:-1, :]).T
-    Ex_half = mix2d(mix_mat * mix2d(Ex_rhs.T[1:-1, :])[1:-1, :]).T
+    Bx = mix2d(mix_mat * mix2d(Bx_rhs[1:-1, :])[1:-1, :])
+    By = mix2d(mix_mat * mix2d(By_rhs.T[1:-1, :])[1:-1, :]).T
+    Ex = mix2d(mix_mat * mix2d(Ex_rhs.T[1:-1, :])[1:-1, :]).T
 
-    return Ex_half, Ey_half, Bx_half, By_half
-
-
-def calculate_Ex_Ey_Bx_By_corrector(
-    grid_step_size, xi_step_size, const, fields_full, ro_beam_full,
-    currents_full, currents_prev, currents_prevprev
-):
-    """
-    Calculate transverse fields as iDST-DCT(mixed_matrix * DST-DCT(RHS.T)).T,
-    with and without transposition depending on the field component.
-    NOTE: density and currents are assumed to be zero on the perimeter
-          (no plasma particles must reach the wall, so the reflection boundary
-           must be closer to the center than the simulation window boundary
-           minus the coarse plasma particle cloud width).
-    """
-    jx_prevprev, jy_prevprev = currents_prevprev.jx, currents_prevprev.jy
-    jx_prev,     jy_prev     = currents_prev.jx,     currents_prev.jy
-    jx_full,     jy_full     = currents_full.jx,     currents_full.jy
-
-    ro_full = currents_full.ro
-    jz_full = currents_full.jz
-
-    # 0. Calculate gradients and RHS.
-    dro_dx, dro_dy = dx_dy(ro_full + ro_beam_full, grid_step_size)
-    djz_dx, djz_dy = dx_dy(jz_full + ro_beam_full, grid_step_size)
-
-    djx_dxi = (- jx_prevprev + 4 * jx_prev - 3 * jx_full) / (2 * xi_step_size)
-    djy_dxi = (- jy_prevprev + 4 * jy_prev - 3 * jy_full) / (2 * xi_step_size)
-
-    # We are solving a Helmholtz equation
-    Ex_rhs = -(dro_dx - djx_dxi - fields_full.Ex)  # -?
-    Ey_rhs = -(dro_dy - djy_dxi - fields_full.Ey)
-    Bx_rhs = +(djz_dy - djy_dxi + fields_full.Bx)
-    By_rhs = -(djz_dx - djx_dxi - fields_full.By)
-
-    # Boundary conditions application (for future reference, ours are zero):
-    # rhs[:, 0] -= bound_bottom[:] * (2 / grid_step_size)
-    # rhs[:, -1] += bound_top[:] * (2 / grid_step_size)
-
-    # 1. Apply our mixed DCT-DST transform to RHS.
-    Ey_f = mix2d(Ey_rhs[1:-1, :])[1:-1, :]
-
-    # 2. Multiply f by the magic matrix.
-    mix_mat = const.field_mixed_matrix
-    Ey_f *= mix_mat
-
-    # 3. Apply our mixed DCT-DST transform again.
-    Ey_full = mix2d(Ey_f)
-
-    # Likewise for other fields:
-    Bx_full = mix2d(mix_mat * mix2d(Bx_rhs[1:-1, :])[1:-1, :])
-    By_full = mix2d(mix_mat * mix2d(By_rhs.T[1:-1, :])[1:-1, :]).T
-    Ex_full = mix2d(mix_mat * mix2d(Ex_rhs.T[1:-1, :])[1:-1, :]).T
-
-    return Ex_full, Ey_full, Bx_full, By_full
+    return Ex, Ey, Bx, By
 
 
 # Pushing particles without any fields (used for initial halfstep estimation) #
@@ -297,25 +243,25 @@ def __init__(self, config: Config):
         self.trick = config.getfloat('field-solver-subtraction-trick')
         self.variant_A = config.getbool('field-solver-variant-A')
 
-    def compute_fields_predictor(
-        self, fields_prev, const, rho_beam, rho_beam_prev, currents_prev,
-        currents_full
+    def compute_fields(
+        self, fields, fields_prev, const, rho_beam_full, rho_beam_prev,
+        currents_prev, currents_full
     ):
         # Looks terrible! TODO: rewrite this function entirely
 
-        Ex_half, Ey_half, Bx_half, By_half = calculate_Ex_Ey_Bx_By_predictor(
-            self.grid_step_size, self.xi_step_size, const, fields_prev,
-            rho_beam, rho_beam_prev, currents_full, currents_prev
+        Ex_half, Ey_half, Bx_half, By_half = calculate_Ex_Ey_Bx_By(
+            self.grid_step_size, self.xi_step_size, const, fields,
+            rho_beam_full, rho_beam_prev, currents_full, currents_prev
         )
 
         Ex_full = 2 * Ex_half - fields_prev.Ex
         Ey_full = 2 * Ey_half - fields_prev.Ey
         Bx_full = 2 * Bx_half - fields_prev.Bx
         By_full = 2 * By_half - fields_prev.By
 
-        Ez_full  = calculate_Ez(self.grid_step_size, const, currents_full)
-        Bz_full  = calculate_Bz(self.grid_step_size, const, currents_full)
-        Phi = fields_prev.Phi #calculate_Phi(const, currents_full)
+        Ez_full = calculate_Ez(self.grid_step_size, const, currents_full)
+        Bz_full = calculate_Bz(self.grid_step_size, const, currents_full)
+        Phi = calculate_Phi(const, currents_full)
 
         fields_full = GPUArrays(
             Ex=Ex_full, Ey=Ey_full, Ez=Ez_full,
@@ -333,25 +279,3 @@ def compute_fields_predictor(
         )
 
         return fields_full, fields_half
-
-    def compute_fields_corrector(
-        self, fields_full, fields_prev, const, rho_beam, currents_prevprev,
-        currents_prev, currents_full
-    ):
-
-        Ex_full, Ey_full, Bx_full, By_full = calculate_Ex_Ey_Bx_By_corrector(
-            self.grid_step_size, self.xi_step_size, const, fields_full,
-            rho_beam, currents_full, currents_prev, currents_prevprev
-        )
-
-        Ez_full  = calculate_Ez(self.grid_step_size, const, currents_full)
-        Bz_full  = calculate_Bz(self.grid_step_size, const, currents_full)
-        Phi_full = calculate_Phi(const, currents_full)
-
-        fields_full = GPUArrays(
-            Ex=Ex_full, Ey=Ey_full, Ez=Ez_full,
-            Bx=Bx_full, By=By_full, Bz=Bz_full,
-            Phi=Phi_full
-        )
-
-        return fields_full, fields_average(fields_full, fields_prev)

lcode2dPy/plasma3d_gpu/solver.py

@@ -19,8 +19,8 @@ def __init__(self, config: Config):
     # *_prevprev = * on the xi step with an index number k - 1
     def step_dxi(
         self, particles_prev: GPUArrays, fields_prev: GPUArrays,
-        currents_prev: GPUArrays, currents_prevprev: GPUArrays,
-        const_arrays: GPUArrays, rho_beam, rho_beam_prev
+        currents_prev: GPUArrays, const_arrays: GPUArrays,
+        rho_beam_full, rho_beam_prev
     ):
         particles_full = self.PMover.move_particles_wo_fields(particles_prev)
 
@@ -32,9 +32,9 @@ def step_dxi(
             particles_full, const_arrays
         )
 
-        fields_full, fields_half= self.FComputer.compute_fields_predictor(
-            fields_prev, const_arrays, rho_beam, rho_beam_prev, currents_prev,
-            currents_full
+        _, fields_half = self.FComputer.compute_fields(
+            fields_prev, fields_prev, const_arrays, rho_beam_full, rho_beam_prev,
+            currents_prev, currents_full
         )
 
 
@@ -45,8 +45,8 @@ def step_dxi(
             particles_full, const_arrays
         )
 
-        fields_full, fields_half = self.FComputer.compute_fields_corrector(
-            fields_full, fields_prev, const_arrays, rho_beam, currents_prevprev,
+        fields_full, fields_half = self.FComputer.compute_fields(
+            fields_half, fields_prev, const_arrays, rho_beam_full, rho_beam_prev,
             currents_prev, currents_full
         )
 
@@ -57,4 +57,4 @@ def step_dxi(
             particles_full, const_arrays
         )
 
-        return particles_full, fields_full, currents_full, currents_prev
+        return particles_full, fields_full, currents_full