Merge branch 'master' into fix_collision_operator_scripts

.github/workflows/documentation.yml

@@ -18,7 +18,10 @@ jobs:
       - uses: julia-actions/cache@v1
       - name: Install dependencies
         run: |
-          pip3 install --user matplotlib
+          # Version 3.9.0 of matplotlib causes an error with PyPlot.jl, so pin
+          # matplotlib version to 3.8.3 until this is fixed. See
+          # https://github.com/JuliaPy/PyPlot.jl/issues/582).
+          pip3 install --user "matplotlib==3.8.3"
       - name: Build and deploy
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} # Authenticate with GitHub Actions token
@@ -33,6 +36,9 @@ jobs:
           version: '1.10'
       - name: Install dependencies
         run: |
-          pip3 install --user matplotlib
+          # Version 3.9.0 of matplotlib causes an error with PyPlot.jl, so pin
+          # matplotlib version to 3.8.3 until this is fixed. See
+          # https://github.com/JuliaPy/PyPlot.jl/issues/582).
+          pip3 install --user "matplotlib==3.8.3"
       - name: Build and deploy
         run: julia --project=docs/ docs/make-pdf.jl

docs/src/moment_kinetic_equations.md

@@ -89,7 +89,7 @@ energy equation over $\tilde{v}_{\|}$ instead of $w_{\|}$,
   \tilde{p}_{\|,s}
   & = \frac{1}{\sqrt{\pi}}\int d\tilde{v}_{\|}\left(\tilde{v}_{\|}
       - \tilde{u}_{s}\right)^{2}\tilde{f}_{s}
-    = \int d\tilde{v}_{\|}\tilde{v}_{\|}^{2}\tilde{f}_{s}
+    = \frac{1}{\sqrt{\pi}}\int d\tilde{v}_{\|}\tilde{v}_{\|}^{2}\tilde{f}_{s}
       - \tilde{n}_{s}\tilde{u}_{s}^{2}\\
 %
   \tilde{q}_{\|,s}
@@ -184,7 +184,7 @@ tildes from here on)
   \frac{\partial u_{i}}{\partial t} + \frac{1}{n_{i}}\frac{\partial p_{\|,i}}{\partial z}
   + u_{i}\frac{\partial u_{i}}{\partial z} + \frac{1}{2}\frac{\partial\phi}{\partial z}
   &= -R_{in}n_{n}\left(u_{i}-u_{n}\right)
-    + R_{\mathrm{ion}}\frac{n_{i}n_{n}}{n_{s}}\left(u_{n}-u_{i}\right)
+    + R_{\mathrm{ion}}n_{i}n_{n}\left(u_{n}-u_{i}\right)
     - \frac{u_{i}}{n_{i}} \int dv_\parallel S_{i}
 \end{align}
 ```
@@ -440,16 +440,16 @@ The derivatives transform as
 \begin{align}
   \left.\frac{\partial f_{s}}{\partial t}\right|_{z,v\|}
   & \rightarrow\left.\frac{\partial f_{s}}{\partial t}\right|_{z,w\|}
-               - \frac{1}{v_{\mathrm{th},s}}\frac{\partial u_{s}}{\partial t}\left.\frac{\partial f_{s}}{\partial w_{\|,s}}\right|_{z,w\|}
-               - \frac{w_{\|,s}}{v_{\mathrm{th},s}}\frac{\partial v_{\mathrm{th},s}}{\partial t}\left.\frac{\partial f_{s}}{\partial w_{\|,s}}\right|_{z,w\|}\\
+               - \frac{1}{v_{\mathrm{th},s}}\frac{\partial u_{s}}{\partial t}\left.\frac{\partial f_{s}}{\partial w_{\|,s}}\right|_{z,t}
+               - \frac{w_{\|,s}}{v_{\mathrm{th},s}}\frac{\partial v_{\mathrm{th},s}}{\partial t}\left.\frac{\partial f_{s}}{\partial w_{\|,s}}\right|_{z,t}\\
 %
   \left.\frac{\partial f_{s}}{\partial z}\right|_{z,v\|}
-  & \rightarrow\left.\frac{\partial f_{s}}{\partial z}\right|_{z,w\|}
-               - \frac{1}{v_{\mathrm{th},s}}\frac{\partial u_{s}}{\partial z}\left.\frac{\partial f_{s}}{\partial w_{\|,s}}\right|_{z,w\|}
-               - \frac{w_{\|,s}}{v_{\mathrm{th},s}}\frac{\partial v_{\mathrm{th},s}}{\partial z}\left.\frac{\partial f_{s}}{\partial w_{\|,s}}\right|_{z,w\|}\\
+  & \rightarrow\left.\frac{\partial f_{s}}{\partial z}\right|_{t,w\|}
+               - \frac{1}{v_{\mathrm{th},s}}\frac{\partial u_{s}}{\partial z}\left.\frac{\partial f_{s}}{\partial w_{\|,s}}\right|_{z,t}
+               - \frac{w_{\|,s}}{v_{\mathrm{th},s}}\frac{\partial v_{\mathrm{th},s}}{\partial z}\left.\frac{\partial f_{s}}{\partial w_{\|,s}}\right|_{z,wt}\\
 %
   \left.\frac{\partial f_{s}}{\partial v_{\|}}\right|_{z,v\|}
-  & \rightarrow\frac{1}{v_{\mathrm{th},s}}\left.\frac{\partial f_{s}}{\partial w_{\|,s}}\right|_{z,w\|}
+  & \rightarrow\frac{1}{v_{\mathrm{th},s}}\left.\frac{\partial f_{s}}{\partial w_{\|,s}}\right|_{z,t}
 \end{align}
 ```
 

