Add calculations of wake potential

Author: kargapolov

lcode2dPy/diagnostics/diagnostics_3d.py

@@ -54,8 +54,8 @@ def dt(self, *parameters):
 
 class Diagnostics_f_xi:
     __allowed_f_xi = ['Ex', 'Ey', 'Ez', 'Bx', 'By', 'Bz', 'ne', 'nb',
-                      'Ez2', 'Bz2', 'nb2']
-                    # 'Phi', 'ni']
+                      'Ez2', 'Bz2', 'nb2', 'Phi']
+                    # 'ni']
                     # TODO: add them and functionality!
     __allowed_f_xi_type = ['numbers', 'pictures', 'both']
     #TODO: add 'pictures' and 'both' and functionality
@@ -125,7 +125,7 @@ def dxi(self, current_time, xi_plasma_layer,
             self.__data['xi'].append(xi_plasma_layer)
 
             for name in self.__f_xi_names:
-                if name in ['Ex', 'Ey', 'Ez', 'Bx', 'By', 'Bz']:
+                if name in ['Ex', 'Ey', 'Ez', 'Bx', 'By', 'Bz', 'Phi']:
                     field = getattr(pl_fields, name)[self.__ax_x, self.__ax_y]
                     self.__data[name].append(field)
 
@@ -172,8 +172,8 @@ def dt(self, *params):
 
 class Diagnostics_colormaps:
     __allowed_colormaps = ['Ex', 'Ey', 'Ez', 'Bx', 'By', 'Bz', 'ne', 'nb',
-                           'px', 'py', 'pz']
-                        # 'Phi', 'ni']
+                           'px', 'py', 'pz', 'Phi']
+                        # 'ni']
                     # TODO: add them and functionality!
     __allowed_colormaps_type = ['numbers']
     #TODO: add 'pictures' and 'both' and functionality
