From 3a308c4e7d85dec2acd5224dcb6665563403f4fe Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Sat, 24 Aug 2024 15:00:14 +0100
Subject: [PATCH 01/19] Add cylindrical flag to coordinate struct to test for
 whether a Jacobian in cylindrical polar coordinates is required for a
 particular dimension. Make vperp and r the two dimensions for which this is
 true. Note that this does not generalise well to other situations with more
 complicated Jacobian factors.

---
 moment_kinetics/src/calculus.jl               |  8 +++---
 moment_kinetics/src/chebyshev.jl              |  4 +--
 moment_kinetics/src/coordinates.jl            | 22 +++++++++------
 moment_kinetics/src/gauss_legendre.jl         | 28 +++++++++----------
 .../test/nonlinear_solver_tests.jl            |  4 +--
 5 files changed, 35 insertions(+), 31 deletions(-)

diff --git a/moment_kinetics/src/calculus.jl b/moment_kinetics/src/calculus.jl
index 2b473a232..b00eb94f1 100644
--- a/moment_kinetics/src/calculus.jl
+++ b/moment_kinetics/src/calculus.jl
@@ -139,7 +139,7 @@ function laplacian_derivative!(d2f, f, coord, spectral; handle_periodic=true)
     # 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
@@ -151,7 +151,7 @@ function laplacian_derivative!(d2f, f, coord, spectral; handle_periodic=true)
     # 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
@@ -240,8 +240,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 a18c2a75e..58b2bb71e 100644
--- a/moment_kinetics/src/coordinates.jl
+++ b/moment_kinetics/src/coordinates.jl
@@ -126,6 +126,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
 
 """
@@ -148,16 +150,18 @@ function define_coordinate(input, parallel_io::Bool=false; run_directory=nothing
     # obtain (local) index mapping from the grid within each element
     # to the full grid
     imin, imax, igrid_full = elemental_to_full_grid_map(input.ngrid, input.nelement_local)
+    # check name of coordinate to determine if radial or vperp cylindrical coordinate
+    cylindrical = (input.name == "vperp") || (input.name == "r")
     # initialise the data used to construct the grid
     # boundaries for each element
-    element_boundaries = set_element_boundaries(input.nelement_global, input.L, input.element_spacing_option, input.name)
+    element_boundaries = set_element_boundaries(input.nelement_global, input.L, input.element_spacing_option, cylindrical)
     # shift and scale factors for each local element
     element_scale, element_shift = set_element_scale_and_shift(input.nelement_global, input.nelement_local, input.irank, element_boundaries)
     # initialize the grid and the integration weights associated with the grid
     # also obtain the Chebyshev theta grid and spacing if chosen as discretization option
     grid, wgts, uniform_grid, radau_first_element = init_grid(input.ngrid,
         input.nelement_local, n_global, n_local, input.irank, input.L, element_scale,
-        element_shift, imin, imax, igrid, input.discretization, input.name)
+        element_shift, imin, imax, igrid, input.discretization, 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
@@ -240,7 +244,7 @@ function define_coordinate(input, parallel_io::Bool=false; run_directory=nothing
         copy(scratch_2d), advection, send_buffer, receive_buffer, input.comm,
         local_io_range, global_io_range, element_scale, element_shift,
         input.element_spacing_option, element_boundaries, radau_first_element,
-        other_nodes, one_over_denominator)
+        other_nodes, one_over_denominator, cylindrical)
 
     if coord.n == 1 && occursin("v", coord.name)
         spectral = null_velocity_dimension_info()
@@ -271,7 +275,7 @@ function define_coordinate(input, parallel_io::Bool=false; run_directory=nothing
     return coord, spectral
 end
 
-function set_element_boundaries(nelement_global, L, element_spacing_option, coord_name)
+function set_element_boundaries(nelement_global, L, element_spacing_option, 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
@@ -301,7 +305,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
@@ -327,7 +331,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
@@ -342,7 +346,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
@@ -358,7 +362,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
@@ -368,7 +372,7 @@ function init_grid(ngrid, nelement_local, n_global, n_local, irank, L, element_s
             grid, wgts = scaled_gauss_legendre_lobatto_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/gauss_legendre.jl b/moment_kinetics/src/gauss_legendre.jl
index 9a0afb6b3..5140de57a 100644
--- a/moment_kinetics/src/gauss_legendre.jl
+++ b/moment_kinetics/src/gauss_legendre.jl
@@ -887,7 +887,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
@@ -941,7 +941,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
@@ -961,7 +961,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
@@ -992,7 +992,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
@@ -1038,7 +1038,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) +
@@ -1062,7 +1062,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
@@ -1105,7 +1105,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
@@ -1127,7 +1127,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 +
@@ -1153,7 +1153,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
@@ -1176,7 +1176,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 + 
@@ -1222,7 +1222,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
@@ -1242,7 +1242,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
@@ -1262,7 +1262,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
@@ -1282,7 +1282,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/test/nonlinear_solver_tests.jl b/moment_kinetics/test/nonlinear_solver_tests.jl
index 23d62acf9..a79060f70 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)
@@ -176,7 +176,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)

From 49093dc2d4c714299bd9eee8045925579a41888b Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Sat, 24 Aug 2024 19:06:11 +0100
Subject: [PATCH 02/19] Initial experiments at making a Poisson solver for the
 spatial domain.

---
 moment_kinetics/src/moment_kinetics.jl |   1 +
 moment_kinetics/src/spatial_poisson.jl | 128 +++++++++++++++++++++++++
 test_scripts/PoissonSolve.jl           |  74 ++++++++++++++
 3 files changed, 203 insertions(+)
 create mode 100644 moment_kinetics/src/spatial_poisson.jl
 create mode 100644 test_scripts/PoissonSolve.jl

diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index 3712ba994..5361ad371 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -48,6 +48,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/spatial_poisson.jl b/moment_kinetics/src/spatial_poisson.jl
new file mode 100644
index 000000000..aea6db09f
--- /dev/null
+++ b/moment_kinetics/src/spatial_poisson.jl
@@ -0,0 +1,128 @@
+"""
+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!
+
+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
+
+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}
+   rhs_dummy::Array{mk_float,1}
+   phi_dummy::Array{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 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)
+   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
+      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 + (im^2)*MN
+      end
+   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,rhs_dummy,phi_dummy)
+end
+
+"""
+function to find the solution to 
+nabla^2 phi = rho in cylindrical polar coordinates
+ phi(r,polar) = the function solved for
+ rho(r,polar) = the source evaluated at the nodal points
+ poisson_arrays = precomputed arrays
+"""
+# for now just support npolar = 1
+# by skipping the FFT
+function spatial_poisson_solve!(phi,rho,poisson_arrays,radial,polar)
+   laplacian_lu_objs = poisson_arrays.laplacian_lu_objs
+   sourcevec = poisson_arrays.sourcevec
+   phi_dummy = poisson_arrays.phi_dummy
+   rhs_dummy = poisson_arrays.rhs_dummy
+   
+   # first FFT rho to hat{rho} appropriate for using the 1D radial operators
+   @. phi = 0.0
+   npolar = polar.n
+   for im in 1:npolar
+      # solve the linear system
+      mul!(rhs_dummy,sourcevec,rho[:,im])
+      lu_object_lhs = laplacian_lu_objs[im]
+      ldiv!(phi_dummy, lu_object_lhs, rhs_dummy)
+      phi[:,im] = phi_dummy
+   end
+   # finally iFFT from hat{phi} to phi 
+   return nothing
+end
+
+end
\ No newline at end of file
diff --git a/test_scripts/PoissonSolve.jl b/test_scripts/PoissonSolve.jl
new file mode 100644
index 000000000..eb4bc86ef
--- /dev/null
+++ b/test_scripts/PoissonSolve.jl
@@ -0,0 +1,74 @@
+export run_poisson_test
+#using Printf
+#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.input_structs: grid_input, advection_input
+using moment_kinetics.coordinates: define_coordinate
+using moment_kinetics.type_definitions: mk_float, mk_int
+using moment_kinetics.spatial_poisson: init_spatial_poisson, spatial_poisson_solve!
+using moment_kinetics.communication
+using moment_kinetics.looping
+
+function run_poisson_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelement_polar=1,ngrid_polar=1,Lpolar=2.0*pi)
+
+   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 
+   #Lvpa = 12.0 #physical box size in reference units 
+   #Lvperp = 6.0 #physical box size in reference units 
+   bc = "" #not required to take a particular value, not used 
+   # fd_option and adv_input not actually used so given values unimportant
+   #discretization = "chebyshev_pseudospectral"
+   discretization = "gausslegendre_pseudospectral"
+   fd_option = "fourth_order_centered"
+   cheb_option = "matrix"
+   adv_input = advection_input("default", 1.0, 0.0, 0.0)
+   nrank = 1
+   irank = 0
+   comm = MPI.COMM_NULL
+   # create the 'input' struct containing input info needed to create a
+   # coordinate
+   element_spacing_option = "uniform"
+   polar_input = grid_input("polar", ngrid_polar, nelement_global_polar, nelement_local_polar, 
+      nrank, irank, Lpolar, discretization, fd_option, cheb_option, bc, adv_input,comm,element_spacing_option)
+   radial_input = grid_input("r", ngrid_radial, nelement_global_radial, nelement_local_radial, 
+      nrank, irank, Lradial, discretization, fd_option, cheb_option, bc, adv_input,comm,element_spacing_option)
+   # create the coordinate struct 'x'
+   println("made inputs")
+  # println("vpa: ngrid: ",ngrid," nelement: ",nelement_local_vpa, " Lvpa: ",Lvpa)
+  # println("vperp: ngrid: ",ngrid," nelement: ",nelement_local_vperp, " Lvperp: ",Lvperp)
+  
+   # Set up MPI
+   initialize_comms!()
+   setup_distributed_memory_MPI(1,1,1,1)
+   polar, polar_spectral = define_coordinate(polar_input)
+   radial, radial_spectral = define_coordinate(radial_input)
+   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()
+  
+   poisson = init_spatial_poisson(radial,polar,radial_spectral)
+   phi = allocate_float(radial.n,polar.n)
+   rho = allocate_float(radial.n,polar.n)
+   @. rho = 1.0
+   spatial_poisson_solve!(phi,rho,poisson,radial,polar)
+   println(phi)
+   finalize_comms!()
+   return nothing
+end
+
+if abspath(PROGRAM_FILE) == @__FILE__
+    using Pkg
+    Pkg.activate(".")
+    
+    run_poisson_test()
+end
\ No newline at end of file