docs/src/zz_boundary_conditions.md

@@ -0,0 +1,6 @@
+`boundary_conditions`
+=====================
+
+```@autodocs
+Modules = [moment_kinetics.boundary_conditions]
+```

docs/src/zz_gyroaverages.md

@@ -0,0 +1,6 @@
+`gyroaverages`
+==============
+
+```@autodocs
+Modules = [moment_kinetics.gyroaverages]
+```

docs/src/zz_lagrange_polynomials.md

@@ -0,0 +1,6 @@
+`lagrange_polynomials`
+======================
+
+```@autodocs
+Modules = [moment_kinetics.lagrange_polynomials]
+```

moment_kinetics/src/chebyshev.jl

@@ -465,7 +465,8 @@ function chebyshev_spectral_derivative!(df,f)
     end
 end
 
-function single_element_interpolate!(result, newgrid, f, imin, imax, coord, chebyshev::chebyshev_base_info)
+function single_element_interpolate!(result, newgrid, f, imin, imax, ielement, coord,
+                                     chebyshev::chebyshev_base_info)
     # Temporary buffer to store Chebyshev coefficients
     cheby_f = chebyshev.df
 

moment_kinetics/src/coordinates.jl

@@ -109,6 +109,10 @@ struct coordinate{T <: AbstractVector{mk_float}}
     element_boundaries::Array{mk_float,1}
     # Does the coordinate use a 'Radau' discretization for the first element?
     radau_first_element::Bool
