Adding in fixes for KE in FFS

Makefile.depend

@@ -468,6 +468,7 @@ response_matrix.o: \
 	 linear_solve.o \
 	 mp.o \
 	 mp_lu_decomposition.o \
+	 parallel_streaming.o \
 	 physics_flags.o \
 	 physics_parameters.o \
 	 run_parameters.o \

extended_zgrid.f90

@@ -16,6 +16,7 @@ module extended_zgrid
    public :: map_from_extended_zgrid
    public :: map_to_iz_ikx_from_izext
 
+   public :: enforce_reality
    private
 
    !> these arrays needed to keep track of connections between different
@@ -433,6 +434,37 @@ subroutine map_from_extended_zgrid(it, ie, iky, gext, g)
 
    end subroutine map_from_extended_zgrid
 
+   subroutine enforce_reality (field)
+
+     use zgrid, only: nzgrid, ntubes
+     use kt_grids, only: naky
+
+     implicit none
+
+     complex, dimension(:, :, -nzgrid:, :), intent(inout) :: field
+     complex, dimension(:), allocatable :: field_ext
+     integer :: ulim, nz_ext
+     integer :: it, ie, iky
+
+     field(1,1,:,:) = 0.0
+
+     do iky = 1, naky
+        do it = 1, ntubes
+           do ie = 1, neigen(iky)
+              nz_ext = nsegments(ie, iky) * nzed_segment + 1
+              allocate (field_ext(nz_ext)); field_ext = 0.0
+              call map_to_extended_zgrid (it, ie, iky, field(iky, :, :, :), field_ext, ulim)
+              call map_from_extended_zgrid (it, ie, iky, field_ext, field(iky,:, :, :))
+              deallocate(field_ext)
+           end do
+        end do
+!        if (periodic(iky)) field(iky,:,nzgrid,:) = field(iky,:,-nzgrid,:)
+     end do
+
+     field(1, :, -nzgrid, :) = field(1, :, nzgrid, :) 
+   end subroutine enforce_reality
+
+
    subroutine map_to_iz_ikx_from_izext(iky, ie, iz_from_izext, ikx_from_izext)
 
       implicit none

fields.fpp

@@ -1210,7 +1210,10 @@ contains
          if (proc0) call time_message(.false., time_field_solve(:, 3), ' int_dv_g')
 
          call get_phi(phi, dist, skip_fsa_local)
-
+
+         ! Enforce periodicity for zonal modes 
+         phi(1, :, nzgrid, :) = phi(1, :, -nzgrid, :)
+
       else if (fphi > epsilon(0.0) .and. include_bpar) then
          if (proc0) call time_message(.false., time_field_solve(:, 3), ' int_dv_g int_dv_g_vperp2')
 
@@ -1292,6 +1295,8 @@ contains
      use kt_grids, only: akx
      use gyro_averages, only: gyro_average
 
+     use extended_zgrid, only: enforce_reality
+     
      implicit none
      complex, dimension(:, :, -nzgrid:, :), intent (in out) :: source
      complex, dimension(:, :, -nzgrid:, :, vmu_lo%llim_proc:), intent(in) :: gold
@@ -1322,9 +1327,11 @@ contains
      
      if (any(gamtot(1, 1, :) < epsilon(0.))) source(1, 1, :, :) = 0.0
      if (akx(1) < epsilon(0.)) then
-         source(1, 1, :, :) = 0.0
-      end if
+        source(1, 1, :, :) = 0.0
+     end if
 
+!     call enforce_reality(source)
+     
      deallocate(source2, gamtot_t) 
      
    end subroutine get_fields_source
@@ -1335,7 +1342,7 @@ contains
      use mp, only: sum_allreduce
      use stella_layouts, only: vmu_lo
      use species, only: spec
-     use zgrid, only: nzgrid
+     use zgrid, only: nzgrid, ntubes
      use kt_grids, only: naky, nakx
      use vpamu_grids, only: integrate_species_ffs
      use gyro_averages, only: gyro_average, j0_B_ffs