From c5d10f2035255e1a695cf2ac2d4647a31f2ffd36 Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Sat, 24 Aug 2024 19:58:55 +0100
Subject: [PATCH 03/19] Working 1D Poisson solve with Dirichlet BC phi=0 on
 radial boundary.

---
 moment_kinetics/src/spatial_poisson.jl | 9 ++++++++-
 test_scripts/PoissonSolve.jl           | 7 +++++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/moment_kinetics/src/spatial_poisson.jl b/moment_kinetics/src/spatial_poisson.jl
index aea6db09f..842d41c61 100644
--- a/moment_kinetics/src/spatial_poisson.jl
+++ b/moment_kinetics/src/spatial_poisson.jl
@@ -73,8 +73,11 @@ function init_spatial_poisson(radial::coordinate, polar::coordinate, radial_spec
           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 + (im^2)*MN
+          @. laplacian[imin:imax,imin:imax,im] += KJ - PP + ((im-1)^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)
@@ -114,9 +117,13 @@ function spatial_poisson_solve!(phi,rho,poisson_arrays,radial,polar)
    # first FFT rho to hat{rho} appropriate for using the 1D radial operators
    @. phi = 0.0
    npolar = polar.n
+   nradial = radial.n
    for im in 1:npolar
       # solve the linear system
+      # form the rhs vector
       mul!(rhs_dummy,sourcevec,rho[:,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)
       phi[:,im] = phi_dummy
diff --git a/test_scripts/PoissonSolve.jl b/test_scripts/PoissonSolve.jl
index eb4bc86ef..3e7a8bbce 100644
--- a/test_scripts/PoissonSolve.jl
+++ b/test_scripts/PoissonSolve.jl
@@ -58,10 +58,17 @@ function run_poisson_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelemen
   
    poisson = init_spatial_poisson(radial,polar,radial_spectral)
    phi = allocate_float(radial.n,polar.n)
+   exact_phi = allocate_float(radial.n,polar.n)
    rho = allocate_float(radial.n,polar.n)
    @. rho = 1.0
+   #@. rho = sinpi(2*radial.grid/radial.L)
    spatial_poisson_solve!(phi,rho,poisson,radial,polar)
    println(phi)
+   println(phi[1])
+   @. exact_phi[:,1] = 0.25*(radial.grid^2 -1)
+   println(exact_phi)
+   println("Maximum error value: ",maximum(abs.(phi- exact_phi)))
+   #println(radial.grid)
    finalize_comms!()
    return nothing
 end

From 4dc5087242c65bb4a2ab02093b022bb0e82a50fe Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Sun, 25 Aug 2024 14:06:13 +0100
Subject: [PATCH 04/19] Add "fourier_pseudospectral" discretisation option for
 near-exact treatment of periodic functions.

---
 moment_kinetics/src/coordinates.jl     |  19 ++
 moment_kinetics/src/fourier.jl         | 313 +++++++++++++++++++++++++
 moment_kinetics/src/moment_kinetics.jl |   1 +
 test_scripts/PoissonSolve.jl           |  30 ++-
 4 files changed, 358 insertions(+), 5 deletions(-)
 create mode 100644 moment_kinetics/src/fourier.jl

diff --git a/moment_kinetics/src/coordinates.jl b/moment_kinetics/src/coordinates.jl
index 58b2bb71e..e5059d6da 100644
--- a/moment_kinetics/src/coordinates.jl
+++ b/moment_kinetics/src/coordinates.jl
@@ -11,6 +11,7 @@ using ..type_definitions: mk_float, mk_int
 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
@@ -265,6 +266,18 @@ function define_coordinate(input, parallel_io::Bool=false; run_directory=nothing
         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 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.
@@ -371,6 +384,12 @@ 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 cylindrical
             # initialize equally spaced grid defined on [0,L]
diff --git a/moment_kinetics/src/fourier.jl b/moment_kinetics/src/fourier.jl
new file mode 100644
index 000000000..aad6b0fa3
--- /dev/null
+++ b/moment_kinetics/src/fourier.jl
@@ -0,0 +1,313 @@
+"""
+"""
+module fourier
+
+export setup_fourier_pseudospectral
+export scaled_fourier_grid
+export fourier_spectral_derivative!
+export fourier_info
+
+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}
+    # Chebyshev spectral coefficients of distribution function f
+    # first dimension contains location within element
+    # second dimension indicates the element
+    f::Array{Complex{mk_float},2}
+    # Chebyshev spectral coefficients of derivative of f
+    df::Array{Complex{mk_float},1}
+    # plan for the complex-to-complex, in-place, forward Fourier transform on Chebyshev-Gauss-Lobatto/Radau grid
+    forward::TForward
+    # plan for the complex-to-complex, in-place, backward Fourier transform on Chebyshev-Gauss-Lobatto/Radau 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
+   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::chebyshev_info)
+
+Fourier transform f to get Chebyshev 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 (cheby.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/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index 5361ad371..b94f0c2fa 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -24,6 +24,7 @@ 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")
diff --git a/test_scripts/PoissonSolve.jl b/test_scripts/PoissonSolve.jl
index 3e7a8bbce..2257d59fb 100644
--- a/test_scripts/PoissonSolve.jl
+++ b/test_scripts/PoissonSolve.jl
@@ -13,6 +13,7 @@ using moment_kinetics.type_definitions: mk_float, mk_int
 using moment_kinetics.spatial_poisson: init_spatial_poisson, spatial_poisson_solve!
 using moment_kinetics.communication
 using moment_kinetics.looping
+using moment_kinetics.calculus: derivative!
 
 function run_poisson_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelement_polar=1,ngrid_polar=1,Lpolar=2.0*pi)
 
@@ -36,7 +37,7 @@ function run_poisson_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelemen
    # coordinate
    element_spacing_option = "uniform"
    polar_input = grid_input("polar", ngrid_polar, nelement_global_polar, nelement_local_polar, 
-      nrank, irank, Lpolar, discretization, fd_option, cheb_option, bc, adv_input,comm,element_spacing_option)
+      nrank, irank, Lpolar, "fourier_pseudospectral", fd_option, cheb_option, "none", adv_input,comm,element_spacing_option)
    radial_input = grid_input("r", ngrid_radial, nelement_global_radial, nelement_local_radial, 
       nrank, irank, Lradial, discretization, fd_option, cheb_option, bc, adv_input,comm,element_spacing_option)
    # create the coordinate struct 'x'
@@ -56,6 +57,25 @@ function run_poisson_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelemen
    
    begin_serial_region()
   
