Electronic Boltzmann made compatible with master

bolt/lib/nonlinear/boundaries.py

@@ -224,10 +224,10 @@ def apply_dirichlet_bcs_f(self, boundary):
 
     return
 
-def apply_mirror_bcs_f(self, boundary):
+def apply_mirror_bcs_f_cartesian(self, boundary):
     """
     Applies mirror boundary conditions along boundary specified 
-    for the distribution function
+    for the distribution function when momentum space is on a cartesian grid
 
     Parameters
     ----------
@@ -310,6 +310,108 @@ def apply_mirror_bcs_f(self, boundary):
 
     return
 
+def apply_mirror_bcs_f_polar2D(self, boundary):
+    """
+    Applies mirror boundary conditions along boundary specified 
+    for the distribution function when momentum space is on a 2D polar grid
+
+    Parameters
+    ----------
+    boundary: str
+              Boundary along which the boundary condition is to be applied.
+    """
+
+    N_g = self.N_ghost
+
+    if(boundary == 'left'):
+        # x-0-x-0-x-0-|-0-x-0-x-0-x-....
+        #   0   1   2   3   4   5
+        # For mirror boundary conditions:
+        # 0 = 5; 1 = 4; 2 = 3;
+        self.f[:, :, :N_g] = af.flip(self.f[:, :, N_g:2 * N_g], 2)
+
+        # For a particle moving with initial momentum at an angle \theta
+        # with the x-axis, a collision with the left boundary changes
+        # the angle of momentum after reflection to (pi - \theta)
+        # To do this, we split the array into to equal halves,
+        # flip each of the halves along the p_theta axis and then
+        # join the two flipped halves together.
+
+        N_theta = self.N_p2
+
+        tmp1 = self._convert_to_p_expanded(self.f)[:, :N_theta/2, :, :]
+        tmp1 = af.flip(tmp1, 1)
+        tmp2 = self._convert_to_p_expanded(self.f)[:, N_theta/2:, :, :]
+        tmp2 = af.flip(tmp2, 1)
+        tmp = af.join(1, tmp1, tmp2)
+
+        self.f[:, :, :N_g] = \
+                self._convert_to_q_expanded(tmp)[:, :, :N_g]
+
+    elif(boundary == 'right'):
+        # ...-x-0-x-0-x-0-|-0-x-0-x-0-x
+        #      -6  -5  -4  -3  -2  -1
+        # For mirror boundary conditions:
+        # -1 = -6; -2 = -5; -3 = -4;
+        self.f[:, :, -N_g:] = af.flip(self.f[:, :, -2 * N_g:-N_g], 2)
+
+        # For a particle moving with initial momentum at an angle \theta
+        # with the x-axis, a collision with the right boundary changes
+        # the angle of momentum after reflection to (pi - \theta)
+        # To do this, we split the array into to equal halves,
+        # flip each of the halves along the p_theta axis and then
+        # join the two flipped halves together.
+
+        N_theta = self.N_p2
+
+        tmp1 = self._convert_to_p_expanded(self.f)[:, :N_theta/2, :, :]
+        tmp1 = af.flip(tmp1, 1)
+        tmp2 = self._convert_to_p_expanded(self.f)[:, N_theta/2:, :, :]
+        tmp2 = af.flip(tmp2, 1)
+        tmp = af.join(1, tmp1, tmp2)
+
+        self.f[:, :, -N_g:] = \
+                self._convert_to_q_expanded(tmp)[:, :, -N_g:]
+
+    elif(boundary == 'bottom'):
+        # x-0-x-0-x-0-|-0-x-0-x-0-x-....
+        #   0   1   2   3   4   5
+        # For mirror boundary conditions:
+        # 0 = 5; 1 = 4; 2 = 3;
+        self.f[:, :, :, :N_g] = af.flip(self.f[:, :, :, N_g:2 * N_g], 3)
+
+        # For a particle moving with initial momentum at an angle \theta
+        # with the x-axis, a collision with the bottom boundary changes
+        # the angle of momentum after reflection to (2*pi - \theta) = (-\theta)
+        # To do this we flip the axis that contains the variation in p_theta
+        self.f[:, :, :, :N_g] = \
+            self._convert_to_q_expanded(af.flip(self._convert_to_p_expanded(self.f), 
+                                                1
+                                               )
+                                       )[:, :, :, :N_g]
+
+    elif(boundary == 'top'):
+        # ...-x-0-x-0-x-0-|-0-x-0-x-0-x
+        #      -6  -5  -4  -3  -2  -1
+        # For mirror boundary conditions:
+        # -1 = -6; -2 = -5; -3 = -4;
+        self.f[:, :, :, -N_g:] = af.flip(self.f[:, :, :, -2 * N_g:-N_g], 3)
+
+        # For a particle moving with initial momentum at an angle \theta
+        # with the x-axis, a collision with the top boundary changes
+        # the angle of momentum after reflection to (2*pi - \theta) = (-\theta)
+        # To do this we flip the axis that contains the variation in p_theta
+        self.f[:, :, :, -N_g:] = \
+            self._convert_to_q_expanded(af.flip(self._convert_to_p_expanded(self.f), 
+                                                1
+                                               )
+                                       )[:, :, :, -N_g:]
+
+    else:
+        raise Exception('Invalid choice for boundary')
+
+    return
+
 def apply_bcs_f(self):
     """
     Applies boundary conditions to the distribution function as specified by 