@@ -1344,11 +1351,14 @@ contains
      use stella_layouts, only: iv_idx, imu_idx, is_idx
      use kt_grids, only: nalpha
      
+     use extended_zgrid, only: enforce_reality
+
      implicit none
      
      complex, dimension(:, :, -nzgrid:, :, vmu_lo%llim_proc:), intent(in) :: g
      complex, dimension(:, :, -nzgrid:), intent(in out) :: source
-     
+
+     complex, dimension (:,:,:,:), allocatable :: source_copy 
      integer :: it, iz, ivmu
      complex, dimension(:, :, :), allocatable :: gyro_g, gyro_g2
      
@@ -1374,6 +1384,12 @@ contains
      !> gather sub-sums from each processor and add them together
      !> store result in phi, which will be further modified below to account for polarization term
      call sum_allreduce(source)
+
+     allocate(source_copy(naky,nakx, -nzgrid:nzgrid, ntubes)) ; source_copy = 0.0
+     source_copy = spread(source, 4, ntubes)
+     call enforce_reality (source_copy)
+     deallocate(source_copy)
+     
      !> no longer need <g>, so deallocate
      deallocate (gyro_g, gyro_g2)
 
@@ -1403,6 +1419,8 @@ contains
       use physics_flags, only: adiabatic_option_switch
       use physics_flags, only: adiabatic_option_fieldlineavg
       use stella_geometry, only: dl_over_b
+
+      use extended_zgrid, only: enforce_reality
       implicit none
 
       complex, dimension(:, :, -nzgrid:, :, vmu_lo%llim_proc:), intent(in) :: g
@@ -1524,6 +1542,9 @@ contains
          phi(1, 1, :, :) = 0.
       end if
 
+      phi(1, 1, :, :) = 0.
+!      call enforce_reality (phi) 
+      
       deallocate (source)
       apar = 0.
       if (include_apar) then
@@ -1537,20 +1558,22 @@ contains
          use mp, only: sum_allreduce
          use stella_layouts, only: vmu_lo
          use species, only: spec
-         use zgrid, only: nzgrid
+         use zgrid, only: nzgrid, ntubes
          use kt_grids, only: naky, nakx
          use vpamu_grids, only: integrate_species_ffs
          use gyro_averages, only: gyro_average, j0_B_ffs
 
          use gyro_averages, only: j0_B_const
          use stella_layouts, only: iv_idx, imu_idx, is_idx
          use kt_grids, only: nalpha
-
+         
+         use extended_zgrid, only: enforce_reality
          implicit none
 
          complex, dimension(:, :, -nzgrid:, :, vmu_lo%llim_proc:), intent(in) :: g
          complex, dimension(:, :, -nzgrid:), intent(in out) :: source
-
+         complex,  dimension(:, :, :, :), allocatable :: source_copy
+         
          integer :: it, iz, ivmu
          complex, dimension(:, :, :), allocatable :: gyro_g
          logical, optional, intent(in) :: implicit_solve
@@ -1582,6 +1605,14 @@ contains
          !> gather sub-sums from each processor and add them together
          !> store result in phi, which will be further modified below to account for polarization term
          call sum_allreduce(source)
+
+         if (present(implicit_solve)) then
+            allocate(source_copy(naky,nakx, -nzgrid:nzgrid, ntubes)) ; source_copy = 0.0
+            source_copy = spread(source, 4, ntubes)
+            call enforce_reality (source_copy)
+            source = source_copy (:,:,:,1) 
+            deallocate(source_copy)
+         end if
          !> no longer need <g>, so deallocate
          deallocate (gyro_g)
 

implicit_solve.f90

@@ -135,8 +135,8 @@ subroutine advance_implicit_terms(g, phi, apar, bpar)
          if (driftkinetic_implicit) then
             call get_source_ffs_itteration (phi_old, g2, phi_source_ffs)
 !!!!!!!            call get_drifts_ffs_itteration (phi_old, g2, drifts_source_ffs)
-!!            phi_source_ffs = phi_source_ffs + drifts_source_ffs
-            phi_source = 0.0
+            !!            phi_source_ffs = phi_source_ffs + drifts_source_ffs
+            phi_source = tupwnd_m * phi
             !> set the g on the RHS to be g from the previous time step  
             !> FFS will have a RHS source term
             !> modify being passes in will make sure that this source is included
