Solver for Poisson's equation in (radial,polar) coordinates.

.github/workflows/test_scripts.yml

@@ -10,7 +10,7 @@ jobs:
       matrix:
         os: [ubuntu-latest, macOS-latest]
       fail-fast: false
-    timeout-minutes: 35
+    timeout-minutes: 40
 
     steps:
       - uses: actions/checkout@v4
@@ -22,4 +22,4 @@ jobs:
         run: |
           touch Project.toml
           julia -O3 --project -e 'import Pkg; Pkg.develop(path="moment_kinetics/"); Pkg.add(["FastGaussQuadrature", "LaTeXStrings", "LegendrePolynomials", "Measures", "MPI", "Plots", "SpecialFunctions"]); Pkg.precompile()'
-          julia -O3 --project -e 'include("test_scripts/2D_FEM_assembly_test.jl"); run_assembly_test(); include("test_scripts/chebyshev_radau_test.jl"); chebyshevradau_test(); include("test_scripts/fkpl_direct_integration_test.jl"); test_rosenbluth_potentials_direct_integration(); include("test_scripts/GaussLobattoLegendre_test.jl"); gausslegendre_test(); include("test_scripts/gyroaverage_test.jl"); gyroaverage_test()'
+          julia -O3 --project -e 'include("test_scripts/2D_FEM_assembly_test.jl"); run_assembly_test(); include("test_scripts/chebyshev_radau_test.jl"); chebyshevradau_test(); include("test_scripts/fkpl_direct_integration_test.jl"); test_rosenbluth_potentials_direct_integration(); include("test_scripts/GaussLobattoLegendre_test.jl"); gausslegendre_test(); include("test_scripts/gyroaverage_test.jl"); gyroaverage_test(); include("test_scripts/PoissonSolve.jl"); run_poisson_test(); run_poisson2D_test()'

moment_kinetics/src/calculus.jl

@@ -147,7 +147,7 @@ function laplacian_derivative!(d2f, f, coord, spectral)
     # First derivative
     elementwise_derivative!(coord, f, spectral)
     derivative_elements_to_full_grid!(coord.scratch3, coord.scratch_2d, coord)
-    if coord.name == "vperp"
+    if coord.cylindrical
         # include the Jacobian
         @. coord.scratch3 *= coord.grid
     end
@@ -159,7 +159,7 @@ function laplacian_derivative!(d2f, f, coord, spectral)
     # Second derivative for element interiors
     elementwise_derivative!(coord, coord.scratch3, spectral)
     derivative_elements_to_full_grid!(d2f, coord.scratch_2d, coord)
-    if coord.name == "vperp"
+    if coord.cylindrical
         # include the Jacobian
         @. d2f /= coord.grid
     end