+   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(radial.n,polar.n)
    exact_phi = allocate_float(radial.n,polar.n)
@@ -63,11 +83,11 @@ function run_poisson_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelemen
    @. rho = 1.0
    #@. rho = sinpi(2*radial.grid/radial.L)
    spatial_poisson_solve!(phi,rho,poisson,radial,polar)
-   println(phi)
-   println(phi[1])
+   #println(phi)
+   #println(phi[1])
    @. exact_phi[:,1] = 0.25*(radial.grid^2 -1)
-   println(exact_phi)
-   println("Maximum error value: ",maximum(abs.(phi- exact_phi)))
+   #println(exact_phi)
+   #println("Maximum error value: ",maximum(abs.(phi- exact_phi)))
    #println(radial.grid)
    finalize_comms!()
    return nothing

From c6fbf131f5e8d50f5d54179979d55f8b6606fd95 Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Sun, 25 Aug 2024 15:55:53 +0100
Subject: [PATCH 05/19] Version of Poisson solver giving physically plausible
 solutions.

---
 moment_kinetics/src/fourier.jl         |  2 +
 moment_kinetics/src/spatial_poisson.jl | 63 ++++++++++++++++++++++----
 test_scripts/PoissonSolve.jl           | 57 +++++++++++++++++++----
 3 files changed, 102 insertions(+), 20 deletions(-)

diff --git a/moment_kinetics/src/fourier.jl b/moment_kinetics/src/fourier.jl
index aad6b0fa3..2864cead2 100644
--- a/moment_kinetics/src/fourier.jl
+++ b/moment_kinetics/src/fourier.jl
@@ -6,6 +6,8 @@ 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
diff --git a/moment_kinetics/src/spatial_poisson.jl b/moment_kinetics/src/spatial_poisson.jl
index 842d41c61..2f2308baa 100644
--- a/moment_kinetics/src/spatial_poisson.jl
+++ b/moment_kinetics/src/spatial_poisson.jl
@@ -16,6 +16,8 @@ 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
 
 struct poisson_arrays 
    # an array of lu objects for solving the mth polar harmonic of Poisson's equation
@@ -25,8 +27,9 @@ struct poisson_arrays
    laplacian::Array{mk_float,3}
    # the matrix that needs to multiply the nodal values of the source function
    sourcevec::Array{mk_float,2}
-   rhs_dummy::Array{mk_float,1}
-   phi_dummy::Array{mk_float,1}
+   rhohat::Array{Complex{mk_float},2}
+   rhs_dummy::Array{Complex{mk_float},1}
+   phi_dummy::Array{Complex{mk_float},1}
 end
 
 """
@@ -44,6 +47,27 @@ function get_imin_imax(coord,iel)
     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
@@ -57,6 +81,7 @@ function init_spatial_poisson(radial::coordinate, polar::coordinate, radial_spec
    nrelement = radial.nelement_global
    laplacian = allocate_float(nrtot,nrtot,npolar)
    sourcevec = allocate_float(nrtot,nrtot)
+   rhohat = allocate_float(nrtot,npolar)
    rhs_dummy = allocate_float(nrtot)
    phi_dummy = allocate_float(nrtot)
    MR = allocate_float(nrgrid,nrgrid)
@@ -67,13 +92,14 @@ function init_spatial_poisson(radial::coordinate, polar::coordinate, radial_spec
    @. 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 + ((im-1)^2)*MN
+          @. laplacian[imin:imax,imin:imax,im] += KJ - PP - ((mwn)^2)*MN
       end
       # set rows for Dirichlet BCs on phi
       laplacian[nrtot,:,im] .= 0.0
@@ -96,7 +122,7 @@ function init_spatial_poisson(radial::coordinate, polar::coordinate, radial_spec
    for im in 1:npolar
       laplacian_lu_objs[im] = lu(laplacian_sparse[im])
    end
-   return poisson_arrays(laplacian_lu_objs,laplacian,sourcevec,rhs_dummy,phi_dummy)
+   return poisson_arrays(laplacian_lu_objs,laplacian,sourcevec,rhohat,rhs_dummy,phi_dummy)
 end
 
 """