@@ -176,7 +176,8 @@ subroutine advance_implicit_terms(g, phi, apar, bpar)
          if (use_deltaphi_for_response_matrix .and. include_bpar) bpar = bpar + bpar_old
 
          if (driftkinetic_implicit) then
-            phi_source = phi
+            !!!!!!phi_source = phi
+            phi_source = tupwnd_m * phi_old + tupwnd_p * phi
          else
 	    ! get time-centered phi
             phi_source = tupwnd_m * phi_old + tupwnd_p * phi
@@ -289,7 +290,7 @@ subroutine update_pdf(mod)
                      ! to the standard zed domain; the mapped pdf is called 'g'
                      call map_from_extended_zgrid(it, ie, iky, pdf2, g(iky, :, :, :, ivmu))
                      deallocate (pdf1, pdf2, phiext, aparext, aparext_new, aparext_old, bparext)
-		               if (present(mod)) deallocate(phiffsext)
+                     if (present(mod)) deallocate(phiffsext)
                   end do
                end do
             end do

init_g.f90

@@ -106,7 +106,8 @@ subroutine ginit(restarted, istep0)
 
       use stella_save, only: init_tstart
       use run_parameters, only: maxwellian_normalization
-
+      use physics_flags, only: full_flux_surface
+
       logical, intent(out) :: restarted
       integer, intent(out) :: istep0
       integer :: istatus
@@ -115,7 +116,11 @@ subroutine ginit(restarted, istep0)
       istep0 = 0
       select case (ginitopt_switch)
       case (ginitopt_default)
-         call ginit_default
+         if(full_flux_surface) then 
+            call ginit_default_ffs
+         else
+            call ginit_default
+         end if
       case (ginitopt_noise)
          call ginit_noise
       case (ginitopt_kpar)
@@ -288,6 +293,132 @@ subroutine ginit_default
 
    end subroutine ginit_default
 
