make zonal diagnostic for all kx

Author: GeorgiaActon

COMPILATION/Makefile.depend

@@ -1395,4 +1395,5 @@ $(OBJ_DIR)/write_diagnostics_to_netcdf.o: \
 	 $(OBJ_DIR)/mp.o \
 	 $(OBJ_DIR)/neasyf.o \
 	 $(OBJ_DIR)/netcdf_utils.o \
+	 $(OBJ_DIR)/parameters_diagnostics.o \
 	 $(OBJ_DIR)/parameters_physics.o 

POST_PROCESSING/post_processing/gzvs_movie.py

@@ -1,6 +1,6 @@
 import numpy as np
 from stella_plots import movie_2d
-from stella_data import gzvs, vpa, zed, ntime, g2nozonal_vs_zvpas, nzed
+from stella_data import gzvs, vpa, zed, ntime, g2nozonal_vs_zvpas, nzed, time
 from stella_data import gzvs_zonal
 from stella_data import outdir
 
@@ -9,22 +9,17 @@
 gmax = np.arange(ntime,dtype=float)
 gmin = np.arange(ntime,dtype=float)
 
-#print('SHAPE=', gzvs_zonal.shape)
-
-#g_zonal = np.zeros((ntime, nvpa,  nzed), dtype=complex)
-#g_zonal = np.array(gzvs_zonal[:, 0, :, :, 0, 0] + 1j*gzvs_zonal[:, 0, :, :, 0, 1], dtype=complex) #gzvs - g2nozonal_vs_zvpas
-
-#print('shape zonal', np.shape(g_zonal))
+#(t, species, vpa, tube, zed, kx, ri)
+g_zonal = np.array(gzvs_zonal[:, 0, :, 0, :, 1, 0] + 1j*gzvs_zonal[:, 0, :, 0, :, 1, 1], dtype=complex) 
 
 for i in range(ntime):
-    gmax[i] = np.absolute(gzvs[i,0,:,:].max())
-    #gmax[i] = np.abs(g_zonal[i,:,:]).max()
+    gmax[i] = np.abs(g_zonal[i,:,:]).max()
+    
+gmin[:] = -gmax[:]
 
-gmin[:] = - gmax[:]
 ylabel = '$v_{\parallel}$'
 xlabel = '$z$'
 title = '$\int d\mu \int d^2 \mathbf{R} g^2$'
 movie_file = outdir + 'gzvs.mp4'
 
-movie_2d(gzvs[:,0,:,:],zed,vpa,gmin,gmax,ntime-1,movie_file,xlabel,ylabel,title,cmp='YlGnBu', abs_flag=False)
-#movie_2d(g_zonal,zed,vpa,gmin,gmax,ntime-1,movie_file,xlabel,ylabel,title,cmp='RdBu', abs_flag=False)
+movie_2d(g_zonal,zed,vpa,gmin,gmax,ntime-1,movie_file,xlabel,ylabel,title,cmp='RdBu', abs_flag=True)

POST_PROCESSING/post_processing/stella_data.py

@@ -27,7 +27,6 @@
 # note stella orders kx as (0, dkx, ..., kx_max, -kx_max, -kx_max+dkx, ..., -dkx)
 nakx_mid = nakx//2+1
 kx = np.concatenate((kx_stella[nakx_mid:],kx_stella[:nakx_mid]))
-
 print('nakx=', nakx)
 # get zed grid
 zed = np.copy(ncfile.variables['zed'][:])

POST_PROCESSING/post_processing/stella_plots.py