@@ -108,27 +134,44 @@ nabla^2 phi = rho in cylindrical polar coordinates
 """
 # for now just support npolar = 1
 # by skipping the FFT
-function spatial_poisson_solve!(phi,rho,poisson_arrays,radial,polar)
+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
    
-   # first FFT rho to hat{rho} appropriate for using the 1D radial operators
-   @. phi = 0.0
    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,rho[:,im])
+      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)
-      phi[:,im] = phi_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
-   # finally iFFT from hat{phi} to phi 
    return nothing
 end
 
diff --git a/test_scripts/PoissonSolve.jl b/test_scripts/PoissonSolve.jl
index 2257d59fb..f98e66c28 100644
--- a/test_scripts/PoissonSolve.jl
+++ b/test_scripts/PoissonSolve.jl
@@ -1,9 +1,10 @@
 export run_poisson_test
 #using Printf
-#using Plots
-#using LaTeXStrings
+using Plots
+#gr()
+using LaTeXStrings
 using MPI
-#using Measures
+using Measures
 #using Dates
 import moment_kinetics
 using moment_kinetics.array_allocation: allocate_float, allocate_shared_float
@@ -15,7 +16,24 @@ using moment_kinetics.communication
 using moment_kinetics.looping
 using moment_kinetics.calculus: derivative!
 
-function run_poisson_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelement_polar=1,ngrid_polar=1,Lpolar=2.0*pi)
+function plot_test_data(func_exact,func_num,func_err,func_name,radial,polar)
+    @views heatmap(polar.grid, radial.grid, func_num[:,:], ylabel=L"r", xlabel=L"\theta", c = :deep, interpolation = :cubic,
+                windowsize = (360,240), margin = 15pt, projection =:polar)
+                outfile = string(func_name*"_num.pdf")
+                savefig(outfile)
+    @views heatmap(polar.grid, radial.grid, func_exact[:,:], ylabel=L"r", xlabel=L"\theta", c = :deep, interpolation = :cubic,
+                windowsize = (360,240), margin = 15pt, projection =:polar)
+                outfile = string(func_name*"_exact.pdf")
+                savefig(outfile)
+    @views heatmap(polar.grid, radial.grid, func_err[:,:], ylabel=L"r", xlabel=L"\theta", c = :deep, interpolation = :cubic,
+                windowsize = (360,240), margin = 15pt, projection =:polar)
+                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=1,Lpolar=2.0*pi,kk=1.0)
 
    nelement_local_polar = nelement_polar # number of elements per rank
    nelement_global_polar = nelement_local_polar # total number of elements 
@@ -79,16 +97,35 @@ function run_poisson_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelemen
    poisson = init_spatial_poisson(radial,polar,radial_spectral)
    phi = allocate_float(radial.n,polar.n)
    exact_phi = allocate_float(radial.n,polar.n)
+   err_phi = allocate_float(radial.n,polar.n)
    rho = allocate_float(radial.n,polar.n)
-   @. rho = 1.0
+   
+   for ipol in 1:polar.n
+      for irad in 1:radial.n
+         exact_phi[irad,ipol] = 0.25*(radial.grid[irad]^2 -1)
+         rho[irad,ipol] = 1.0
+      end
+   end
+   #@. rho = sinpi(2*radial.grid/radial.L)
+   spatial_poisson_solve!(phi,rho,poisson,radial,polar,polar_spectral)
+   #println(phi)
+   #println(exact_phi)
+   println("Maximum error value Test rho=1 : ",maximum(abs.(phi- exact_phi)))
+   
+   for ipol in 1:polar.n
+      for irad in 1:radial.n
+         exact_phi[irad,ipol] = 0.25*(radial.grid[irad]^2 -1)*cos(2.0*pi*kk*polar.grid[ipol]/polar.L)
+         rho[irad,ipol] = cos(2.0*kk*pi*polar.grid[ipol]/polar.L)
+      end
+   end
+   
    #@. rho = sinpi(2*radial.grid/radial.L)
-   spatial_poisson_solve!(phi,rho,poisson,radial,polar)
+   spatial_poisson_solve!(phi,rho,poisson,radial,polar,polar_spectral)
+   @. err_phi = abs(phi - exact_phi)
    #println(phi)
-   #println(phi[1])
-   @. exact_phi[:,1] = 0.25*(radial.grid^2 -1)
    #println(exact_phi)
-   #println("Maximum error value: ",maximum(abs.(phi- exact_phi)))
-   #println(radial.grid)
+   println("Maximum error value Test rho = cos(2 pi kk P/L): ",maximum(err_phi))
+   plot_test_data(exact_phi,phi,err_phi,"phi",radial,polar)
    finalize_comms!()
    return nothing
 end

From 84003d224e493a8c2d00cda5e68d3bb9e1f689f3 Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Mon, 26 Aug 2024 17:22:04 +0100
Subject: [PATCH 06/19] Correct test for kk>=1 to demonstrate correct
 functioning of solver for non-trivial functions of polar angle (with phi = 0
 boundary conditions only).

---
 test_scripts/PoissonSolve.jl | 48 +++++++++++++++---------------------
 1 file changed, 20 insertions(+), 28 deletions(-)

diff --git a/test_scripts/PoissonSolve.jl b/test_scripts/PoissonSolve.jl
index f98e66c28..09ea67cd5 100644
--- a/test_scripts/PoissonSolve.jl
+++ b/test_scripts/PoissonSolve.jl
@@ -1,7 +1,5 @@
 export run_poisson_test
-#using Printf
 using Plots
-#gr()
 using LaTeXStrings
 using MPI
 using Measures
@@ -17,34 +15,30 @@ 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(polar.grid, radial.grid, func_num[:,:], ylabel=L"r", xlabel=L"\theta", c = :deep, interpolation = :cubic,
-                windowsize = (360,240), margin = 15pt, projection =:polar)
+    @views heatmap(polar.grid, radial.grid, func_num[:,:], ylabel=L"r", xlabel=L"\theta", c = :deep, interpolation = :cubic, #, projection =:polar
+                windowsize = (360,240), margin = 15pt)
                 outfile = string(func_name*"_num.pdf")
                 savefig(outfile)
-    @views heatmap(polar.grid, radial.grid, func_exact[:,:], ylabel=L"r", xlabel=L"\theta", c = :deep, interpolation = :cubic,
-                windowsize = (360,240), margin = 15pt, projection =:polar)
+    @views heatmap(polar.grid, radial.grid, func_exact[:,:], ylabel=L"r", xlabel=L"\theta", c = :deep, interpolation = :cubic, #, projection =:polar
+                windowsize = (360,240), margin = 15pt)
                 outfile = string(func_name*"_exact.pdf")
                 savefig(outfile)
-    @views heatmap(polar.grid, radial.grid, func_err[:,:], ylabel=L"r", xlabel=L"\theta", c = :deep, interpolation = :cubic,
-                windowsize = (360,240), margin = 15pt, projection =:polar)
+    @views heatmap(polar.grid, radial.grid, func_err[:,:], ylabel=L"r", xlabel=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=1,Lpolar=2.0*pi,kk=1.0)
+function run_poisson_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelement_polar=1,ngrid_polar=1,Lpolar=2.0*pi,kk::mk_int=1)
 
    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 
-   #Lvpa = 12.0 #physical box size in reference units 
-   #Lvperp = 6.0 #physical box size in reference units 
-   bc = "" #not required to take a particular value, not used 
+   bc = "none" #not required to take a particular value, not used, set to "none" to avoid extra BC impositions 
    # fd_option and adv_input not actually used so given values unimportant
-   #discretization = "chebyshev_pseudospectral"
-   discretization = "gausslegendre_pseudospectral"
    fd_option = "fourth_order_centered"
    cheb_option = "matrix"
    adv_input = advection_input("default", 1.0, 0.0, 0.0)
@@ -55,13 +49,10 @@ function run_poisson_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelemen
    # coordinate
    element_spacing_option = "uniform"
    polar_input = grid_input("polar", ngrid_polar, nelement_global_polar, nelement_local_polar, 
-      nrank, irank, Lpolar, "fourier_pseudospectral", fd_option, cheb_option, "none", adv_input,comm,element_spacing_option)
+      nrank, irank, Lpolar, "fourier_pseudospectral", fd_option, cheb_option, bc, adv_input,comm,element_spacing_option)
    radial_input = grid_input("r", ngrid_radial, nelement_global_radial, nelement_local_radial, 
-      nrank, irank, Lradial, discretization, fd_option, cheb_option, bc, adv_input,comm,element_spacing_option)
-   # create the coordinate struct 'x'
+      nrank, irank, Lradial, "gausslegendre_pseudospectral", fd_option, cheb_option, bc, adv_input,comm,element_spacing_option)
    println("made inputs")
-  # println("vpa: ngrid: ",ngrid," nelement: ",nelement_local_vpa, " Lvpa: ",Lvpa)
-  # println("vperp: ngrid: ",ngrid," nelement: ",nelement_local_vperp, " Lvperp: ",Lvperp)
   
    # Set up MPI
    initialize_comms!()
@@ -106,25 +97,26 @@ function run_poisson_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelemen
          rho[irad,ipol] = 1.0
       end
    end
-   #@. rho = sinpi(2*radial.grid/radial.L)
    spatial_poisson_solve!(phi,rho,poisson,radial,polar,polar_spectral)
-   #println(phi)
-   #println(exact_phi)
    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
+   
    for ipol in 1:polar.n
       for irad in 1:radial.n
-         exact_phi[irad,ipol] = 0.25*(radial.grid[irad]^2 -1)*cos(2.0*pi*kk*polar.grid[ipol]/polar.L)
-         rho[irad,ipol] = cos(2.0*kk*pi*polar.grid[ipol]/polar.L)
+         exact_phi[irad,ipol] = (1.0 - radial.grid[irad])*(radial.grid[irad]^kk)*cos(2.0*pi*kk*polar.grid[ipol]/polar.L)
+         rho[irad,ipol] = (kk^2 - (kk+1)^2)*(radial.grid[irad]^(kk-1))*cos(2.0*kk*pi*polar.grid[ipol]/polar.L)
       end
    end
    
-   #@. rho = sinpi(2*radial.grid/radial.L)
    spatial_poisson_solve!(phi,rho,poisson,radial,polar,polar_spectral)
    @. err_phi = abs(phi - exact_phi)
-   #println(phi)
-   #println(exact_phi)
-   println("Maximum error value Test rho = cos(2 pi kk P/L): ",maximum(err_phi))
+   println("Maximum error value Test rho = (kk^2 - (kk+1)^2) * cos(2 pi kk P/L) * r^(kk-1): ",maximum(err_phi))
    plot_test_data(exact_phi,phi,err_phi,"phi",radial,polar)
    finalize_comms!()
    return nothing

From 24a493fd968e9719b72ff4a643b7d2de7e3c7a02 Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Mon, 26 Aug 2024 17:41:18 +0100
Subject: [PATCH 07/19] Add some sensible warnings, avoid plots by default.

---
 test_scripts/PoissonSolve.jl | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/test_scripts/PoissonSolve.jl b/test_scripts/PoissonSolve.jl
index 09ea67cd5..7e472dfc6 100644
--- a/test_scripts/PoissonSolve.jl
+++ b/test_scripts/PoissonSolve.jl
@@ -31,7 +31,7 @@ function plot_test_data(func_exact,func_num,func_err,func_name,radial,polar)
 end
 
 
-function run_poisson_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelement_polar=1,ngrid_polar=1,Lpolar=2.0*pi,kk::mk_int=1)
+function run_poisson_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelement_polar=1,ngrid_polar=1,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 
@@ -106,6 +106,9 @@ function run_poisson_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelemen
    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 ipol in 1:polar.n
       for irad in 1:radial.n
@@ -117,7 +120,9 @@ function run_poisson_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelemen
    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))
-   plot_test_data(exact_phi,phi,err_phi,"phi",radial,polar)
+   if plot_results
+      plot_test_data(exact_phi,phi,err_phi,"phi",radial,polar)
+   end
    finalize_comms!()
    return nothing
 end

From 92e9d9ae59ea13b682a7ef408d315f7d40c21213 Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Mon, 26 Aug 2024 17:46:40 +0100
Subject: [PATCH 08/19] Correct comments.

---
 moment_kinetics/src/fourier.jl | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/moment_kinetics/src/fourier.jl b/moment_kinetics/src/fourier.jl
index 2864cead2..d881dcc69 100644
--- a/moment_kinetics/src/fourier.jl
+++ b/moment_kinetics/src/fourier.jl
@@ -24,15 +24,15 @@ struct fourier_info{TForward <: FFTW.cFFTWPlan, TBackward <: AbstractFFTs.Scaled
     # 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}
-    # Chebyshev spectral coefficients of distribution function f
+    # Fourier spectral coefficients of distribution function f
     # first dimension contains location within element
     # second dimension indicates the element
     f::Array{Complex{mk_float},2}
-    # Chebyshev spectral coefficients of derivative of f
+    # 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 Chebyshev-Gauss-Lobatto/Radau grid
+    # 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 Chebyshev-Gauss-Lobatto/Radau grid
+    # 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
@@ -203,6 +203,7 @@ function elementwise_derivative!(coord, ff, fourier::fourier_info)
        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
@@ -211,9 +212,9 @@ function elementwise_derivative!(coord, ff, fourier::fourier_info)
 end
 
 """
