Precompute three quantities needed for momentum and energy conservation: z0, w0, s0 (z0le, w0le, s0le if use_le_layout chosen in the input file defined) !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! Now get z0 (first form) $OMP PARALLEL DO DEFAULT(none) & $OMP PRIVATE(iglo, ik, ie, il, is, isgn) & $OMP SHARED(g_lo, conservative, z0tmp, code_dt, vns, vpdiff, speed, aj0, vpa) & $OMP SCHEDULE(static) !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! Now get v0z0 $OMP PARALLEL DO DEFAULT(none) & $OMP PRIVATE(iglo, isgn) & $OMP SHARED(g_lo, vpa, gtmp, aj0, z0tmp) & $OMP COLLAPSE(2) & $OMP SCHEDULE(static) !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! Redefine z0 = z0 / (1 + v0z0) $OMP PARALLEL DO DEFAULT(none) & $OMP PRIVATE(iglo, ik, it, is, isgn) & $OMP SHARED(g_lo, z0tmp, dtmp) & $OMP SCHEDULE(static) !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! du == int (E nu_s f_0); du = du(z, kx, ky, s) duinv = 1/du
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! Now get s0 (first form) $OMP PARALLEL DO DEFAULT(none) & $OMP PRIVATE(iglo, ik, it, ie, il, is, isgn, ig) & $OMP SHARED(g_lo, conservative, s0tmp, vns, vpdiff, speed, aj0, code_dt, duinv, vpa) & $OMP SCHEDULE(static) !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! Now get v0s0 $OMP PARALLEL DO DEFAULT(none) & $OMP PRIVATE(iglo, isgn) & $OMP SHARED(g_lo, gtmp, vpa, aj0, s0tmp) & $OMP COLLAPSE(2) & $OMP SCHEDULE(static) !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! Redefine s0 = s0 - v0s0 * z0 / (1 + v0z0) $OMP PARALLEL DO DEFAULT(none) & $OMP PRIVATE(iglo, ik, it, is, isgn) & $OMP SHARED(g_lo, s0tmp, dtmp, z0tmp) & $OMP SCHEDULE(static) !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! Now get v1s0 $OMP PARALLEL DO DEFAULT(none) & $OMP PRIVATE(iglo, ik, ie, il, is, isgn) & $OMP SHARED(g_lo, conservative, gtmp, vns, speed, vpdiff, aj0, s0tmp, vpa) & $OMP SCHEDULE(static) !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! Redefine s0 = s0 / (1 + v0s0) $OMP PARALLEL DO DEFAULT(none) & $OMP PRIVATE(iglo, ik, it, is, isgn) & $OMP SHARED(g_lo, s0tmp, dtmp) & $OMP SCHEDULE(static) !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! Now get w0 $OMP PARALLEL DO DEFAULT(none) & $OMP PRIVATE(iglo, ik, it, ie, il, is, isgn, ig) & $OMP SHARED(g_lo, ntgrid, forbid, w0tmp, vns, energy, al , aj1, code_dt, & $OMP spec, kperp2, duinv, bmag, conserve_forbid_zero) & $OMP SCHEDULE(static) !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! Now get v0w0 $OMP PARALLEL DO DEFAULT(none) & $OMP PRIVATE(iglo, isgn) & $OMP SHARED(g_lo, gtmp, vpa, aj0, w0tmp) & $OMP COLLAPSE(2) & $OMP SCHEDULE(static) !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! Redefine w0 = w0 - v0w0 * z0 / (1 + v0z0) (this is w1 from MAB notes) $OMP PARALLEL DO DEFAULT(none) & $OMP PRIVATE(iglo, ik, it, is, isgn) & $OMP SHARED(g_lo, w0tmp, z0tmp, dtmp) & $OMP SCHEDULE(static) !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! Get v1w1 $OMP PARALLEL DO DEFAULT(none) & $OMP PRIVATE(iglo, ik, ie, il, is, isgn) & $OMP SHARED(g_lo, conservative, gtmp, vns, speed, vpdiff, aj0, w0tmp, vpa) & $OMP SCHEDULE(static) !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! Redefine w0 = w1 - v1w1 * s1 / (1 + v1s1) (this is w2 from MAB notes) $OMP PARALLEL DO DEFAULT(none) & $OMP PRIVATE(iglo, ik, it, is, isgn) & $OMP SHARED(g_lo, w0tmp, s0tmp, dtmp) & $OMP SCHEDULE(static) !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! Get v2w2 $OMP PARALLEL DO DEFAULT(none) & $OMP PRIVATE(iglo, ik, ie, il, is, isgn) & $OMP SHARED(g_lo, gtmp, vns, energy, al, aj1, w0tmp) & $OMP SCHEDULE(static) !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! Redefine w0 = w2 / (1 + v2w2) $OMP PARALLEL DO DEFAULT(none) & $OMP PRIVATE(iglo, ik, it, is, isgn, ig) & $OMP SHARED(g_lo, w0tmp, dtmp) & $OMP SCHEDULE(static)
subroutine init_lorentz_conserve
use gs2_layouts, only: g_lo, ie_idx, is_idx, ik_idx, il_idx, it_idx
use species, only: nspec, spec, is_electron_species
use kt_grids, only: kperp2, naky, ntheta0
use theta_grid, only: ntgrid, bmag
use le_grids, only: energy => energy_maxw, speed => speed_maxw, al, &
integrate_moment, negrid, forbid
use gs2_time, only: code_dt, tunits
use dist_fn_arrays, only: aj0, aj1, vpa
use le_grids, only: g2le
use gs2_layouts, only: le_lo
use redistribute, only: gather, scatter
use array_utils, only: zero_array
implicit none
complex, dimension (1,1,1) :: dum1, dum2
real, dimension (:,:,:), allocatable :: gtmp
real, dimension (:,:,:,:), allocatable :: duinv, dtmp, vns
integer :: ie, il, ik, is, isgn, iglo, it, ig
complex, dimension (:,:,:), allocatable :: ctmp, z_big
complex, dimension(:,:,:), allocatable :: s0tmp, w0tmp, z0tmp
logical, parameter :: all_procs = .true.
if(use_le_layout) then
allocate (ctmp(nxi_lim, negrid+1, le_lo%llim_proc:le_lo%ulim_alloc))
! We need to initialise ctmp as it is used as receiving buffer in
! g2le redistribute, which doesn't populate all elements
call zero_array(ctmp)
end if
dum1 = 0. ; dum2 = 0.
allocate(s0tmp(-ntgrid:ntgrid,2,g_lo%llim_proc:g_lo%ulim_alloc))
allocate(w0tmp(-ntgrid:ntgrid,2,g_lo%llim_proc:g_lo%ulim_alloc))
allocate(z0tmp(-ntgrid:ntgrid,2,g_lo%llim_proc:g_lo%ulim_alloc))
allocate (gtmp(-ntgrid:ntgrid,2,g_lo%llim_proc:g_lo%ulim_alloc))
allocate (duinv(-ntgrid:ntgrid, ntheta0, naky, nspec))
allocate (dtmp(-ntgrid:ntgrid, ntheta0, naky, nspec))
allocate (vns(naky,negrid,nspec,3))
call zero_array(duinv)
call zero_array(dtmp)
vns(:,:,:,1) = vnmult(1)*vnew_D
vns(:,:,:,2) = vnmult(1)*vnew_s
vns(:,:,:,3) = 0.0
if (drag) then
do is = 1, nspec
if (.not. is_electron_species(spec(is))) cycle
do ik = 1, naky
vns(ik,:,is,3) = vnmult(1)*spec(is)%vnewk*tunits(ik)/energy**1.5
end do
end do
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! Now get z0 (first form)
!$OMP PARALLEL DO DEFAULT(none) &
!$OMP PRIVATE(iglo, ik, ie, il, is, isgn) &
!$OMP SHARED(g_lo, conservative, z0tmp, code_dt, vns, vpdiff, speed, aj0, vpa) &
!$OMP SCHEDULE(static)
do iglo = g_lo%llim_proc, g_lo%ulim_proc
ik = ik_idx(g_lo,iglo)
ie = ie_idx(g_lo,iglo)
il = il_idx(g_lo,iglo)
is = is_idx(g_lo,iglo)
do isgn = 1, 2
! u0 = -2 nu_D^{ei} vpa J0 dt f0
if (conservative) then
z0tmp(:,isgn,iglo) = -2.0*code_dt*vns(ik,ie,is,3)*vpdiff(:,isgn,il) &
* speed(ie)*aj0(:,iglo)
else
z0tmp(:,isgn,iglo) = -2.0*code_dt*vns(ik,ie,is,3)*vpa(:,isgn,iglo)*aj0(:,iglo)
end if
end do
end do
!$OMP END PARALLEL DO
call zero_out_passing_hybrid_electrons(z0tmp)
if(use_le_layout) then
call gather (g2le, z0tmp, ctmp)
call solfp_lorentz (ctmp)
call scatter (g2le, ctmp, z0tmp) ! z0 is redefined below
else
call solfp_lorentz (z0tmp,dum1,dum2) ! z0 is redefined below
end if
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! Now get v0z0
!$OMP PARALLEL DO DEFAULT(none) &
!$OMP PRIVATE(iglo, isgn) &
!$OMP SHARED(g_lo, vpa, gtmp, aj0, z0tmp) &
!$OMP COLLAPSE(2) &
!$OMP SCHEDULE(static)
do iglo = g_lo%llim_proc, g_lo%ulim_proc
do isgn = 1, 2
! v0 = vpa J0 f0
gtmp(:,isgn,iglo) = vpa(:,isgn,iglo) * aj0(:,iglo) * real(z0tmp(:,isgn,iglo))
end do
end do
!$OMP END PARALLEL DO
call integrate_moment (gtmp, dtmp, all_procs) ! v0z0
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! Redefine z0 = z0 / (1 + v0z0)
!$OMP PARALLEL DO DEFAULT(none) &
!$OMP PRIVATE(iglo, ik, it, is, isgn) &
!$OMP SHARED(g_lo, z0tmp, dtmp) &
!$OMP SCHEDULE(static)
do iglo = g_lo%llim_proc, g_lo%ulim_proc
it = it_idx(g_lo,iglo)
ik = ik_idx(g_lo,iglo)
is = is_idx(g_lo,iglo)
do isgn = 1, 2
z0tmp(:,isgn,iglo) = z0tmp(:,isgn,iglo) / (1.0 + dtmp(:,it,ik,is))
end do
end do
!$OMP END PARALLEL DO
else
!If drag is false vns(...,3) is zero and hence z0tmp is zero here.
call zero_array(z0tmp)
end if
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! du == int (E nu_s f_0); du = du(z, kx, ky, s)
! duinv = 1/du
if (conservative) then
!$OMP PARALLEL DO DEFAULT(none) &
!$OMP PRIVATE(iglo, ik, ie, il, is, isgn) &
!$OMP SHARED(g_lo, gtmp, vns, vpa, vpdiff, speed) &
!$OMP SCHEDULE(static)
do iglo = g_lo%llim_proc, g_lo%ulim_proc
ik = ik_idx(g_lo,iglo)
ie = ie_idx(g_lo,iglo)
il = il_idx(g_lo,iglo)
is = is_idx(g_lo,iglo)
do isgn = 1, 2
gtmp(:,isgn,iglo) = vns(ik,ie,is,1)*vpa(:,isgn,iglo) &
* vpdiff(:,isgn,il)*speed(ie)
end do
end do
!$OMP END PARALLEL DO
else
!$OMP PARALLEL DO DEFAULT(none) &
!$OMP PRIVATE(iglo, ik, ie, is, isgn) &
!$OMP SHARED(g_lo, gtmp, vpa, vns) &
!$OMP SCHEDULE(static)
do iglo = g_lo%llim_proc, g_lo%ulim_proc
ik = ik_idx(g_lo,iglo)
ie = ie_idx(g_lo,iglo)
is = is_idx(g_lo,iglo)
do isgn = 1, 2
gtmp(:,isgn,iglo) = vns(ik,ie,is,1)*vpa(:,isgn,iglo)**2
end do
end do
!$OMP END PARALLEL DO
end if
call integrate_moment (gtmp, duinv, all_procs) ! not 1/du yet
! Could replace this with OpenMP using an explicit loop. TAG
where (abs(duinv) > epsilon(0.0)) ! necessary b/c some species may have vnewk=0
!duinv=0 iff vnew=0 so ok to keep duinv=0.
duinv = 1./duinv ! now it is 1/du
end where
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! Now get s0 (first form)
!$OMP PARALLEL DO DEFAULT(none) &
!$OMP PRIVATE(iglo, ik, it, ie, il, is, isgn, ig) &
!$OMP SHARED(g_lo, conservative, s0tmp, vns, vpdiff, speed, aj0, code_dt, duinv, vpa) &
!$OMP SCHEDULE(static)
do iglo = g_lo%llim_proc, g_lo%ulim_proc
it = it_idx(g_lo,iglo)
ik = ik_idx(g_lo,iglo)
ie = ie_idx(g_lo,iglo)
il = il_idx(g_lo,iglo)
is = is_idx(g_lo,iglo)
do isgn = 1, 2
! u1 = -3 nu_s vpa dt J0 f_0 / du
if (conservative) then
s0tmp(:,isgn,iglo) = -vns(ik,ie,is,1)*vpdiff(:,isgn,il)*speed(ie) &
* aj0(:,iglo)*code_dt*duinv(:,it,ik,is)
else
s0tmp(:,isgn,iglo) = -vns(ik,ie,is,1)*vpa(:,isgn,iglo) &
* aj0(:,iglo)*code_dt*duinv(:,it,ik,is)
end if
end do
end do
!$OMP END PARALLEL DO
call zero_out_passing_hybrid_electrons(s0tmp)
if(use_le_layout) then
call gather (g2le, s0tmp, ctmp)
call solfp_lorentz (ctmp)
call scatter (g2le, ctmp, s0tmp) ! s0
else
call solfp_lorentz (s0tmp,dum1,dum2) ! s0
end if
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! Now get v0s0
!$OMP PARALLEL DO DEFAULT(none) &
!$OMP PRIVATE(iglo, isgn) &
!$OMP SHARED(g_lo, gtmp, vpa, aj0, s0tmp) &
!$OMP COLLAPSE(2) &
!$OMP SCHEDULE(static)
do iglo = g_lo%llim_proc, g_lo%ulim_proc
do isgn = 1, 2
! v0 = vpa J0 f0
gtmp(:,isgn,iglo) = vpa(:,isgn,iglo) * aj0(:,iglo) * real(s0tmp(:,isgn,iglo))
end do
end do
!OMP END PARALLEL DO
call integrate_moment (gtmp, dtmp, all_procs) ! v0s0
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! Redefine s0 = s0 - v0s0 * z0 / (1 + v0z0)
!$OMP PARALLEL DO DEFAULT(none) &
!$OMP PRIVATE(iglo, ik, it, is, isgn) &
!$OMP SHARED(g_lo, s0tmp, dtmp, z0tmp) &
!$OMP SCHEDULE(static)
do iglo = g_lo%llim_proc, g_lo%ulim_proc
ik = ik_idx(g_lo,iglo)
it = it_idx(g_lo,iglo)
is = is_idx(g_lo,iglo)
do isgn=1,2
s0tmp(:,isgn,iglo) = s0tmp(:,isgn,iglo) - dtmp(:,it,ik,is) * z0tmp(:,isgn,iglo)
end do
end do
!$OMP END PARALLEL DO
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! Now get v1s0
!$OMP PARALLEL DO DEFAULT(none) &
!$OMP PRIVATE(iglo, ik, ie, il, is, isgn) &
!$OMP SHARED(g_lo, conservative, gtmp, vns, speed, vpdiff, aj0, s0tmp, vpa) &
!$OMP SCHEDULE(static)
do iglo = g_lo%llim_proc, g_lo%ulim_proc
ik = ik_idx(g_lo,iglo)
ie = ie_idx(g_lo,iglo)
il = il_idx(g_lo,iglo)
is = is_idx(g_lo,iglo)
do isgn = 1, 2
! v1 = nu_D vpa J0
if (conservative) then
gtmp(:,isgn,iglo) = vns(ik,ie,is,1)*speed(ie)*vpdiff(:,isgn,il) &
* aj0(:,iglo) * real(s0tmp(:,isgn,iglo))
else
gtmp(:,isgn,iglo) = vns(ik,ie,is,1)*vpa(:,isgn,iglo)*aj0(:,iglo) &
* real(s0tmp(:,isgn,iglo))
end if
end do
end do
!$OMP END PARALLEL DO
call integrate_moment (gtmp, dtmp, all_procs) ! v1s0
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! Redefine s0 = s0 / (1 + v0s0)
!$OMP PARALLEL DO DEFAULT(none) &
!$OMP PRIVATE(iglo, ik, it, is, isgn) &
!$OMP SHARED(g_lo, s0tmp, dtmp) &
!$OMP SCHEDULE(static)
do iglo = g_lo%llim_proc, g_lo%ulim_proc
ik = ik_idx(g_lo,iglo)
it = it_idx(g_lo,iglo)
is = is_idx(g_lo,iglo)
do isgn=1,2
s0tmp(:,isgn,iglo) = s0tmp(:,isgn,iglo) / (1.0 + dtmp(:,it,ik,is))
end do
end do
!$OMP END PARALLEL DO
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! Now get w0
!$OMP PARALLEL DO DEFAULT(none) &
!$OMP PRIVATE(iglo, ik, it, ie, il, is, isgn, ig) &
!$OMP SHARED(g_lo, ntgrid, forbid, w0tmp, vns, energy, al , aj1, code_dt, &
!$OMP spec, kperp2, duinv, bmag, conserve_forbid_zero) &
!$OMP SCHEDULE(static)
do iglo = g_lo%llim_proc, g_lo%ulim_proc
it = it_idx(g_lo,iglo)
ik = ik_idx(g_lo,iglo)
ie = ie_idx(g_lo,iglo)
il = il_idx(g_lo,iglo)
is = is_idx(g_lo,iglo)
do isgn = 1, 2
do ig=-ntgrid,ntgrid
! u2 = -3 dt J1 vperp vus a f0 / du
if ( .not. (forbid(ig,il) .and. conserve_forbid_zero) ) then
! Note no conservative branch here, is that right?
! Note: energy * al * smz^2 * kperp2 / bmag is alpha^2 where
! alpha is the argument to the Bessel function, i.e. aj1 = J1(alpha) / alpha
! This appears to leave us with alpha * J1(alpha) whilst Barnes' paper
! only includes terms with J1(alpha). Note that alpha = vperp kperp smz/B
! so alpha * J1(alpha) = (kperp * smz / B) * vperp J1
w0tmp(ig, isgn, iglo) = -vns(ik, ie, is, 1) * energy(ie) * al(il) &
* aj1(ig, iglo) * code_dt * spec(is)%smz**2 * kperp2(ig, it, ik) * &
duinv(ig, it, ik, is) / bmag(ig)
else
w0tmp(ig,isgn,iglo) = 0.
endif
end do
end do
end do
!$OMP END PARALLEL DO
call zero_out_passing_hybrid_electrons(w0tmp)
if(use_le_layout) then
call gather (g2le, w0tmp, ctmp)
call solfp_lorentz (ctmp)
call scatter (g2le, ctmp, w0tmp)
else
call solfp_lorentz (w0tmp,dum1,dum2)
end if
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! Now get v0w0
!$OMP PARALLEL DO DEFAULT(none) &
!$OMP PRIVATE(iglo, isgn) &
!$OMP SHARED(g_lo, gtmp, vpa, aj0, w0tmp) &
!$OMP COLLAPSE(2) &
!$OMP SCHEDULE(static)
do iglo = g_lo%llim_proc, g_lo%ulim_proc
do isgn = 1, 2
! v0 = vpa J0 f0
gtmp(:,isgn,iglo) = vpa(:,isgn,iglo) * aj0(:,iglo) * real(w0tmp(:,isgn,iglo))
end do
end do
!$OMP END PARALLEL DO
call integrate_moment (gtmp, dtmp, all_procs) ! v0w0
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! Redefine w0 = w0 - v0w0 * z0 / (1 + v0z0) (this is w1 from MAB notes)
!$OMP PARALLEL DO DEFAULT(none) &
!$OMP PRIVATE(iglo, ik, it, is, isgn) &
!$OMP SHARED(g_lo, w0tmp, z0tmp, dtmp) &
!$OMP SCHEDULE(static)
do iglo = g_lo%llim_proc, g_lo%ulim_proc
ik = ik_idx(g_lo,iglo)
it = it_idx(g_lo,iglo)
is = is_idx(g_lo,iglo)
do isgn=1,2
w0tmp(:,isgn,iglo) = w0tmp(:,isgn,iglo) - z0tmp(:,isgn,iglo)*dtmp(:,it,ik,is)
end do
end do
!$OMP END PARALLEL DO
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! Get v1w1
!$OMP PARALLEL DO DEFAULT(none) &
!$OMP PRIVATE(iglo, ik, ie, il, is, isgn) &
!$OMP SHARED(g_lo, conservative, gtmp, vns, speed, vpdiff, aj0, w0tmp, vpa) &
!$OMP SCHEDULE(static)
do iglo = g_lo%llim_proc, g_lo%ulim_proc
ik = ik_idx(g_lo,iglo)
ie = ie_idx(g_lo,iglo)
il = il_idx(g_lo,iglo)
is = is_idx(g_lo,iglo)
do isgn = 1, 2
! v1 = nud vpa J0 f0
if (conservative) then
gtmp(:,isgn,iglo) = vns(ik,ie,is,1)*speed(ie)*vpdiff(:,isgn,il) &
* aj0(:,iglo) * real(w0tmp(:,isgn,iglo))
else
gtmp(:,isgn,iglo) = vns(ik,ie,is,1)*vpa(:,isgn,iglo)*aj0(:,iglo) &
* real(w0tmp(:,isgn,iglo))
end if
end do
end do
!$OMP END PARALLEL DO
call integrate_moment (gtmp, dtmp, all_procs) ! v1w1
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! Redefine w0 = w1 - v1w1 * s1 / (1 + v1s1) (this is w2 from MAB notes)
!$OMP PARALLEL DO DEFAULT(none) &
!$OMP PRIVATE(iglo, ik, it, is, isgn) &
!$OMP SHARED(g_lo, w0tmp, s0tmp, dtmp) &
!$OMP SCHEDULE(static)
do iglo = g_lo%llim_proc, g_lo%ulim_proc
ik = ik_idx(g_lo,iglo)
it = it_idx(g_lo,iglo)
is = is_idx(g_lo,iglo)
do isgn=1,2
w0tmp(:,isgn,iglo) = w0tmp(:,isgn,iglo) - s0tmp(:,isgn,iglo)*dtmp(:,it,ik,is)
end do
end do
!$OMP END PARALLEL DO
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! Get v2w2
!$OMP PARALLEL DO DEFAULT(none) &
!$OMP PRIVATE(iglo, ik, ie, il, is, isgn) &
!$OMP SHARED(g_lo, gtmp, vns, energy, al, aj1, w0tmp) &
!$OMP SCHEDULE(static)
do iglo = g_lo%llim_proc, g_lo%ulim_proc
ik = ik_idx(g_lo,iglo)
ie = ie_idx(g_lo,iglo)
il = il_idx(g_lo,iglo)
is = is_idx(g_lo,iglo)
do isgn = 1, 2
! v2 = nud vperp J1 f0
! Note : aj1 = J1(alpha) / alpha, where we have
! alpha^2 = energy * al * smz^2 * kperp2 / bmag = vperp^2 smz^2 kperp2 / bmag^2
! so energy * al * aj1 = (B / [kperp2 * smz^2]) alpha^2 aj1
! = (B / [kperp2 * smz^2]) alpha J1 = vperp J1 / (kperp * smz). As w0tmp
! appears to have an extra factor of kperp * smz / B this _may_ work out to
! (vperp J1 / B) * ... Is there an extra factor of 1 / B here?
gtmp(:, isgn, iglo) = vns(ik, ie, is, 1) * energy(ie) * al(il) * aj1(:, iglo) &
* real(w0tmp(:, isgn, iglo))
end do
end do
!$OMP END PARALLEL DO
call integrate_moment (gtmp, dtmp, all_procs) ! v2w2
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! Redefine w0 = w2 / (1 + v2w2)
!$OMP PARALLEL DO DEFAULT(none) &
!$OMP PRIVATE(iglo, ik, it, is, isgn, ig) &
!$OMP SHARED(g_lo, w0tmp, dtmp) &
!$OMP SCHEDULE(static)
do iglo = g_lo%llim_proc, g_lo%ulim_proc
ik = ik_idx(g_lo,iglo)
it = it_idx(g_lo,iglo)
is = is_idx(g_lo,iglo)
do isgn=1,2
w0tmp(:,isgn,iglo) = w0tmp(:,isgn,iglo) / (1.0 + dtmp(:,it,ik,is))
end do
end do
!$OMP END PARALLEL DO
deallocate (gtmp, duinv, dtmp, vns)
if(use_le_layout) then
allocate (z_big(-ntgrid:ntgrid, 2, g_lo%llim_proc:g_lo%ulim_alloc))
! first set s0le, w0le & z0le
z_big = cmplx(real(s0tmp), real(w0tmp))
! get rid of z0, s0, w0 now that we've converted to z0le, s0le, w0le
if (allocated(s0tmp)) deallocate(s0tmp)
if (allocated(w0tmp)) deallocate(w0tmp)
call gather (g2le, z_big, ctmp)
if (.not. allocated(s0le)) then
allocate (s0le(nxi_lim, negrid+1, le_lo%llim_proc:le_lo%ulim_alloc))
allocate (w0le(nxi_lim, negrid+1, le_lo%llim_proc:le_lo%ulim_alloc))
end if
s0le = real(ctmp)
w0le = aimag(ctmp)
! set z0le
call gather (g2le, z0tmp, ctmp)
if (allocated(z0tmp)) deallocate(z0tmp)
if (.not. allocated(z0le)) allocate (z0le(nxi_lim, negrid+1, le_lo%llim_proc:le_lo%ulim_alloc))
z0le = real(ctmp)
deallocate (ctmp, z_big)
else
!Only need the real components (imaginary part should be zero, just
!use complex arrays to allow reuse of existing integrate routines etc.)
if (.not. allocated(s0)) then
allocate (s0(-ntgrid:ntgrid,2,g_lo%llim_proc:g_lo%ulim_alloc))
s0=real(s0tmp)
deallocate(s0tmp)
endif
if (.not. allocated(w0)) then
allocate (w0(-ntgrid:ntgrid,2,g_lo%llim_proc:g_lo%ulim_alloc))
w0=real(w0tmp)
deallocate(w0tmp)
endif
if (.not. allocated(z0)) then
allocate (z0(-ntgrid:ntgrid,2,g_lo%llim_proc:g_lo%ulim_alloc))
z0=real(z0tmp)
deallocate(z0tmp)
endif
end if
end subroutine init_lorentz_conserve