@@ -332,14 +434,26 @@ def apply_bcs_f(self):
             apply_dirichlet_bcs_f(self, 'left')
 
         elif(self.boundary_conditions.in_q1_left == 'mirror'):
-            apply_mirror_bcs_f(self, 'left')            
+            if (self.physical_system.params.p_space_grid == 'cartesian'):
+                apply_mirror_bcs_f_cartesian(self, 'left')            
+            elif (self.physical_system.params.p_space_grid == 'polar2D'):
+                apply_mirror_bcs_f_polar2D(self, 'left')
+            else :
+                raise NotImplementedError('Unsupported coordinate system in p_space')
 
         elif(self.boundary_conditions.in_q1_left == 'mirror+dirichlet'):
-            apply_mirror_bcs_f(self, 'left')            
+            if (self.physical_system.params.p_space_grid == 'cartesian'):
+                apply_mirror_bcs_f_cartesian(self, 'left')
+            elif (self.physical_system.params.p_space_grid == 'polar2D'):
+                apply_mirror_bcs_f_polar2D(self, 'left')
+            else :
+                raise NotImplementedError('Unsupported coordinate system in p_space')
             apply_dirichlet_bcs_f(self, 'left')
 
         # This is automatically handled by the PETSc function globalToLocal()
-        elif(self.boundary_conditions.in_q1_left == 'periodic'):
+        elif(   self.boundary_conditions.in_q1_left == 'periodic'
+             or self.boundary_conditions.in_q1_left == 'none' # no ghost zones (1D)
+            ):
             pass
 
         elif(self.boundary_conditions.in_q1_left == 'shearing-box'):
@@ -355,14 +469,26 @@ def apply_bcs_f(self):
             apply_dirichlet_bcs_f(self, 'right')
 
         elif(self.boundary_conditions.in_q1_right == 'mirror'):
-            apply_mirror_bcs_f(self, 'right')
+            if (self.physical_system.params.p_space_grid == 'cartesian'):
+                apply_mirror_bcs_f_cartesian(self, 'right')
+            elif (self.physical_system.params.p_space_grid == 'polar2D'):
+                apply_mirror_bcs_f_polar2D(self, 'right')
+            else:
+                raise NotImplementedError('Unsupported coordinate system in p_space')
 
         elif(self.boundary_conditions.in_q1_right == 'mirror+dirichlet'):
-            apply_mirror_bcs_f(self, 'right')            
+            if (self.physical_system.params.p_space_grid == 'cartesian'):
+                apply_mirror_bcs_f_cartesian(self, 'right')
+            elif (self.physical_system.params.p_space_grid == 'polar2D'):
+                apply_mirror_bcs_f_polar2D(self, 'right')
+            else:
+                raise NotImplementedError('Unsupported coordinate system in p_space')
             apply_dirichlet_bcs_f(self, 'right')
-        
+
         # This is automatically handled by the PETSc function globalToLocal()
-        elif(self.boundary_conditions.in_q1_right == 'periodic'):
+        elif(   self.boundary_conditions.in_q1_right == 'periodic'
+             or self.boundary_conditions.in_q1_right == 'none' # no ghost zones (1D)
+            ):
             pass
 
         elif(self.boundary_conditions.in_q1_right == 'shearing-box'):
@@ -378,14 +504,26 @@ def apply_bcs_f(self):
             apply_dirichlet_bcs_f(self, 'bottom')
 
         elif(self.boundary_conditions.in_q2_bottom == 'mirror'):