+   subroutine ginit_default_ffs
+
+     use constants, only: zi
+     use species, only: spec
+     use zgrid, only: nzgrid, zed, ntubes, delzed
+     use kt_grids, only: naky, nakx, ikx_max, naky_all, ny 
+     use kt_grids, only: zed0, akx
+     use kt_grids, only: reality, zonal_mode
+     use vpamu_grids, only: nvpa, nmu
+     use vpamu_grids, only: vpa
+     use vpamu_grids, only: maxwell_vpa, maxwell_mu, maxwell_fac, maxwell_mu_avg
+     use dist_fn_arrays, only: gvmu, gnew
+     use stella_layouts, only: kxkyz_lo, iz_idx, ikx_idx, iky_idx, is_idx
+     use stella_layouts, only: imu_idx, iv_idx
+     use extended_zgrid, only: map_to_iz_ikx_from_izext
+     use redistribute, only: scatter
+     use dist_redistribute, only: kxkyz2vmu
+     use stella_layouts, only: vmu_lo
+     use extended_zgrid, only: map_from_extended_zgrid
+     use extended_zgrid, only: nsegments, nzed_segment
+     use extended_zgrid, only: neigen, enforce_reality
+     use constants, only: pi
+
+     use kt_grids, only: swap_kxky, swap_kxky_back
+     use stella_transforms, only: transform_ky2y, transform_y2ky
+     implicit none
+
+     complex, dimension (:), allocatable :: phiext
+     complex, dimension(naky, nakx, -nzgrid:nzgrid, ntubes) :: phi
+     complex, dimension(:,:,:,:), allocatable :: g0x
+     complex, dimension(:,:,:), allocatable :: phiy
+     complex, dimension (:,:), allocatable :: g_swap
+
+     integer :: nz_ext
+     real, dimension (:), allocatable :: zed_ext
+     logical :: right
+     integer :: ie, it, j
+     integer :: imu,is, iv, ivmu
+     integer :: iz, iky, ikx, ia, iy
+     integer :: ikxkyz
+
+     integer, dimension(:), allocatable :: iz_from_izext
+     right = .not. left
+
+     it = 1
+     do iky = 1, naky
+        do it = 1, ntubes
+           do ie = 1, neigen(iky)
+              nz_ext = nsegments(ie, iky) * nzed_segment + 1
+              allocate (phiext(nz_ext))
+              allocate (zed_ext(nz_ext))
+              zed_ext = [((2 * pi * j) / (nz_ext + 1), j=-nz_ext/2 , nz_ext/2)]
+              do iz = 1, nz_ext
+                 phiext(iz) = exp(-(zed_ext(iz)/width0)**2) * cmplx(1.0, 1.0)
+              end do
+              call map_from_extended_zgrid(it, ie, iky, phiext, phi(iky, :, :, :))
+              deallocate(phiext, zed_ext)
+           end do
+        end do
+     end do
+
+     ! if (chop_side) then
+     !    if (left) phi(:, :, :-1, :) = 0.0
+     !    if (right) phi(:, :, 1:, :) = 0.0
+     ! end if
+
+     ! do iz = -nzgrid, nzgrid
+     !    phi(:, :, iz,1) = exp(-((zed(iz)) / width0)**2) * cmplx(1.0, 1.0)
+     ! end do
+     phi(1, 1, :, :) = 0.0
+
+     ! if (zonal_mode(1)) then
+     !    if (abs(akx(1)) < epsilon(0.0)) then
+     !       phi(1, 1, :, :) = 0.0
+     !    end if
+
+     !    if (reality) then
+     !       do ikx = 1, nakx - ikx_max
+     !          phi(1, nakx - ikx + 1, :, :) = conjg(phi(1, ikx + 1, :, :))
+     !       end do
+     !    end if
+     ! end if
+
+     allocate (g_swap(naky_all, ikx_max))
+     allocate (phiy(ny, ikx_max, -nzgrid:nzgrid)) ; phiy = 0.0
+     do iz = -nzgrid, nzgrid
+        call swap_kxky(phi(:, :, iz, 1), g_swap)
+        call transform_ky2y(g_swap, phiy(:, :, iz))
+     end do
+
+     allocate (g0x(ny, ikx_max, -nzgrid:nzgrid, vmu_lo%llim_proc:vmu_lo%ulim_alloc)); g0x = 0.0
+     do ivmu = vmu_lo%llim_proc, vmu_lo%ulim_proc
+        iv = iv_idx(vmu_lo, ivmu)
+        imu = imu_idx(vmu_lo, ivmu)
+        is = is_idx(vmu_lo, ivmu)
+        do iy = 1, ny
+           do ikx = 1, ikx_max
+              do iz = -nzgrid, nzgrid 
+                 g0x(iy, ikx, iz, ivmu) = phiinit * phiy(iy, ikx, iz) / abs(spec(is)%z) &
+                      * (den0 + 2.0 * zi * vpa(iv) * upar0) &
+!                      * maxwell_mu(iy, iz, imu, is) &
+                      * maxwell_vpa(iv, is) * maxwell_fac(is)
+              end do
+           end do
+        end do
+     end do
+
+     do ivmu = vmu_lo%llim_proc, vmu_lo%ulim_proc
+        do iz = -nzgrid, nzgrid
+           call transform_y2ky(g0x(:, :, iz, ivmu), g_swap)
+           call swap_kxky_back(g_swap, gnew(:, :, iz, 1, ivmu))
+        end do
+     end do
+     gnew = spread(gnew(:,:,:,1,:), 4, ntubes)
+     gnew (1,1,:,:,:) = 0.0
+     do ivmu = vmu_lo%llim_proc, vmu_lo%ulim_proc
+        call enforce_reality (gnew(:, :, :, :, ivmu))
+     end do
+     deallocate (g_swap)
+
+     gvmu = 0.
+     call scatter(kxkyz2vmu, gnew, gvmu)
+
+     deallocate(phiy, g0x)
+
+   end subroutine ginit_default_ffs
    ! initialize two kys and kx=0
 !   subroutine ginit_nltest