-    elementwise_derivative!(coord, ff, adv_fac, spectral::chebyshev_info)
+    elementwise_derivative!(coord, ff, adv_fac, spectral::fourier_info)
 
-Fourier transform f to get Chebyshev spectral coefficients and use them to calculate f'.
+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
@@ -278,7 +279,7 @@ function fourier_forward_transform!(fhat, fext, ff, transform, imidm, imidp, n)
             fext[j] = complex(ff[j-imidm],0.0) 
         end
     end
-    # perform the forward, complex-to-complex FFT in-place (cheby.fext is overwritten)
+    # perform the forward, complex-to-complex FFT in-place (fext is overwritten)
     transform*fext
     # set out data
     fhat .= fext

From 0281f092ad9ef40763f732a7d6817c38a2c3cfa8 Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Fri, 6 Sep 2024 15:10:53 +0100
Subject: [PATCH 09/19] Improve documentation

---
 moment_kinetics/src/spatial_poisson.jl | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/moment_kinetics/src/spatial_poisson.jl b/moment_kinetics/src/spatial_poisson.jl
index 2f2308baa..6189c539d 100644
--- a/moment_kinetics/src/spatial_poisson.jl
+++ b/moment_kinetics/src/spatial_poisson.jl
@@ -126,11 +126,23 @@ function init_spatial_poisson(radial::coordinate, polar::coordinate, radial_spec
 end
 
 """
-function to find the solution to 
+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(r,polar) = the function solved for
  rho(r,polar) = 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

From aad7e274a2645577c1120eae6e622aecd3ed489c Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Fri, 6 Sep 2024 15:13:06 +0100
Subject: [PATCH 10/19] Automatic test (running in serial only) for Poisson
 solver in 2D cylindrical coordinates (radial,polar), as opposed to
 (radial,vertical) coordinates used for the velocity space solvers in the
 Fokker-Planck operator.

---
 .../test/poisson_radial_polar_tests.jl        | 177 ++++++++++++++++++
 moment_kinetics/test/runtests.jl              |   1 +
 2 files changed, 178 insertions(+)
 create mode 100644 moment_kinetics/test/poisson_radial_polar_tests.jl

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..9d0552782
--- /dev/null
+++ b/moment_kinetics/test/poisson_radial_polar_tests.jl
@@ -0,0 +1,177 @@
+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.input_structs: grid_input, advection_input
+using moment_kinetics.coordinates: define_coordinate
+using moment_kinetics.type_definitions: mk_float, mk_int
+using moment_kinetics.spatial_poisson: init_spatial_poisson, spatial_poisson_solve!
+using moment_kinetics.communication
+using moment_kinetics.looping
+using moment_kinetics.calculus: derivative!
+
+function run_test()
+    @testset "poisson_radial_polar_tests" begin
+        println("poisson_radial_polar_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_test(nelement_radial=nelement_radial,ngrid_radial=ngrid_radial,Lradial=Lradial,ngrid_polar=ngrid_polar,kk=kk,
+                              atol_phi1=3.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_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-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_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-13,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_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-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_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-13,atol_phi2=2.0e-11)
+        end
+    end
+end
+
+function run_poisson_radial_polar_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 
+   # fd_option and adv_input not actually used so given values unimportant
+   fd_option = "fourth_order_centered"
+   cheb_option = "matrix"
+   adv_input = advection_input("default", 1.0, 0.0, 0.0)
+   nrank = 1
+   irank = 0
+   comm = MPI.COMM_NULL
+   # create the 'input' struct containing input info needed to create a
+   # coordinate
+   element_spacing_option = "uniform"
+   polar_input = grid_input("polar", ngrid_polar, nelement_global_polar, nelement_local_polar, 
+      nrank, irank, Lpolar, "fourier_pseudospectral", fd_option, cheb_option, bc, adv_input,comm,element_spacing_option)
+   radial_input = grid_input("r", ngrid_radial, nelement_global_radial, nelement_local_radial, 
+      nrank, irank, Lradial, "gausslegendre_pseudospectral", fd_option, cheb_option, bc, adv_input,comm,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(polar_input, ignore_MPI=true)
+   radial, radial_spectral = define_coordinate(radial_input)
+   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(radial.n,polar.n)
+       exact_phi = allocate_float(radial.n,polar.n)
+       err_phi = allocate_float(radial.n,polar.n)
+       rho = allocate_float(radial.n,polar.n)
+       
+       for ipol in 1:polar.n
+          for irad in 1:radial.n
+             exact_phi[irad,ipol] = 0.25*(radial.grid[irad]^2 -1)
+             rho[irad,ipol] = 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 ipol in 1:polar.n
+              for irad in 1:radial.n
+                 exact_phi[irad,ipol] = (1.0 - radial.grid[irad])*(radial.grid[irad]^kk)*cos(2.0*pi*kk*polar.grid[ipol]/polar.L)
+                 rho[irad,ipol] = (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
+
+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
 

From 2b7d5e1eb7842cc765e00c864e52cc665be53c7d Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Mon, 16 Sep 2024 08:36:03 +0100
Subject: [PATCH 11/19] Relax error values for MacOS tests.

---
 moment_kinetics/test/poisson_radial_polar_tests.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/moment_kinetics/test/poisson_radial_polar_tests.jl b/moment_kinetics/test/poisson_radial_polar_tests.jl
index 9d0552782..9377406a5 100644
--- a/moment_kinetics/test/poisson_radial_polar_tests.jl
+++ b/moment_kinetics/test/poisson_radial_polar_tests.jl
@@ -26,7 +26,7 @@ function run_test()
         kk=1
         @testset "nelement_radial $nelement_radial ngrid_radial $ngrid_radial Lradial $Lradial ngrid_polar $ngrid_polar kk $kk" begin
              run_poisson_radial_polar_test(nelement_radial=nelement_radial,ngrid_radial=ngrid_radial,Lradial=Lradial,ngrid_polar=ngrid_polar,kk=kk,
-                              atol_phi1=3.0e-14)
+                              atol_phi1=4.0e-14)
         end
         nelement_radial=5
         ngrid_radial=5
@@ -35,7 +35,7 @@ function run_test()
         kk=3
         @testset "nelement_radial $nelement_radial ngrid_radial $ngrid_radial Lradial $Lradial ngrid_polar $ngrid_polar kk $kk" begin
              run_poisson_radial_polar_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-14,atol_phi2=1.0e-14)
+                              atol_fourier=1.0e-14,atol_phi1=4.0e-14,atol_phi2=1.0e-14)
         end
         nelement_radial=5
         ngrid_radial=9
@@ -44,7 +44,7 @@ function run_test()
         kk=3
         @testset "nelement_radial $nelement_radial ngrid_radial $ngrid_radial Lradial $Lradial ngrid_polar $ngrid_polar kk $kk" begin
              run_poisson_radial_polar_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-13,atol_phi2=1.0e-14)
+                              atol_fourier=1.0e-14,atol_phi1=3.0e-12,atol_phi2=1.0e-14)
         end
         nelement_radial=10
         ngrid_radial=5
@@ -53,7 +53,7 @@ function run_test()
         kk=3
         @testset "nelement_radial $nelement_radial ngrid_radial $ngrid_radial Lradial $Lradial ngrid_polar $ngrid_polar kk $kk" begin
              run_poisson_radial_polar_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-13,atol_phi2=1.0e-14)
+                              atol_fourier=1.0e-14,atol_phi1=9.0e-13,atol_phi2=1.0e-14)
         end
         nelement_radial=40
         ngrid_radial=5

From 88ca919b0ac1e674a257caf4e77c21cba8d4a830 Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Mon, 16 Sep 2024 09:03:42 +0100
Subject: [PATCH 12/19] Relax error values for MacOS tests.

---
 moment_kinetics/test/poisson_radial_polar_tests.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/moment_kinetics/test/poisson_radial_polar_tests.jl b/moment_kinetics/test/poisson_radial_polar_tests.jl
index 9377406a5..30ac76f20 100644
--- a/moment_kinetics/test/poisson_radial_polar_tests.jl
+++ b/moment_kinetics/test/poisson_radial_polar_tests.jl
@@ -62,7 +62,7 @@ function run_test()
         kk=5
         @testset "nelement_radial $nelement_radial ngrid_radial $ngrid_radial Lradial $Lradial ngrid_polar $ngrid_polar kk $kk" begin
              run_poisson_radial_polar_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-13,atol_phi2=2.0e-11)
+                              atol_fourier=1.0e-14,atol_phi1=9.0e-13,atol_phi2=2.0e-11)
         end
     end
 end

From 02c13a85dbf567252eaf54eeac87ec4a0a9542da Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Mon, 16 Sep 2024 12:04:41 +0100
Subject: [PATCH 13/19] Reorder arguments in phi and rho so that phi =
 phi(polar,r) and rho = rho(polar,r).

---
 moment_kinetics/src/spatial_poisson.jl        | 14 ++++-----
 .../test/poisson_radial_polar_tests.jl        | 24 +++++++--------
 test_scripts/PoissonSolve.jl                  | 30 +++++++++----------
 3 files changed, 34 insertions(+), 34 deletions(-)

diff --git a/moment_kinetics/src/spatial_poisson.jl b/moment_kinetics/src/spatial_poisson.jl
index 6189c539d..9020029b3 100644
--- a/moment_kinetics/src/spatial_poisson.jl
+++ b/moment_kinetics/src/spatial_poisson.jl
@@ -81,7 +81,7 @@ function init_spatial_poisson(radial::coordinate, polar::coordinate, radial_spec
    nrelement = radial.nelement_global
    laplacian = allocate_float(nrtot,nrtot,npolar)
    sourcevec = allocate_float(nrtot,nrtot)
-   rhohat = allocate_float(nrtot,npolar)
+   rhohat = allocate_float(npolar,nrtot)
    rhs_dummy = allocate_float(nrtot)
    phi_dummy = allocate_float(nrtot)
    MR = allocate_float(nrgrid,nrgrid)
@@ -133,8 +133,8 @@ nabla^2 phi = (1/r)d/dr(r dphi/dr) + (1/r^2)d^2 phi/dpolar^2
 
 The arguments are 
  
- phi(r,polar) = the function solved for
- rho(r,polar) = the source evaluated at the nodal points
+ 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
@@ -159,7 +159,7 @@ function spatial_poisson_solve!(phi,rho,poisson_arrays,radial,polar,polar_spectr
    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)
+         @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)
@@ -168,18 +168,18 @@ function spatial_poisson_solve!(phi,rho,poisson_arrays,radial,polar,polar_spectr
    for im in 1:npolar
       # solve the linear system
       # form the rhs vector
-      mul!(rhs_dummy,sourcevec,rhohat[:,im])
+      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
+      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)
+         @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)
diff --git a/moment_kinetics/test/poisson_radial_polar_tests.jl b/moment_kinetics/test/poisson_radial_polar_tests.jl
index 30ac76f20..92ca7dc65 100644
--- a/moment_kinetics/test/poisson_radial_polar_tests.jl
+++ b/moment_kinetics/test/poisson_radial_polar_tests.jl
@@ -125,15 +125,15 @@ function run_poisson_radial_polar_test(; nelement_radial=5,ngrid_radial=5,Lradia
        
        if print_to_screen println("Test Poisson") end
        poisson = init_spatial_poisson(radial,polar,radial_spectral)
-       phi = allocate_float(radial.n,polar.n)
-       exact_phi = allocate_float(radial.n,polar.n)
-       err_phi = allocate_float(radial.n,polar.n)
-       rho = allocate_float(radial.n,polar.n)
+       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 ipol in 1:polar.n
-          for irad in 1:radial.n
-             exact_phi[irad,ipol] = 0.25*(radial.grid[irad]^2 -1)
-             rho[irad,ipol] = 1.0
+       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)
@@ -152,10 +152,10 @@ function run_poisson_radial_polar_test(; nelement_radial=5,ngrid_radial=5,Lradia
               error("ERROR: polar coordinate assumed angle for definition of the following test - set Lpolar = 2.0*pi")
            end
            
-           for ipol in 1:polar.n
-              for irad in 1:radial.n
-                 exact_phi[irad,ipol] = (1.0 - radial.grid[irad])*(radial.grid[irad]^kk)*cos(2.0*pi*kk*polar.grid[ipol]/polar.L)
-                 rho[irad,ipol] = (kk^2 - (kk+1)^2)*(radial.grid[irad]^(kk-1))*cos(2.0*kk*pi*polar.grid[ipol]/polar.L)
+           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
            
diff --git a/test_scripts/PoissonSolve.jl b/test_scripts/PoissonSolve.jl
index 7e472dfc6..3d15347a4 100644
--- a/test_scripts/PoissonSolve.jl
+++ b/test_scripts/PoissonSolve.jl
@@ -15,15 +15,15 @@ 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(polar.grid, radial.grid, func_num[:,:], ylabel=L"r", xlabel=L"\theta", c = :deep, interpolation = :cubic, #, projection =: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(polar.grid, radial.grid, func_exact[:,:], ylabel=L"r", xlabel=L"\theta", c = :deep, interpolation = :cubic, #, projection =:polar
+    @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(polar.grid, radial.grid, func_err[:,:], ylabel=L"r", xlabel=L"\theta", c = :deep, interpolation = :cubic, #, projection =:polar
+    @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)
@@ -86,15 +86,15 @@ function run_poisson_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelemen
    
    println("Test Poisson")
    poisson = init_spatial_poisson(radial,polar,radial_spectral)
-   phi = allocate_float(radial.n,polar.n)
-   exact_phi = allocate_float(radial.n,polar.n)
-   err_phi = allocate_float(radial.n,polar.n)
-   rho = allocate_float(radial.n,polar.n)
+   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 ipol in 1:polar.n
-      for irad in 1:radial.n
-         exact_phi[irad,ipol] = 0.25*(radial.grid[irad]^2 -1)
-         rho[irad,ipol] = 1.0
+   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)
@@ -110,10 +110,10 @@ function run_poisson_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelemen
       error("ERROR: polar coordinate assumed angle for definition of the following test - set Lpolar = 2.0*pi")
    end
    
-   for ipol in 1:polar.n
-      for irad in 1:radial.n
-         exact_phi[irad,ipol] = (1.0 - radial.grid[irad])*(radial.grid[irad]^kk)*cos(2.0*pi*kk*polar.grid[ipol]/polar.L)
-         rho[irad,ipol] = (kk^2 - (kk+1)^2)*(radial.grid[irad]^(kk-1))*cos(2.0*kk*pi*polar.grid[ipol]/polar.L)
+   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
    

From c12f46572ad7ab9e3930bd4e0411d55ee4a4ce4c Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Mon, 16 Sep 2024 16:14:17 +0100
Subject: [PATCH 14/19] Extract sparse matrix worker functions from
 fokker_planck_calculus.jl to enable them to be reused for a different set of
 coordinates.

---
 moment_kinetics/src/fokker_planck_calculus.jl | 59 +------------
 moment_kinetics/src/moment_kinetics.jl        |  1 +
 .../src/sparse_matrix_functions.jl            | 88 +++++++++++++++++++
 3 files changed, 92 insertions(+), 56 deletions(-)
 create mode 100644 moment_kinetics/src/sparse_matrix_functions.jl

diff --git a/moment_kinetics/src/fokker_planck_calculus.jl b/moment_kinetics/src/fokker_planck_calculus.jl
index f11dcf43c..88dfe7f66 100644
--- a/moment_kinetics/src/fokker_planck_calculus.jl
+++ b/moment_kinetics/src/fokker_planck_calculus.jl
@@ -41,6 +41,9 @@ using ..communication
 using ..communication: MPISharedArray, global_rank
 using ..lagrange_polynomials: lagrange_poly, lagrange_poly_optimised
 using ..looping
+using ..sparse_matrix_functions: icsc_func, allocate_sparse_matrix_constructor,
+                                 assemble_constructor_data!, assign_constructor_data!,
+                                 sparse_matrix_constructor, create_sparse_matrix
 using moment_kinetics.gauss_legendre: get_QQ_local!
 using Dates
 using SpecialFunctions: ellipk, ellipe
@@ -926,62 +929,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
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index b94f0c2fa..01f983350 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -22,6 +22,7 @@ include("array_allocation.jl")
 include("interpolation.jl")
 include("calculus.jl")
 include("lagrange_polynomials.jl")
+include("sparse_matrix_functions.jl")
 include("clenshaw_curtis.jl")
 include("gauss_legendre.jl")
 include("fourier.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..24bb05bcc
--- /dev/null
+++ b/moment_kinetics/src/sparse_matrix_functions.jl
@@ -0,0 +1,88 @@
+"""
+module for storing sparse matrix worker functions
+"""
+module sparse_matrix_functions
+
+export icsc_func
+export allocate_sparse_matrix_constructor
+export assemble_constructor_data!
+export assign_constructor_data!
+export sparse_matrix_constructor
+export create_sparse_matrix
+
+using ..type_definitions: mk_float, mk_int
+using SparseArrays: sparse
+
+"""
+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
+
+"""
+"""
+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

From fc435776cf4c0dc574b94c51cee149a9e40204fa Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Mon, 16 Sep 2024 16:29:54 +0100
Subject: [PATCH 15/19] Split out more indexing functions from
 fokker_planck_calculus.jl.

---
 moment_kinetics/src/fokker_planck_calculus.jl | 37 ++----------------
 moment_kinetics/src/moment_kinetics.jl        |  3 +-
 .../src/sparse_matrix_functions.jl            | 38 ++++++++++++++++++-
 3 files changed, 42 insertions(+), 36 deletions(-)

diff --git a/moment_kinetics/src/fokker_planck_calculus.jl b/moment_kinetics/src/fokker_planck_calculus.jl
index 88dfe7f66..83dcadb45 100644
--- a/moment_kinetics/src/fokker_planck_calculus.jl
+++ b/moment_kinetics/src/fokker_planck_calculus.jl
@@ -41,9 +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, allocate_sparse_matrix_constructor,
+using ..sparse_matrix_functions: icsc_func, ic_func, allocate_sparse_matrix_constructor,
                                  assemble_constructor_data!, assign_constructor_data!,
-                                 sparse_matrix_constructor, create_sparse_matrix
+                                 sparse_matrix_constructor, create_sparse_matrix,
+                                 get_global_compound_index
 using moment_kinetics.gauss_legendre: get_QQ_local!
 using Dates
 using SpecialFunctions: ellipk, ellipe
@@ -884,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)
 