-            apply_mirror_bcs_f(self, 'bottom')            
+            if (self.physical_system.params.p_space_grid =='cartesian'):
+                apply_mirror_bcs_f_cartesian(self, 'bottom')
+            elif (self.physical_system.params.p_space_grid == 'polar2D'):
+                apply_mirror_bcs_f_polar2D(self, 'bottom')
+            else:
+                raise NotImplementedError('Unsupported coordinate system in p_space')
 
         elif(self.boundary_conditions.in_q2_bottom == 'mirror+dirichlet'):
-            apply_mirror_bcs_f(self, 'bottom')            
+            if (self.physical_system.params.p_space_grid == 'cartesian'):
+                apply_mirror_bcs_f_cartesian(self, 'bottom')
+            elif (self.physical_system.params.p_space_grid == 'polar2D'):
+                apply_mirror_bcs_f_polar2D(self, 'bottom')
+            else:
+                raise NotImplementedError('Unsupported coordinate system in p_space')
             apply_dirichlet_bcs_f(self, 'bottom')
 
         # This is automatically handled by the PETSc function globalToLocal()
-        elif(self.boundary_conditions.in_q2_bottom == 'periodic'):
+        elif(   self.boundary_conditions.in_q2_bottom == 'periodic'
+             or self.boundary_conditions.in_q2_bottom == 'none' # no ghost zones (1D)
+            ):
             pass
 
         elif(self.boundary_conditions.in_q2_bottom == 'shearing-box'):
@@ -401,14 +539,26 @@ def apply_bcs_f(self):
             apply_dirichlet_bcs_f(self, 'top')
 
         elif(self.boundary_conditions.in_q2_top == 'mirror'):
-            apply_mirror_bcs_f(self, 'top')
+            if (self.physical_system.params.p_space_grid == 'cartesian'):
+                apply_mirror_bcs_f_cartesian(self, 'top')
+            elif (self.physical_system.params.p_space_grid == 'polar2D'):
+                apply_mirror_bcs_f_polar2D(self, 'top')
+            else:
+                raise NotImplementedError('Unsupported coordinate system in p_space')
 
         elif(self.boundary_conditions.in_q2_top == 'mirror+dirichlet'):
-            apply_mirror_bcs_f(self, 'top')            
+            if (self.physical_system.params.p_space_grid == 'cartesian'):
+                apply_mirror_bcs_f_cartesian(self, 'top')            
+            elif (self.physical_system.params.p_space_grid == 'polar2D'):
+                apply_mirror_bcs_f_polar2D(self, 'top')
+            else:
+                raise NotImplementedError('Unsupported coordinate system in p_space')
             apply_dirichlet_bcs_f(self, 'top')
 
         # This is automatically handled by the PETSc function globalToLocal()
-        elif(self.boundary_conditions.in_q2_top == 'periodic'):
+        elif(   self.boundary_conditions.in_q2_top == 'periodic'
+             or self.boundary_conditions.in_q2_top == 'none' # no ghost zones (1D)
+            ):
             pass
 
         elif(self.boundary_conditions.in_q2_top == 'shearing-box'):

bolt/lib/nonlinear/communicate.py

@@ -2,6 +2,8 @@
 # -*- coding: utf-8 -*-
 
 import arrayfire as af
