Micro-optimise a few electron-advance loops

Removes some allocations, and speeds up a few loops. The speed-up for a
given loop is substantial, but these loops were not a large fraction of
the run-time, so the overall effect will be minimal.

Also adds a few extra timers to ensure all functions in
`electron_kinetic_equation_euler_update!()` are timed.

moment_kinetics/src/electron_fluid_equations.jl

@@ -132,6 +132,7 @@ function calculate_electron_moments!(scratch, pdf, moments, composition, collisi
     else
         pdf_electron = pdf.electron
     end
+    electron_ppar = scratch.electron_ppar
     calculate_electron_density!(scratch.electron_density, moments.electron.dens_updated,
                                 scratch.density)
     calculate_electron_upar_from_charge_conservation!(
@@ -141,20 +142,20 @@ function calculate_electron_moments!(scratch, pdf, moments, composition, collisi
                                        kinetic_electrons_with_temperature_equation)
         begin_r_z_region()
         @loop_r_z ir iz begin
-            scratch.electron_ppar[iz,ir] = 0.5 * composition.me_over_mi *
-                                           scratch.electron_density[iz,ir] *
-                                           moments.electron.vth[iz,ir]^2
+            electron_ppar[iz,ir] = 0.5 * composition.me_over_mi *
+                                   scratch.electron_density[iz,ir] *
+                                   moments.electron.vth[iz,ir]^2
         end
         moments.electron.ppar_updated[] = true
     end
-    update_electron_vth_temperature!(moments, scratch.electron_ppar,
-                                     scratch.electron_density, composition)
-    calculate_electron_qpar!(moments.electron, pdf_electron, scratch.electron_ppar,
+    update_electron_vth_temperature!(moments, electron_ppar, scratch.electron_density,
+                                     composition)
+    calculate_electron_qpar!(moments.electron, pdf_electron, electron_ppar,
                              scratch.electron_upar, scratch.upar,
                              collisions.electron_fluid.nu_ei, composition.me_over_mi,
                              composition.electron_physics, vpa)
     if composition.electron_physics == braginskii_fluid
-        electron_fluid_qpar_boundary_condition!(scratch.electron_ppar,
+        electron_fluid_qpar_boundary_condition!(electron_ppar,
                                                 scratch.electron_upar,
                                                 scratch.electron_density,
                                                 moments.electron, z)
@@ -869,7 +870,7 @@ function calculate_electron_qpar_from_pdf!(qpar, ppar, vth, pdf, vpa)
     begin_r_z_region()
     ivperp = 1
     @loop_r_z ir iz begin
-        @views qpar[iz, ir] = 2*ppar[iz,ir]*vth[iz,ir]*integrate_over_vspace(pdf[:, ivperp, iz, ir], vpa.grid.^3, vpa.wgts)
+        @views qpar[iz, ir] = 2*ppar[iz,ir]*vth[iz,ir]*integrate_over_vspace(pdf[:, ivperp, iz, ir], vpa.grid, 3, vpa.wgts)
     end
 end
 

moment_kinetics/src/electron_kinetic_equation.jl

@@ -4775,18 +4775,20 @@ Note that this function operates on a single point in `r`, given by `ir`, and `f
                    composition, external_source_settings.electron, num_diss_params, z, ir)
 
         if ion_dt !== nothing
-            # Add source term to turn steady state solution into a backward-Euler update of
-            # electron_ppar with the ion timestep `ion_dt`.
-            ppar_previous_ion_step = moments.electron.ppar
-            begin_z_region()
-            @loop_z iz begin
-                # At this point, ppar_out = ppar_in + dt*RHS(ppar_in). Here we add a
-                # source/damping term so that in the steady state of the electron
-                # pseudo-timestepping iteration,
-                #   RHS(ppar) - (ppar - ppar_previous_ion_step) / ion_dt = 0,
-                # resulting in a backward-Euler step (as long as the pseudo-timestepping
-                # loop converges).
-                ppar_out[iz] += -dt * (ppar_in[iz] - ppar_previous_ion_step[iz,ir]) / ion_dt
+            @timeit global_timer "electron_ppar_ion_dt" begin
+                # Add source term to turn steady state solution into a backward-Euler
+                # update of electron_ppar with the ion timestep `ion_dt`.
+                ppar_previous_ion_step = moments.electron.ppar
+                begin_z_region()
+                @loop_z iz begin
+                    # At this point, ppar_out = ppar_in + dt*RHS(ppar_in). Here we add a
+                    # source/damping term so that in the steady state of the electron
+                    # pseudo-timestepping iteration,
+                    #   RHS(ppar) - (ppar - ppar_previous_ion_step) / ion_dt = 0,
+                    # resulting in a backward-Euler step (as long as the
+                    # pseudo-timestepping loop converges).
+                    ppar_out[iz] += -dt * (ppar_in[iz] - ppar_previous_ion_step[iz,ir]) / ion_dt
+                end
             end
         end
     end
@@ -5521,16 +5523,16 @@ function add_source_term!(source_term, vpa, z, dvth_dz)
     return nothing
 end
 
-function add_dissipation_term!(pdf_out, pdf_in, scratch_dummy, z_spectral, z, vpa,
-                               vpa_spectral, num_diss_params, dt)
+@timeit global_timer add_dissipation_term!(pdf_out, pdf_in, scratch_dummy, z_spectral, z,
+                                           vpa, vpa_spectral, num_diss_params, dt) = begin
     if num_diss_params.electron.vpa_dissipation_coefficient ≤ 0.0
         return nothing
     end
 
     begin_z_vperp_region()
     @loop_z_vperp iz ivperp begin
         @views second_derivative!(vpa.scratch, pdf_in[:,ivperp,iz], vpa, vpa_spectral)
-        @. pdf_out[:,ivperp,iz] += dt * num_diss_params.electron.vpa_dissipation_coefficient * vpa.scratch
+        @views @. pdf_out[:,ivperp,iz] += dt * num_diss_params.electron.vpa_dissipation_coefficient * vpa.scratch
     end
     return nothing
 end
@@ -5803,13 +5805,14 @@ function calculate_pdf_dot_prefactor!(pdf_dot_prefactor, ppar, vth, dens, ddens_
 end
 
 # add contribution to the residual coming from the term proporational to the pdf
-function add_contribution_from_pdf_term!(pdf_out, pdf_in, ppar, dens, upar, moments, vpa,
-                                         z, dt, electron_source_settings, ir)
+@timeit global_timer add_contribution_from_pdf_term!(pdf_out, pdf_in, ppar, dens, upar,
+                                                     moments, vpa, z, dt,
+                                                     electron_source_settings, ir) = begin
     vth = @view moments.electron.vth[:,ir]
     ddens_dz = @view moments.electron.ddens_dz[:,ir]
     dvth_dz = @view moments.electron.dvth_dz[:,ir]
     dqpar_dz = @view moments.electron.dqpar_dz[:,ir]
-    begin_z_vperp_vpa_region()
+    begin_z_vperp_region()
     @loop_z iz begin
         this_dqpar_dz = dqpar_dz[iz]
         this_ppar = ppar[iz]
@@ -5818,11 +5821,11 @@ function add_contribution_from_pdf_term!(pdf_out, pdf_in, ppar, dens, upar, mome
         this_dens = dens[iz]
         this_dvth_dz = dvth_dz[iz]
         this_vth = vth[iz]
-        @loop_vperp_vpa ivperp ivpa begin
-            pdf_out[ivpa,ivperp,iz] +=
-                dt * (-0.5 * this_dqpar_dz / this_ppar - vpa[ivpa] * this_vth *
+        @loop_vperp ivperp begin
+            @views @. pdf_out[:,ivperp,iz] +=
+                dt * (-0.5 * this_dqpar_dz / this_ppar - vpa * this_vth *
                       (this_ddens_dz / this_dens - this_dvth_dz / this_vth)) *
-                pdf_in[ivpa,ivperp,iz]
+                pdf_in[:,ivperp,iz]
         end
     end
 
@@ -5832,11 +5835,11 @@ function add_contribution_from_pdf_term!(pdf_out, pdf_in, ppar, dens, upar, mome
             @views source_momentum_amplitude = moments.electron.external_source_momentum_amplitude[:, ir, index]
             @views source_pressure_amplitude = moments.electron.external_source_pressure_amplitude[:, ir, index]
             @loop_z iz begin
-                term = dt * (1.5 * source_density_amplitude[iz,ir] / dens[iz,ir] -
-                            (0.5 * source_pressure_amplitude[iz,ir] +
-                            source_momentum_amplitude[iz,ir]) / ppar[iz,ir])
-                @loop_vperp_vpa ivperp ivpa begin
-                    pdf_out[ivpa,ivperp,iz,ir] -= term * pdf_in[ivpa,ivperp,iz,ir]
+                term = dt * (1.5 * source_density_amplitude[iz] / dens[iz] -
+                            (0.5 * source_pressure_amplitude[iz] +
+                            source_momentum_amplitude[iz]) / ppar[iz])
+                @loop_vperp ivperp begin
+                    @views @. pdf_out[:,ivperp,iz] -= term * pdf_in[:,ivperp,iz]
                 end
             end
         end

moment_kinetics/src/electron_vpa_advection.jl

@@ -61,10 +61,10 @@ function update_electron_speed_vpa!(advect, density, upar, ppar, moments, vpa,
     dppar_dz = @view moments.electron.dppar_dz[:,ir]
     dqpar_dz = @view moments.electron.dqpar_dz[:,ir]
     dvth_dz = @view moments.electron.dvth_dz[:,ir]
-    speed = advect.speed
+    speed = @view advect.speed[:,:,:,ir]
     # calculate the advection speed in wpa
     @loop_z_vperp_vpa iz ivperp ivpa begin
-        speed[ivpa,ivperp,iz,ir] = ((vth[iz] * dppar_dz[iz] + vpa[ivpa] * dqpar_dz[iz])
+        speed[ivpa,ivperp,iz] = ((vth[iz] * dppar_dz[iz] + vpa[ivpa] * dqpar_dz[iz])
                                     / (2 * ppar[iz]) - vpa[ivpa]^2 * dvth_dz[iz])
     end
 
@@ -80,8 +80,8 @@ function update_electron_speed_vpa!(advect, density, upar, ppar, moments, vpa,
                             2.0 * upar[iz] * source_momentum_amplitude[iz]) /
                         ppar[iz] +
                     0.5 * source_density_amplitude[iz] / density[iz]
-                @loop_vperp_vpa ivperp ivpa begin
-                    speed[ivpa,ivperp,iz,ir] += term1 + vpa[ivpa] * term2_over_vpa
+                @loop_vperp ivperp begin
+                    @. speed[:,ivperp,iz] += term1 + vpa * term2_over_vpa
                 end
             end
         end

moment_kinetics/src/external_sources.jl

@@ -989,14 +989,13 @@ Note that this function operates on a single point in `r`, given by `ir`, and `p
         this_upar = electron_upar[iz]
         this_prefactor = dt * this_vth / electron_density[iz] * vth_factor *
                          source_amplitude[iz]
-        @loop_vperp_vpa ivperp ivpa begin
+        @loop_vperp ivperp begin
             # Factor of 1/sqrt(π) (for 1V) or 1/π^(3/2) (for 2V/3V) is absorbed by the
             # normalisation of F
             vperp_unnorm = vperp_grid[ivperp] * this_vth
-            vpa_unnorm = vpa_grid[ivpa] * this_vth + this_upar
-            pdf_out[ivpa,ivperp,iz] +=
+            @. pdf_out[:,ivperp,iz] +=
                 this_prefactor *
-                exp(-(vperp_unnorm^2 + vpa_unnorm^2) * me_over_mi / source_T)
+                exp(-(vperp_unnorm^2 + (vpa_grid * this_vth + this_upar)^2) * me_over_mi / source_T)
         end
     end