Merge branch 'master' into mms_bugfixes_and_docs

.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/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/calculus.jl

@@ -14,6 +14,8 @@ using ..communication: block_rank
 using ..communication: _block_synchronize
 using ..looping
 
+using LinearAlgebra
+
 """
     elementwise_derivative!(coord, f, adv_fac, spectral)
     elementwise_derivative!(coord, f, spectral)
@@ -27,17 +29,6 @@ Result is stored in coord.scratch_2d.
 """
 function elementwise_derivative! end
 
-"""
-    elementwise_second_derivative!(coord, f, spectral)
-
-Generic function for element-by-element second derivatives.
-
-Note: no upwinding versions of second deriatives.
-
-Result is stored in coord.scratch_2d.
-"""
-function elementwise_second_derivative! end
-
 """
     derivative!(df, f, coord, adv_fac, spectral)
 
@@ -73,7 +64,7 @@ function derivative!(df, f, coord, spectral::Union{null_spatial_dimension_info,
     return nothing
 end
 
-function second_derivative!(d2f, f, coord, spectral)
+function second_derivative!(d2f, f, coord, spectral; handle_periodic=true)
     # computes d^2f / d(coord)^2
     # For spectral element methods, calculate second derivative by applying first
     # derivative twice, with special treatment for element boundaries
@@ -152,40 +143,54 @@ is an input.
 """
 function mass_matrix_solve! end
 
-"""
-Apply 'K-matrix' as part of a weak-form second derivative
-"""
-function elementwise_apply_Kmat! end
-
-function second_derivative!(d2f, f, coord, spectral::weak_discretization_info)
+function second_derivative!(d2f, f, coord, spectral::weak_discretization_info; handle_periodic=true)
     # obtain the RHS of numerical weak-form of the equation 
     # g = d^2 f / d coord^2, which is 
     # M * g = K * f, with M the mass matrix and K an appropriate stiffness matrix
     # by multiplying by basis functions and integrating by parts    
-    elementwise_apply_Kmat!(coord, f, spectral)
-    # map the RHS vector K * f from the elemental grid to the full grid;
-    # at element boundaries, use the average of K * f from neighboring elements.
-    derivative_elements_to_full_grid!(coord.scratch, coord.scratch_2d, coord)
+    mul!(coord.scratch, spectral.K_matrix, f)
+
+    if handle_periodic && coord.bc == "periodic"
+        if coord.nrank > 1
+            error("second_derivative!() cannot handle periodic boundaries for a "
+                  * "distributed coordinate")
+        end
+
+        coord.scratch[1] = 0.5 * (coord.scratch[1] + coord.scratch[end])
+        coord.scratch[end] = coord.scratch[1]
+    end
+
     # solve weak form matrix problem M * g = K * f to obtain g = d^2 f / d coord^2
+    if coord.nrank > 1
+        error("mass_matrix_solve!() does not support a "
+              * "distributed coordinate")
+    end
     mass_matrix_solve!(d2f, coord.scratch, spectral)
 end
 
-"""
-Apply 'L-matrix' as part of a weak-form Laplacian derivative
-"""
-function elementwise_apply_Lmat! end
-
 function laplacian_derivative!(d2f, f, coord, spectral::weak_discretization_info)
     # for coord.name 'vperp' obtain the RHS of numerical weak-form of the equation 
     # g = (1/coord) d/d coord ( coord  d f / d coord ), which is 
     # M * g = K * f, with M the mass matrix, and K an appropriate stiffness matrix,
     # by multiplying by basis functions and integrating by parts.
     # for all other coord.name, do exactly the same as second_derivative! above.
-    elementwise_apply_Lmat!(coord, f, spectral)
-    # map the RHS vector K * f from the elemental grid to the full grid;
-    # at element boundaries, use the average of K * f from neighboring elements.
-    derivative_elements_to_full_grid!(coord.scratch, coord.scratch_2d, coord)
+    mul!(coord.scratch, spectral.L_matrix, f)
+
+    if handle_periodic && coord.bc == "periodic"
+        if coord.nrank > 1
+            error("second_derivative!() cannot handle periodic boundaries for a "
+                  * "distributed coordinate")
+        end
+
+        coord.scratch[1] = 0.5 * (coord.scratch[1] + coord.scratch[end])
+        coord.scratch[end] = coord.scratch[1]
+    end
+
     # solve weak form matrix problem M * g = K * f to obtain g = d^2 f / d coord^2