+from bolt.lib.utils.af_petsc_conversion import af_to_petsc_glob_array
+from bolt.lib.utils.af_petsc_conversion import petsc_local_array_to_af
 
 def communicate_f(self):
     """
@@ -14,43 +16,19 @@ def communicate_f(self):
     if(self.performance_test_flag == True):
         tic = af.time()
 
-    # Obtaining start coordinates for the local zone
-    # Additionally, we also obtain the size of the local zone
-    ((i_q1_start, i_q2_start), (N_q1_local, N_q2_local)) = self._da_f.getCorners()
-
-    N_g = self.N_ghost
-
-    # Assigning the local array only when Dirichlet
-    # boundary conditions are applied. This is needed since
-    # only inflowing characteristics are to be changed by 
-    # the apply boundary conditions function.
-
-    if(   self.boundary_conditions.in_q1_left   == 'dirichlet'
-       or self.boundary_conditions.in_q1_right  == 'dirichlet' 
-       or self.boundary_conditions.in_q2_bottom == 'dirichlet' 
-       or self.boundary_conditions.in_q2_top    == 'dirichlet' 
-      ):
-        af.flat(self.f).to_ndarray(self._local_f_array)
-
-    # Global value is non-inclusive of the ghost-zones:
-    af.flat(self.f[:, :, N_g:-N_g, N_g:-N_g]).to_ndarray(self._glob_f_array)
+    # Transfer data from af.Array to PETSc.Vec
+    af_to_petsc_glob_array(self, self.f, self._glob_f_array)
 
     # The following function takes care of interzonal communications
     # Additionally, it also automatically applies periodic BCs when necessary
     self._da_f.globalToLocal(self._glob_f, self._local_f)
 
     # Converting back from PETSc.Vec to af.Array:
-    f_flattened = af.to_array(self._local_f_array)
-    self.f      = af.moddims(f_flattened,
-                               self.N_p1 
-                             * self.N_p2 
-                             * self.N_p3,
-                             self.N_species,
-                             N_q1_local + 2 * N_g,
-                             N_q2_local + 2 * N_g
-                            )
-
-    af.eval(self.f)
+    self.f = petsc_local_array_to_af(self, 
+                                     self.N_p1*self.N_p2*self.N_p3,
+                                     self.N_species,
+                                     self._local_f_array
+                                    )
 
     if(self.performance_test_flag == True):
         af.sync()
@@ -59,7 +37,6 @@ def communicate_f(self):
 
     return
 
-
 def communicate_fields(self, on_fdtd_grid = False):
     """
     Used in communicating the values at the boundary zones for each of
@@ -72,46 +49,26 @@ def communicate_fields(self, on_fdtd_grid = False):
     if(self.performance_test_flag == True):
         tic = af.time()
 
-    # Obtaining start coordinates for the local zone
-    # Additionally, we also obtain the size of the local zone
-    ((i_q1_start, i_q2_start), (N_q1_local, N_q2_local)) = self._da_fields.getCorners()
-
-    N_g = self.N_g
-
     # Assigning the values of the af.Array 
     # fields quantities to the PETSc.Vec:
     if(on_fdtd_grid is True):
-        flattened_global_EM_fields_array = \
-            af.flat(self.yee_grid_EM_fields[:, :, N_g:-N_g, N_g:-N_g])
-        flattened_global_EM_fields_array.to_ndarray(self._glob_fields_array)
-
+        tmp_array = self.yee_grid_EM_fields
     else:
-        flattened_global_EM_fields_array = \
-            af.flat(self.cell_centered_EM_fields[:, :, N_g:-N_g, N_g:-N_g])
-        flattened_global_EM_fields_array.to_ndarray(self._glob_fields_array)
+        tmp_array = self.cell_centered_EM_fields
+
+    af_to_petsc_glob_array(self, tmp_array, self._glob_fields_array)
 
     # Takes care of boundary conditions and interzonal communications:
     self._da_fields.globalToLocal(self._glob_fields, self._local_fields)
 
     # Converting back to af.Array
-    if(on_fdtd_grid is True):
-
-        self.yee_grid_EM_fields = af.moddims(af.to_array(self._local_fields_array),
-                                             6, 1, N_q1_local + 2 * N_g,
-                                             N_q2_local + 2 * N_g
-                                            )
-
-        af.eval(self.yee_grid_EM_fields)
+    tmp_array = petsc_local_array_to_af(self, 6, 1, self._local_fields_array)
 
+    if(on_fdtd_grid is True):
+        self.yee_grid_EM_fields = tmp_array
     else:
-
-        self.cell_centered_EM_fields = af.moddims(af.to_array(self._local_fields_array),
-                                                  6, 1, N_q1_local + 2 * N_g,
-                                                  N_q2_local + 2 * N_g
-                                                 )
+        self.cell_centered_EM_fields = tmp_array
 
-        af.eval(self.cell_centered_EM_fields)
-
     if(self.performance_test_flag == True):
         af.sync()
         toc = af.time()

bolt/lib/nonlinear/compute_moments.py

@@ -38,10 +38,12 @@ def compute_moments(self, moment_name, f=None):
 
     if(f is None):
         f = self.f
+
+    #measure = (4./(2.*np.pi*self.physical_system.params.h_bar)**2) * self.dp3 * self.dp2 * self.dp1
 
     moment = af.broadcast(getattr(self.physical_system.moments, 
                                   moment_name
-                                 ), f, p1, p2, p3, self.dp3 * self.dp2 * self.dp1
+                                 ), f, p1, p2, p3, self.physical_system.params.integral_measure
                          )
 
     af.eval(moment)