@@ -260,7 +260,7 @@ def dxi(self, current_time, xi_plasma_layer,
             self.__data['xi'].append(xi_plasma_layer)
 
             for name in self.__colormaps_names:
-                if name in ['Ex', 'Ey', 'Ez', 'Bx', 'By', 'Bz']:
+                if name in ['Ex', 'Ey', 'Ez', 'Bx', 'By', 'Bz', 'Phi']:
                     val = getattr(pl_fields, name)[
                         self.__steps//2, self.__r_f:self.__r_t]
                     self.__data[name].append(val)
@@ -362,7 +362,7 @@ def dt(self, current_time,
                     px=pl_particles.px, py=pl_particles.py, pz=pl_particles.pz,
                     Ex=pl_fields.Ex, Ey=pl_fields.Ey, Ez=pl_fields.Ez,
                     Bx=pl_fields.Bx, By=pl_fields.By, Bz=pl_fields.Bz,
-                    ro=pl_currents.ro,
+                    Phi=pl_fields.Phi, ro=pl_currents.ro,
                     jx=pl_currents.jx, jy=pl_currents.jy, jz=pl_currents.jz)
 
 

lcode2dPy/plasma3d/data.py

@@ -12,6 +12,7 @@
     ('Bx', _float_array),
     ('By', _float_array),
     ('Bz', _float_array),
+    ('Phi', _float_array),
 ]
 
 
@@ -24,6 +25,7 @@ def __init__(self, n_cells: int) -> None:
         self.Bx = np.zeros((n_cells, n_cells), dtype=np.float64)
         self.By = np.zeros((n_cells, n_cells), dtype=np.float64)
         self.Bz = np.zeros((n_cells, n_cells), dtype=np.float64)
+        self.Phi = np.zeros((n_cells, n_cells), dtype=np.float64)
 
     # Average operation is necessary on intermediate steps to have
     # better approximations of all fields
@@ -35,6 +37,7 @@ def average(self, other):
         fields.Bx = (self.Bx + other.Bx) / 2
         fields.By = (self.By + other.By) / 2
         fields.Bz = (self.Bz + other.Bz) / 2
+        fields.Phi = (self.Phi + other.Phi) / 2
         return fields
 
     def copy(self):
@@ -45,6 +48,7 @@ def copy(self):
         new_fields.Bx = np.copy(self.Bx)
         new_fields.By = np.copy(self.By)
         new_fields.Bz = np.copy(self.Bz)
+        new_fields.Phi = np.copy(self.Phi)
         return new_fields
 
 
@@ -113,4 +117,4 @@ def __init__(self, ro_initial, dirichlet_matrix,
         self.ro_initial = np.copy(ro_initial)
         self.dirichlet_matrix = np.copy(dirichlet_matrix)
         self.field_mixed_matrix = np.copy(field_mixed_matrix)
-        self.neumann_matrix = np.copy(neumann_matrix)
+        self.neumann_matrix = np.copy(neumann_matrix)

lcode2dPy/plasma3d/fields.py

@@ -7,7 +7,7 @@
 from lcode2dPy.plasma3d.data import Fields, Currents, Const_Arrays
 
 
-# Solving Laplace equation with Dirichlet boundary conditions (Ez) #
+# Solving Laplace equation with Dirichlet boundary conditions (Ez and Phi) #
 
 def calculate_Ez(grid_step_size, const: Const_Arrays, currents: Currents):
     """
@@ -32,12 +32,30 @@ def calculate_Ez(grid_step_size, const: Const_Arrays, currents: Currents):
     #    unnormalized DST-Type1 is its own inverse, up to a factor 2(N+1)
     #    and we take all scaling matters into account with a single factor
     #    hidden inside dirichlet_matrix.
-    
+
     Ez_inner = scipy.fftpack.dstn(f, type=1)
     Ez = np.pad(Ez_inner, 1, 'constant', constant_values=0)
     return Ez
 
 
+def calculate_Phi(const: Const_Arrays, currents: Currents):
+    """
+    Calculates Phi as iDST2D(dirichlet_matrix * DST2D(-ro + jz)).
+    """
+    ro, jz = currents.ro, currents.jz
+
+    rhs_inner = (ro - jz)[1:-1, 1:-1]
+
+    f = scipy.fftpack.dstn(rhs_inner, type=1)
+
+    f *= const.dirichlet_matrix
+
+    Phi_inner = scipy.fftpack.dstn(f, type=1)
+    Phi = np.pad(Phi_inner, 1, 'constant', constant_values=0)
+
+    return Phi
+
+
 # Solving Laplace or Helmholtz equation with mixed boundary conditions #
 
 def mix2d(a):
@@ -178,15 +196,15 @@ def compute_fields(self, flds: Fields, flds_prev: Fields,
                        curr_prev: Currents, curr: Currents):
         # Looks terrible! TODO: rewrite this function entirely
         new_flds = Fields((flds.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)
-                                                               
+
         if self.variant_A:
             new_flds.Ex = 2 * new_flds.Ex - flds_prev.Ex
             new_flds.Ey = 2 * new_flds.Ey - flds_prev.Ey
@@ -195,5 +213,7 @@ def compute_fields(self, flds: Fields, flds_prev: Fields,
 
         new_flds.Ez = calculate_Ez(self.grid_step_size, const, curr)
         new_flds.Bz = calculate_Bz(self.grid_step_size, const, curr)
-        
-        return new_flds, new_flds.average(flds_prev)
+
+        new_flds.Phi = calculate_Phi(const, curr)
+
+        return new_flds, new_flds.average(flds_prev)

lcode2dPy/plasma3d/initialization.py

@@ -173,6 +173,7 @@ def load_plasma(config: Config, path_to_plasmastate: str):
              state['Ex'], state['Ey'], state['Ez']
         fields.Bx, fields.By, fields.Bz =\
              state['Bx'], state['By'], state['Bz']
+        fields.Phi = state['Phi']
 
         particles.x_offt, particles.x_offt =\
             state['x_offt'], state['y_offt']

lcode2dPy/plasma3d_gpu/data.py

@@ -83,5 +83,6 @@ def fields_average(fields1: GPUArrays, fields2: GPUArrays):
             Ez = (fields1.Ez + fields2.Ez) / 2,
             Bx = (fields1.Bx + fields2.Bx) / 2,
             By = (fields1.By + fields2.By) / 2,
-            Bz = (fields1.Bz + fields2.Bz) / 2
-        )
+            Bz = (fields1.Bz + fields2.Bz) / 2,
+            Phi = (fields1.Phi + fields2.Phi) / 2
+        )

lcode2dPy/plasma3d_gpu/fields.py

@@ -5,7 +5,7 @@
 from lcode2dPy.plasma3d_gpu.data import GPUArrays, fields_average
 
 
-# Solving Laplace equation with Dirichlet boundary conditions (Ez) #
+# Solving Laplace equation with Dirichlet boundary conditions (Ez and Phi) #
 
 def dst2d(a):
     """
@@ -56,6 +56,24 @@ def calculate_Ez(grid_step_size, const, currents):
     return Ez
 
 
+def calculate_Phi(const, currents):
+    """
+    Calculates Phi as iDST2D(dirichlet_matrix * DST2D(-ro + jz)).
+    """
+    ro, jz = currents.ro, currents.jz
+
+    rhs_inner = (ro - jz)[1:-1, 1:-1]
+
+    f = dst2d(rhs_inner)
+
+    f *= const.dirichlet_matrix
+
+    Phi_inner = dst2d(f)
+    Phi = cp.pad(Phi_inner, 1, 'constant', constant_values=0)
+
+    return Phi
+
+
 # Solving Laplace or Helmholtz equation with mixed boundary conditions #
 
 def mix2d(a):
@@ -230,7 +248,7 @@ def compute_fields(self, fields, fields_prev, const,
                             self.grid_step_size, self.xi_step_size,
                             self.trick, self.variant_A, const,
                             fields, rho_beam, currents, currents_prev)
-                                            
+
         if self.variant_A:
             Ex = 2 * Ex - fields_prev.Ex
             Ey = 2 * Ey - fields_prev.Ey
@@ -240,7 +258,10 @@ def compute_fields(self, fields, fields_prev, const,
         Ez = calculate_Ez(self.grid_step_size, const, currents)
         Bz = calculate_Bz(self.grid_step_size, const, currents)
 
+        Phi = calculate_Phi(const, currents)
+
         new_fields = GPUArrays(Ex=Ex, Ey=Ey, Ez=Ez,
-                               Bx=Bx, By=By, Bz=Bz)
+                               Bx=Bx, By=By, Bz=Bz,
+                               Phi=Phi)
 
         return new_fields, fields_average(new_fields, fields_prev)

lcode2dPy/plasma3d_gpu/initialization.py

@@ -163,7 +163,8 @@ def zeros():
         return cp.zeros((grid_steps, grid_steps), dtype=cp.float64)
 
     fields = GPUArrays(Ex=zeros(), Ey=zeros(), Ez=zeros(),
-                       Bx=zeros(), By=zeros(), Bz=zeros())
+                       Bx=zeros(), By=zeros(), Bz=zeros(),
+                       Phi=zeros())
 
     particles = GPUArrays(x_init=x_init, y_init=y_init,
                           x_offt=x_offt, y_offt=y_offt,
@@ -184,7 +185,8 @@ def load_plasma(config: Config, path_to_plasmastate: str):
     with np.load(file=path_to_plasmastate) as state:
         fields = GPUArrays(Ex=state['Ex'], Ey=state['Ey'],
                            Ez=state['Ez'], Bx=state['Bx'],
-                           By=state['By'], Bz=state['Bz'])
+                           By=state['By'], Bz=state['Bz'],
+                           Phi=state['Phi'])
 
         particles = GPUArrays(x_init=particles.x_init, y_init=particles.y_init,
                               q=particles.q, m=particles.m,