@@ -1128,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/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index 01f983350..b642b21d3 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -22,7 +22,7 @@ include("array_allocation.jl")
 include("interpolation.jl")
 include("calculus.jl")
 include("lagrange_polynomials.jl")
-include("sparse_matrix_functions.jl")
+
 include("clenshaw_curtis.jl")
 include("gauss_legendre.jl")
 include("fourier.jl")
@@ -35,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")
diff --git a/moment_kinetics/src/sparse_matrix_functions.jl b/moment_kinetics/src/sparse_matrix_functions.jl
index 24bb05bcc..6e2f429c5 100644
--- a/moment_kinetics/src/sparse_matrix_functions.jl
+++ b/moment_kinetics/src/sparse_matrix_functions.jl
@@ -3,15 +3,18 @@ module for storing sparse matrix worker functions
 """
 module sparse_matrix_functions
 
-export icsc_func
+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
@@ -38,6 +41,39 @@ function icsc_func(ix_local::mk_int,ixp_local::mk_int,
     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

From eab874259cf9bbd8a64beb998dca9d31ac3c0581 Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Tue, 17 Sep 2024 16:02:25 +0100
Subject: [PATCH 16/19] Initial Poisson solver for (polar,radial) problem using
 2D Gauss-Legendre tensor product elements (not using Fourier transform).

---
 moment_kinetics/src/spatial_poisson.jl | 235 ++++++++++++++++++++++++-
 test_scripts/PoissonSolve.jl           |  85 +++++++++
 2 files changed, 318 insertions(+), 2 deletions(-)

diff --git a/moment_kinetics/src/spatial_poisson.jl b/moment_kinetics/src/spatial_poisson.jl
index 9020029b3..b314fdab7 100644
--- a/moment_kinetics/src/spatial_poisson.jl
+++ b/moment_kinetics/src/spatial_poisson.jl
@@ -8,6 +8,8 @@ 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!
@@ -18,7 +20,12 @@ 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}
@@ -187,4 +194,228 @@ function spatial_poisson_solve!(phi,rho,poisson_arrays,radial,polar,polar_spectr
    return nothing
 end
 
-end
\ No newline at end of file
+# 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)
+   npolar = polar.n
+   nradial = radial.n
+   LP2D_sparse, MR2D_sparse = init_sparse_laplacian2D(radial,polar,radial_spectral,polar_spectral)
+   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 
+"""
+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
+                                    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[ipolar_local,ipolarp_local]*
+                                             (KJradial[iradial_local,iradialp_local] - 
+                                              PPradial[iradial_local,iradialp_local]) +
+                                             KKpolar[ipolar_local,ipolarp_local]*
+                                             MNradial[iradial_local,iradialp_local]))    
+                                end
+                            end # Laplcian 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
+                                    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[ipolar_local,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
+
+"""
+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/test_scripts/PoissonSolve.jl b/test_scripts/PoissonSolve.jl
index 3d15347a4..0bb7af443 100644
--- a/test_scripts/PoissonSolve.jl
+++ b/test_scripts/PoissonSolve.jl
@@ -10,6 +10,7 @@ using moment_kinetics.input_structs: grid_input, advection_input
 using moment_kinetics.coordinates: define_coordinate
 using moment_kinetics.type_definitions: mk_float, mk_int
 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!
@@ -127,9 +128,93 @@ function run_poisson_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelemen
    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,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 
+   # fd_option and adv_input not actually used so given values unimportant
+   fd_option = "fourth_order_centered"
+   cheb_option = "matrix"
+   adv_input = advection_input("default", 1.0, 0.0, 0.0)
+   nrank = 1
+   irank = 0
+   comm = MPI.COMM_NULL
+   # create the 'input' struct containing input info needed to create a
+   # coordinate
+   element_spacing_option = "uniform"
+   polar_input = grid_input("polar", ngrid_polar, nelement_global_polar, nelement_local_polar, 
+      nrank, irank, Lpolar, "gausslegendre_pseudospectral", fd_option, cheb_option, bc, adv_input,comm,element_spacing_option)
+   radial_input = grid_input("r", ngrid_radial, nelement_global_radial, nelement_local_radial, 
+      nrank, irank, Lradial, "gausslegendre_pseudospectral", fd_option, cheb_option, bc, adv_input,comm,element_spacing_option)
+   println("made inputs")
+  
+   # Set up MPI
+   initialize_comms!()
+   setup_distributed_memory_MPI(1,1,1,1)
+   polar, polar_spectral = define_coordinate(polar_input)
+   radial, radial_spectral = define_coordinate(radial_input)
+   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()
+  
+   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): ",phi[1,:])
+      println("phi(theta=pi): ",phi[end,:])
+   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)
+         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_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) * r^(kk-1): ",maximum(err_phi))
+   if plot_results
+      plot_test_data(exact_phi,phi,err_phi,"phi_2",radial,polar)
+   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

From 17f61558d0d491131fc2a88066530b0398ec34fa Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Wed, 18 Sep 2024 09:21:44 +0100
Subject: [PATCH 17/19] Correction to make nelement_polar = 1 solve work for
 arbitrary phase of periodic function -- make sure data for lower row assembly
 comes from correct location.

---
 moment_kinetics/src/spatial_poisson.jl | 16 ++++++++++++----
 test_scripts/PoissonSolve.jl           | 12 +++++++-----
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/moment_kinetics/src/spatial_poisson.jl b/moment_kinetics/src/spatial_poisson.jl
index b314fdab7..4b4f23c82 100644
--- a/moment_kinetics/src/spatial_poisson.jl
+++ b/moment_kinetics/src/spatial_poisson.jl
@@ -321,6 +321,10 @@ function init_sparse_laplacian2D(radial::coordinate, polar::coordinate,
                                              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,
@@ -328,13 +332,13 @@ function init_sparse_laplacian2D(radial::coordinate, polar::coordinate,
                                            ielement_radial,
                                            ngrid_radial,nelement_radial)
                                     assemble_constructor_data!(LP2D,icsc_bc,ic_global,icp_global_bc,
-                                            (MMpolar[ipolar_local,ipolarp_local]*
+                                            (MMpolar[ngrid_polar,ipolarp_local]*
                                              (KJradial[iradial_local,iradialp_local] - 
                                               PPradial[iradial_local,iradialp_local]) +
-                                             KKpolar[ipolar_local,ipolarp_local]*
+                                             KKpolar[ngrid_polar,ipolarp_local]*
                                              MNradial[iradial_local,iradialp_local]))    
                                 end
-                            end # Laplcian assembly
+                            end # Laplacian assembly
                             
                             # mass matrices for RHS of matrix equation (enforce periodicity only)
                             if upper_boundary_row_polar
@@ -363,6 +367,10 @@ function init_sparse_laplacian2D(radial::coordinate, polar::coordinate,
                                              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,
@@ -370,7 +378,7 @@ function init_sparse_laplacian2D(radial::coordinate, polar::coordinate,
                                            ielement_radial,
                                            ngrid_radial,nelement_radial)
                                     assemble_constructor_data!(MR2D,icsc_bc,ic_global,icp_global_bc,
-                                            (MMpolar[ipolar_local,ipolarp_local]*
+                                            (MMpolar[ngrid_polar,ipolarp_local]*
                                              MRradial[iradial_local,iradialp_local]))    
                                 end # mass matrix assembly
                             end
diff --git a/test_scripts/PoissonSolve.jl b/test_scripts/PoissonSolve.jl
index 0bb7af443..3e595cec2 100644
--- a/test_scripts/PoissonSolve.jl
+++ b/test_scripts/PoissonSolve.jl
@@ -128,7 +128,7 @@ function run_poisson_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelemen
    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,plot_results=false)
+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=0.0,plot_results=false)
 
    nelement_local_polar = nelement_polar # number of elements per rank
    nelement_global_polar = nelement_local_polar # total number of elements 
@@ -180,8 +180,8 @@ function run_poisson2D_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelem
    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): ",phi[1,:])
-      println("phi(theta=pi): ",phi[end,:])
+      println("phi(theta=-pi,r=0): ",phi[1,1])
+      println("phi(theta=pi,r=0): ",phi[end,1])
    end
    
    if kk < 1
@@ -196,8 +196,8 @@ function run_poisson2D_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelem
    
    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)
+         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
    
@@ -206,6 +206,8 @@ function run_poisson2D_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelem
    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_2",radial,polar)
+      println("phi(theta=-pi,r=0): ",phi[1,1])
+      println("phi(theta=pi,r=0): ",phi[end,1])
    end
    finalize_comms!()
    return nothing

From 8acf3dc4cfad634b0453d55da4262fbde67a763d Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Wed, 18 Sep 2024 12:20:46 +0100
Subject: [PATCH 18/19] Add assembly routine for 2D Poisson equation that
 constructs a dense matrix using ic_global and icp_global indices only (no
 icsc index). This assembly is slower for larger resolution due to memory
 usage, but can be done correctly with existing indexing functions. Update
 automatic tests to test the 2D Poisson solve as well as the 1D+Fourier
 decomposition version.

---
 moment_kinetics/src/spatial_poisson.jl        | 162 +++++++++++++++++-
 .../test/poisson_radial_polar_tests.jl        | 159 ++++++++++++++++-
 test_scripts/PoissonSolve.jl                  |   2 +-
 3 files changed, 311 insertions(+), 12 deletions(-)

diff --git a/moment_kinetics/src/spatial_poisson.jl b/moment_kinetics/src/spatial_poisson.jl
index 4b4f23c82..26054a0aa 100644
--- a/moment_kinetics/src/spatial_poisson.jl
+++ b/moment_kinetics/src/spatial_poisson.jl
@@ -216,10 +216,15 @@ Poisson's equation problem
 """
 function init_spatial_poisson2D(radial::coordinate, polar::coordinate, 
                                 radial_spectral::weak_discretization_info,
-                                polar_spectral::weak_discretization_info)
+                                polar_spectral::weak_discretization_info;
+                                use_sparse_init=false)
    npolar = polar.n
    nradial = radial.n
