Merge pull request #144 from mabarnes/fix-source-terms-z-derivative-bug

source_terms!(), source_terms_neutral!() use distributed MPI derivatives
mabarnes · Nov 17, 2023 · 7142fb9 · 7142fb9
2 parents 10e6799 + 14925e4
commit 7142fb9
Show file tree

Hide file tree

Showing 4 changed files with 94 additions and 96 deletions.
diff --git a/src/source_terms.jl b/src/source_terms.jl
@@ -23,8 +23,9 @@ function source_terms!(pdf_out, fvec_in, moments, vpa, z, r, dt, spectral, compo
             @views source_terms_evolve_ppar_no_collisions!(
                 pdf_out[:,:,:,:,is], fvec_in.pdf[:,:,:,:,is], fvec_in.density[:,:,is],
                 fvec_in.upar[:,:,is], fvec_in.ppar[:,:,is], moments.charged.vth[:,:,is],
-                moments.charged.qpar[:,:,is], moments, z, r, dt, spectral,
-                ion_source_settings)
+                moments.charged.qpar[:,:,is], moments.charged.ddens_dz[:,:,is],
+                moments.charged.dvth_dz[:,:,is], moments.charged.dqpar_dz[:,:,is],
+                moments, z, r, dt, spectral, ion_source_settings)
             if composition.n_neutral_species > 0
                 if abs(collisions.charge_exchange) > 0.0 || abs(collisions.ionization) > 0.0
                     @views source_terms_evolve_ppar_collisions!(
@@ -40,26 +41,26 @@ function source_terms!(pdf_out, fvec_in, moments, vpa, z, r, dt, spectral, compo
         @loop_s is begin
             @views source_terms_evolve_density!(
                 pdf_out[:,:,:,:,is], fvec_in.pdf[:,:,:,:,is], fvec_in.density[:,:,is],
-                fvec_in.upar[:,:,is], moments, z, r, dt, spectral, ion_source_settings)
+                fvec_in.upar[:,:,is], moments.charged.ddens_dz[:,:,is],
+                moments.charged.dupar_dz[:,:,is], moments, z, r, dt, spectral,
+                ion_source_settings)
         end
     end
     return nothing
 end
 
 """
 """
-function source_terms_evolve_density!(pdf_out, pdf_in, dens, upar, moments, z, r, dt,
-                                      spectral, ion_source_settings)
+function source_terms_evolve_density!(pdf_out, pdf_in, dens, upar, ddens_dz, dupar_dz,
+                                      moments, z, r, dt, spectral, ion_source_settings)
     # update the density
     nvpa = size(pdf_out, 1)
-    @loop_r ir begin
-        # calculate d(n*upar)/dz
-        @views @. z.scratch = dens[:,ir]*upar[:,ir]
-        derivative!(z.scratch, z.scratch, z, spectral)
-        @views @. z.scratch *= dt/dens[:,ir]
-        #derivative!(z.scratch, z.scratch, z, -upar, spectral)
-        @loop_z_vperp_vpa iz ivperp ivpa begin
-            pdf_out[ivpa,ivperp,iz,ir] += pdf_in[ivpa,ivperp,iz,ir]*z.scratch[iz]
+    @loop_r_z ir iz begin
+        # calculate dt * d(n*upar)/dz / n
+        factor = dt * (dens[iz,ir] * dupar_dz[iz,ir] + upar[iz,ir] * ddens_dz[iz,ir]) /
+                 dens[iz,ir]
+        @loop_vperp_vpa ivperp ivpa begin
+            pdf_out[ivpa,ivperp,iz,ir] += pdf_in[ivpa,ivperp,iz,ir] * factor
         end
     end
 
@@ -81,25 +82,17 @@ update the evolved pdf to account for the collisionless source terms in the kine
 arising due to the re-normalization of the pdf as g = f * vth / n
 """
 function source_terms_evolve_ppar_no_collisions!(pdf_out, pdf_in, dens, upar, ppar, vth,
-                                                 qpar, moments, z, r, dt, spectral,
+                                                 qpar, ddens_dz, dvth_dz, dqpar_dz,
+                                                 moments, z, r, dt, spectral,
                                                  ion_source_settings)
     nvpa = size(pdf_out, 1)
-    @loop_r ir begin
-        # calculate dn/dz
-        derivative!(z.scratch, view(dens,:,ir), z, spectral)
-        # update the pdf to account for the density gradient contribution to the source
-        @views @. z.scratch *= upar[:,ir]/dens[:,ir]
-        # calculate dvth/dz
-        derivative!(z.scratch2, view(vth,:,ir), z, spectral)
-        # update the pdf to account for the -g*upar/vth * dvth/dz contribution to the source
-        @views @. z.scratch -= z.scratch2*upar[:,ir]/vth[:,ir]
-        # calculate dqpar/dz
-        derivative!(z.scratch2, view(qpar,:,ir), z, spectral)
-        # update the pdf to account for the parallel heat flux contribution to the source
-        @views @. z.scratch -= 0.5*z.scratch2/ppar[:,ir]
+    @loop_r_z ir iz begin
+        factor = dt * (ddens_dz[iz,ir] * upar[iz,ir] / dens[iz,ir] -
+                       dvth_dz[iz,ir] * upar[iz,ir] / vth[iz,ir] -
+                       0.5 * dqpar_dz[iz,ir] / ppar[iz,ir])
 
-        @loop_z_vperp_vpa iz ivperp ivpa begin
-            pdf_out[ivpa,ivperp,iz,ir] += dt*pdf_in[ivpa,ivperp,iz,ir]*z.scratch[iz]
+        @loop_vperp_vpa ivperp ivpa begin
+            pdf_out[ivpa,ivperp,iz,ir] += pdf_in[ivpa,ivperp,iz,ir] * factor
         end
     end
 
@@ -157,8 +150,9 @@ function source_terms_neutral!(pdf_out, fvec_in, moments, vpa, z, r, dt, spectra
                 pdf_out[:,:,:,:,:,isn], fvec_in.pdf_neutral[:,:,:,:,:,isn],
                 fvec_in.density_neutral[:,:,isn], fvec_in.uz_neutral[:,:,isn],
                 fvec_in.pz_neutral[:,:,isn], moments.neutral.vth[:,:,isn],
-                moments.neutral.qz[:,:,isn], moments, z, r, dt, spectral,
-                neutral_source_settings)
+                moments.neutral.qz[:,:,isn], moments.neutral.ddens_dz[:,:,isn],
+                moments.neutral.dvth_dz[:,:,isn], moments.neutral.dqz_dz[:,:,isn],
+                moments, z, r, dt, spectral, neutral_source_settings)
             if abs(collisions.charge_exchange) > 0.0 || abs(collisions.ionization) > 0.0
                 @views source_terms_evolve_ppar_collisions_neutral!(
                     pdf_out[:,:,:,:,:,isn], fvec_in.pdf_neutral[:,:,:,:,:,isn],
@@ -172,27 +166,27 @@ function source_terms_neutral!(pdf_out, fvec_in, moments, vpa, z, r, dt, spectra
         @loop_sn isn begin
             @views source_terms_evolve_density_neutral!(
                 pdf_out[:,:,:,:,:,isn], fvec_in.pdf_neutral[:,:,:,:,:,isn],
-                fvec_in.density_neutral[:,:,isn], fvec_in.uz_neutral[:,:,isn], moments, z,
-                r, dt, spectral, neutral_source_settings)
+                fvec_in.density_neutral[:,:,isn], fvec_in.uz_neutral[:,:,isn],
+                moments.neutral.ddens_dz[:,:,isn], moments.neutral.duz_dz[:,:,isn],
+                moments, z, r, dt, spectral, neutral_source_settings)
         end
     end
     return nothing
 end
 
 """
 """
-function source_terms_evolve_density_neutral!(pdf_out, pdf_in, dens, upar, moments, z, r,
-                                              dt, spectral, neutral_source_settings)
+function source_terms_evolve_density_neutral!(pdf_out, pdf_in, dens, upar, ddens_dz,
+                                              dupar_dz, moments, z, r, dt, spectral,
+                                              neutral_source_settings)
     # update the density
     nvpa = size(pdf_out, 1)
-    @loop_r ir begin
-        # calculate d(n*upar)/dz
-        @views @. z.scratch = dens[:,ir]*upar[:,ir]
-        derivative!(z.scratch, z.scratch, z, spectral)
-        @views @. z.scratch *= dt/dens[:,ir]
-        #derivative!(z.scratch, z.scratch, z, -upar, spectral)
-        @loop_z_vzeta_vr_vz iz ivzeta ivr ivz begin
-            pdf_out[ivz,ivr,ivzeta,iz,ir] += pdf_in[ivz,ivr,ivzeta,iz,ir]*z.scratch[iz]
+    @loop_r_z ir iz begin
+        # calculate dt * d(n*upar)/dz / n
+        factor = dt * (dens[iz,ir] * dupar_dz[iz,ir] + upar[iz,ir] * ddens_dz[iz,ir]) /
+                 dens[iz,ir]
+        @loop_vzeta_vr_vz ivzeta ivr ivz begin
+            pdf_out[ivz,ivr,ivzeta,iz,ir] += pdf_in[ivz,ivr,ivzeta,iz,ir] * factor
         end
     end
 
@@ -214,26 +208,16 @@ update the evolved pdf to account for the collisionless source terms in the kine
 arising due to the re-normalization of the pdf as g = f * vth / n
 """
 function source_terms_evolve_ppar_no_collisions_neutral!(pdf_out, pdf_in, dens, upar,
-                                                         ppar, vth, qpar, moments, z, r,
+                                                         ppar, vth, qpar, ddens_dz,
+                                                         dvth_dz, dqpar_dz, moments, z, r,
                                                          dt, spectral,
                                                          neutral_source_settings)
     nvpa = size(pdf_out, 1)
-    @loop_r ir begin
-        # calculate dn/dz
-        derivative!(z.scratch, view(dens,:,ir), z, spectral)
-        # update the pdf to account for the density gradient contribution to the source
-        @views @. z.scratch *= upar[:,ir]/dens[:,ir]
-        # calculate dvth/dz
-        derivative!(z.scratch2, view(vth,:,ir), z, spectral)
-        # update the pdf to account for the -g*upar/vth * dvth/dz contribution to the source
-        @views @. z.scratch -= z.scratch2*upar[:,ir]/vth[:,ir]
-        # calculate dqpar/dz
-        derivative!(z.scratch2, view(qpar,:,ir), z, spectral)
-        # update the pdf to account for the parallel heat flux contribution to the source
-        @views @. z.scratch -= 0.5*z.scratch2/ppar[:,ir]
-
-        @loop_z_vzeta_vr_vz iz ivzeta ivr ivz begin
-            pdf_out[ivz,ivr,ivzeta,iz,ir] += dt*pdf_in[ivz,ivr,ivzeta,iz,ir]*z.scratch[iz]
+    @loop_r_z ir iz begin
+        factor = dt * (ddens_dz[iz,ir] * upar[iz,ir] / dens[iz,ir] - dvth_dz[iz,ir] *
+                       upar[iz,ir] / vth[iz,ir] - 0.5 * dqpar_dz[iz,ir] / ppar[iz,ir])
+        @loop_vzeta_vr_vz ivzeta ivr ivz begin
+            pdf_out[ivz,ivr,ivzeta,iz,ir] += pdf_in[ivz,ivr,ivzeta,iz,ir] * factor
         end
     end
 

diff --git a/src/velocity_moments.jl b/src/velocity_moments.jl
@@ -88,6 +88,8 @@ struct moments_charged_substruct
     # v_norm_fac accounts for this: it is vth if using the above definition for the parallel velocity,
     # and it is one otherwise
     v_norm_fac::Union{MPISharedArray{mk_float,3},Nothing}
+    # this is the z-derivative of the particle density
+    ddens_dz::Union{MPISharedArray{mk_float,3},Nothing}
     # this is the upwinded z-derivative of the particle density
     ddens_dz_upwind::Union{MPISharedArray{mk_float,3},Nothing}
     # this is the second-z-derivative of the particle density
@@ -170,6 +172,8 @@ struct moments_neutral_substruct
     # and it is one otherwise
     v_norm_fac::MPISharedArray{mk_float,3}
     # this is the z-derivative of the particle density
+    ddens_dz::Union{MPISharedArray{mk_float,3},Nothing}
+    # this is the z-derivative of the particle density
     ddens_dz_upwind::Union{MPISharedArray{mk_float,3},Nothing}
     # this is the second-z-derivative of the particle density
     d2dens_dz2::Union{MPISharedArray{mk_float,3},Nothing}
@@ -244,8 +248,10 @@ function create_moments_charged(nz, nr, n_species, evolve_density, evolve_upar,
     end
 
     if evolve_density
+        ddens_dz = allocate_shared_float(nz, nr, n_species)
         ddens_dz_upwind = allocate_shared_float(nz, nr, n_species)
     else
+        ddens_dz = nothing
         ddens_dz_upwind = nothing
     end
     if evolve_density &&
@@ -329,10 +335,11 @@ function create_moments_charged(nz, nr, n_species, evolve_density, evolve_upar,
         parallel_flow_updated, parallel_pressure, parallel_pressure_updated,perpendicular_pressure,
         parallel_heat_flux, parallel_heat_flux_updated, thermal_speed, 
         chodura_integral_lower, chodura_integral_upper, v_norm_fac,
-        ddens_dz_upwind, d2dens_dz2, dupar_dz, dupar_dz_upwind, d2upar_dz2, dppar_dz,
-        dppar_dz_upwind, d2ppar_dz2, dqpar_dz, dvth_dz, external_source_amplitude,
-        external_source_density_amplitude, external_source_momentum_amplitude,
-        external_source_pressure_amplitude, external_source_controller_integral)
+        ddens_dz, ddens_dz_upwind, d2dens_dz2, dupar_dz, dupar_dz_upwind, d2upar_dz2,
+        dppar_dz, dppar_dz_upwind, d2ppar_dz2, dqpar_dz, dvth_dz,
+        external_source_amplitude, external_source_density_amplitude,
+        external_source_momentum_amplitude, external_source_pressure_amplitude,
+        external_source_controller_integral)
 end
 
 # neutral particles have natural mean velocities 
@@ -379,8 +386,10 @@ function create_moments_neutral(nz, nr, n_species, evolve_density, evolve_upar,
     qz_updated .= false
 
     if evolve_density
+        ddens_dz = allocate_shared_float(nz, nr, n_species)
         ddens_dz_upwind = allocate_shared_float(nz, nr, n_species)
     else
+        ddens_dz = nothing
         ddens_dz_upwind = nothing
     end
     if evolve_density &&
@@ -462,9 +471,9 @@ function create_moments_neutral(nz, nr, n_species, evolve_density, evolve_upar,
     # return struct containing arrays needed to update moments
     return moments_neutral_substruct(density, density_updated, uz, uz_updated, ur,
         ur_updated, uzeta, uzeta_updated, pz, pz_updated, pr, pr_updated, pzeta,
-        pzeta_updated, ptot, qz, qz_updated, vth, v_norm_fac, ddens_dz_upwind, d2dens_dz2,
-        duz_dz, duz_dz_upwind, d2uz_dz2, dpz_dz, dpz_dz_upwind, d2pz_dz2, dqz_dz, dvth_dz,
-        external_source_amplitude, external_source_density_amplitude,
+        pzeta_updated, ptot, qz, qz_updated, vth, v_norm_fac, ddens_dz, ddens_dz_upwind,
+        d2dens_dz2, duz_dz, duz_dz_upwind, d2uz_dz2, dpz_dz, dpz_dz_upwind, d2pz_dz2,
+        dqz_dz, dvth_dz, external_source_amplitude, external_source_density_amplitude,
         external_source_momentum_amplitude, external_source_pressure_amplitude,
         external_source_controller_integral)
 end
@@ -932,6 +941,8 @@ function calculate_moment_derivatives!(moments, scratch, scratch_dummy, z, z_spe
         buffer_r_5 = @view scratch_dummy.buffer_rs_5[:,is]
         buffer_r_6 = @view scratch_dummy.buffer_rs_6[:,is]
         if moments.evolve_density
+            @views derivative_z!(moments.charged.ddens_dz[:,:,is], density, buffer_r_1,
+                                 buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
             # Upwinded using upar as advection velocity, to be used in continuity equation
             @loop_r_z ir iz begin
                 dummy_zr[iz,ir] = -upar[iz,ir]
@@ -1465,6 +1476,8 @@ function calculate_moment_derivatives_neutral!(moments, scratch, scratch_dummy,
         buffer_r_5 = @view scratch_dummy.buffer_rsn_5[:,isn]
         buffer_r_6 = @view scratch_dummy.buffer_rsn_6[:,isn]
         if moments.evolve_density
+            @views derivative_z!(moments.neutral.ddens_dz[:,:,isn], density, buffer_r_1,
+                                 buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
             # Upwinded using upar as advection velocity, to be used in continuity equation
             @loop_r_z ir iz begin
                 dummy_zr[iz,ir] = -uz[iz,ir]

diff --git a/test/harrisonthompson.jl b/test/harrisonthompson.jl
@@ -252,21 +252,21 @@ function runtests()
         end
         @testset "Chebyshev split 1" begin
             run_test(test_input_chebyshev_split1, 3.e-2, 3.e-3,
-                     [-0.8089566414811241, -0.6619131770360419, -0.4308291908103036,
-                      -0.2958203456849356, -0.1934418964299151, -0.14925142403473443,
-                      -0.1197751269800481, -0.12060862943337616, -0.1134210587901286,
-                      -0.12060862943337315, -0.11977512698004258, -0.14925142403473443,
-                      -0.1934418964299073, -0.29582034568493026, -0.43082919081031584,
-                      -0.6619131770360454, -0.8089566414811092], 5.0e-9, 1.e-15)
+                     [-0.808956646073449, -0.6619131832543625, -0.4308291868843453,
+                      -0.295820339728472, -0.19344190006125275, -0.1492514208442407,
+                      -0.11977511930743077, -0.12060863604650167, -0.11342106824862994,
+                      -0.12060863604649866, -0.11977511930742626, -0.14925142084423915,
+                      -0.1934419000612479, -0.295820339728463, -0.4308291868843545,
+                      -0.6619131832543678, -0.808956646073442], 5.0e-9, 1.e-15)
         end
         @testset "Chebyshev split 2" begin
-            run_test(test_input_chebyshev_split2, 4.e-2, 3.e-3,
-                     [-0.782235614136118, -0.6326283459102813, -0.3904371921294129,
-                      -0.2608763603062158, -0.16526267868436909, -0.11064467678165119,
-                      -0.08500523595307348, -0.07891517979585916, -0.0749662287596728,
-                      -0.07891517979568231, -0.08500523595075139, -0.11064467678548737,
-                      -0.1652626786833772, -0.26087636031024986, -0.3904371921199402,
-                      -0.6326283459130052, -0.7822356141643002], 5.0e-9, 1.e-15)
+            run_test(test_input_chebyshev_split2, 5.e-2, 3.e-3,
+                     [-0.7667804422571606, -0.6128777083267765, -0.39031953439035494,
+                      -0.27326504140885904, -0.15311275955907663, -0.11567486122959246,
+                      -0.09226471519174786, -0.07477085120501512, -0.07206945428218994,
+                      -0.07477085120545898, -0.09226471518828984, -0.11567486123016281,
+                      -0.15311275955613904, -0.273265041412353, -0.3903195344134153,
+                      -0.612877708320375, -0.766780442235556], 5.0e-9, 1.e-15)
         end
         # The 'split 3' test is pretty badly resolved, but don't want to increase
         # run-time!

diff --git a/test/recycling_fraction_tests.jl b/test/recycling_fraction_tests.jl
@@ -202,25 +202,26 @@ function runtests()
         end
         @testset "Split 1" begin
             run_test(test_input_split1,
-                     [-0.05479385362313513, -0.017535505402861598, -0.001471824975450105,
-                      0.0016368028167570122, 0.002097476960724949, 0.0013534479492176561,
-                      0.0013561386240175473, 0.003653749578548698, 0.006719973898101642,
-                      0.014855703778346607, 0.03452400416817504, 0.03590889139838701,
-                      0.022290971814958923, 0.0074469188093168, 0.005048816543623428,
-                      0.0016968659971732512, 0.0013266650464258426, 0.0017028444376377347,
-                      0.0025344688564808227, -0.00018702558321354208,
-                      -0.009661131226016069, -0.048448362188207916])
+                     [-0.054793853738618496, -0.017535475032013862,
+                      -0.0014718402826481662, 0.0016368065803215382, 0.002097475822421603,
+                      0.001353447830403315, 0.001356138437924921, 0.0036537497347573,
+                      0.006719973928112565, 0.014855703760316889, 0.03452400419220982,
+                      0.03590889137214591, 0.022290971843531463, 0.007446918804323913,
+                      0.005048816472156039, 0.0016968661957691385, 0.0013266658105610114,
+                      0.0017028442360018413, 0.002534466861251151,
+                      -0.00018703865529355897, -0.009661145065079906,
+                      -0.0484483682752969])
         end
         @testset "Split 2" begin
             run_test(test_input_split2,
-                     [-0.055555691893139914, -0.020145180423245236, 0.001182129287268245,
-                      0.002193138222766261, 0.001944121337024971, 0.0011789363569829913,
-                      0.0013514285031174852, 0.0037355306549218033, 0.00672369539198614,
-                      0.014826903274076, 0.03454936283464004, 0.03587040869161285,
-                      0.022277731128342124, 0.007403053124167948, 0.005121534871576419,
-                      0.0017463615798537773, 0.001145272466843452, 0.0014049855368865622,
-                      0.002275541369035829, 0.0016780270704040375, -0.008381042083345462,
-                      -0.05005526304897398])
+                     [-0.05555568198447252, -0.020145183717956348, 0.001182118478411508,
+                      0.002193148323751635, 0.0019441188563940751, 0.0011789368818662881,
+                      0.0013514249605048384, 0.003735531583031493, 0.006723696092974834,
+                      0.014826903180374499, 0.03454936277756109, 0.03587040875737859,
+                      0.022277731154827392, 0.007403052912240603, 0.00512153431160143,
+                      0.0017463637584066217, 0.0011452779397062784, 0.0014049872146431029,
+                      0.0022755389057580316, 0.0016780234234311344, -0.008381041468024259,
+                      -0.05005526194222513])
         end
         @testset "Split 3" begin
             run_test(test_input_split3,