diff --git a/.github/workflows/test_scripts.yml b/.github/workflows/test_scripts.yml index 98eede2b0..5ab223ca8 100644 --- a/.github/workflows/test_scripts.yml +++ b/.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()' diff --git a/moment_kinetics/src/calculus.jl b/moment_kinetics/src/calculus.jl index 47036e077..034cee4b6 100644 --- a/moment_kinetics/src/calculus.jl +++ b/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 " diff --git a/moment_kinetics/src/chebyshev.jl b/moment_kinetics/src/chebyshev.jl index ad8139120..6f28331ac 100644 --- a/moment_kinetics/src/chebyshev.jl +++ b/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] diff --git a/moment_kinetics/src/coordinates.jl b/moment_kinetics/src/coordinates.jl index a49c52d02..b76a50df0 100644 --- a/moment_kinetics/src/coordinates.jl +++ b/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 diff --git a/moment_kinetics/src/fokker_planck_calculus.jl b/moment_kinetics/src/fokker_planck_calculus.jl index f11dcf43c..83dcadb45 100644 --- a/moment_kinetics/src/fokker_planck_calculus.jl +++ b/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 diff --git a/moment_kinetics/src/fourier.jl b/moment_kinetics/src/fourier.jl new file mode 100644 index 000000000..d881dcc69 --- /dev/null +++ b/moment_kinetics/src/fourier.jl @@ -0,0 +1,316 @@ +""" +""" +module fourier + +export setup_fourier_pseudospectral +export scaled_fourier_grid +export fourier_spectral_derivative! +export fourier_info +export fourier_forward_transform! +export fourier_backward_transform! + +using FFTW +using MPI +using ..type_definitions: mk_float, mk_int +using ..array_allocation: allocate_float, allocate_complex +import ..calculus: elementwise_derivative! +using ..communication +using ..moment_kinetics_structs: discretization_info + +""" +Fourier pseudospectral discretization +""" +struct fourier_info{TForward <: FFTW.cFFTWPlan, TBackward <: AbstractFFTs.ScaledPlan} + # fext is an array for storing f(z) on the extended domain needed + # to perform complex-to-complex FFT using the fact that f(theta) is even in theta + fext::Array{Complex{mk_float},1} + # Fourier spectral coefficients of distribution function f + # first dimension contains location within element + # second dimension indicates the element + f::Array{Complex{mk_float},2} + # Fourier spectral coefficients of derivative of f + df::Array{Complex{mk_float},1} + # plan for the complex-to-complex, in-place, forward Fourier transform on uniform grid + forward::TForward + # plan for the complex-to-complex, in-place, backward Fourier transform on uniform grid + # backward_transform::FFTW.cFFTWPlan + backward::TBackward + # midpoint integer, highest integer of postive physical wavenumbers + imidm::mk_int + imidp::mk_int +end + + +""" +create arrays needed for explicit Fourier pseudospectral treatment +and create the plans for the forward and backward fast Fourier transforms +""" +function setup_fourier_pseudospectral(coord, run_directory; ignore_MPI=false) + # First set up the FFTW plans on the (global) root process, then save the 'FFTW + # wisdom' and load it on all other processes, to ensure that we use the exact same + # FFT algorithms on all processes for consistency. + if run_directory === nothing + if global_size[] != 1 && !ignore_MPI + error("run_directory is required by setup_fourier_pseudospectral() when " + * "running in parallel, in order to save FFTW wisdom.") + end + wisdom_filename = nothing + else + wisdom_filename = joinpath(run_directory, "fftw_wisdom_fourier.save") + end + + # When using FFTW.WISDOM_ONLY, the flag should be combined with the flag that was + # originally used to generate the 'wisdom' otherwise if the original flag was 'lower + # effort' (i.e. was FFTW.ESTIMATE) then the default (FFTW.MEASURE) will be used + # instead. + base_flag = FFTW.MEASURE + + function this_barrier() + if !ignore_MPI + # Normal case, all processors are creating the coordinate + MPI.Barrier(comm_world) + elseif run_directory !== nothing && comm_inter_block[] != MPI.COMM_NULL + # ignore_MPI=true was passed, but non-null communicator exists. This happens + # in calls from load_restart_coordinates(), which is only called on + # block_rank[]==0. + MPI.Barrier(comm_inter_block[]) + else + # Should be serial (e.g. used in post-processing), so no Barrier + end + end + + if global_rank[] != 0 + # Wait for rank-0 + this_barrier() + if wisdom_filename !== nothing + # Load wisdom + FFTW.import_wisdom(wisdom_filename) + # Flags can be combined with a bitwise-or operation `|`. + fftw_flags = base_flag | FFTW.WISDOM_ONLY + else + fftw_flags = base_flag + end + else + fftw_flags = base_flag + end + + fourier_spectral = setup_fourier_pseudospectral_grids(coord, fftw_flags) + + if global_rank[] == 0 + if wisdom_filename !== nothing + FFTW.export_wisdom(wisdom_filename) + end + this_barrier() + end + + # Ensure root does not start modifying 'wisdom file' while other processes are still + # reading it - root waits here for all other processes. + this_barrier() + + return fourier_spectral +end + +function setup_fourier_pseudospectral_grids(coord, fftw_flags) + ngrid_fft = coord.ngrid + # create array for f on extended [0,2π] domain in theta = ArcCos[z] + fext = allocate_complex(ngrid_fft) + # create arrays for storing spectral coefficients of f and f' + fhat = allocate_complex(coord.ngrid, coord.nelement_local) + dhat = allocate_complex(coord.ngrid) + # setup the plans for the forward and backward Fourier transforms + forward_transform = plan_fft!(fext, flags=fftw_flags) + backward_transform = plan_ifft!(fext, flags=fftw_flags) + if mod(ngrid_fft,2) == 0 + imidm = mk_int((ngrid_fft/2) - 1 + 1) + imidp = mk_int((ngrid_fft/2) - 1 + 1) + else + imidm = mk_int(((ngrid_fft - 1)/2)) + imidp = mk_int(((ngrid_fft - 1)/2) + 1) + end + # return a structure containing the information needed to carry out + # a 1D Fourier transform + return fourier_info(fext, fhat, dhat, forward_transform, backward_transform, imidm, imidp) +end + + +""" +initialize uniform Fourier grid scaled to interval [-box_length/2, box_length/2] +we no longer pass the box_length to this function, but instead pass precomputed +arrays element_scale and element_shift that are needed to compute the grid. + +ngrid -- number of points per element (including boundary points) +nelement_local -- number of elements in the local (distributed memory MPI) grid +n -- total number of points in the local grid (excluding duplicate points) +element_scale -- the scale factor in the transform from the coordinates + where the element limits are -1, 1 to the coordinate where + the limits are Aj = coord.grid[imin[j]-1] and Bj = coord.grid[imax[j]] + element_scale = 0.5*(Bj - Aj) +element_shift -- the centre of the element in the extended grid coordinate + element_shift = 0.5*(Aj + Bj) +imin -- the array of minimum indices of each element on the extended grid. + By convention, the duplicated points are not included, so for element index j > 1 + the lower boundary point is actually imin[j] - 1 +imax -- the array of maximum indices of each element on the extended grid. +""" +function scaled_fourier_grid(ngrid, nelement_local, n, + element_scale, element_shift, imin, imax) + # this uniform grid goes from -1 to 1(-) + uniform_grid = uniformpoints(ngrid) + # create array for the full grid + grid = allocate_float(n) + + if nelement_local > 1 + error("Fourier grids should have nelement = 1") + end + @inbounds for j ∈ 1:1 + scale_factor = element_scale[j] + shift = element_shift[j] + # apply the scale factor and shift + grid[imin[j]:imax[j]] .= (uniform_grid[1:ngrid] * scale_factor) .+ shift + + end + # get the weights + uniform_wgts = allocate_float(n) + @. uniform_wgts = 2/n + wgts = uniform_wgts*element_scale[1] + + return grid, wgts +end + +""" + elementwise_derivative!(coord, ff, fourier::fourier_info) + +Fourier transform f to get Fourier spectral coefficients and use them to calculate f'. +""" +function elementwise_derivative!(coord, ff, fourier::fourier_info) + df = coord.scratch_2d + # define local variable nelement for convenience + nelement = coord.nelement_local + # check array bounds + @boundscheck nelement == size(df,2) && coord.ngrid == size(df,1) || throw(BoundsError(df)) + # note that one must multiply by a coordinate transform factor 1/element_scale[j] + # for each element j to get derivative on the extended grid + + # note that one must multiply by 1/element_scale[j] get derivative + # in scaled coordinate on element j + + j = 1 # the first and only element + imin = coord.imin[j] + # imax is the maximum index on the full grid for this (jth) element + imax = coord.imax[j] + #println("elementwise_derivative!: ",ff) + @views fourier_derivative_single_element!(df[:,j], ff[imin:imax], + fourier.f[:,j], fourier.df, fourier.fext, fourier.forward, fourier.backward, fourier.imidm, fourier.imidp, coord) + # and multiply by scaling factor needed to go + # from Fourier z coordinate to actual z + # factor 2 because Fourier grid normally runs from [0,1), but here we defined it [-1,1). + for i ∈ 1:coord.ngrid + df[i,j] /= 2.0*coord.element_scale[j] + end + + return nothing +end + +""" + elementwise_derivative!(coord, ff, adv_fac, spectral::fourier_info) + +Fourier transform f to get Fourier spectral coefficients and use them to calculate f'. + +Note: Fourier derivative does not make use of upwinding information. This function +is only provided for compatibility +""" +function elementwise_derivative!(coord, ff, adv_fac, spectral::fourier_info) + return elementwise_derivative!(coord, ff, spectral) +end + +""" +""" +function fourier_derivative_single_element!(df, ff, fhat, dfhat, fext, + forward, backward, imidm, imidp, coord) + # calculate the fourier coefficients of the real-space function ff and return + fourier_forward_transform!(fhat, fext, ff, forward, imidm, imidp, coord.ngrid) + # calculate the fourier coefficients of the derivative of ff with respect to coord.grid + fourier_spectral_derivative!(dfhat, fhat, imidm, imidp) + # inverse fourier transform to get df/dcoord + fourier_backward_transform!(df, fext, dfhat, backward, imidm, imidp, coord.ngrid) +end + +""" +use Fourier basis to compute the first derivative of f +""" +function fourier_spectral_derivative!(df,f,imidm,imidp) + m = length(f) + @boundscheck m == length(df) || throw(BoundsError(df)) + @inbounds begin + for i in 1:imidm + df[i] = 2*pi*im*(i-1)*f[i] + end + for i in imidm+1:m + df[i] = 2*pi*im*((i-1)-m)*f[i] + end + end +end + +""" +returns the uniform grid points [-1,1) on an n point grid +""" +function uniformpoints(n) + grid = allocate_float(n) + nfac = 1/n + @inbounds begin + for j ∈ 1:n + grid[j] = -1.0 + 2.0*(j-1)*nfac + end + end + return grid +end + + +""" +""" +function fourier_forward_transform!(fhat, fext, ff, transform, imidm, imidp, n) + # put ff into fft order + @inbounds begin + # first, fill in values for f into complex array function + for j in 1:imidm + fext[j] = complex(ff[j+imidp],0.0) + end + for j in imidm+1:n + fext[j] = complex(ff[j-imidm],0.0) + end + end + # perform the forward, complex-to-complex FFT in-place (fext is overwritten) + transform*fext + # set out data + fhat .= fext + return nothing +end + +""" +""" +function fourier_backward_transform!(ff, fext, fhat, transform, imidm, imidp, n) + fext .= fhat + # perform the backward, complex-to-complex FFT in-place (fext is overwritten) + transform*fext + @inbounds begin + # fill in entries for ff + # put ff out of fft order + @inbounds begin + # fill in values for f into complex array function + # normalisation appears to be handled by the transform + for j in 1:imidp + ff[j] = real(fext[j+imidm]) + end + for j in imidp+1:n + ff[j] = real(fext[j-imidp]) + end + end + end + return nothing +end + + + + + +end diff --git a/moment_kinetics/src/gauss_legendre.jl b/moment_kinetics/src/gauss_legendre.jl index 7d4b5f2e1..c496faa03 100644 --- a/moment_kinetics/src/gauss_legendre.jl +++ b/moment_kinetics/src/gauss_legendre.jl @@ -902,7 +902,7 @@ function setup_global_weak_form_matrix!(QQ_global::Array{mk_float,2}, if dirichlet_bc # Make matrix diagonal for first/last grid points so it does not change the values # there - if !(coord.name == "vperp") + if !(coord.cylindrical) # modify lower endpoint if not a radial/cylindrical coordinate if coord.irank == 0 QQ_global[1,:] .= 0.0 @@ -978,7 +978,7 @@ function get_MM_local!(QQ,ielement, scale_factor = coord.element_scale[ielement] shift_factor = coord.element_shift[ielement] - if coord.name == "vperp" # assume integrals of form int^infty_0 (.) vperp d vperp + if coord.cylindrical # assume integrals of form int^infty_0 (.) vperp d vperp # extra scale and shift factors required because of vperp in integral if ielement > 1 || coord.irank > 0 # lobatto points @. QQ = (shift_factor*lobatto.M0 + scale_factor*lobatto.M1)*scale_factor @@ -998,7 +998,7 @@ function get_SS_local!(QQ,ielement, scale_factor = coord.element_scale[ielement] shift_factor = coord.element_shift[ielement] - if coord.name == "vperp" # assume integrals of form int^infty_0 (.) vperp d vperp + if coord.cylindrical # assume integrals of form int^infty_0 (.) vperp d vperp # extra scale and shift factors required because of vperp in integral if ielement > 1 || coord.irank > 0 # lobatto points @. QQ = shift_factor*lobatto.S0 + scale_factor*lobatto.S1 @@ -1029,7 +1029,7 @@ function get_KK_local!(QQ,ielement, nelement = coord.nelement_local scale_factor = coord.element_scale[ielement] shift_factor = coord.element_shift[ielement] - if coord.name == "vperp" # assume integrals of form int^infty_0 (.) vperp d vperp + if coord.cylindrical # assume integrals of form int^infty_0 (.) vperp d vperp # extra scale and shift factors required because of vperp in integral # P0 factors make this a d^2 / dvperp^2 rather than (1/vperp) d ( vperp d (.) / d vperp) if ielement > 1 || coord.irank > 0 # lobatto points @@ -1075,7 +1075,7 @@ function get_KJ_local!(QQ,ielement, scale_factor = scale_factor_func(coord.L,coord.nelement_global) shift_factor = shift_factor_func(coord.L,coord.nelement_global,coord.nelement_local,coord.irank,ielement) + 0.5*coord.L - if coord.name == "vperp" # assume integrals of form int^infty_0 (.) vperp d vperp + if coord.cylindrical # assume integrals of form int^infty_0 (.) vperp d vperp # extra scale and shift factors required because of vperp^2 in integral if ielement > 1 || coord.irank > 0 # lobatto points @. QQ = (lobatto.K0*((shift_factor^2)/scale_factor) + @@ -1099,7 +1099,7 @@ function get_LL_local!(QQ,ielement, nelement = coord.nelement_local scale_factor = coord.element_scale[ielement] shift_factor = coord.element_shift[ielement] - if coord.name == "vperp" # assume integrals of form int^infty_0 (.) vperp d vperp + if coord.cylindrical # assume integrals of form int^infty_0 (.) vperp d vperp # extra scale and shift factors required because of vperp in integral # (1/vperp) d ( vperp d (.) / d vperp) if ielement > 1 || coord.irank > 0 # lobatto points @@ -1144,7 +1144,7 @@ function get_MN_local!(QQ,ielement, coord) scale_factor = coord.element_scale[ielement] - if coord.name == "vperp" # assume integrals of form int^infty_0 (.) vperp d vperp + if coord.cylindrical # assume integrals of form int^infty_0 (.) vperp d vperp # extra scale and shift factors required because of vperp in integral if ielement > 1 || coord.irank > 0 # lobatto points @. QQ = lobatto.M0*scale_factor @@ -1166,7 +1166,7 @@ function get_MR_local!(QQ,ielement, scale_factor = coord.element_scale[ielement] shift_factor = coord.element_shift[ielement] - if coord.name == "vperp" # assume integrals of form int^infty_0 (.) vperp d vperp + if coord.cylindrical # assume integrals of form int^infty_0 (.) vperp d vperp # extra scale and shift factors required because of vperp in integral if ielement > 1 || coord.irank > 0 # lobatto points @. QQ = (lobatto.M0*shift_factor^2 + @@ -1192,7 +1192,7 @@ function get_PP_local!(QQ,ielement, scale_factor = coord.element_scale[ielement] shift_factor = coord.element_shift[ielement] - if coord.name == "vperp" # assume integrals of form int^infty_0 (.) vperp d vperp + if coord.cylindrical # assume integrals of form int^infty_0 (.) vperp d vperp # extra scale and shift factors required because of vperp in integral if ielement > 1 || coord.irank > 0 # lobatto points @. QQ = lobatto.P0*shift_factor + lobatto.P1*scale_factor @@ -1215,7 +1215,7 @@ function get_PU_local!(QQ,ielement, scale_factor = coord.element_scale[ielement] shift_factor = coord.element_shift[ielement] - if coord.name == "vperp" # assume integrals of form int^infty_0 (.) vperp d vperp + if coord.cylindrical # assume integrals of form int^infty_0 (.) vperp d vperp # extra scale and shift factors required because of vperp in integral if ielement > 1 || coord.irank > 0 # lobatto points @. QQ = (lobatto.P0*shift_factor^2 + @@ -1261,7 +1261,7 @@ function get_YY0_local!(QQ,ielement, scale_factor = coord.element_scale[ielement] shift_factor = coord.element_shift[ielement] - if coord.name == "vperp" # assume integrals of form int^infty_0 (.) vperp d vperp + if coord.cylindrical # assume integrals of form int^infty_0 (.) vperp d vperp # extra scale and shift factors required because of vperp in integral if ielement > 1 || coord.irank > 0 # lobatto points @. QQ = (shift_factor*lobatto.Y00 + scale_factor*lobatto.Y01)*scale_factor @@ -1281,7 +1281,7 @@ function get_YY1_local!(QQ,ielement, scale_factor = coord.element_scale[ielement] shift_factor = coord.element_shift[ielement] - if coord.name == "vperp" # assume integrals of form int^infty_0 (.) vperp d vperp + if coord.cylindrical # assume integrals of form int^infty_0 (.) vperp d vperp # extra scale and shift factors required because of vperp in integral if ielement > 1 || coord.irank > 0 # lobatto points @. QQ = shift_factor*lobatto.Y10 + scale_factor*lobatto.Y11 @@ -1301,7 +1301,7 @@ function get_YY2_local!(QQ,ielement, scale_factor = coord.element_scale[ielement] shift_factor = coord.element_shift[ielement] - if coord.name == "vperp" # assume integrals of form int^infty_0 (.) vperp d vperp + if coord.cylindrical # assume integrals of form int^infty_0 (.) vperp d vperp # extra scale and shift factors required because of vperp in integral if ielement > 1 || coord.irank > 0 # lobatto points @. QQ = (shift_factor/scale_factor)*lobatto.Y20 + lobatto.Y21 @@ -1321,7 +1321,7 @@ function get_YY3_local!(QQ,ielement, scale_factor = coord.element_scale[ielement] shift_factor = coord.element_shift[ielement] - if coord.name == "vperp" # assume integrals of form int^infty_0 (.) vperp d vperp + if coord.cylindrical # assume integrals of form int^infty_0 (.) vperp d vperp # extra scale and shift factors required because of vperp in integral if ielement > 1 || coord.irank > 0 # lobatto points @. QQ = shift_factor*lobatto.Y30 + scale_factor*lobatto.Y31 diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl index 86810a66f..a2f8a9d32 100644 --- a/moment_kinetics/src/moment_kinetics.jl +++ b/moment_kinetics/src/moment_kinetics.jl @@ -22,8 +22,10 @@ include("array_allocation.jl") include("interpolation.jl") include("calculus.jl") include("lagrange_polynomials.jl") + include("clenshaw_curtis.jl") include("gauss_legendre.jl") +include("fourier.jl") include("chebyshev.jl") include("finite_differences.jl") include("quadrature.jl") @@ -33,6 +35,7 @@ include("input_structs.jl") include("runge_kutta.jl") include("reference_parameters.jl") include("coordinates.jl") +include("sparse_matrix_functions.jl") include("nonlinear_solvers.jl") include("file_io.jl") include("geo.jl") @@ -48,6 +51,7 @@ include("moment_constraints.jl") include("fokker_planck_test.jl") include("fokker_planck_calculus.jl") include("fokker_planck.jl") +include("spatial_poisson.jl") include("boundary_conditions.jl") include("advection.jl") include("vpa_advection.jl") diff --git a/moment_kinetics/src/sparse_matrix_functions.jl b/moment_kinetics/src/sparse_matrix_functions.jl new file mode 100644 index 000000000..6e2f429c5 --- /dev/null +++ b/moment_kinetics/src/sparse_matrix_functions.jl @@ -0,0 +1,124 @@ +""" +module for storing sparse matrix worker functions +""" +module sparse_matrix_functions + +export icsc_func, ic_func +export allocate_sparse_matrix_constructor +export assemble_constructor_data! +export assign_constructor_data! +export sparse_matrix_constructor +export create_sparse_matrix +export get_global_compound_index + + +using ..type_definitions: mk_float, mk_int +using SparseArrays: sparse +using ..coordinates: coordinate + +""" +function that returns the sparse matrix index +used to directly construct the nonzero entries +of a 2D assembled sparse matrix + +Assume that this indexing is for an array A that +in a 2D view is indexed by A[ix,iy] +""" +function icsc_func(ix_local::mk_int,ixp_local::mk_int, + ielement_x::mk_int, + ngrid_x::mk_int,nelement_x::mk_int, + iy_local::mk_int,iyp_local::mk_int, + ielement_y::mk_int, + ngrid_y::mk_int,nelement_y::mk_int) + ntot_x = (nelement_x - 1)*(ngrid_x^2 - 1) + ngrid_x^2 + #ntot_y = (nelement_y - 1)*(ngrid_y^2 - 1) + ngrid_y^2 + + icsc_x = ((ixp_local - 1) + (ix_local - 1)*ngrid_x + + (ielement_x - 1)*(ngrid_x^2 - 1)) + icsc_y = ((iyp_local - 1) + (iy_local - 1)*ngrid_y + + (ielement_y - 1)*(ngrid_y^2 - 1)) + icsc = 1 + icsc_x + ntot_x*icsc_y + return icsc +end + +""" + ic_func(ix::mk_int,iy::mk_int,nx::mk_int) + +Get the 'linear index' corresponding to `ix` and `iy`. Defined so that the linear +index corresponds to the underlying layout in memory of a 2d array indexed by +`[ix,iy]`, i.e. for a 2d array `f2d`: +* `size(f2d) == (x.n, y.n)` +* For a reference to `f2d` that is reshaped to a vector (a 1d array) `f1d = vec(f2d)` than + for any `ix` and `iy` it is true that `f1d[ic_func(ix,iy)] == + f2d[ix,iy]`. +""" +function ic_func(ix::mk_int,iy::mk_int,nx::mk_int) + return ix + nx*(iy-1) +end + +""" + get_global_compound_index(x,y,ielement_x,ielement_y,ix_local,iy_local) + +For local (within the single element specified by `ielement_x` and `ielement_y`) +indices `ix_local` and `iy_local`, get the global index in the 'linear-indexed' 2d +space of size `(x.n, y.n)` (as returned by [`ic_func`](@ref)). +""" +function get_global_compound_index(x::coordinate,y::coordinate, + ielement_x::mk_int,ielement_y::mk_int, + ix_local::mk_int,iy_local::mk_int) + # global indices on the grids + ix_global = x.igrid_full[ix_local,ielement_x] + iy_global = y.igrid_full[iy_local,ielement_y] + # global compound index + ic_global = ic_func(ix_global,iy_global,x.n) + return ic_global +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 + +end diff --git a/moment_kinetics/src/spatial_poisson.jl b/moment_kinetics/src/spatial_poisson.jl new file mode 100644 index 000000000..26054a0aa --- /dev/null +++ b/moment_kinetics/src/spatial_poisson.jl @@ -0,0 +1,585 @@ +""" +module for solving Poisson's equation on a spatial cylindrical domain +Take the approach of solving 1D radial equation in a cylinder, which is +translationally symmetric in the axial direction, using +Fourier series for the polar coordinate. +""" +module spatial_poisson + +export init_spatial_poisson +export spatial_poisson_solve! +export init_spatial_poisson2D +export spatial_poisson2D_solve! + +using ..type_definitions: mk_float, mk_int +using ..gauss_legendre: get_QQ_local! +using ..coordinates: coordinate +using ..array_allocation: allocate_float +using SparseArrays: sparse, AbstractSparseArray +using LinearAlgebra: ldiv!, mul!, LU, lu +using SuiteSparse +using ..fourier: fourier_info, fourier_forward_transform!, fourier_backward_transform! +using ..moment_kinetics_structs: null_spatial_dimension_info +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_structs: weak_discretization_info + +struct poisson_arrays + # an array of lu objects for solving the mth polar harmonic of Poisson's equation + laplacian_lu_objs::Array{SuiteSparse.UMFPACK.UmfpackLU{mk_float,mk_int},1} + # an Array of 2D arrays (one 1D operator on r for + # each Fourier component in the polar coordinate + laplacian::Array{mk_float,3} + # the matrix that needs to multiply the nodal values of the source function + sourcevec::Array{mk_float,2} + rhohat::Array{Complex{mk_float},2} + rhs_dummy::Array{Complex{mk_float},1} + phi_dummy::Array{Complex{mk_float},1} +end + +""" +duplicate function, should be moved to coordinates +""" +function get_imin_imax(coord,iel) + j = iel + if j > 1 + k = 1 + else + k = 0 + end + imin = coord.imin[j] - k + imax = coord.imax[j] + return imin, imax +end + + +function mwavenumber(polar,im) + npolar = polar.n + if mod(npolar,2) == 0 + imid = mk_int((npolar/2)) + else + imid = mk_int(((npolar-1)/2)) + end + mwvc = (2.0*pi/polar.L) + if npolar > 1 + if im < imid+1 + mwvc *= (im-1) + else + mwvc *= ((im-1)-npolar) + end + else + mwvc *= (im-1) + end + return mwvc +end + +""" +function to initialise the arrays needed for the weak-form +Poisson's equation problem + radial = a coordinate struct of a cylindrical=true coordinate + polar = a coordinate struct of a periodic coordinate with length = 2pi +""" +function init_spatial_poisson(radial::coordinate, polar::coordinate, radial_spectral) + npolar = polar.n + nrtot = radial.n + nrgrid = radial.ngrid + nrelement = radial.nelement_global + laplacian = allocate_float(nrtot,nrtot,npolar) + sourcevec = allocate_float(nrtot,nrtot) + rhohat = allocate_float(npolar,nrtot) + rhs_dummy = allocate_float(nrtot) + phi_dummy = allocate_float(nrtot) + MR = allocate_float(nrgrid,nrgrid) + MN = allocate_float(nrgrid,nrgrid) + KJ = allocate_float(nrgrid,nrgrid) + PP = allocate_float(nrgrid,nrgrid) + # initialise to zero + @. laplacian = 0.0 + @. sourcevec = 0.0 + for im in 1:npolar + mwn = mwavenumber(polar,im) + for irel in 1:nrelement + imin, imax = get_imin_imax(radial,irel) + get_QQ_local!(MN,irel,radial_spectral.lobatto,radial_spectral.radau,radial,"N") + get_QQ_local!(KJ,irel,radial_spectral.lobatto,radial_spectral.radau,radial,"J") + get_QQ_local!(PP,irel,radial_spectral.lobatto,radial_spectral.radau,radial,"P") + # assemble the Laplacian matrix + @. laplacian[imin:imax,imin:imax,im] += KJ - PP - ((mwn)^2)*MN + end + # set rows for Dirichlet BCs on phi + laplacian[nrtot,:,im] .= 0.0 + laplacian[nrtot,nrtot,im] = 1.0 + end + for irel in 1:nrelement + imin, imax = get_imin_imax(radial,irel) + get_QQ_local!(MR,irel,radial_spectral.lobatto,radial_spectral.radau,radial,"R") + # assemble the weak-form source vector that should multiply the node values of the source + @. sourcevec[imin:imax,imin:imax] += MR + end + + # construct the sparse Laplacian matrix + laplacian_sparse = Array{AbstractSparseArray{mk_float,mk_int,2},1}(undef,npolar) + for im in 1:npolar + laplacian_sparse[im] = sparse(laplacian[:,:,im]) + end + # construct LU objects + laplacian_lu_objs = Array{SuiteSparse.UMFPACK.UmfpackLU{mk_float,mk_int},1}(undef,npolar) + for im in 1:npolar + laplacian_lu_objs[im] = lu(laplacian_sparse[im]) + end + return poisson_arrays(laplacian_lu_objs,laplacian,sourcevec,rhohat,rhs_dummy,phi_dummy) +end + +""" +Function to find the solution to + +nabla^2 phi = rho in cylindrical polar coordinates +nabla^2 phi = (1/r)d/dr(r dphi/dr) + (1/r^2)d^2 phi/dpolar^2 + +The arguments are + + phi(polar,r) = the function solved for + rho(polar,r) = the source evaluated at the nodal points + poisson_arrays = precomputed arrays + radial = coordinate + polar = coordinate + polar_spectral = fourier_info + +The function uses a 1D Fourier transform to convert the 2D Poisson's +equation into M 1D ODEs, which are solved using 1D elemental weak-form matrices. +The Fourier transform reconstructs the solution. +""" +# for now just support npolar = 1 +# by skipping the FFT +function spatial_poisson_solve!(phi,rho,poisson_arrays,radial,polar,polar_spectral::Union{fourier_info,null_spatial_dimension_info}) + laplacian_lu_objs = poisson_arrays.laplacian_lu_objs + sourcevec = poisson_arrays.sourcevec + phi_dummy = poisson_arrays.phi_dummy + rhs_dummy = poisson_arrays.rhs_dummy + rhohat = poisson_arrays.rhohat + #phihat = poisson_arrays.phihat + + npolar = polar.n + nradial = radial.n + if npolar > 1 + # first FFT rho to hat{rho} appropriate for using the 1D radial operators + for irad in 1:nradial + @views fourier_forward_transform!(rhohat[:,irad], polar_spectral.fext, rho[:,irad], polar_spectral.forward, polar_spectral.imidm, polar_spectral.imidp, polar.ngrid) + end + else + @. rhohat = complex(rho,0.0) + end + + for im in 1:npolar + # solve the linear system + # form the rhs vector + mul!(rhs_dummy,sourcevec,rhohat[im,:]) + # set the Dirichlet BC phi = 0 + rhs_dummy[nradial] = 0.0 + lu_object_lhs = laplacian_lu_objs[im] + ldiv!(phi_dummy, lu_object_lhs, rhs_dummy) + rhohat[im,:] = phi_dummy + end + + if npolar > 1 + # finally iFFT from hat{phi} to phi + for irad in 1:nradial + @views fourier_backward_transform!(phi[:,irad], polar_spectral.fext, rhohat[:,irad], polar_spectral.backward, polar_spectral.imidm, polar_spectral.imidp, polar.ngrid) + end + else + @. phi = real(rhohat) + end + return nothing +end + +# functions associated with a 2D (polar,radial) solver +# using a 2D sparse matrix +struct poisson2D_arrays{M <: AbstractSparseArray{mk_float,mk_int,N} where N} + # an array of lu objects for solving the mth polar harmonic of Poisson's equation + laplacian_lu_obj::SuiteSparse.UMFPACK.UmfpackLU{mk_float,mk_int} + # an Array of 2D arrays (one 1D operator on r for + # each Fourier component in the polar coordinate + laplacian2D::M + # the matrix that needs to multiply the nodal values of the source function + massmatrix2D::M + # dummy array for forming RHS in LU solve + rhspr::Array{mk_float,2} +end + +""" +function to initialise the arrays needed for the 2D weak-form +Poisson's equation problem + radial = a coordinate struct of a cylindrical=true coordinate + polar = a coordinate struct of a periodic polar coordinate with length = 2pi +""" +function init_spatial_poisson2D(radial::coordinate, polar::coordinate, + radial_spectral::weak_discretization_info, + polar_spectral::weak_discretization_info; + use_sparse_init=false) + npolar = polar.n + nradial = radial.n + if use_sparse_init + LP2D_sparse, MR2D_sparse = init_sparse_laplacian2D(radial,polar,radial_spectral,polar_spectral) + else + LP2D_sparse, MR2D_sparse = init_laplacian2D(radial,polar,radial_spectral,polar_spectral) + end + laplacian_lu = lu(LP2D_sparse) + rhspr = allocate_float(npolar,nradial) + return poisson2D_arrays(laplacian_lu,LP2D_sparse,MR2D_sparse,rhspr) +end + +""" +This function makes sparse matrices representing the 2D Laplacian operator. +The compound index runs over the radial and polar coordinates. +We use sparse matrix constructor functions to build the matrix with minimal +memory usage, avoiding global matrices in the compound index. +We assume that the polar coordinate has Jacobian = 1 +so that we do not need to label elemental matrices when +assembling terms from the boundary conditions. + +This function is currently broken due to incorrect sparse indexing functions +which do not handle the periodic boundary conditions correctly. +""" +function init_sparse_laplacian2D(radial::coordinate, polar::coordinate, + radial_spectral::weak_discretization_info, + polar_spectral::weak_discretization_info) + nradial = radial.n + npolar = polar.n + + ngrid_radial = radial.ngrid + ngrid_polar = polar.ngrid + nelement_radial = radial.nelement_local + nelement_polar = polar.nelement_local + + ntot_polar = (polar.nelement_local - 1)*(polar.ngrid^2 - 1) + polar.ngrid^2 + ntot_radial = (radial.nelement_local - 1)*(radial.ngrid^2 - 1) + radial.ngrid^2 + nsparse = ntot_polar*ntot_radial + + MR2D = allocate_sparse_matrix_constructor(nsparse) + # Laplacian matrix + LP2D = allocate_sparse_matrix_constructor(nsparse) + # local dummy arrays + MMpolar = Array{mk_float,2}(undef,ngrid_polar,ngrid_polar) + KKpolar = Array{mk_float,2}(undef,ngrid_polar,ngrid_polar) + MNradial = Array{mk_float,2}(undef,ngrid_radial,ngrid_radial) + MRradial = Array{mk_float,2}(undef,ngrid_radial,ngrid_radial) + KJradial = Array{mk_float,2}(undef,ngrid_radial,ngrid_radial) + PPradial = Array{mk_float,2}(undef,ngrid_radial,ngrid_radial) + # loop over elements to carry out assembly + for ielement_radial in 1:nelement_radial + # get local radial 1D matrices + get_QQ_local!(MNradial,ielement_radial,radial_spectral.lobatto,radial_spectral.radau,radial,"N") + get_QQ_local!(MRradial,ielement_radial,radial_spectral.lobatto,radial_spectral.radau,radial,"R") + get_QQ_local!(KJradial,ielement_radial,radial_spectral.lobatto,radial_spectral.radau,radial,"J") + get_QQ_local!(PPradial,ielement_radial,radial_spectral.lobatto,radial_spectral.radau,radial,"P") + for ielement_polar in 1:nelement_polar + # get local polar 1D matrices + get_QQ_local!(MMpolar,ielement_polar,polar_spectral.lobatto,polar_spectral.radau,polar,"M") + get_QQ_local!(KKpolar,ielement_polar,polar_spectral.lobatto,polar_spectral.radau,polar,"K") + # loop over grid points + for iradialp_local in 1:ngrid_radial + for iradial_local in 1:ngrid_radial + for ipolarp_local in 1:ngrid_polar + for ipolar_local in 1:ngrid_polar + ic_global = get_global_compound_index(polar,radial,ielement_polar,ielement_radial,ipolar_local,iradial_local) + icp_global = get_global_compound_index(polar,radial,ielement_polar,ielement_radial,ipolarp_local,iradialp_local) + icsc = icsc_func(ipolar_local,ipolarp_local,ielement_polar, + ngrid_polar,nelement_polar, + iradial_local,iradialp_local, + ielement_radial, + ngrid_radial,nelement_radial) + # boundary condition possibilities + lower_boundary_row_polar = (ielement_polar == 1 && ipolar_local == 1) + upper_boundary_row_polar = (ielement_polar == polar.nelement_local && ipolar_local == polar.ngrid) + lower_boundary_row_radial = (ielement_radial == 1 && iradial_local == 1) + upper_boundary_row_radial = (ielement_radial == radial.nelement_local && iradial_local == radial.ngrid) + + # Laplacian assembly + if upper_boundary_row_radial + # Dirichlet boundary condition on outer boundary + if iradialp_local == radial.ngrid && ipolar_local == ipolarp_local + assign_constructor_data!(LP2D,icsc,ic_global,icp_global,1.0) + else + assign_constructor_data!(LP2D,icsc,ic_global,icp_global,0.0) + end + elseif upper_boundary_row_polar + # ensure periodicity by setting this row to 0 apart from appropriately placed 1, -1. + if ipolarp_local == ngrid_polar && iradial_local == iradialp_local + # assign -1 to this point using loop icsc, ic, icp + assign_constructor_data!(LP2D,icsc,ic_global,icp_global,-1.0) + # find other end of polar domain for this ipolar_local + # and set ipolarp_local = ielement_polar = 1. + icp_global_bc = get_global_compound_index(polar,radial,1,ielement_radial,1,iradialp_local) + icsc_bc = icsc_func(ipolar_local,1,1, + ngrid_polar,nelement_polar, + iradial_local,iradialp_local, + ielement_radial, + ngrid_radial,nelement_radial) + # assign 1 to this point using ic from loop and updated icsc, icp + assign_constructor_data!(LP2D,icsc_bc,ic_global,icp_global_bc,1.0) + else + assign_constructor_data!(LP2D,icsc,ic_global,icp_global,0.0) + end + else + # carry out assembly + # of Laplacian matrix data + assemble_constructor_data!(LP2D,icsc,ic_global,icp_global, + (MMpolar[ipolar_local,ipolarp_local]* + (KJradial[iradial_local,iradialp_local] - + PPradial[iradial_local,iradialp_local]) + + KKpolar[ipolar_local,ipolarp_local]* + MNradial[iradial_local,iradialp_local])) + if lower_boundary_row_polar + # assemble the contributions from the upper_boundary_row_polar location + # take data from + # ielement_polar = nelement_polar + # ipolar_local = ngrid_polar + # but ensure that data is put into correct target index (ic_global is unchanged from loop) + icp_global_bc = get_global_compound_index(polar,radial,nelement_polar,ielement_radial,ipolarp_local,iradialp_local) + icsc_bc = icsc_func(ipolar_local,ipolarp_local,nelement_polar, + ngrid_polar,nelement_polar, + iradial_local,iradialp_local, + ielement_radial, + ngrid_radial,nelement_radial) + assemble_constructor_data!(LP2D,icsc_bc,ic_global,icp_global_bc, + (MMpolar[ngrid_polar,ipolarp_local]* + (KJradial[iradial_local,iradialp_local] - + PPradial[iradial_local,iradialp_local]) + + KKpolar[ngrid_polar,ipolarp_local]* + MNradial[iradial_local,iradialp_local])) + end + end # Laplacian assembly + + # mass matrices for RHS of matrix equation (enforce periodicity only) + if upper_boundary_row_polar + # ensure periodicity by setting this row to 0 apart from appropriately placed 1, -1. + if ipolarp_local == ngrid_polar && iradial_local == iradialp_local + # assign -1 to this point using loop icsc, ic, icp + assign_constructor_data!(MR2D,icsc,ic_global,icp_global,-1.0) + # find other end of polar domain for this ipolar_local + # and set ipolarp_local = ielement_polar = 1. + icp_global_bc = get_global_compound_index(polar,radial,1,ielement_radial,1,iradialp_local) + icsc_bc = icsc_func(ipolar_local,1,1, + ngrid_polar,nelement_polar, + iradial_local,iradialp_local, + ielement_radial, + ngrid_radial,nelement_radial) + # assign 1 to this point using ic from loop and updated icsc, icp + assign_constructor_data!(MR2D,icsc_bc,ic_global,icp_global_bc,1.0) + else + assign_constructor_data!(MR2D,icsc,ic_global,icp_global,0.0) + end + else + # carry out assembly + # of mass matrix data + assemble_constructor_data!(MR2D,icsc,ic_global,icp_global, + (MMpolar[ipolar_local,ipolarp_local]* + MRradial[iradial_local,iradialp_local])) + if lower_boundary_row_polar + # assemble the contributions from the upper_boundary_row_polar location + # take data from + # ielement_polar = nelement_polar + # ipolar_local = ngrid_polar + # but ensure that data is put into correct target index (ic_global is unchanged from loop) + icp_global_bc = get_global_compound_index(polar,radial,nelement_polar,ielement_radial,ipolarp_local,iradialp_local) + icsc_bc = icsc_func(ipolar_local,ipolarp_local,nelement_polar, + ngrid_polar,nelement_polar, + iradial_local,iradialp_local, + ielement_radial, + ngrid_radial,nelement_radial) + assemble_constructor_data!(MR2D,icsc_bc,ic_global,icp_global_bc, + (MMpolar[ngrid_polar,ipolarp_local]* + MRradial[iradial_local,iradialp_local])) + end # mass matrix assembly + end + end + end + end + end + end + end + LP2D_sparse = create_sparse_matrix(LP2D) + MR2D_sparse = create_sparse_matrix(MR2D) + + return LP2D_sparse, MR2D_sparse +end + +""" +This function makes sparse matrices representing the 2D Laplacian operator. +The compound index runs over the radial and polar coordinates. +We use global matrices in the compound index. +We assume that the polar coordinate has Jacobian = 1 +so that we do not need to label elemental matrices when +assembling terms from the boundary conditions. +""" +function init_laplacian2D(radial::coordinate, polar::coordinate, + radial_spectral::weak_discretization_info, + polar_spectral::weak_discretization_info) + nradial = radial.n + npolar = polar.n + + nc_global = nradial*npolar + ngrid_radial = radial.ngrid + ngrid_polar = polar.ngrid + nelement_radial = radial.nelement_local + nelement_polar = polar.nelement_local + + MR2D = allocate_float(nc_global,nc_global) + @. MR2D = 0.0 + # Laplacian matrix + LP2D = allocate_float(nc_global,nc_global) + @. LP2D = 0.0 + # local dummy arrays + MMpolar = allocate_float(ngrid_polar,ngrid_polar) + KKpolar = allocate_float(ngrid_polar,ngrid_polar) + MNradial = allocate_float(ngrid_radial,ngrid_radial) + MRradial = allocate_float(ngrid_radial,ngrid_radial) + KJradial = allocate_float(ngrid_radial,ngrid_radial) + PPradial = allocate_float(ngrid_radial,ngrid_radial) + # loop over elements to carry out assembly + for ielement_radial in 1:nelement_radial + # get local radial 1D matrices + get_QQ_local!(MNradial,ielement_radial,radial_spectral.lobatto,radial_spectral.radau,radial,"N") + get_QQ_local!(MRradial,ielement_radial,radial_spectral.lobatto,radial_spectral.radau,radial,"R") + get_QQ_local!(KJradial,ielement_radial,radial_spectral.lobatto,radial_spectral.radau,radial,"J") + get_QQ_local!(PPradial,ielement_radial,radial_spectral.lobatto,radial_spectral.radau,radial,"P") + for ielement_polar in 1:nelement_polar + # get local polar 1D matrices + get_QQ_local!(MMpolar,ielement_polar,polar_spectral.lobatto,polar_spectral.radau,polar,"M") + get_QQ_local!(KKpolar,ielement_polar,polar_spectral.lobatto,polar_spectral.radau,polar,"K") + # loop over grid points + for iradialp_local in 1:ngrid_radial + for iradial_local in 1:ngrid_radial + for ipolarp_local in 1:ngrid_polar + for ipolar_local in 1:ngrid_polar + ic_global = get_global_compound_index(polar,radial,ielement_polar,ielement_radial,ipolar_local,iradial_local) + icp_global = get_global_compound_index(polar,radial,ielement_polar,ielement_radial,ipolarp_local,iradialp_local) + # boundary condition possibilities + lower_boundary_row_polar = (ielement_polar == 1 && ipolar_local == 1) + upper_boundary_row_polar = (ielement_polar == polar.nelement_local && ipolar_local == polar.ngrid) + lower_boundary_row_radial = (ielement_radial == 1 && iradial_local == 1) + upper_boundary_row_radial = (ielement_radial == radial.nelement_local && iradial_local == radial.ngrid) + + # Laplacian assembly + if upper_boundary_row_radial + # Dirichlet boundary condition on outer boundary + if iradialp_local == radial.ngrid && ipolar_local == ipolarp_local + LP2D[ic_global,icp_global] = 1.0 + else + LP2D[ic_global,icp_global] = 0.0 + end + elseif upper_boundary_row_polar + # ensure periodicity by setting this row to 0 apart from appropriately placed 1, -1. + if ipolarp_local == ngrid_polar && iradial_local == iradialp_local + # assign -1 to this point using loop icsc, ic, icp + LP2D[ic_global,icp_global] = -1.0 + # find other end of polar domain for this ipolar_local + # and set ipolarp_local = ielement_polar = 1. + icp_global_bc = get_global_compound_index(polar,radial,1,ielement_radial,1,iradialp_local) + # assign 1 to this point using ic from loop and updated icsc, icp + LP2D[ic_global,icp_global_bc] = 1.0 + else + LP2D[ic_global,icp_global] = 0.0 + end + else + # carry out assembly + # of Laplacian matrix data + LP2D[ic_global,icp_global] +=(MMpolar[ipolar_local,ipolarp_local]* + (KJradial[iradial_local,iradialp_local] - + PPradial[iradial_local,iradialp_local]) + + KKpolar[ipolar_local,ipolarp_local]* + MNradial[iradial_local,iradialp_local]) + if lower_boundary_row_polar + # assemble the contributions from the upper_boundary_row_polar location + # take data from + # ielement_polar = nelement_polar + # ipolar_local = ngrid_polar + # but ensure that data is put into correct target index (ic_global is unchanged from loop) + # no elemental index on polar matrices here => assume elemental matrices independent of element index + icp_global_bc = get_global_compound_index(polar,radial,nelement_polar,ielement_radial,ipolarp_local,iradialp_local) + LP2D[ic_global,icp_global_bc] +=(MMpolar[ngrid_polar,ipolarp_local]* + (KJradial[iradial_local,iradialp_local] - + PPradial[iradial_local,iradialp_local]) + + KKpolar[ngrid_polar,ipolarp_local]* + MNradial[iradial_local,iradialp_local]) + end + end # Laplacian assembly + + # mass matrices for RHS of matrix equation (enforce periodicity only) + if upper_boundary_row_polar + # ensure periodicity by setting this row to 0 apart from appropriately placed 1, -1. + if ipolarp_local == ngrid_polar && iradial_local == iradialp_local + # assign -1 to this point using loop icsc, ic, icp + MR2D[ic_global,icp_global] = -1.0 + # find other end of polar domain for this ipolar_local + # and set ipolarp_local = ielement_polar = 1. + icp_global_bc = get_global_compound_index(polar,radial,1,ielement_radial,1,iradialp_local) + # assign 1 to this point using ic from loop and updated icsc, icp + MR2D[ic_global,icp_global_bc] = 1.0 + else + MR2D[ic_global,icp_global] = 0.0 + end + else + # carry out assembly + # of mass matrix data + MR2D[ic_global,icp_global] +=(MMpolar[ipolar_local,ipolarp_local]* + MRradial[iradial_local,iradialp_local]) + if lower_boundary_row_polar + # assemble the contributions from the upper_boundary_row_polar location + # take data from + # ielement_polar = nelement_polar + # ipolar_local = ngrid_polar + # but ensure that data is put into correct target index (ic_global is unchanged from loop) + # no elemental index on polar matrices here => assume elemental matrices independent of element index + icp_global_bc = get_global_compound_index(polar,radial,nelement_polar,ielement_radial,ipolarp_local,iradialp_local) + MR2D[ic_global,icp_global_bc] +=(MMpolar[ngrid_polar,ipolarp_local]* + MRradial[iradial_local,iradialp_local]) + end # mass matrix assembly + end + end + end + end + end + end + end + LP2D_sparse = sparse(LP2D) + MR2D_sparse = sparse(MR2D) + + return LP2D_sparse, MR2D_sparse +end + +""" +function for carrying out a 2D Poisson solve, solving the equation +(1/r) d( r d phi / d r) d r + (1/r^2)d^2 phi/ d polar^2 = rho +with + phi = an Array[polar,radial] representing phi(polar,r) + rho = an Array[polar,radial] representing rho(polar,r) + poisson = a struct of type poisson2D_arrays + +The vec() function is used to get 1D views of the 2D arrays of shape Array{mk_float,2}(npolar,nradial) +phi has assumed zero boundary conditions at r = L, and we assume r spans (0,L]. +""" +function spatial_poisson2D_solve!(phi::Array{mk_float,2},rho::Array{mk_float,2}, + poisson::poisson2D_arrays) + laplacian_lu_obj = poisson.laplacian_lu_obj + massmatrix2D = poisson.massmatrix2D + rhspr = poisson.rhspr + # get data into the compound index format + # by making 1D views of the array with vec() + rhoc = vec(rho) + rhsc = vec(rhspr) + phic = vec(phi) + # form RHS of LP2D . phic = MR2D . rhoc + mul!(rhsc,massmatrix2D,rhoc) + # enforce the boundary conditions using the original 2D reference to the array + # phi(r=L) = 0.0 + @. rhspr[:,end] = 0.0 + # solve system with LU decomposition + ldiv!(phic,laplacian_lu_obj,rhsc) + return nothing +end + + +end \ No newline at end of file diff --git a/moment_kinetics/test/calculus_tests.jl b/moment_kinetics/test/calculus_tests.jl index 58fc00f1f..5ae798d5f 100644 --- a/moment_kinetics/test/calculus_tests.jl +++ b/moment_kinetics/test/calculus_tests.jl @@ -8,6 +8,7 @@ using moment_kinetics.calculus: laplacian_derivative! using MPI using Random +using LinearAlgebra: mul!, ldiv! function runtests() @testset "calculus" verbose=use_verbose begin @@ -1565,6 +1566,297 @@ function runtests() norm=maxabs_norm) end end + + @testset "GaussLegendre pseudospectral cylindrical laplacian ODE solve, zero" verbose=false begin + @testset "$nelement $ngrid" for (nelement, ngrid, rtol) ∈ + ( + (1, 8, 5.e-2), + (1, 9, 3.e-2), + (1, 10, 4.e-3), + (1, 11, 3.e-3), + (1, 12, 6.e-4), + (1, 13, 4.e-4), + (1, 14, 2.e-4), + (1, 15, 5.e-5), + (1, 16, 3.e-5), + (1, 17, 9.e-6), + + (2, 6, 8.e-3), + (2, 7, 4.e-3), + (2, 8, 4.e-4), + (2, 9, 2.e-4), + (2, 10, 5.e-5), + (2, 11, 1.e-5), + (2, 12, 5.e-6), + (2, 13, 4.e-7), + (2, 14, 4.e-7), + (2, 15, 5.e-8), + (2, 16, 2.e-8), + (2, 17, 5.e-9), + + (3, 6, 2.e-3), + (3, 7, 2.e-4), + (3, 8, 5.e-5), + (3, 9, 3.e-6), + (3, 10, 2.e-6), + (3, 11, 3.e-7), + (3, 12, 6.e-8), + (3, 13, 9.e-9), + (3, 14, 2.e-9), + (3, 15, 4.e-10), + (3, 16, 3.e-11), + (3, 17, 1.e-11), + + (4, 5, 1.e-3), + (4, 6, 8.e-5), + (4, 7, 3.e-5), + (4, 8, 3.e-6), + (4, 9, 6.e-7), + (4, 10, 8.e-8), + (4, 11, 2.e-8), + (4, 12, 3.e-9), + (4, 13, 3.e-10), + (4, 14, 5.e-11), + (4, 15, 3.e-12), + (4, 16, 2.e-12), + (4, 17, 1.e-12), + + (5, 5, 4.e-4), + (5, 6, 3.e-5), + (5, 7, 6.e-6), + (5, 8, 4.e-7), + (5, 9, 9.e-8), + (5, 10, 4.e-9), + (5, 11, 2.e-9), + (5, 12, 8.e-11), + (5, 13, 3.e-11), + (5, 14, 1.e-12), + (5, 15, 4.e-13), + (5, 16, 4.e-13), + (5, 17, 2.e-12), + ) + + # define inputs needed for the test + L = 6.0 + bc = "zero" + element_spacing_option = "uniform" + # create the coordinate struct 'x' + # This test runs effectively in serial, so implicitly uses + # `ignore_MPI=true` to avoid errors due to communicators not being fully + # set up. + x, spectral = define_test_coordinate("vperp"; ngrid=ngrid, + nelement=nelement, L=L, + discretization="gausslegendre_pseudospectral", + bc=bc, + element_spacing_option=element_spacing_option, + collision_operator_dim=false) + expected_f = @. exp(-x.grid^2) + d2f = @. 4.0*(x.grid^2 - 1.0)*exp(-x.grid^2) + # create array for the numerical solution + f = similar(expected_f) + # create array for RHS vector b + b = similar(expected_f) + # solve for f + mul!(b,spectral.mass_matrix,d2f) + # Dirichlet zero BC at upper endpoint + b[end] = 0.0 + # solve ODE + ldiv!(f,spectral.L_matrix_lu,b) + #err = maximum(abs.(f.-expected_f)) + #println("$nelement $ngrid $err") + @test isapprox(f, expected_f, rtol=rtol, atol=1.e-10, + norm=maxabs_norm) + end + end + + @testset "GaussLegendre pseudospectral second derivative ODE solve, zero" verbose=false begin + @testset "$nelement $ngrid" for (nelement, ngrid, rtol) ∈ + ( + (1, 23, 1.e-2), + (1, 24, 5.e-3), + (1, 25, 3.e-3), + (1, 26, 2.e-3), + (1, 27, 1.e-3), + (1, 28, 6.e-4), + (1, 29, 5.e-4), + (1, 30, 3.e-4), + (1, 31, 2.e-4), + (1, 32, 8.e-5), + + (2, 9, 3.e-2), + (2, 10, 2.e-2), + (2, 11, 4.e-3), + (2, 12, 1.e-3), + (2, 13, 4.e-4), + (2, 14, 2.e-4), + (2, 15, 9.e-5), + (2, 16, 3.e-5), + (2, 17, 7.e-6), + + (3, 9, 2.e-2), + (3, 10, 2.e-3), + (3, 11, 6.e-4), + (3, 12, 2.e-4), + (3, 13, 7.e-5), + (3, 14, 1.e-5), + (3, 15, 9.e-6), + (3, 16, 9.e-7), + (3, 17, 1.e-6), + + (4, 7, 3.e-3), + (4, 8, 8.e-4), + (4, 9, 2.e-4), + (4, 10, 4.e-5), + (4, 11, 2.e-5), + (4, 12, 4.e-6), + (4, 13, 9.e-7), + (4, 14, 4.e-7), + (4, 15, 3.e-8), + (4, 16, 3.e-8), + (4, 17, 4.e-9), + + (5, 8, 2.e-4), + (5, 9, 7.e-5), + (5, 10, 7.e-6), + (5, 11, 4.e-6), + (5, 12, 3.e-7), + (5, 13, 3.e-7), + (5, 14, 9.e-9), + (5, 15, 1.e-8), + (5, 16, 3.e-10), + (5, 17, 4.e-10), + ) + + # define inputs needed for the test + L = 12.0 + bc = "zero" + element_spacing_option = "uniform" + # create the coordinate struct 'x' + # This test runs effectively in serial, so implicitly uses + # `ignore_MPI=true` to avoid errors due to communicators not being fully + # set up. + x, spectral = define_test_coordinate("vpa"; ngrid=ngrid, + nelement=nelement, L=L, + discretization="gausslegendre_pseudospectral", + bc=bc, + element_spacing_option=element_spacing_option, + collision_operator_dim=false) + expected_f = @. exp(-x.grid^2) + d2f = @. (4.0*x.grid^2 - 2.0)*exp(-x.grid^2) + # create array for the numerical solution + f = similar(expected_f) + # create array for RHS vector b + b = similar(expected_f) + # solve for f + mul!(b,spectral.mass_matrix,d2f) + # Dirichlet zero BC at lower and upper endpoint + b[1] = 0.0 + b[end] = 0.0 + # solve ODE + ldiv!(f,spectral.L_matrix_lu,b) + #err = maximum(abs.(f.-expected_f)) + #maxfe = maximum(expected_f) + #maxf = maximum(f) + #minf = minimum(f) + #println("$nelement $ngrid $err $maxfe $maxf $minf") + @test isapprox(f, expected_f, rtol=rtol, atol=1.e-10, + norm=maxabs_norm) + end + end + + @testset "GaussLegendre pseudospectral second derivative ODE solve, periodic" verbose=false begin + @testset "$nelement $ngrid" for (nelement, ngrid, rtol) ∈ + ( + (1, 10, 6.e-6), + (1, 11, 2.e-6), + (1, 12, 8.e-8), + (1, 13, 2.e-8), + (1, 14, 1.e-9), + (1, 15, 2.e-10), + (1, 16, 9.e-12), + (1, 17, 2.e-12), + + (2, 9, 5.e-8), + (2, 10, 6.e-9), + (2, 11, 3.e-10), + (2, 12, 3.e-11), + (2, 13, 7.e-13), + (2, 14, 1.e-13), + (2, 15, 1.e-13), + (2, 16, 1.e-13), + (2, 17, 1.e-13), + + (3, 9, 2.e-9), + (3, 10, 7.e-11), + (3, 11, 4.e-12), + (3, 12, 3.e-13), + (3, 13, 3.e-13), + (3, 14, 1.e-13), + (3, 15, 1.e-13), + (3, 16, 3.e-13), + (3, 17, 1.e-13), + + (4, 7, 5.e-8), + (4, 8, 3.e-9), + (4, 9, 8.e-11), + (4, 10, 4.e-12), + (4, 11, 3.e-13), + (4, 12, 1.e-13), + (4, 13, 3.e-13), + (4, 14, 1.e-13), + (4, 15, 1.e-13), + (4, 16, 3.e-13), + (4, 17, 3.e-13), + + (5, 8, 5.e-10), + (5, 9, 9.e-12), + (5, 10, 5.e-13), + (5, 11, 3.e-13), + (5, 12, 1.e-13), + (5, 13, 5.e-13), + (5, 14, 1.e-13), + (5, 15, 2.e-13), + (5, 16, 5.e-13), + (5, 17, 5.e-13), + ) + + # define inputs needed for the test + L = 1.0 + bc = "periodic" + element_spacing_option = "uniform" + phase = pi/3.0 + # create the coordinate struct 'x' + # This test runs effectively in serial, so implicitly uses + # `ignore_MPI=true` to avoid errors due to communicators not being fully + # set up. + x, spectral = define_test_coordinate("vpa"; ngrid=ngrid, + nelement=nelement, L=L, + discretization="gausslegendre_pseudospectral", + bc=bc, + element_spacing_option=element_spacing_option, + collision_operator_dim=false) + expected_f = @. sin((2.0*pi*x.grid/x.L) + phase) + d2f = @. -((2.0*pi/x.L)^2)*sin((2.0*pi*x.grid/x.L)+phase) + # create array for the numerical solution + f = similar(expected_f) + # create array for RHS vector b + b = similar(expected_f) + # solve for f + mul!(b,spectral.mass_matrix,d2f) + # Dirichlet zero BC at lower endpoint, periodic bc at upper endpoint + b[1] = sin((2.0*pi*x.grid[1]/x.L) + phase) # fixes constant piece of solution + b[end] = 0.0 # makes sure periodicity is enforced + # solve ODE + ldiv!(f,spectral.L_matrix_lu,b) + #err = maximum(abs.(f.-expected_f)) + #maxfe = maximum(expected_f) + #maxf = maximum(f) + #minf = minimum(f) + #println("$nelement $ngrid $err $maxfe $maxf $minf") + @test isapprox(f, expected_f, rtol=rtol, atol=1.e-10, + norm=maxabs_norm) + end + end end end diff --git a/moment_kinetics/test/fokker_planck_tests.jl b/moment_kinetics/test/fokker_planck_tests.jl index 8f2dafbd7..267a2f32b 100644 --- a/moment_kinetics/test/fokker_planck_tests.jl +++ b/moment_kinetics/test/fokker_planck_tests.jl @@ -9,7 +9,7 @@ using moment_kinetics.communication using moment_kinetics.looping using moment_kinetics.array_allocation: allocate_float, allocate_shared_float using moment_kinetics.coordinates: define_coordinate -using moment_kinetics.type_definitions: mk_float, mk_int +using moment_kinetics.type_definitions: mk_float, mk_int, OptionsDict using moment_kinetics.velocity_moments: get_density, get_upar, get_ppar, get_pperp, get_pressure using moment_kinetics.fokker_planck: init_fokker_planck_collisions_weak_form, fokker_planck_collision_operator_weak_form! diff --git a/moment_kinetics/test/nonlinear_solver_tests.jl b/moment_kinetics/test/nonlinear_solver_tests.jl index 806f399e3..cd3e0769d 100644 --- a/moment_kinetics/test/nonlinear_solver_tests.jl +++ b/moment_kinetics/test/nonlinear_solver_tests.jl @@ -64,7 +64,7 @@ function linear_test() zeros(mk_float, 0), MPI.COMM_NULL, 1:n, 1:n, zeros(mk_float, 0), zeros(mk_float, 0), "", zeros(mk_float, 0), false, zeros(mk_float, 0, 0, 0), - zeros(mk_float, 0, 0)) + zeros(mk_float, 0, 0), false) coords = NamedTuple(c => the_coord for c ∈ coord_names) function rhs_func!(residual, x) @@ -177,7 +177,7 @@ function nonlinear_test() zeros(mk_float, 0), MPI.COMM_NULL, 1:n, 1:n, zeros(mk_float, 0), zeros(mk_float, 0), "", zeros(mk_float, 0), false, zeros(mk_float, 0, 0, 0), - zeros(mk_float, 0, 0)) + zeros(mk_float, 0, 0), false) coords = NamedTuple(c => the_coord for c ∈ coord_names) function rhs_func!(residual, x) diff --git a/moment_kinetics/test/poisson_radial_polar_tests.jl b/moment_kinetics/test/poisson_radial_polar_tests.jl new file mode 100644 index 000000000..c1012cad0 --- /dev/null +++ b/moment_kinetics/test/poisson_radial_polar_tests.jl @@ -0,0 +1,314 @@ +module poisson_radial_polar_tests +# currently only runs on a single core + +include("setup.jl") + +export run_poisson_radial_polar_test +export run_test +using MPI +import moment_kinetics +using moment_kinetics.array_allocation: allocate_float, allocate_shared_float +using moment_kinetics.coordinates: define_coordinate +using moment_kinetics.type_definitions: mk_float, mk_int, OptionsDict +using moment_kinetics.spatial_poisson: init_spatial_poisson, spatial_poisson_solve! +using moment_kinetics.spatial_poisson: init_spatial_poisson2D, spatial_poisson2D_solve! +using moment_kinetics.communication +using moment_kinetics.looping +using moment_kinetics.calculus: derivative! + +function run_test() + @testset "poisson_radial_polar_fourier_tests" begin + println("poisson_radial_polar_fourier_tests") + nelement_radial=5 + ngrid_radial=5 + Lradial=1.0 + ngrid_polar=1 + kk=1 + @testset "nelement_radial $nelement_radial ngrid_radial $ngrid_radial Lradial $Lradial ngrid_polar $ngrid_polar kk $kk" begin + run_poisson_radial_polar_fourier_test(nelement_radial=nelement_radial,ngrid_radial=ngrid_radial,Lradial=Lradial,ngrid_polar=ngrid_polar,kk=kk, + atol_phi1=4.0e-14) + end + nelement_radial=5 + ngrid_radial=5 + Lradial=1.0 + ngrid_polar=8 + kk=3 + @testset "nelement_radial $nelement_radial ngrid_radial $ngrid_radial Lradial $Lradial ngrid_polar $ngrid_polar kk $kk" begin + run_poisson_radial_polar_fourier_test(nelement_radial=nelement_radial,ngrid_radial=ngrid_radial,Lradial=Lradial,ngrid_polar=ngrid_polar,kk=kk, + atol_fourier=1.0e-14,atol_phi1=4.0e-14,atol_phi2=1.0e-14) + end + nelement_radial=5 + ngrid_radial=9 + Lradial=1.0 + ngrid_polar=8 + kk=3 + @testset "nelement_radial $nelement_radial ngrid_radial $ngrid_radial Lradial $Lradial ngrid_polar $ngrid_polar kk $kk" begin + run_poisson_radial_polar_fourier_test(nelement_radial=nelement_radial,ngrid_radial=ngrid_radial,Lradial=Lradial,ngrid_polar=ngrid_polar,kk=kk, + atol_fourier=1.0e-14,atol_phi1=3.0e-12,atol_phi2=1.0e-14) + end + nelement_radial=10 + ngrid_radial=5 + Lradial=1.0 + ngrid_polar=8 + kk=3 + @testset "nelement_radial $nelement_radial ngrid_radial $ngrid_radial Lradial $Lradial ngrid_polar $ngrid_polar kk $kk" begin + run_poisson_radial_polar_fourier_test(nelement_radial=nelement_radial,ngrid_radial=ngrid_radial,Lradial=Lradial,ngrid_polar=ngrid_polar,kk=kk, + atol_fourier=1.0e-14,atol_phi1=9.0e-13,atol_phi2=1.0e-14) + end + nelement_radial=40 + ngrid_radial=5 + Lradial=1.0 + ngrid_polar=16 + kk=5 + @testset "nelement_radial $nelement_radial ngrid_radial $ngrid_radial Lradial $Lradial ngrid_polar $ngrid_polar kk $kk" begin + run_poisson_radial_polar_fourier_test(nelement_radial=nelement_radial,ngrid_radial=ngrid_radial,Lradial=Lradial,ngrid_polar=ngrid_polar,kk=kk, + atol_fourier=1.0e-14,atol_phi1=9.0e-13,atol_phi2=2.0e-11) + end + end + + @testset "poisson_radial_polar_tests" begin + println("poisson_radial_polar_tests") + nelement_radial=5 + ngrid_radial=5 + Lradial=1.0 + nelement_polar=5 + ngrid_polar=5 + kk=1 + phase=pi/3.0 + @testset "nelement_radial $nelement_radial ngrid_radial $ngrid_radial Lradial $Lradial nelement_polar $nelement_polar ngrid_polar $ngrid_polar kk $kk phase $phase" begin + run_poisson_radial_polar_test(nelement_radial=nelement_radial,ngrid_radial=ngrid_radial,Lradial=Lradial,nelement_polar=nelement_polar,ngrid_polar=ngrid_polar,kk=kk,phase=phase, + atol_phi1=7.0e-14,atol_phi2=9.0e-7) + end + nelement_radial=5 + ngrid_radial=5 + Lradial=1.0 + nelement_polar=10 + ngrid_polar=9 + kk=3 + phase=pi/3.0 + @testset "nelement_radial $nelement_radial ngrid_radial $ngrid_radial Lradial $Lradial nelement_polar $nelement_polar ngrid_polar $ngrid_polar kk $kk phase $phase" begin + run_poisson_radial_polar_test(nelement_radial=nelement_radial,ngrid_radial=ngrid_radial,Lradial=Lradial,nelement_polar=nelement_polar,ngrid_polar=ngrid_polar,kk=kk,phase=phase, + atol_phi1=3.0e-13,atol_phi2=5.0e-11) + end + nelement_radial=5 + ngrid_radial=9 + Lradial=1.0 + nelement_polar=10 + ngrid_polar=9 + kk=3 + phase=pi/3.0 + @testset "nelement_radial $nelement_radial ngrid_radial $ngrid_radial Lradial $Lradial nelement_polar $nelement_polar ngrid_polar $ngrid_polar kk $kk phase $phase" begin + run_poisson_radial_polar_test(nelement_radial=nelement_radial,ngrid_radial=ngrid_radial,Lradial=Lradial,nelement_polar=nelement_polar,ngrid_polar=ngrid_polar,kk=kk,phase=phase, + atol_phi1=3.0e-12,atol_phi2=5.0e-11) + end + nelement_radial=10 + ngrid_radial=5 + Lradial=1.0 + nelement_polar=10 + ngrid_polar=9 + kk=3 + phase=pi/3.0 + @testset "nelement_radial $nelement_radial ngrid_radial $ngrid_radial Lradial $Lradial nelement_polar $nelement_polar ngrid_polar $ngrid_polar kk $kk phase $phase" begin + run_poisson_radial_polar_test(nelement_radial=nelement_radial,ngrid_radial=ngrid_radial,Lradial=Lradial,nelement_polar=nelement_polar,ngrid_polar=ngrid_polar,kk=kk,phase=phase, + atol_phi1=9.0e-13,atol_phi2=5.0e-11) + end + nelement_radial=40 + ngrid_radial=5 + Lradial=1.0 + nelement_polar=20 + ngrid_polar=9 + kk=5 + phase=pi/3.0 + @testset "nelement_radial $nelement_radial ngrid_radial $ngrid_radial Lradial $Lradial nelement_polar $nelement_polar ngrid_polar $ngrid_polar kk $kk phase $phase" begin + run_poisson_radial_polar_test(nelement_radial=nelement_radial,ngrid_radial=ngrid_radial,Lradial=Lradial,nelement_polar=nelement_polar,ngrid_polar=ngrid_polar,kk=kk,phase=phase, + atol_phi1=9.0e-13,atol_phi2=2.0e-11) + end + end +end + +function run_poisson_radial_polar_fourier_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelement_polar=1,ngrid_polar=1,Lpolar=2.0*pi,kk::mk_int=1, + atol_fourier=1.0e-14,atol_phi1=1.0e-14,atol_phi2=1.0e-14,print_to_screen=false) + + nelement_local_polar = nelement_polar # number of elements per rank + nelement_global_polar = nelement_local_polar # total number of elements + nelement_local_radial = nelement_radial # number of elements per rank + nelement_global_radial = nelement_local_radial # total number of elements + bc = "none" #not required to take a particular value, not used, set to "none" to avoid extra BC impositions + GL_discretization = "gausslegendre_pseudospectral" + F_discretization = "fourier_pseudospectral" + element_spacing_option = "uniform" + coords_input = OptionsDict( + "radial"=>OptionsDict("ngrid"=>ngrid_radial, "nelement"=>nelement_global_radial, + "nelement_local"=>nelement_local_radial, "L"=>Lradial, + "discretization"=>GL_discretization, "bc"=>bc, + "element_spacing_option"=>element_spacing_option), + "polar"=>OptionsDict("ngrid"=>ngrid_polar, "nelement"=>nelement_global_polar, + "nelement_local"=>nelement_local_polar, "L"=>Lpolar, + "discretization"=>F_discretization, "bc"=>bc, + "element_spacing_option"=>element_spacing_option), + ) + # Set up MPI + initialize_comms!() + setup_distributed_memory_MPI(1,1,1,1) + # ignore MPI here to avoid FFTW wisdom problems, test runs on a single core below + polar, polar_spectral = define_coordinate(coords_input, "polar", ignore_MPI=true) + radial, radial_spectral = define_coordinate(coords_input, "radial", ignore_MPI=true) + looping.setup_loop_ranges!(block_rank[], block_size[]; + s=1, sn=1, + r=radial.n, z=polar.n, vperp=1, vpa=1, + vzeta=1, vr=1, vz=1) + + begin_serial_region() + @serial_region begin # run tests purely in serial for now + if ngrid_polar > 1 + if print_to_screen println("Test fourier_pseudospectral") end + ff = allocate_float(polar.n) + df = allocate_float(polar.n) + df_exact = allocate_float(polar.n) + df_err = allocate_float(polar.n) + for i in 1:polar.n + ff[i] = sin(2.0*pi*polar.grid[i]/polar.L) + df_exact[i] = cos(2.0*pi*polar.grid[i]/polar.L)*(2.0*pi/polar.L) + end + derivative!(df,ff,polar,polar_spectral) + @. df_err = abs(df - df_exact) + max_df_err = maximum(df_err) + if print_to_screen println("maximum(df_err): ",max_df_err) end + @test max_df_err < atol_fourier + wgts_err = abs( 1.0 - (sum(polar.wgts)/polar.L) ) + if print_to_screen println("wgts err: ",wgts_err) end + @test wgts_err < atol_fourier + end + + if print_to_screen println("Test Poisson") end + poisson = init_spatial_poisson(radial,polar,radial_spectral) + phi = allocate_float(polar.n,radial.n) + exact_phi = allocate_float(polar.n,radial.n) + err_phi = allocate_float(polar.n,radial.n) + rho = allocate_float(polar.n,radial.n) + + for irad in 1:radial.n + for ipol in 1:polar.n + exact_phi[ipol,irad] = 0.25*(radial.grid[irad]^2 -1) + rho[ipol,irad] = 1.0 + end + end + spatial_poisson_solve!(phi,rho,poisson,radial,polar,polar_spectral) + max_phi_err = maximum(abs.(phi- exact_phi)) + if print_to_screen println("Maximum error value Test rho=1 : ",max_phi_err) end + @test max_phi_err < atol_phi1 + + if ngrid_polar > 1 + if kk < 1 + error("ERROR: kk >=1 required for test") + end + if !(mod(kk,1) == 0) + error("ERROR: kk integer required for test") + end + if abs(polar.L - 2.0*pi) >1.0e-14 + error("ERROR: polar coordinate assumed angle for definition of the following test - set Lpolar = 2.0*pi") + end + + for irad in 1:radial.n + for ipol in 1:polar.n + exact_phi[ipol,irad] = (1.0 - radial.grid[irad])*(radial.grid[irad]^kk)*cos(2.0*pi*kk*polar.grid[ipol]/polar.L) + rho[ipol,irad] = (kk^2 - (kk+1)^2)*(radial.grid[irad]^(kk-1))*cos(2.0*kk*pi*polar.grid[ipol]/polar.L) + end + end + + spatial_poisson_solve!(phi,rho,poisson,radial,polar,polar_spectral) + @. err_phi = abs(phi - exact_phi) + max_phi_err = maximum(err_phi) + if print_to_screen println("Maximum error value Test rho = (kk^2 - (kk+1)^2) * cos(2 pi kk P/L) * r^(kk-1): ",max_phi_err) end + @test max_phi_err < atol_phi2 + end + end + finalize_comms!() + return nothing +end + +function run_poisson_radial_polar_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelement_polar=5,ngrid_polar=5,Lpolar=2.0*pi,kk::mk_int=1,phase::mk_float=0.0, + atol_phi1=1.0e-14,atol_phi2=1.0e-14,print_to_screen=false) + + nelement_local_polar = nelement_polar # number of elements per rank + nelement_global_polar = nelement_local_polar # total number of elements + nelement_local_radial = nelement_radial # number of elements per rank + nelement_global_radial = nelement_local_radial # total number of elements + bc = "none" #not required to take a particular value, finite element assembly takes care of BCs, so not used, set to "none" to avoid extra BC impositions + discretization = "gausslegendre_pseudospectral" + element_spacing_option = "uniform" + coords_input = OptionsDict( + "radial"=>OptionsDict("ngrid"=>ngrid_radial, "nelement"=>nelement_global_radial, + "nelement_local"=>nelement_local_radial, "L"=>Lradial, + "discretization"=>discretization, "bc"=>bc, + "element_spacing_option"=>element_spacing_option), + "polar"=>OptionsDict("ngrid"=>ngrid_polar, "nelement"=>nelement_global_polar, + "nelement_local"=>nelement_local_polar, "L"=>Lpolar, + "discretization"=>discretization, "bc"=>bc, + "element_spacing_option"=>element_spacing_option), + ) + # Set up MPI + initialize_comms!() + setup_distributed_memory_MPI(1,1,1,1) + # ignore MPI here to avoid FFTW wisdom problems, test runs on a single core below + polar, polar_spectral = define_coordinate(coords_input, "polar", ignore_MPI=true) + radial, radial_spectral = define_coordinate(coords_input, "radial", ignore_MPI=true) + looping.setup_loop_ranges!(block_rank[], block_size[]; + s=1, sn=1, + r=radial.n, z=polar.n, vperp=1, vpa=1, + vzeta=1, vr=1, vz=1) + + begin_serial_region() + @serial_region begin # run tests purely in serial for now + if print_to_screen println("Test Poisson2D") end + poisson = init_spatial_poisson2D(radial,polar,radial_spectral,polar_spectral) + phi = allocate_float(polar.n,radial.n) + exact_phi = allocate_float(polar.n,radial.n) + err_phi = allocate_float(polar.n,radial.n) + rho = allocate_float(polar.n,radial.n) + + for irad in 1:radial.n + for ipol in 1:polar.n + exact_phi[ipol,irad] = 0.25*(radial.grid[irad]^2 -1) + rho[ipol,irad] = 1.0 + end + end + spatial_poisson2D_solve!(phi,rho,poisson) + max_phi_err = maximum(abs.(phi- exact_phi)) + if print_to_screen println("Maximum error value Test rho=1 : ",max_phi_err) end + @test max_phi_err < atol_phi1 + + if ngrid_polar > 1 + if kk < 1 + error("ERROR: kk >=1 required for test") + end + if !(mod(kk,1) == 0) + error("ERROR: kk integer required for test") + end + if abs(polar.L - 2.0*pi) >1.0e-14 + error("ERROR: polar coordinate assumed angle for definition of the following test - set Lpolar = 2.0*pi") + end + + for irad in 1:radial.n + for ipol in 1:polar.n + exact_phi[ipol,irad] = (1.0 - radial.grid[irad])*(radial.grid[irad]^kk)*cos(2.0*pi*kk*polar.grid[ipol]/polar.L + phase) + rho[ipol,irad] = (kk^2 - (kk+1)^2)*(radial.grid[irad]^(kk-1))*cos(2.0*kk*pi*polar.grid[ipol]/polar.L + phase) + end + end + + spatial_poisson2D_solve!(phi,rho,poisson) + @. err_phi = abs(phi - exact_phi) + max_phi_err = maximum(err_phi) + if print_to_screen println("Maximum error value Test rho = (kk^2 - (kk+1)^2) * cos(2 pi kk P/L + phase) * r^(kk-1): ",max_phi_err) end + @test max_phi_err < atol_phi2 + end + end + finalize_comms!() + return nothing +end + +end #poisson_radial_polar_tests + +using .poisson_radial_polar_tests + +poisson_radial_polar_tests.run_test() diff --git a/moment_kinetics/test/runtests.jl b/moment_kinetics/test/runtests.jl index 73d688f06..cddd0ccea 100644 --- a/moment_kinetics/test/runtests.jl +++ b/moment_kinetics/test/runtests.jl @@ -20,6 +20,7 @@ function runtests() include(joinpath(@__DIR__, "fokker_planck_tests.jl")) include(joinpath(@__DIR__, "fokker_planck_time_evolution_tests.jl")) include(joinpath(@__DIR__, "gyroaverage_tests.jl")) + include(joinpath(@__DIR__, "poisson_radial_polar_tests.jl")) end end diff --git a/test_scripts/PoissonSolve.jl b/test_scripts/PoissonSolve.jl new file mode 100644 index 000000000..a6f6135a6 --- /dev/null +++ b/test_scripts/PoissonSolve.jl @@ -0,0 +1,218 @@ +export run_poisson_test +using Plots +using LaTeXStrings +using MPI +using Measures +#using Dates +import moment_kinetics +using moment_kinetics.array_allocation: allocate_float, allocate_shared_float +using moment_kinetics.coordinates: define_coordinate +using moment_kinetics.type_definitions: mk_float, mk_int, OptionsDict +using moment_kinetics.spatial_poisson: init_spatial_poisson, spatial_poisson_solve! +using moment_kinetics.spatial_poisson: init_spatial_poisson2D, spatial_poisson2D_solve! +using moment_kinetics.communication +using moment_kinetics.looping +using moment_kinetics.calculus: derivative! + +function plot_test_data(func_exact,func_num,func_err,func_name,radial,polar) + @views heatmap(radial.grid, polar.grid, func_num[:,:], xlabel=L"r", ylabel=L"\theta", c = :deep, interpolation = :cubic, #, projection =:polar + windowsize = (360,240), margin = 15pt) + outfile = string(func_name*"_num.pdf") + savefig(outfile) + @views heatmap(radial.grid, polar.grid, func_exact[:,:], xlabel=L"r", ylabel=L"\theta", c = :deep, interpolation = :cubic, #, projection =:polar + windowsize = (360,240), margin = 15pt) + outfile = string(func_name*"_exact.pdf") + savefig(outfile) + @views heatmap(radial.grid, polar.grid, func_err[:,:], xlabel=L"r", ylabel=L"\theta", c = :deep, interpolation = :cubic, #, projection =:polar + windowsize = (360,240), margin = 15pt) + outfile = string(func_name*"_err.pdf") + savefig(outfile) + return nothing +end + + +function run_poisson_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelement_polar=1,ngrid_polar=5,Lpolar=2.0*pi,kk::mk_int=1,plot_results=false) + + nelement_local_polar = nelement_polar # number of elements per rank + nelement_global_polar = nelement_local_polar # total number of elements + nelement_local_radial = nelement_radial # number of elements per rank + nelement_global_radial = nelement_local_radial # total number of elements + bc = "none" #not required to take a particular value, not used, set to "none" to avoid extra BC impositions + GL_discretization = "gausslegendre_pseudospectral" + F_discretization = "fourier_pseudospectral" + element_spacing_option = "uniform" + coords_input = OptionsDict( + "radial"=>OptionsDict("ngrid"=>ngrid_radial, "nelement"=>nelement_global_radial, + "nelement_local"=>nelement_local_radial, "L"=>Lradial, + "discretization"=>GL_discretization, "bc"=>bc, + "element_spacing_option"=>element_spacing_option), + "polar"=>OptionsDict("ngrid"=>ngrid_polar, "nelement"=>nelement_global_polar, + "nelement_local"=>nelement_local_polar, "L"=>Lpolar, + "discretization"=>F_discretization, "bc"=>bc, + "element_spacing_option"=>element_spacing_option), + ) + # Set up MPI + initialize_comms!() + setup_distributed_memory_MPI(1,1,1,1) + # ignore MPI here to avoid FFTW wisdom problems, test runs on a single core below + polar, polar_spectral = define_coordinate(coords_input, "polar", ignore_MPI=true) + radial, radial_spectral = define_coordinate(coords_input, "radial", ignore_MPI=true) + looping.setup_loop_ranges!(block_rank[], block_size[]; + s=1, sn=1, + r=radial.n, z=polar.n, vperp=1, vpa=1, + vzeta=1, vr=1, vz=1) + + begin_serial_region() + @serial_region begin + println("Test fourier_pseudospectral") + ff = allocate_float(polar.n) + df = allocate_float(polar.n) + df_exact = allocate_float(polar.n) + df_err = allocate_float(polar.n) + for i in 1:polar.n + ff[i] = sin(2.0*pi*polar.grid[i]/polar.L) + df_exact[i] = cos(2.0*pi*polar.grid[i]/polar.L)*(2.0*pi/polar.L) + end + #println("ff: ",ff) + derivative!(df,ff,polar,polar_spectral) + #println("polar.grid: ",polar.grid) + @. df_err = abs(df - df_exact) + println("maximum(df_err): ",maximum(df_err)) + #println("df_exact: ",df_exact) + #println(polar.wgts) + println("wgts err: ",abs( 1.0 - (sum(polar.wgts)/polar.L) )) + + println("Test Poisson") + poisson = init_spatial_poisson(radial,polar,radial_spectral) + phi = allocate_float(polar.n,radial.n) + exact_phi = allocate_float(polar.n,radial.n) + err_phi = allocate_float(polar.n,radial.n) + rho = allocate_float(polar.n,radial.n) + + for irad in 1:radial.n + for ipol in 1:polar.n + exact_phi[ipol,irad] = 0.25*(radial.grid[irad]^2 -1) + rho[ipol,irad] = 1.0 + end + end + spatial_poisson_solve!(phi,rho,poisson,radial,polar,polar_spectral) + println("Maximum error value Test rho=1 : ",maximum(abs.(phi- exact_phi))) + + if kk < 1 + error("ERROR: kk >=1 required for test") + end + if !(mod(kk,1) == 0) + error("ERROR: kk integer required for test") + end + if abs(polar.L - 2.0*pi) >1.0e-14 + error("ERROR: polar coordinate assumed angle for definition of the following test - set Lpolar = 2.0*pi") + end + + for irad in 1:radial.n + for ipol in 1:polar.n + exact_phi[ipol,irad] = (1.0 - radial.grid[irad])*(radial.grid[irad]^kk)*cos(2.0*pi*kk*polar.grid[ipol]/polar.L) + rho[ipol,irad] = (kk^2 - (kk+1)^2)*(radial.grid[irad]^(kk-1))*cos(2.0*kk*pi*polar.grid[ipol]/polar.L) + end + end + + spatial_poisson_solve!(phi,rho,poisson,radial,polar,polar_spectral) + @. err_phi = abs(phi - exact_phi) + println("Maximum error value Test rho = (kk^2 - (kk+1)^2) * cos(2 pi kk P/L) * r^(kk-1): ",maximum(err_phi)) + if plot_results + plot_test_data(exact_phi,phi,err_phi,"phi",radial,polar) + end + end + finalize_comms!() + return nothing +end + +function run_poisson2D_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelement_polar=5,ngrid_polar=5,Lpolar=2.0*pi,kk::mk_int=1,phase::mk_float=pi/3.0,plot_results=false) + + nelement_local_polar = nelement_polar # number of elements per rank + nelement_global_polar = nelement_local_polar # total number of elements + nelement_local_radial = nelement_radial # number of elements per rank + nelement_global_radial = nelement_local_radial # total number of elements + bc = "none" #not required to take a particular value, not used, set to "none" to avoid extra BC impositions + discretization = "gausslegendre_pseudospectral" + element_spacing_option = "uniform" + coords_input = OptionsDict( + "radial"=>OptionsDict("ngrid"=>ngrid_radial, "nelement"=>nelement_global_radial, + "nelement_local"=>nelement_local_radial, "L"=>Lradial, + "discretization"=>discretization, "bc"=>bc, + "element_spacing_option"=>element_spacing_option), + "polar"=>OptionsDict("ngrid"=>ngrid_polar, "nelement"=>nelement_global_polar, + "nelement_local"=>nelement_local_polar, "L"=>Lpolar, + "discretization"=>discretization, "bc"=>bc, + "element_spacing_option"=>element_spacing_option), + ) + # Set up MPI + initialize_comms!() + setup_distributed_memory_MPI(1,1,1,1) + # ignore MPI here to avoid FFTW wisdom problems, test runs on a single core below + polar, polar_spectral = define_coordinate(coords_input, "polar", ignore_MPI=true) + radial, radial_spectral = define_coordinate(coords_input, "radial", ignore_MPI=true) + looping.setup_loop_ranges!(block_rank[], block_size[]; + s=1, sn=1, + r=radial.n, z=polar.n, vperp=1, vpa=1, + vzeta=1, vr=1, vz=1) + + begin_serial_region() + @serial_region begin + println("Test Poisson") + poisson = init_spatial_poisson2D(radial,polar,radial_spectral,polar_spectral) + phi = allocate_float(polar.n,radial.n) + exact_phi = allocate_float(polar.n,radial.n) + err_phi = allocate_float(polar.n,radial.n) + rho = allocate_float(polar.n,radial.n) + + for irad in 1:radial.n + for ipol in 1:polar.n + exact_phi[ipol,irad] = 0.25*(radial.grid[irad]^2 -1) + rho[ipol,irad] = 1.0 + end + end + spatial_poisson2D_solve!(phi,rho,poisson) + println("Maximum error value Test rho=1 : ",maximum(abs.(phi- exact_phi))) + if plot_results + plot_test_data(exact_phi,phi,err_phi,"phi_1",radial,polar) + println("phi(theta=-pi,r=0): ",phi[1,1]) + println("phi(theta=pi,r=0): ",phi[end,1]) + end + + if kk < 1 + error("ERROR: kk >=1 required for test") + end + if !(mod(kk,1) == 0) + error("ERROR: kk integer required for test") + end + if abs(polar.L - 2.0*pi) >1.0e-14 + error("ERROR: polar coordinate assumed angle for definition of the following test - set Lpolar = 2.0*pi") + end + + for irad in 1:radial.n + for ipol in 1:polar.n + exact_phi[ipol,irad] = (1.0 - radial.grid[irad])*(radial.grid[irad]^kk)*cos(2.0*pi*kk*polar.grid[ipol]/polar.L + phase) + rho[ipol,irad] = (kk^2 - (kk+1)^2)*(radial.grid[irad]^(kk-1))*cos(2.0*kk*pi*polar.grid[ipol]/polar.L + phase) + end + end + + spatial_poisson2D_solve!(phi,rho,poisson) + @. err_phi = abs(phi - exact_phi) + println("Maximum error value Test rho = (kk^2 - (kk+1)^2) * cos(2 pi kk P/L + phase) * r^(kk-1): ",maximum(err_phi)) + if plot_results + plot_test_data(exact_phi,phi,err_phi,"phi_2",radial,polar) + println("phi(theta=-pi,r=0): ",phi[1,1]) + println("phi(theta=pi,r=0): ",phi[end,1]) + end + end + finalize_comms!() + return nothing +end + +if abspath(PROGRAM_FILE) == @__FILE__ + using Pkg + Pkg.activate(".") + + run_poisson_test() + run_poisson2D_test() +end \ No newline at end of file