-   LP2D_sparse, MR2D_sparse = init_sparse_laplacian2D(radial,polar,radial_spectral,polar_spectral)
+   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)
@@ -228,7 +233,14 @@ end
 """
 This function makes sparse matrices representing the 2D Laplacian operator.
 The compound index runs over the radial and polar coordinates. 
-We use 
+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,
@@ -394,6 +406,150 @@ function init_sparse_laplacian2D(radial::coordinate, polar::coordinate,
     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
diff --git a/moment_kinetics/test/poisson_radial_polar_tests.jl b/moment_kinetics/test/poisson_radial_polar_tests.jl
index 92ca7dc65..69d61aef9 100644
--- a/moment_kinetics/test/poisson_radial_polar_tests.jl
+++ b/moment_kinetics/test/poisson_radial_polar_tests.jl
@@ -12,20 +12,21 @@ using moment_kinetics.input_structs: grid_input, advection_input
 using moment_kinetics.coordinates: define_coordinate
 using moment_kinetics.type_definitions: mk_float, mk_int
 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_tests" begin
-        println("poisson_radial_polar_tests")
+    @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_test(nelement_radial=nelement_radial,ngrid_radial=ngrid_radial,Lradial=Lradial,ngrid_polar=ngrid_polar,kk=kk,
+             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
@@ -34,7 +35,7 @@ function run_test()
         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_test(nelement_radial=nelement_radial,ngrid_radial=ngrid_radial,Lradial=Lradial,ngrid_polar=ngrid_polar,kk=kk,
+             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
@@ -43,7 +44,7 @@ function run_test()
         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_test(nelement_radial=nelement_radial,ngrid_radial=ngrid_radial,Lradial=Lradial,ngrid_polar=ngrid_polar,kk=kk,
+             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
@@ -52,7 +53,7 @@ function run_test()
         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_test(nelement_radial=nelement_radial,ngrid_radial=ngrid_radial,Lradial=Lradial,ngrid_polar=ngrid_polar,kk=kk,
+             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
@@ -61,13 +62,72 @@ function run_test()
         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_test(nelement_radial=nelement_radial,ngrid_radial=ngrid_radial,Lradial=Lradial,ngrid_polar=ngrid_polar,kk=kk,
+             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=4.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=7.0e-14,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_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelement_polar=1,ngrid_polar=1,Lpolar=2.0*pi,kk::mk_int=1,
+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
@@ -170,6 +230,89 @@ function run_poisson_radial_polar_test(; nelement_radial=5,ngrid_radial=5,Lradia
    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, not used, set to "none" to avoid extra BC impositions 
+   # fd_option and adv_input not actually used so given values unimportant
+   fd_option = "fourth_order_centered"
+   cheb_option = "matrix"
+   adv_input = advection_input("default", 1.0, 0.0, 0.0)
+   nrank = 1
+   irank = 0
+   comm = MPI.COMM_NULL
+   # create the 'input' struct containing input info needed to create a
+   # coordinate
+   element_spacing_option = "uniform"
+   polar_input = grid_input("polar", ngrid_polar, nelement_global_polar, nelement_local_polar, 
+      nrank, irank, Lpolar, "gausslegendre_pseudospectral", fd_option, cheb_option, bc, adv_input,comm,element_spacing_option)
+   radial_input = grid_input("r", ngrid_radial, nelement_global_radial, nelement_local_radial, 
+      nrank, irank, Lradial, "gausslegendre_pseudospectral", fd_option, cheb_option, bc, adv_input,comm,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(polar_input, ignore_MPI=true)
+   radial, radial_spectral = define_coordinate(radial_input)
+   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
diff --git a/test_scripts/PoissonSolve.jl b/test_scripts/PoissonSolve.jl
index 3e595cec2..1cf455dff 100644
--- a/test_scripts/PoissonSolve.jl
+++ b/test_scripts/PoissonSolve.jl
@@ -203,7 +203,7 @@ function run_poisson2D_test(; nelement_radial=5,ngrid_radial=5,Lradial=1.0,nelem
    
    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) * r^(kk-1): ",maximum(err_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])

From 79d24e73ec58313d4aff0a9bdf48a6da62eca75c Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Wed, 18 Sep 2024 12:50:38 +0100
Subject: [PATCH 19/19] Relax error tolerances for MacOS.

---
 moment_kinetics/test/poisson_radial_polar_tests.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/moment_kinetics/test/poisson_radial_polar_tests.jl b/moment_kinetics/test/poisson_radial_polar_tests.jl
index 69d61aef9..379edaf2e 100644
--- a/moment_kinetics/test/poisson_radial_polar_tests.jl
+++ b/moment_kinetics/test/poisson_radial_polar_tests.jl
@@ -78,7 +78,7 @@ function run_test()
         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=4.0e-14,atol_phi2=9.0e-7)
+                              atol_phi1=7.0e-14,atol_phi2=9.0e-7)
         end
         nelement_radial=5
         ngrid_radial=5
@@ -89,7 +89,7 @@ function run_test()
         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=5.0e-11)
+                              atol_phi1=3.0e-13,atol_phi2=5.0e-11)
         end
         nelement_radial=5
         ngrid_radial=9