+    # 'Other' nodes where the j'th Lagrange polynomial (which is 1 at x[j]) is equal to 0
+    other_nodes::Array{mk_float,3}
+    # One over the denominators of the Lagrange polynomials
+    one_over_denominator::Array{mk_float,2}
 end
 
 """
@@ -190,13 +194,37 @@ function define_coordinate(input, parallel_io::Bool=false; run_directory=nothing
         local_io_range = 1 : n_local-1
         global_io_range = input.irank*(n_local-1)+1 : (input.irank+1)*(n_local-1)
     end
+
+    # Precompute some values for Lagrange polynomial evaluation
+    other_nodes = allocate_float(input.ngrid-1, input.ngrid, input.nelement_local)
+    one_over_denominator = allocate_float(input.ngrid, input.nelement_local)
+    for ielement ∈ 1:input.nelement_local
+        if ielement == 1
+            this_imin = imin[ielement]
+        else
+            this_imin = imin[ielement] - 1
+        end
+        this_imax = imax[ielement]
+        this_grid = grid[this_imin:this_imax]
+        for j ∈ 1:input.ngrid
+            @views other_nodes[1:j-1,j,ielement] .= this_grid[1:j-1]
+            @views other_nodes[j:end,j,ielement] .= this_grid[j+1:end]
+
+            if input.ngrid == 1
+                one_over_denominator[j,ielement] = 1.0
+            else
+                one_over_denominator[j,ielement] = 1.0 / prod(this_grid[j] - n for n ∈ @view other_nodes[:,j,ielement])
+            end
+        end
+    end
+
     coord = coordinate(input.name, n_global, n_local, input.ngrid,
         input.nelement_global, input.nelement_local, input.nrank, input.irank, input.L, grid,
         cell_width, igrid, ielement, imin, imax, igrid_full, input.discretization, input.fd_option, input.cheb_option,
         input.bc, wgts, uniform_grid, duniform_dgrid, scratch, copy(scratch), copy(scratch), scratch_shared, scratch_shared2,
         scratch_2d, copy(scratch_2d), advection, send_buffer, receive_buffer, input.comm,
         local_io_range, global_io_range, element_scale, element_shift, input.element_spacing_option,
-        element_boundaries, radau_first_element)
+        element_boundaries, radau_first_element, other_nodes, one_over_denominator)
 
     if coord.n == 1 && occursin("v", coord.name)
         spectral = null_velocity_dimension_info()

moment_kinetics/src/fokker_planck_calculus.jl

@@ -39,7 +39,7 @@ using ..array_allocation: allocate_float, allocate_shared_float
 using ..calculus: derivative!
 using ..communication
 using ..communication: MPISharedArray, global_rank
-using ..lagrange_polynomials: lagrange_poly
+using ..lagrange_polynomials: lagrange_poly, lagrange_poly_optimised
 using ..looping
 using moment_kinetics.gauss_legendre: get_QQ_local!
 using Dates
@@ -646,12 +646,16 @@ end
 # `ellipe(k) = \int^{\pi/2}\_0 \frac{1}{\sqrt{ 1 - m \sin^2(\theta)}} d \theta`
 
 function local_element_integration!(G0_weights,G1_weights,H0_weights,H1_weights,H2_weights,H3_weights,
-                            nquad_vpa,ielement_vpa,vpa_nodes,vpa, # info about primed vpa grids
-                            nquad_vperp,ielement_vperp,vperp_nodes,vperp, # info about primed vperp grids
+                            nquad_vpa,ielement_vpa,vpa, # info about primed vpa grids
+                            nquad_vperp,ielement_vperp,vperp, # info about primed vperp grids
                             x_vpa, w_vpa, x_vperp, w_vperp, # points and weights for primed (source) grids
                             vpa_val, vperp_val) # values and indices for unprimed (field) grids
     for igrid_vperp in 1:vperp.ngrid
+        vperp_other_nodes = @view vperp.other_nodes[:,igrid_vperp,ielement_vperp]
+        vperp_one_over_denominator = vperp.one_over_denominator[igrid_vperp,ielement_vperp]
         for igrid_vpa in 1:vpa.ngrid
+            vpa_other_nodes = @view vpa.other_nodes[:,igrid_vpa,ielement_vpa]
+            vpa_one_over_denominator = vpa.one_over_denominator[igrid_vpa,ielement_vpa]
             # get grid index for point on full grid  
             ivpap = vpa.igrid_full[igrid_vpa,ielement_vpa]   
             ivperpp = vperp.igrid_full[igrid_vperp,ielement_vperp]   
@@ -679,8 +683,12 @@ function local_element_integration!(G0_weights,G1_weights,H0_weights,H1_weights,
                     H_elliptic_integral_factor = 2.0*ellipk_mm/(pi*prefac)
                     H1_elliptic_integral_factor = -(2.0/(pi*prefac))*( (mm-2.0)*(ellipk_mm/mm) + (2.0*ellipe_mm/mm) )
                     H2_elliptic_integral_factor = (2.0/(pi*prefac))*( (3.0*mm^2 - 8.0*mm + 8.0)*(ellipk_mm/(3.0*mm^2)) + (4.0*mm - 8.0)*ellipe_mm/(3.0*mm^2) )
-                    lagrange_poly_vpa = lagrange_poly(igrid_vpa,vpa_nodes,x_kvpa)
-                    lagrange_poly_vperp = lagrange_poly(igrid_vperp,vperp_nodes,x_kvperp)
+                    lagrange_poly_vpa = lagrange_poly_optimised(vpa_other_nodes,
+                                                                vpa_one_over_denominator,
+                                                                x_kvpa)
+                    lagrange_poly_vperp = lagrange_poly_optimised(vperp_other_nodes,
+                                                                  vperp_one_over_denominator,
+                                                                  x_kvperp)
 
                     (G0_weights[ivpap,ivperpp] += 
                         lagrange_poly_vpa*lagrange_poly_vperp*
@@ -741,8 +749,8 @@ function loop_over_vpa_elements!(G0_weights,G1_weights,H0_weights,H1_weights,H2_
         vpa_min, vpa_max = vpa_nodes[1], vpa_nodes[end]
         nquad_vpa = get_scaled_x_w_no_divergences!(x_vpa, w_vpa, x_legendre, w_legendre, vpa_min, vpa_max)
         local_element_integration!(G0_weights,G1_weights,H0_weights,H1_weights,H2_weights,H3_weights,
-                    nquad_vpa,ielement_vpap,vpa_nodes,vpa,
-                    nquad_vperp,ielement_vperpp,vperp_nodes,vperp,
+                    nquad_vpa,ielement_vpap,vpa,
+                    nquad_vperp,ielement_vperpp,vperp,
                     x_vpa, w_vpa, x_vperp, w_vperp, 
                     vpa_val, vperp_val)
     end
@@ -755,8 +763,8 @@ function loop_over_vpa_elements!(G0_weights,G1_weights,H0_weights,H1_weights,H2_
         #nquad_vpa = get_scaled_x_w_no_divergences!(x_vpa, w_vpa, x_legendre, w_legendre, vpa_min, vpa_max)
         nquad_vpa = get_scaled_x_w_with_divergences!(x_vpa, w_vpa, x_legendre, w_legendre, x_laguerre, w_laguerre, vpa_min, vpa_max, vpa_nodes, igrid_vpa, vpa_val)
         local_element_integration!(G0_weights,G1_weights,H0_weights,H1_weights,H2_weights,H3_weights,
-                    nquad_vpa,ielement_vpap,vpa_nodes,vpa,
-                    nquad_vperp,ielement_vperpp,vperp_nodes,vperp,
+                    nquad_vpa,ielement_vpap,vpa,
+                    nquad_vperp,ielement_vperpp,vperp,
                     x_vpa, w_vpa, x_vperp, w_vperp, 
                     vpa_val, vperp_val)
     end
@@ -767,8 +775,8 @@ function loop_over_vpa_elements!(G0_weights,G1_weights,H0_weights,H1_weights,H2_
         vpa_min, vpa_max = vpa_nodes[1], vpa_nodes[end]
         nquad_vpa = get_scaled_x_w_no_divergences!(x_vpa, w_vpa, x_legendre, w_legendre, vpa_min, vpa_max)
         local_element_integration!(G0_weights,G1_weights,H0_weights,H1_weights,H2_weights,H3_weights,
-                    nquad_vpa,ielement_vpap,vpa_nodes,vpa,
-                    nquad_vperp,ielement_vperpp,vperp_nodes,vperp,
+                    nquad_vpa,ielement_vpap,vpa,
+                    nquad_vperp,ielement_vperpp,vperp,
                     x_vpa, w_vpa, x_vperp, w_vperp, 
                     vpa_val, vperp_val)
 
@@ -788,8 +796,8 @@ function loop_over_vpa_elements_no_divergences!(G0_weights,G1_weights,H0_weights
         vpa_min, vpa_max = vpa_nodes[1], vpa_nodes[end]
         nquad_vpa = get_scaled_x_w_no_divergences!(x_vpa, w_vpa, x_legendre, w_legendre, vpa_min, vpa_max)
         local_element_integration!(G0_weights,G1_weights,H0_weights,H1_weights,H2_weights,H3_weights,
-                    nquad_vpa,ielement_vpap,vpa_nodes,vpa,
-                    nquad_vperp,ielement_vperpp,vperp_nodes,vperp,
+                    nquad_vpa,ielement_vpap,vpa,
+                    nquad_vperp,ielement_vperpp,vperp,
                     x_vpa, w_vpa, x_vperp, w_vperp, 
                     vpa_val, vperp_val)
 
@@ -2366,6 +2374,25 @@ function interpolate_2D_vspace!(pdf_out,pdf_in,vpa,vperp,scalefac)
     end
     return nothing
 end
+# Alternative version that should be faster - to be tested
+#function interpolate_2D_vspace!(pdf_out, pdf_in, vpa, vpa_spectral, vperp, vperp_spectral,
+#                                scalefac, pdf_buffer)
+#    newgrid_vperp = vperp.scratch .= scalefac .* vperp.grid
+#    newgrid_vpa = vpa.scratch .= scalefac .* vpa.grid
+#
+#    begin_anyv_vpa_region()
+#    @loop_vpa ivpa begin
+#        @views interpolate_to_grid_1d!(pdf_buffer[ivpa,:], newgrid_vperp,
+#                                       pdf_in[ivpa,:], vperp, vperp_spectral)
+#    end
+#
+#    begin_anyv_vperp_region()
+#    @loop_vperp ivperp begin
+#        @views interpolate_to_grid_1d!(pdf_out[:,ivperp], newgrid_vpa,
+#                                       pdf_buffer[:,ivperp], vpa, vpa_spectral)
+
+#    end
+#end
 
 """
 function to find the element in which x sits