@@ -249,8 +249,8 @@ function laplacian_derivative!(d2f, f, coord, spectral::weak_discretization_info
     # for all other coord.name, do exactly the same as second_derivative! above.
     mul!(coord.scratch3, spectral.L_matrix, f)
 
-    if coord.bc == "periodic" && coord.name == "vperp"
-        error("laplacian_derivative!() cannot handle periodic boundaries for vperp")
+    if coord.bc == "periodic" && coord.cylindrical
+        error("laplacian_derivative!() cannot handle periodic boundaries for cylindrical coordinates like vperp")
     elseif coord.bc == "periodic"
         if coord.nrank > 1
             error("laplacian_derivative!() cannot handle periodic boundaries for a "

moment_kinetics/src/chebyshev.jl

@@ -291,7 +291,7 @@ function elementwise_derivative!(coord, ff, chebyshev::chebyshev_info)
         imin = coord.imin[j]-k
         # imax is the maximum index on the full grid for this (jth) element
         imax = coord.imax[j]        
-        if coord.name == "vperp" && coord.irank == 0 # differentiate this element with the Radau scheme
+        if coord.cylindrical && coord.irank == 0 # differentiate this element with the Radau scheme
             @views mul!(df[:,j],chebyshev.radau.Dmat[:,:],ff[imin:imax])
         else #differentiate using the Lobatto scheme
             @views mul!(df[:,j],chebyshev.lobatto.Dmat[:,:],ff[imin:imax])
@@ -324,7 +324,7 @@ function elementwise_derivative!(coord, ff, chebyshev::chebyshev_info)
         # at element boundaries (see below for further explanation)
         k = 0
         j = 1 # the first element
-        if coord.name == "vperp" && coord.irank == 0 # differentiate this element with the Radau scheme
+        if coord.cylindrical && coord.irank == 0 # differentiate this element with the Radau scheme
             imin = coord.imin[j]-k
             # imax is the maximum index on the full grid for this (jth) element
             imax = coord.imax[j]

moment_kinetics/src/coordinates.jl

@@ -11,6 +11,7 @@ using ..type_definitions: mk_float, mk_int, OptionsDict
 using ..array_allocation: allocate_float, allocate_shared_float, allocate_int
 using ..calculus: derivative!
 using ..chebyshev: scaled_chebyshev_grid, scaled_chebyshev_radau_grid, setup_chebyshev_pseudospectral
+using ..fourier: scaled_fourier_grid, setup_fourier_pseudospectral
 using ..communication
 using ..finite_differences: finite_difference_info
 using ..gauss_legendre: scaled_gauss_legendre_lobatto_grid, scaled_gauss_legendre_radau_grid, setup_gausslegendre_pseudospectral
@@ -126,6 +127,8 @@ struct coordinate{T <: AbstractVector{mk_float}}
     other_nodes::Array{mk_float,3}
     # One over the denominators of the Lagrange polynomials
     one_over_denominator::Array{mk_float,2}
+    # flag to determine if the coordinate is cylindrical
+    cylindrical::Bool
 end
 
 """
@@ -255,12 +258,15 @@ function define_coordinate(coord_input::NamedTuple; parallel_io::Bool=false,
     # to the full grid
     imin, imax, igrid_full = elemental_to_full_grid_map(coord_input.ngrid,
                                                         coord_input.nelement_local)
+    # check name of coordinate to determine if radial or vperp cylindrical coordinate
+    cylindrical = (coord_input.name == "vperp") || (coord_input.name == "radial")
     # initialise the data used to construct the grid
     # boundaries for each element
     element_boundaries = set_element_boundaries(coord_input.nelement,
                                                 coord_input.L,
                                                 coord_input.element_spacing_option,
-                                                coord_input.name)
+                                                coord_input.name,
+                                                cylindrical)
     # shift and scale factors for each local element
     element_scale, element_shift =
         set_element_scale_and_shift(coord_input.nelement, coord_input.nelement_local,
@@ -270,7 +276,7 @@ function define_coordinate(coord_input::NamedTuple; parallel_io::Bool=false,
     grid, wgts, uniform_grid, radau_first_element =
         init_grid(coord_input.ngrid, coord_input.nelement_local, n_global, n_local,
                   irank, coord_input.L, element_scale, element_shift, imin, imax, igrid,
-                  coord_input.discretization, coord_input.name)
+                  coord_input.discretization, coord_input.name, cylindrical)
     # 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
@@ -358,7 +364,7 @@ function define_coordinate(coord_input::NamedTuple; parallel_io::Bool=false,
         scratch_shared2, scratch_2d, copy(scratch_2d), advection, send_buffer,
         receive_buffer, comm, local_io_range, global_io_range, element_scale,
         element_shift, coord_input.element_spacing_option, element_boundaries,
-        radau_first_element, other_nodes, one_over_denominator)
+        radau_first_element, other_nodes, one_over_denominator, cylindrical)
 
     if coord.n == 1 && occursin("v", coord.name)
         spectral = null_velocity_dimension_info()
@@ -379,6 +385,18 @@ function define_coordinate(coord_input::NamedTuple; parallel_io::Bool=false,
         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)
+    elseif coord_input.discretization == "fourier_pseudospectral"
+        if !(coord.bc == "none")
+         error("fourier_pseudospectral option requires bc='none' (periodicity enforced by basis functions, not be explicit assignment)")
+        end
+        if !(coord.nelement_global == 1)
+         error("fourier_pseudospectral option requires nelement=1")
+        end        
+        # create arrays needed for explicit GaussLegendre pseudospectral treatment in this
+        # coordinate and create the matrices for differentiation
+        spectral = setup_fourier_pseudospectral(coord, run_directory; ignore_MPI=ignore_MPI)
+        # obtain the local derivatives of the uniform grid with respect to the used grid
+        derivative!(coord.duniform_dgrid, coord.uniform_grid, coord, spectral)
     else
         # finite_difference_info is just a type so that derivative methods, etc., dispatch
         # to the finite difference versions, it does not contain any information.
@@ -419,7 +437,7 @@ function define_test_coordinate(name; collision_operator_dim=true, kwargs...)
                                   ignore_MPI=true)
 end
 
-function set_element_boundaries(nelement_global, L, element_spacing_option, coord_name)
+function set_element_boundaries(nelement_global, L, element_spacing_option, coord_name, coord_cylindrical)
     # set global element boundaries between [-L/2,L/2]
     element_boundaries = allocate_float(nelement_global+1)
     if element_spacing_option == "sqrt" && nelement_global > 3
@@ -476,7 +494,7 @@ function set_element_boundaries(nelement_global, L, element_spacing_option, coor
     else 
         println("ERROR: element_spacing_option: ",element_spacing_option, " not supported")
     end
-    if coord_name == "vperp"
+    if coord_cylindrical
         #shift so that the range of element boundaries is [0,L]
         for j in 1:nelement_global+1
             element_boundaries[j] += L/2.0
@@ -502,7 +520,7 @@ end
 setup a grid with n_global grid points on the interval [-L/2,L/2]
 """
 function init_grid(ngrid, nelement_local, n_global, n_local, irank, L, element_scale, element_shift,
-                   imin, imax, igrid, discretization, name)
+                   imin, imax, igrid, discretization, name, cylindrical)
     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
@@ -517,7 +535,7 @@ function init_grid(ngrid, nelement_local, n_global, n_local, irank, L, element_s
             wgts[1] = 1.0
         end
     elseif discretization == "chebyshev_pseudospectral"
-        if name == "vperp"
+        if cylindrical
             # initialize chebyshev grid defined on [-L/2,L/2]
             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
@@ -533,7 +551,7 @@ function init_grid(ngrid, nelement_local, n_global, n_local, irank, L, element_s
             grid, wgts = scaled_chebyshev_grid(ngrid, nelement_local, n_local, element_scale, element_shift, imin, imax)
         end
     elseif discretization == "gausslegendre_pseudospectral"
-        if name == "vperp"
+        if cylindrical
             # use a radau grid for the 1st element near the origin
             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
@@ -542,8 +560,14 @@ function init_grid(ngrid, nelement_local, n_global, n_local, irank, L, element_s
         else
             grid, wgts = scaled_gauss_legendre_lobatto_grid(ngrid, nelement_local, n_local, element_scale, element_shift, imin, imax)
         end
+    elseif discretization == "fourier_pseudospectral"
+        if cylindrical
+            error("fourier_pseudospectral is inappropriate for cylindrical radial coordinates")
+        else
+            grid, wgts = scaled_fourier_grid(ngrid, nelement_local, n_local, element_scale, element_shift, imin, imax)
+        end
     elseif discretization == "finite_difference"
-        if name == "vperp"
+        if cylindrical
             # initialize equally spaced grid defined on [0,L]
             grid = uniform_grid_shifted
             # use composite Simpson's rule to obtain integration weights associated with this coordinate

moment_kinetics/src/fokker_planck_calculus.jl

@@ -41,6 +41,10 @@ using ..communication
 using ..communication: MPISharedArray, global_rank
 using ..lagrange_polynomials: lagrange_poly, lagrange_poly_optimised
 using ..looping
+using ..sparse_matrix_functions: icsc_func, ic_func, allocate_sparse_matrix_constructor,
+                                 assemble_constructor_data!, assign_constructor_data!,
+                                 sparse_matrix_constructor, create_sparse_matrix,
+                                 get_global_compound_index
 using moment_kinetics.gauss_legendre: get_QQ_local!
 using Dates
 using SpecialFunctions: ellipk, ellipe
@@ -881,22 +885,6 @@ function loop_over_vperp_vpa_elements_no_divergences!(G0_weights,G1_weights,H0_w
     return nothing
 end 
 
-
-"""
-    ic_func(ivpa::mk_int,ivperp::mk_int,nvpa::mk_int)
-
-Get the 'linear index' corresponding to `ivpa` and `ivperp`. Defined so that the linear
-index corresponds to the underlying layout in memory of a 2d array indexed by
-`[ivpa,ivperp]`, i.e. for a 2d array `f2d`:
-* `size(f2d) == (vpa.n, vperp.n)`
-* For a reference to `f2d` that is reshaped to a vector (a 1d array) `f1d = vec(f2d)` than
-  for any `ivpa` and `ivperp` it is true that `f1d[ic_func(ivpa,ivperp)] ==
-  f2d[ivpa,ivperp]`.
-"""
-function ic_func(ivpa::mk_int,ivperp::mk_int,nvpa::mk_int)
-    return ivpa + nvpa*(ivperp-1)
-end
-
 """
     ivperp_func(ic::mk_int,nvpa::mk_int)
 
@@ -926,62 +914,6 @@ function ivpa_func(ic::mk_int,nvpa::mk_int)
     return ivpa
 end
 
-# function that returns the sparse matrix index
-# used to directly construct the nonzero entries
-# of a 2D assembled sparse matrix
-function icsc_func(ivpa_local::mk_int,ivpap_local::mk_int,
-                   ielement_vpa::mk_int,
-                   ngrid_vpa::mk_int,nelement_vpa::mk_int,
-                   ivperp_local::mk_int,ivperpp_local::mk_int,
-                   ielement_vperp::mk_int,
-                   ngrid_vperp::mk_int,nelement_vperp::mk_int)
-    ntot_vpa = (nelement_vpa - 1)*(ngrid_vpa^2 - 1) + ngrid_vpa^2
-    #ntot_vperp = (nelement_vperp - 1)*(ngrid_vperp^2 - 1) + ngrid_vperp^2
-
-    icsc_vpa = ((ivpap_local - 1) + (ivpa_local - 1)*ngrid_vpa +
-                (ielement_vpa - 1)*(ngrid_vpa^2 - 1))
-    icsc_vperp = ((ivperpp_local - 1) + (ivperp_local - 1)*ngrid_vperp + 
-                    (ielement_vperp - 1)*(ngrid_vperp^2 - 1))
-    icsc = 1 + icsc_vpa + ntot_vpa*icsc_vperp
-    return icsc
-end
-
-struct sparse_matrix_constructor
-    # the Ith row
-    II::Array{mk_float,1}
-    # the Jth column
-    JJ::Array{mk_float,1}
-    # the data S[I,J]
-    SS::Array{mk_float,1}
-end
-
-function allocate_sparse_matrix_constructor(nsparse::mk_int)
-    II = Array{mk_int,1}(undef,nsparse)
-    @. II = 0
-    JJ = Array{mk_int,1}(undef,nsparse)
-    @. JJ = 0
-    SS = Array{mk_float,1}(undef,nsparse)
-    @. SS = 0.0
-    return sparse_matrix_constructor(II,JJ,SS)
-end
-
-function assign_constructor_data!(data::sparse_matrix_constructor,icsc::mk_int,ii::mk_int,jj::mk_int,ss::mk_float)
-    data.II[icsc] = ii
-    data.JJ[icsc] = jj
-    data.SS[icsc] = ss
-    return nothing
-end
-function assemble_constructor_data!(data::sparse_matrix_constructor,icsc::mk_int,ii::mk_int,jj::mk_int,ss::mk_float)
-    data.II[icsc] = ii
-    data.JJ[icsc] = jj
-    data.SS[icsc] += ss
-    return nothing
-end
-
-function create_sparse_matrix(data::sparse_matrix_constructor)
-    return sparse(data.II,data.JJ,data.SS)
-end
-
 function allocate_boundary_data(vpa,vperp)
     # The following velocity-space-sized buffer arrays are used to evaluate the
     # collision operator for a single species at a single spatial point. They are
@@ -1181,22 +1113,6 @@ function test_boundary_data(func,func_exact,func_name,vpa,vperp,buffer_vpa,buffe
     return max_err
 end
 
-"""
-    get_global_compound_index(vpa,vperp,ielement_vpa,ielement_vperp,ivpa_local,ivperp_local)
-
-For local (within the single element specified by `ielement_vpa` and `ielement_vperp`)
-indices `ivpa_local` and `ivperp_local`, get the global index in the 'linear-indexed' 2d
-space of size `(vperp.n, vpa.n)` (as returned by [`ic_func`](@ref)).
-"""
-function get_global_compound_index(vpa,vperp,ielement_vpa,ielement_vperp,ivpa_local,ivperp_local)
-    # global indices on the grids
-    ivpa_global = vpa.igrid_full[ivpa_local,ielement_vpa]
-    ivperp_global = vperp.igrid_full[ivperp_local,ielement_vperp]
-    # global compound index
-    ic_global = ic_func(ivpa_global,ivperp_global,vpa.n)
-    return ic_global
-end
-
 function enforce_zero_bc!(fvpavperp,vpa,vperp;impose_BC_at_zero_vperp=false)
     # lower vpa boundary
     @loop_vperp ivperp begin