+    if coord.nrank > 1
+        error("mass_matrix_solve!() does not support a "
+              * "distributed coordinate")
+    end
     mass_matrix_solve!(d2f, coord.scratch, spectral)
 end
 

moment_kinetics/src/chebyshev.jl

@@ -18,7 +18,7 @@ using ..array_allocation: allocate_float, allocate_complex
 using ..clenshaw_curtis: clenshawcurtisweights
 import ..calculus: elementwise_derivative!
 using ..communication
-import ..interpolation: interpolate_to_grid_1d!
+import ..interpolation: single_element_interpolate!
 using ..moment_kinetics_structs: discretization_info
 
 """
@@ -465,88 +465,8 @@ function chebyshev_spectral_derivative!(df,f)
     end
 end
 
-"""
-Interpolation from a regular grid to a 1d grid with arbitrary spacing
-
-Arguments
----------
-result : Array{mk_float, 1}
-    Array to be overwritten with the result of the interpolation
-new_grid : Array{mk_float, 1}
-    Grid of points to interpolate `coord` to
-f : Array{mk_float}
-    Field to be interpolated
-coord : coordinate
-    `coordinate` struct giving the coordinate along which f varies
-chebyshev : chebyshev_info
-    struct containing information for Chebyshev transforms
-
-Note that this routine does not support Gauss-Chebyshev-Radau elements
-"""
-function interpolate_to_grid_1d!(result, newgrid, f, coord, chebyshev::chebyshev_info)
-    # define local variable nelement for convenience
-    nelement = coord.nelement_local
-    # check array bounds
-    @boundscheck nelement == size(chebyshev.lobatto.f,2) || throw(BoundsError(chebyshev.lobatto.f))
-
-    n_new = size(newgrid)[1]
-    # Find which points belong to which element.
-    # kstart[j] contains the index of the first point in newgrid that is within element
-    # j, and kstart[nelement+1] is n_new if the last point is within coord.grid, or the
-    # index of the first element outside coord.grid otherwise.
-    # Assumes points in newgrid are sorted.
-    # May not be the most efficient algorithm.
-    # Find start/end points for each element, storing their indices in kstart
-    kstart = Vector{mk_int}(undef, nelement+1)
-    # set the starting index by finding the start of coord.grid
-    kstart[1] = searchsortedfirst(newgrid, coord.grid[1])
-    # check to see if any of the newgrid points are to the left of the first grid point
-    for j ∈ 1:kstart[1]-1
-        # if the new grid location is outside the bounds of the original grid,
-        # extrapolate f with Gaussian-like decay beyond the domain
-        result[j] = f[1] * exp(-(coord.grid[1] - newgrid[j])^2)
-    end
-    @inbounds for j ∈ 1:nelement
-        # Search from kstart[j] to try to speed up the sort, but means result of
-        # searchsortedfirst() is offset by kstart[j]-1 from the beginning of newgrid.
-        kstart[j+1] = kstart[j] - 1 + @views searchsortedfirst(newgrid[kstart[j]:end], coord.grid[coord.imax[j]])
-    end
-
-    # First element includes both boundary points, while all others have only one (to
-    # avoid duplication), so calculate the first element outside the loop.
-    if kstart[1] < kstart[2]
-        imin = coord.imin[1]
-        imax = coord.imax[1]
-        kmin = kstart[1]
-        kmax = kstart[2] - 1
-        @views chebyshev_interpolate_single_element!(result[kmin:kmax],
-                                                     newgrid[kmin:kmax],
-                                                     f[imin:imax],
-                                                     imin, imax, coord, chebyshev.lobatto)
-    end
-    @inbounds for j ∈ 2:nelement
-        kmin = kstart[j]
-        kmax = kstart[j+1] - 1
-        if kmin <= kmax
-            imin = coord.imin[j] - 1
-            imax = coord.imax[j]
-            @views chebyshev_interpolate_single_element!(result[kmin:kmax],
-                                                         newgrid[kmin:kmax],
-                                                         f[imin:imax],
-                                                         imin, imax, coord, chebyshev.lobatto)
-        end
-    end
-
-    for k ∈ kstart[nelement+1]:n_new
-        result[k] = f[end] * exp(-(newgrid[k] - coord.grid[end])^2)
-    end
-
-    return nothing
-end
-
-"""
-"""
-function chebyshev_interpolate_single_element!(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

@@ -107,6 +107,12 @@ struct coordinate{T <: AbstractVector{mk_float}}
     element_spacing_option::String
     # list of element boundaries
     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
 
 """