moment_kinetics/src/gauss_legendre.jl

@@ -29,7 +29,7 @@ using ..type_definitions: mk_float, mk_int
 using ..array_allocation: allocate_float
 import ..calculus: elementwise_derivative!, mass_matrix_solve!
 import ..interpolation: single_element_interpolate!
-using ..lagrange_polynomials: lagrange_poly
+using ..lagrange_polynomials: lagrange_poly_optimised
 using ..moment_kinetics_structs: weak_discretization_info
 
 
@@ -308,11 +308,23 @@ function elementwise_derivative!(coord, ff, adv_fac, spectral::gausslegendre_inf
     return elementwise_derivative!(coord, ff, spectral)
 end
 
-function single_element_interpolate!(result, newgrid, f, imin, imax, coord, gausslegendre::gausslegendre_base_info)
-    for i ∈ 1:length(newgrid)
-        result[i] = 0.0
-        for j ∈ 1:coord.ngrid
-            @views result[i] += f[j] * lagrange_poly(j, coord.grid[imin:imax], newgrid[i])
+function single_element_interpolate!(result, newgrid, f, imin, imax, ielement, coord,
+                                     gausslegendre::gausslegendre_base_info)
+    n_new = length(newgrid)
+
+    i = 1
+    other_nodes = @view coord.other_nodes[:,i,ielement]
+    one_over_denominator = coord.one_over_denominator[i,ielement]
+    this_f = f[i]
+    for j ∈ 1:n_new
+        result[j] = this_f * lagrange_poly_optimised(other_nodes, one_over_denominator, newgrid[j])
+    end
+    for i ∈ 2:coord.ngrid
+        other_nodes = @view coord.other_nodes[:,i,ielement]
+        one_over_denominator = coord.one_over_denominator[i,ielement]
+        this_f = f[i]
+        for j ∈ 1:n_new
+            result[j] += this_f * lagrange_poly_optimised(other_nodes, one_over_denominator, newgrid[j])
         end
     end
 

moment_kinetics/src/interpolation.jl

@@ -99,7 +99,7 @@ function interpolate_to_grid_1d!(result, newgrid, f, coord, spectral)
         kmin = kstart[1]
         kmax = kstart[2] - 1
         @views single_element_interpolate!(result[kmin:kmax], newgrid[kmin:kmax],
-                                           f[imin:imax], imin, imax, coord,
+                                           f[imin:imax], imin, imax, 1, coord,
                                            first_element_spectral)
     end
     @inbounds for j ∈ 2:nelement
@@ -109,7 +109,7 @@ function interpolate_to_grid_1d!(result, newgrid, f, coord, spectral)
             imin = coord.imin[j] - 1
             imax = coord.imax[j]
             @views single_element_interpolate!(result[kmin:kmax], newgrid[kmin:kmax],
-                                               f[imin:imax], imin, imax, coord,
+                                               f[imin:imax], imin, imax, j, coord,
                                                spectral.lobatto)
         end
     end

moment_kinetics/src/lagrange_polynomials.jl

@@ -8,7 +8,7 @@ their being scattered (and possibly duplicated) in other modules.
 """
 module lagrange_polynomials
 
-export lagrange_poly
+export lagrange_poly, lagrange_poly_optimised
 
 """
 Lagrange polynomial
@@ -33,4 +33,21 @@ function lagrange_poly(j,x_nodes,x)
     return poly
 end
 
+"""
+    lagrange_poly_optimised(other_nodes, one_over_denominator, x)
+
+Optimised version of Lagrange polynomial calculation, making use of pre-calculated quantities.
+
+`other_nodes` is a vector of the grid points in this element where this Lagrange
+polynomial is zero (the other nodes than the one where it is 1).
+
+`one_over_denominator` is `1/prod(x0 - n for n ∈ other_nodes)` where `x0` is the grid
+point where this Lagrange polynomial is 1.
+
+`x` is the point to evaluate the Lagrange polynomial at.
+"""
+function lagrange_poly_optimised(other_nodes, one_over_denominator, x)
+    return prod(x - n for n ∈ other_nodes) * one_over_denominator
+end
+
 end