From 250a324fa0d31e471ee5199aedca57e3ad176a04 Mon Sep 17 00:00:00 2001
From: Michael Hardman <michael.hardman@physics.ox.ac.uk>
Date: Wed, 6 Dec 2023 12:40:23 +0000
Subject: [PATCH] Address issue
 https://github.com/mabarnes/moment_kinetics/issues/111 by incorporating the
 true geometric factor in the ExB drift terms. Note that this means that ExB
 drifts will only be present if bzeta > 0, so the pitch input parameter must
 be set to less that 1. This changes significantly the usage of the code so
 far, which assumed that bzeta ~= 1, whilst allowing an input parameter to set
 bzed.

---
 src/geo.jl         | 7 +++++--
 src/r_advection.jl | 6 +++++-
 src/z_advection.jl | 6 +++++-
 3 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/src/geo.jl b/src/geo.jl
index c1bc4169f..8f12407f0 100644
--- a/src/geo.jl
+++ b/src/geo.jl
@@ -44,7 +44,8 @@ bzeta::Array{mk_float,2}
 dBdz::Array{mk_float,2}
 # d Bmag d r
 dBdr::Array{mk_float,2}
-
+# jacobian =  grad r x grad z . grad zeta
+jacobian::Array{mk_float,2}
 end
 
 """
@@ -63,6 +64,7 @@ function init_magnetic_geometry(geometry_input_data::geometry_input,z,r)
     bzeta = allocate_float(nz,nr)
     dBdr = allocate_float(nz,nr)
     dBdz = allocate_float(nz,nr)
+    jacobian = allocate_float(nz,nr)
     
     option = geometry_input_data.option
     rhostar = geometry_input_data.rhostar
@@ -77,6 +79,7 @@ function init_magnetic_geometry(geometry_input_data::geometry_input,z,r)
                 Bzeta[iz,ir] = Bmag[iz,ir]*bzeta[iz,ir]
                 dBdr[iz,ir] = 0.0
                 dBdz[iz,ir] = 0.0
+                jacobian[iz,ir] = 1.0
             end
         end
     else 
@@ -84,7 +87,7 @@ function init_magnetic_geometry(geometry_input_data::geometry_input,z,r)
     end
 
     geometry = geometric_coefficients(geometry_input_data, rhostar,
-               Bzed,Bzeta,Bmag,bzed,bzeta,dBdz,dBdr)
+               Bzed,Bzeta,Bmag,bzed,bzeta,dBdz,dBdr,jacobian)
     return geometry
 end
 
diff --git a/src/r_advection.jl b/src/r_advection.jl
index bb18a7edc..f9c858aa4 100644
--- a/src/r_advection.jl
+++ b/src/r_advection.jl
@@ -95,9 +95,13 @@ function update_speed_r!(advect, upar, vth, fields, evolve_upar, evolve_ppar, vp
     if r.advection.option == "default" && r.n > 1
         Bmag = geometry.Bmag
         ExBfac = 0.5*geometry.rhostar
+        bzeta = geometry.bzeta
+        jacobian = geometry.jacobian
+        geofac = r.scratch
         @inbounds begin
             @loop_z_vperp_vpa iz ivperp ivpa begin
-                @views @. advect.speed[:,ivpa,ivperp,iz] = ExBfac*fields.Ez[iz,:]/Bmag[iz,:]
+                @. geofac = bzeta[iz,:]*jacobian[iz,:]/Bmag[iz,:]
+                @views @. advect.speed[:,ivpa,ivperp,iz] = ExBfac*geofac*fields.Ez[iz,:]
             end
         end
     elseif r.advection.option == "default" && r.n == 1
diff --git a/src/z_advection.jl b/src/z_advection.jl
index 21a30e1cd..3b70d9ce0 100644
--- a/src/z_advection.jl
+++ b/src/z_advection.jl
@@ -94,10 +94,14 @@ function update_speed_z!(advect, upar, vth, evolve_upar, evolve_ppar, fields, vp
         # bzed = B_z/B only used for z.advection.option == "default"
         bzed = geometry.bzed
         Bmag = geometry.Bmag
+        bzeta = geometry.bzeta
+        jacobian = geometry.jacobian
         ExBfac = -0.5*geometry.rhostar
+        geofac = z.scratch
         @inbounds begin
             @loop_r_vperp_vpa ir ivperp ivpa begin
-                @. @views advect.speed[:,ivpa,ivperp,ir] = vpa.grid[ivpa]*bzed[:,ir] + ExBfac*fields.Er[:,ir]/Bmag[:,ir]
+                @. geofac = bzeta[:,ir]*jacobian[:,ir]/Bmag[:,ir]
+                @. @views advect.speed[:,ivpa,ivperp,ir] = vpa.grid[ivpa]*bzed[:,ir] + ExBfac*geofac*fields.Er[:,ir]
             end
             if evolve_ppar
                 @loop_r_vperp_vpa ir ivperp ivpa begin