@@ -115,7 +121,7 @@ setup the coordinate grid, and populate the coordinate structure
 containing all of this information
 """
 function define_coordinate(input, parallel_io::Bool=false; run_directory=nothing,
-                           ignore_MPI=false, init_YY::Bool=true)
+                           ignore_MPI=false, collision_operator_dim::Bool=true)
     # total number of grid points is ngrid for the first element
     # plus ngrid-1 unique points for each additional element due
     # to the repetition of a point at the element boundary
@@ -136,8 +142,9 @@ function define_coordinate(input, parallel_io::Bool=false; run_directory=nothing
     element_scale, element_shift = set_element_scale_and_shift(input.nelement_global, input.nelement_local, input.irank, element_boundaries)
     # initialize the grid and the integration weights associated with the grid
     # also obtain the Chebyshev theta grid and spacing if chosen as discretization option
-    grid, wgts, uniform_grid = init_grid(input.ngrid, input.nelement_local, n_global, n_local, input.irank, input.L, element_scale, element_shift,
-        imin, imax, igrid, input.discretization, input.name)
+    grid, wgts, uniform_grid, radau_first_element = init_grid(input.ngrid,
+        input.nelement_local, n_global, n_local, input.irank, input.L, element_scale,
+        element_shift, imin, imax, igrid, input.discretization, input.name)
     # calculate the widths of the cells between neighboring grid points
     cell_width = grid_spacing(grid, n_local)
     # duniform_dgrid is the local derivative of the uniform grid with respect to
@@ -187,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)
+        element_boundaries, radau_first_element, other_nodes, one_over_denominator)
 
     if coord.n == 1 && occursin("v", coord.name)
         spectral = null_velocity_dimension_info()
@@ -211,7 +242,7 @@ function define_coordinate(input, parallel_io::Bool=false; run_directory=nothing
     elseif input.discretization == "gausslegendre_pseudospectral"
         # create arrays needed for explicit GaussLegendre pseudospectral treatment in this
         # coordinate and create the matrices for differentiation
-        spectral = setup_gausslegendre_pseudospectral(coord,init_YY=init_YY)
+        spectral = setup_gausslegendre_pseudospectral(coord, collision_operator_dim=collision_operator_dim)
         # obtain the local derivatives of the uniform grid with respect to the used grid
         derivative!(coord.duniform_dgrid, coord.uniform_grid, coord, spectral)
     else
@@ -283,6 +314,7 @@ function init_grid(ngrid, nelement_local, n_global, n_local, irank, L, element_s
                    imin, imax, igrid, discretization, name)
     uniform_grid = equally_spaced_grid(n_global, n_local, irank, L)
     uniform_grid_shifted = equally_spaced_grid_shifted(n_global, n_local, irank, L)
+    radau_first_element = false
     if n_global == 1
         grid = allocate_float(n_local)
         grid[1] = 0.0
@@ -299,6 +331,7 @@ function init_grid(ngrid, nelement_local, n_global, n_local, irank, L, element_s
             grid, wgts = scaled_chebyshev_radau_grid(ngrid, nelement_local, n_local, element_scale, element_shift, imin, imax, irank)
             wgts = 2.0 .* wgts .* grid # to include 2 vperp in jacobian of integral
                                        # see note above on normalisation
+            radau_first_element = true
         else
             # initialize chebyshev grid defined on [-L/2,L/2]
             # with n grid points chosen to facilitate
@@ -314,6 +347,7 @@ function init_grid(ngrid, nelement_local, n_global, n_local, irank, L, element_s
             grid, wgts = scaled_gauss_legendre_radau_grid(ngrid, nelement_local, n_local, element_scale, element_shift, imin, imax, irank)
             wgts = 2.0 .* wgts .* grid # to include 2 vperp in jacobian of integral
                                        # see note above on normalisation
+            radau_first_element = true
         else
             grid, wgts = scaled_gauss_legendre_lobatto_grid(ngrid, nelement_local, n_local, element_scale, element_shift, imin, imax)
         end
@@ -336,7 +370,7 @@ function init_grid(ngrid, nelement_local, n_global, n_local, irank, L, element_s
         error("discretization option '$discretization' unrecognized")
     end
     # return the locations of the grid points
-    return grid, wgts, uniform_grid
+    return grid, wgts, uniform_grid, radau_first_element
 end
 
 """