@@ -83,8 +83,8 @@ def movie_2d(z, xin, yin, zmin, zmax, nframes, outfile,
         z_even = np.sign(np.real(z_even_full)) * np.abs(z_even_full)
         z_odd  = np.sign(np.real(z_odd_full)) * np.abs(z_odd_full)
     else:
-        z_even = np.real(z_even_full)#0.5*(z + np.flip(z, axis=1)))
-        z_odd  = np.real(z_odd_full) #0.5*(z - np.flip(z, axis=1)))
+        z_even = np.real(z_even_full)
+        z_odd  = np.real(z_odd_full) 
 
     # Ensure frames list respects step
     frames = list(range(0, nframes, step))

STELLA_CODE/calculations/calculations_velocity_integrals.f90

@@ -58,6 +58,7 @@ module calculations_velocity_integrals
    interface integrate_mu
       module procedure integrate_mu_local
       module procedure integrate_mu_nonlocal
+      module procedure integrate_mu_nonlocal_complex
    end interface
 
    ! Integrations over the parallel velocity
@@ -145,6 +146,49 @@ subroutine integrate_mu_nonlocal(iz, g, total)
       if (nproc > 1) call sum_reduce(total, 0)
 
    end subroutine integrate_mu_nonlocal
+
+
+   !---------------- Each processor has a number of ivmu points ----------------
+   subroutine integrate_mu_nonlocal_complex(iz, g, total)
+
+      use mp, only: nproc, sum_allreduce
+      use parallelisation_layouts, only: vmu_lo
+      use parallelisation_layouts, only: is_idx, imu_idx, iv_idx
+
+      implicit none
+
+      ! Arguments
+      integer, intent(in) :: iz
+      complex, dimension(vmu_lo%llim_proc:), intent(in) :: g
+      complex, dimension(:, :), intent(out) :: total
+
+      ! Local variables
+      integer :: is, imu, iv, ivmu, ia
+      
+      !-------------------------------------------------------------------------
+
+      ! Initialise sum
+      total = 0.
+      ! Assume we only have a single flux tube
+      ia = 1
+      
+      ! Iterate over the i[vpa, mu, s] points
+      do ivmu = vmu_lo%llim_proc, vmu_lo%ulim_proc
+      
+         ! Obtain the [ivpa, imu, is] indices from [ivmus]
+         is = is_idx(vmu_lo, ivmu)
+         imu = imu_idx(vmu_lo, ivmu)
+         iv = iv_idx(vmu_lo, ivmu)
+         
+         ! Integrate over mu
+         total(iv, is) = total(iv, is) + wgts_mu(ia, iz, imu) * g(ivmu)
+         
+      end do
+
+      ! Each processor has a few [ivmu] points, so sum all calculations
+      call sum_allreduce(total)
+
+   end subroutine integrate_mu_nonlocal_complex
 
 !###############################################################################
 !############################## INTEGRALS OVER VPA #############################
@@ -291,7 +335,6 @@ subroutine integrate_vmu_vmulo_complex(g, weights, total)
             total(:, :, iz, :, is) = total(:, :, iz, :, is) + &
                wgts_mu(ia, iz, imu) * wgts_vpa(iv) * g(:, :, iz, :, ivmu) * weights(is)
          end do
-         
       end do
 
       ! Each processor has a few [ivmu] points, so sum all calculations

STELLA_CODE/diagnostics/diagnostics_distribution.f90

@@ -128,7 +128,7 @@ subroutine write_distribution_to_netcdf_file(nout, timer)
       ! Dimensions
       use grids_kxky, only: nakx, naky
       use grids_velocity, only: nvpa, nmu
-      use grids_z, only: nztot, ntubes
+      use grids_z, only: nztot, ntubes, nzgrid
       use grids_species, only: nspec
 
       ! Write to netcdf file
@@ -176,7 +176,6 @@ subroutine write_distribution_to_netcdf_file(nout, timer)
 
       use parameters_diagnostics, only: write_g_vs_zvpas_zonal
       use parameters_diagnostics, only: write_free_energy
-      use parameters_diagnostics, only: number_zonals_kxs, zonal_iks
 
       ! Calculations
       use calculations_tofrom_ghf, only: g_to_f, g_to_h
@@ -193,7 +192,7 @@ subroutine write_distribution_to_netcdf_file(nout, timer)
       real, dimension(:, :, :, :), allocatable :: g2_vs_zvpas, g2_vs_zmus, g2nozonal_vs_zvpas, g2nozonal_vs_zmus
 
       ! Zonals
-      real, dimension(:, :, :, :, :, :), allocatable :: g_vs_zvpas_zonal
+      complex, dimension(:, :, :, :, :), allocatable :: g_vs_zvpas_zonal
       real, dimension(:), allocatable :: free_energy_vs_kx
 
       real, dimension(:, :, :, :, :), allocatable :: g2_vs_zkykxs, g2_vs_zvpamus, g2nozonal_vs_zvpamus
@@ -222,7 +221,7 @@ subroutine write_distribution_to_netcdf_file(nout, timer)
       if (write_g2_vs_kxkyzs) allocate (g2_vs_zkykxs(ntubes, nztot, naky, nakx, nspec))
 
       ! Zonal diagnostics
-      if (write_g_vs_zvpas_zonal) allocate (g_vs_zvpas_zonal(2, number_zonals_kxs, ntubes, nztot, nvpa, nspec))
+      if (write_g_vs_zvpas_zonal) allocate (g_vs_zvpas_zonal(nakx, -nzgrid:nzgrid, ntubes, nvpa, nspec))
       if (write_free_energy) allocate (free_energy_vs_kx(nakx)) 
 
       ! Redistribute the data from <gnew>(ky,kx,z,tube,i[vpa,mu,s]) to <gvmu>(vpa,mu,i[kx,ky,z,s]) for |g|^2(z,kx,ky,s)
@@ -248,9 +247,7 @@ subroutine write_distribution_to_netcdf_file(nout, timer)
          if (write_g2_vs_zvpamus .and. proc0) call write_g2nozonal_vs_zvpamus_nc(nout, g2nozonal_vs_zvpamus)
 
          ! Zonals
-         if (write_g_vs_zvpas_zonal .and. proc0) then 
-            call write_g_vs_zvpas_zonal_nc(nout, g_vs_zvpas_zonal)
-         end if 
+         if (write_g_vs_zvpas_zonal .and. proc0) call write_g_vs_zvpas_zonal_nc(nout, g_vs_zvpas_zonal)
 
          if (write_free_energy) then 
             call g_to_h(gnew, phi, bpar, fphi)
@@ -380,7 +377,6 @@ subroutine calculate_distribution(g_vs_kykxztube, g_vs_vpamuikxkyzs, g2_vs_tzmus
       use parameters_diagnostics, only: write_g2_vs_zvpamus
 
       use parameters_diagnostics, only: write_g_vs_zvpas_zonal
-      use parameters_diagnostics, only: number_zonals_kxs, zonal_iks
 
       implicit none
 
@@ -392,14 +388,14 @@ subroutine calculate_distribution(g_vs_kykxztube, g_vs_vpamuikxkyzs, g2_vs_tzmus
       real, dimension(:, :, :), intent(out) :: g2_vs_vpamus, g2nozonal_vs_vpamus
 
       ! Zonals
-      real, dimension(:, :,  :, :, :, :), intent(out) :: g_vs_zvpas_zonal
+      complex, dimension(:, -nzgrid:, :, :, :), intent(out) :: g_vs_zvpas_zonal
 
       ! Arrays needed to perform calculations
       real, dimension(:, :, :), allocatable :: g2_vs_ztubeivmus, g2nozonal_vs_ztubeivmus 
       real, dimension(:, :), allocatable :: g2_vs_ztube, g2_vs_vpamu
 
       ! Zonals
-      real, dimension(:, :, :, :, :), allocatable :: g_vs_zivmus_zonal
+      complex, dimension(:, :, :, :), allocatable :: g_vs_zivmus_zonal
 
       integer :: ivmus, ia, iz, it, ikx, iky, izp, is, iv, imu, ikxkyzs
       integer :: j, ikx_plot
@@ -414,8 +410,8 @@ subroutine calculate_distribution(g_vs_kykxztube, g_vs_vpamuikxkyzs, g2_vs_tzmus
       if (write_g2_vs_kxkyzs) then; allocate (g2_vs_vpamu(nvpa, nmu)); g2_vs_vpamu = 0.; end if
          allocate (g2_vs_ztube(-nzgrid:nzgrid, ntubes)); g2_vs_ztube = 0.
 
-      if (write_g_vs_zvpas_zonal .and. number_zonals_kxs > 0) then
-         allocate (g_vs_zivmus_zonal(2, number_zonals_kxs, -nzgrid:nzgrid, ntubes, vmu_lo%llim_proc:vmu_lo%ulim_alloc))
+      if (write_g_vs_zvpas_zonal) then
+         allocate (g_vs_zivmus_zonal(nakx, -nzgrid:nzgrid, ntubes, vmu_lo%llim_proc:vmu_lo%ulim_alloc))
          g_vs_zivmus_zonal = 0.0
       end if 
 
@@ -464,28 +460,19 @@ subroutine calculate_distribution(g_vs_kykxztube, g_vs_vpamuikxkyzs, g2_vs_tzmus
 
          ! Zonals
          if (write_g_vs_zvpas_zonal) then
-            if (number_zonals_kxs == 1) then 
+            if (nakx == 1) then 
                if (abs(akx(1)) < epsilon(0.)) then
                   ikx_plot = 2
                else
                   ikx_plot = 1
                end if
-               g_vs_zivmus_zonal(1, 1, :, :, ivmus) = real(g_vs_kykxztube(1, ikx_plot, :, :, ivmus))
-               g_vs_zivmus_zonal(2, 1, :, :, ivmus) = aimag(g_vs_kykxztube(1, ikx_plot, :, :, ivmus))
+               g_vs_zivmus_zonal(1, :, :, ivmus) = g_vs_kykxztube(1, ikx_plot, :, :, ivmus)
 
             else 
 
                ikx_plot = 0
 
-               do ikx = 1, nakx
-                  if (any(zonal_iks == ikx)) then
-                     ikx_plot = ikx_plot + 1
-
-                     g_vs_zivmus_zonal(1, ikx_plot, :, :, ivmus) = real(g_vs_kykxztube(1, ikx, :, :, ivmus))
-                     g_vs_zivmus_zonal(2, ikx_plot, :, :, ivmus) = aimag(g_vs_kykxztube(1, ikx, :, :, ivmus))
-
-                  end if
-               end do
+               g_vs_zivmus_zonal(:, :, :, ivmus) = g_vs_kykxztube(1, :, :, :, ivmus)
             end if 
 
          end if 
@@ -517,25 +504,28 @@ subroutine calculate_distribution(g_vs_kykxztube, g_vs_vpamuikxkyzs, g2_vs_tzmus
       ! The velocity parallel integration takes |g|^2(ivmus) and returns int dvpa |g|^2(ivmus) = |g|^2(mu,s)
       ! There is a sum_reduce() inside of the velocity integration, so <g2_vs_tzmus> holds the total sum
       if (debug) write (*, *) 'diagnostics::diagnostics_distribution::calculate_distribution::write_g2_vs_zvpas .or. write_g2_vs_zmus'
-      if (write_g2_vs_zvpas .or. write_g2_vs_zmus .or. write_g_vs_zvpas_zonal) then
+      if (write_g2_vs_zvpas .or. write_g2_vs_zmus) then
          do it = 1, ntubes
             do iz = -nzgrid, nzgrid
                izp = iz + nzgrid + 1
                if (write_g2_vs_zvpas) call integrate_mu(iz, g2_vs_ztubeivmus(iz, it, :), g2_vs_zvpas(it, izp, :, :))
-               if (write_g_vs_zvpas_zonal) then 
-                  do ikx = 1, number_zonals_kxs
-                     do j = 1, 2
-                        call integrate_mu(iz, g_vs_zivmus_zonal(j, ikx, iz, it, :), g_vs_zvpas_zonal(j, ikx, it, izp, :, :))
-                     end do 
-                  end do
-               end if 
                if (write_g2_vs_zmus) call integrate_vpa(g2_vs_ztubeivmus(iz, it, :), g2_vs_tzmus(it, izp, :, :))
                if (write_g2_vs_zvpas) call integrate_mu(iz, g2nozonal_vs_ztubeivmus(iz, it, :), g2nozonal_vs_zvpas(it, izp, :, :))
                if (write_g2_vs_zmus) call integrate_vpa(g2nozonal_vs_ztubeivmus(iz, it, :), g2nozonal_vs_tzmus(it, izp, :, :))
             end do
          end do
       end if
 
+      if (write_g_vs_zvpas_zonal) then
+         do it = 1, ntubes
+            do iz = -nzgrid, nzgrid
+               do ikx = 1, nakx
+                  call integrate_mu(iz, g_vs_zivmus_zonal(ikx, iz, it, :), g_vs_zvpas_zonal(ikx, iz, it, :, :))
+               end do
+            end do
+         end do
+      end if 
+
       ! For the field line average normalise to account for contributions from multiple flux tubes
       ! in a flux tube train, and sum the values on all processors and to send them to <proc0>
       if (debug) write (*, *) 'diagnostics::diagnostics_distribution::calculate_distribution::sum_all_reduce'
@@ -552,6 +542,10 @@ subroutine calculate_distribution(g_vs_kykxztube, g_vs_vpamuikxkyzs, g2_vs_tzmus
          if (nproc > 1) call sum_reduce(g2_vs_zkykxs, 0)
       end if
 
+      if (write_g_vs_zvpas_zonal) then
+         if (nproc > 1) call sum_reduce(g_vs_zvpas_zonal, 0) 
+      end if 
+
       ! Deallocate local arrays
       if (write_g2_vs_zvpas) deallocate (g2nozonal_vs_ztubeivmus)
       if (write_g2_vs_zvpas) deallocate (g2_vs_ztubeivmus)

STELLA_CODE/diagnostics/write_diagnostics_to_netcdf.fpp

@@ -305,7 +305,6 @@ contains
       use grids_z, only: nzgrid, ntubes, zed
       use grids_velocity, only: nvpa, vpa, nmu, mu
       use grids_species, only: nspec
-      use parameters_diagnostics, only: number_zonals_kxs, zonal_ks
       ! Geometry
       use parameters_physics, only: rhostar
       use geometry, only: geo_surf, dxdpsi, q_as_x
@@ -320,7 +319,6 @@ contains
       integer :: ix
       real :: nmesh  ! Total number of mesh points
 
-
       !-------------------------------------------------------------------------
       !                                DIMENSIONS                               
       !-------------------------------------------------------------------------
@@ -337,8 +335,6 @@ contains
       call neasyf_dim(file_id, "t", unlimited=.true., long_name="Time", units="a_ref/v_ref")
       call neasyf_dim(file_id, "ri", dim_size=2, long_name="Complex components", units="(real, imaginary)")
 
-      call neasyf_dim(file_id, "zonal_ks", values=zonal_ks)
-
       ! Dimensions for various string variables
       call neasyf_dim(file_id, "char10", dim_size=10, dimid=char10_dim)
       call neasyf_dim(file_id, "char200", dim_size=200)
@@ -958,12 +954,18 @@ contains
    end subroutine write_g2_vs_zvpas_nc
 
    subroutine write_g_vs_zvpas_zonal_nc (nout, g_vs_zvpas_zonal)
-     implicit none
-     integer, intent(in) :: nout
-     real, dimension(:, :, :, :, :, :), intent(in) :: g_vs_zvpas_zonal
-     call neasyf_write(ncid, "g_vs_zvpas_zonal", g_vs_zvpas_zonal, &
-          dim_names=[character(len=8)::"ri", "zonal_ks", "tube", "zed", "vpa", "species", "t"], start=[1, 1, 1, 1, 1, 1, nout], &
+      implicit none
+      integer, intent(in) :: nout
+      complex, dimension(:, :, :, :, :), intent(in) :: g_vs_zvpas_zonal
+#ifdef NETCDF
+      character(*), dimension(*), parameter :: dims = [character(7)::"ri", "kx", "zed", "tube", "vpa", "species", "t"]
+      integer, dimension(7) :: start
+      start = [1, 1, 1, 1, 1, 1, nout]
+
+      call netcdf_write_complex(ncid, "g_vs_zvpas_zonal", g_vs_zvpas_zonal, &
+          dim_names=dims, start=start, &
           long_name="Guiding center Zonal distribution function for min kx, ky = 0, averaged over (mu)")
+#endif
    end subroutine write_g_vs_zvpas_zonal_nc
 
       subroutine write_free_energy_nc(nout, free_energy_vs_kx)

STELLA_CODE/parameters/parameters_diagnostics.f90

@@ -269,9 +269,8 @@ subroutine broadcast_parameters
 
          call broadcast(write_g_vs_zvpas_zonal)
          call broadcast(write_free_energy)
-         call broadcast(number_zonals_kxs)
-         call broadcast(number_zonals_kxs)
-         
+         call broadcast(number_zonals_kxs) 
+
          call broadcast(write_radial_fluxes)
          call broadcast(write_radial_moments)
          call broadcast(write_fluxes_kxkyz)

STELLA_CODE/read_namelists_from_input_file/namelist_diagnostics.f90

@@ -737,7 +737,7 @@ subroutine read_input_file_diagnostics_zonal
          implicit none
          integer :: ios
          integer, parameter :: max_kxs = 100
-         real(kind=8) :: zonal_kxs(max_kxs)
+         real :: zonal_kxs(max_kxs)
          integer :: j
 
          namelist /diagnostics_zonal/ write_g_vs_zvpas_zonal, write_free_energy, &