Start writing automated tests for MMS; renormalise manufactured distribution functions

src/chebyshev.jl

@@ -322,7 +322,7 @@ with all grid points for Clenshaw-Curtis quadrature
 """
 function clenshaw_curtis_weights(ngrid, nelement, n, imin, imax, scale_factor)
     # create array containing the integration weights
-    wgts = zeros(n)
+    wgts = zeros(mk_float, n)
     # calculate the modified Chebshev moments of the first kind
     μ = chebyshevmoments(ngrid)
     @inbounds begin

src/coordinates.jl

@@ -7,15 +7,15 @@ export equally_spaced_grid
 
 using ..type_definitions: mk_float, mk_int
 using ..array_allocation: allocate_float, allocate_int
-using ..file_io: open_output_file
-using ..chebyshev: scaled_chebyshev_grid
+using ..calculus: derivative!
+using ..chebyshev: scaled_chebyshev_grid, setup_chebyshev_pseudospectral
 using ..quadrature: composite_simpson_weights
 using ..input_structs: advection_input
 
 """
 structure containing basic information related to coordinates
 """
-struct coordinate
+struct coordinate{Tadvection<:Union{advection_input,Nothing}}
     # name is the name of the variable associated with this coordiante
     name::String
     # n is the total number of grid points associated with this coordinate
@@ -60,7 +60,7 @@ struct coordinate
     # nelement entries
     scratch_2d::Array{mk_float,2}
     # struct containing advection speed options/inputs
-    advection::advection_input
+    advection::Tadvection
 end
 
 """
@@ -95,10 +95,25 @@ function define_coordinate(input, composition=nothing)
     # struct containing the advection speed options/inputs for this coordinate
     advection = input.advection
 
-    return coordinate(input.name, n, input.ngrid, input.nelement, input.L, grid,
+    coord = coordinate(input.name, n, input.ngrid, input.nelement, input.L, grid,
         cell_width, igrid, ielement, imin, imax, input.discretization, input.fd_option,
         input.bc, wgts, uniform_grid, duniform_dgrid, scratch, copy(scratch),
         scratch_2d, advection)
+
+    if input.discretization == "chebyshev_pseudospectral" && coord.ngrid > 1
+        # create arrays needed for explicit Chebyshev pseudospectral treatment in this
+        # coordinate and create the plans for the forward and backward fast Chebyshev
+        # transforms
+        spectral = setup_chebyshev_pseudospectral(coord)
+        # obtain the local derivatives of the uniform grid with respect to the used grid
+        derivative!(coord.duniform_dgrid, coord.uniform_grid, coord, spectral)
+    else
+        # create dummy Bool variable to return in place of the above struct
+        spectral = false
+        coord.duniform_dgrid .= 1.0
+    end
+
+    return coord, spectral
 end
 
 """

src/file_io.jl

@@ -10,6 +10,7 @@ export write_data_to_binary
 
 using NCDatasets
 using ..communication: _block_synchronize
+using ..coordinates: coordinate
 using ..looping
 using ..moment_kinetics_structs: scratch_pdf, em_fields_struct
 using ..type_definitions: mk_float, mk_int
@@ -142,6 +143,10 @@ function define_dimensions!(fid, nvz, nvr, nvzeta, nvpa, nvperp, nz, nr, n_speci
     end
     # define the time dimension, with an expandable size (denoted by Inf)
     defDim(fid, "ntime", Inf)
+    # define a length-1 dimension for storing strings. Don't know why they cannot be
+    # stored as scalars, maybe did not find the right function/method, maybe a missing
+    # feature or bug in NCDatasets.jl?
+    defDim(fid, "str_dim", 1)
 
     return nothing
 end
@@ -152,97 +157,71 @@ end
 Define static (i.e. time-independent) variables for an output file.
 """
 function define_static_variables!(fid,vz,vr,vzeta,vpa,vperp,z,r,composition,collisions)
-    # create and write the "r" variable to file
-    varname = "r"
-    attributes = Dict("description" => "radial coordinate")
-    dims = ("nr",)
-    vartype = mk_float
-    var = defVar(fid, varname, vartype, dims, attrib=attributes)
-    var[:] = r.grid
-    # create and write the "r_wgts" variable to file
-    varname = "r_wgts"
-    attributes = Dict("description" => "integration weights for radial coordinate")
-    vartype = mk_float
-    var = defVar(fid, varname, vartype, dims, attrib=attributes)
-    var[:] = r.wgts
-    # create and write the "z" variable to file
-    varname = "z"
-    attributes = Dict("description" => "parallel coordinate")
-    dims = ("nz",)
-    vartype = mk_float
-    var = defVar(fid, varname, vartype, dims, attrib=attributes)
-    var[:] = z.grid
-    # create and write the "z_wgts" variable to file
-    varname = "z_wgts"
-    attributes = Dict("description" => "integration weights for parallel coordinate")
-    vartype = mk_float
-    var = defVar(fid, varname, vartype, dims, attrib=attributes)
-    var[:] = z.wgts
-    # create and write the "vperp" variable to file
-    varname = "vperp"
-    attributes = Dict("description" => "parallel velocity")
-    dims = ("nvperp",)
-    vartype = mk_float
-    var = defVar(fid, varname, vartype, dims, attrib=attributes)
-    var[:] = vperp.grid
-    # create and write the "vperp_wgts" variable to file
-    varname = "vperp_wgts"
-    attributes = Dict("description" => "integration weights for parallel velocity coordinate")
-    vartype = mk_float
-    var = defVar(fid, varname, vartype, dims, attrib=attributes)
-    var[:] = vperp.wgts
-    # create and write the "vpa" variable to file
-    varname = "vpa"
-    attributes = Dict("description" => "parallel velocity")
-    dims = ("nvpa",)
-    vartype = mk_float
-    var = defVar(fid, varname, vartype, dims, attrib=attributes)
-    var[:] = vpa.grid
-    # create and write the "vpa_wgts" variable to file
-    varname = "vpa_wgts"
-    attributes = Dict("description" => "integration weights for parallel velocity coordinate")
-    vartype = mk_float
-    var = defVar(fid, varname, vartype, dims, attrib=attributes)
-    var[:] = vpa.wgts
-    # create and write the "vzeta" variable to file
-    varname = "vzeta"
-    attributes = Dict("description" => "parallel velocity")
-    dims = ("nvzeta",)
-    vartype = mk_float
-    var = defVar(fid, varname, vartype, dims, attrib=attributes)
-    var[:] = vzeta.grid
-    # create and write the "vzeta_wgts" variable to file
-    varname = "vzeta_wgts"
-    attributes = Dict("description" => "integration weights for parallel velocity coordinate")
-    vartype = mk_float
-    var = defVar(fid, varname, vartype, dims, attrib=attributes)
-    var[:] = vzeta.wgts
-    # create and write the "vr" variable to file
-    varname = "vr"
-    attributes = Dict("description" => "parallel velocity")
-    dims = ("nvr",)
-    vartype = mk_float
-    var = defVar(fid, varname, vartype, dims, attrib=attributes)
-    var[:] = vr.grid
-    # create and write the "vr_wgts" variable to file
-    varname = "vr_wgts"
-    attributes = Dict("description" => "integration weights for parallel velocity coordinate")
-    vartype = mk_float
-    var = defVar(fid, varname, vartype, dims, attrib=attributes)
-    var[:] = vr.wgts
-    # create and write the "vz" variable to file
-    varname = "vz"
-    attributes = Dict("description" => "parallel velocity")
-    dims = ("nvz",)
-    vartype = mk_float
-    var = defVar(fid, varname, vartype, dims, attrib=attributes)
-    var[:] = vz.grid
-    # create and write the "vz_wgts" variable to file
-    varname = "vz_wgts"
-    attributes = Dict("description" => "integration weights for parallel velocity coordinate")
-    vartype = mk_float
-    var = defVar(fid, varname, vartype, dims, attrib=attributes)
-    var[:] = vz.wgts
+    function save_coordinate(coord::coordinate, description::String)
+        # Create and write the grid for coord
+        varname = coord.name
+        attributes = Dict("description" => description)
+        dims = ("n$(coord.name)",)
+        vartype = mk_float
+        var = defVar(fid, varname, vartype, dims, attrib=attributes)
+        var[:] = coord.grid
+
+        # create and write the weights for coord
+        varname = "$(coord.name)_wgts"
+        attributes = Dict("description" => "integration weights for $(coord.name) coordinate")
+        vartype = mk_float
+        var = defVar(fid, varname, vartype, dims, attrib=attributes)
+        var[:] = coord.wgts
+
+        # create and write ngrid for coord
+        varname = "$(coord.name)_ngrid"
+        attributes = Dict("description" => "ngrid for $(coord.name) coordinate")
+        dims = ()
+        vartype = mk_int
+        var = defVar(fid, varname, vartype, dims, attrib=attributes)
+        var[:] = coord.ngrid
+
+        # create and write nelement for coord
+        varname = "$(coord.name)_nelement"
+        attributes = Dict("description" => "nelement for $(coord.name) coordinate")
+        dims = ()
+        vartype = mk_int
+        var = defVar(fid, varname, vartype, dims, attrib=attributes)
+        var[:] = coord.nelement
+
+        # create and write discretization for coord
+        varname = "$(coord.name)_discretization"
+        attributes = Dict("description" => "discretization for $(coord.name) coordinate")
+        dims = ("str_dim",)
+        vartype = String
+        var = defVar(fid, varname, vartype, dims, attrib=attributes)
+        var[:] = coord.discretization
+
+        # create and write fd_option for coord
+        varname = "$(coord.name)_fd_option"
+        attributes = Dict("description" => "fd_option for $(coord.name) coordinate")
+        dims = ("str_dim",)
+        vartype = String
+        var = defVar(fid, varname, vartype, dims, attrib=attributes)
+        var[:] = coord.fd_option
+
+        # create and write bc for coord
+        varname = "$(coord.name)_bc"
+        attributes = Dict("description" => "bc for $(coord.name) coordinate")
+        dims = ("str_dim",)
+        vartype = String
+        var = defVar(fid, varname, vartype, dims, attrib=attributes)
+        var[:] = coord.bc
+    end
+
+    # create and write the coordinate variables
+    save_coordinate(r, "radial coordinate")
+    save_coordinate(z, "parallel coordinate")
+    save_coordinate(vperp, "perpendicular velocity")
+    save_coordinate(vpa, "parallel velocity")
+    save_coordinate(vzeta, "neutral binormal velocity")
+    save_coordinate(vr, "neutral radial velocity")
+    save_coordinate(vz, "neutral parallel velocity")
     # create and write the "T_e" variable to file
     varname = "T_e"
     attributes = Dict("description" => "electron temperature")

src/initial_conditions.jl

@@ -184,24 +184,18 @@ function init_pdf!(pdf, moments, vz, vr, vzeta, vpa, vperp, z, r, n_ion_species,
 end
 
 function init_pdf_moments_manufactured_solns!(pdf, moments, vz, vr, vzeta, vpa, vperp, z, r, n_ion_species, n_neutral_species, geometry,composition)
-    manufactured_solns_list = manufactured_solutions(r.L,z.L,r.bc,z.bc,geometry,composition,r.n) 
+    manufactured_solns_list = manufactured_solutions(r.L,z.L,vperp.L,vpa.L,vzeta.L,vr.L,vz.L,r.bc,z.bc,geometry,composition,r.n) 
     dfni_func = manufactured_solns_list.dfni_func
     densi_func = manufactured_solns_list.densi_func
     dfnn_func = manufactured_solns_list.dfnn_func
     densn_func = manufactured_solns_list.densn_func
     #nb manufactured functions not functions of species
-    for is in 1:n_ion_species
-        for ir in 1:r.n
-            for iz in 1:z.n
-                moments.charged.dens[iz,ir,is] = densi_func(z.grid[iz],r.grid[ir],0.0)
-                for ivperp in 1:vperp.n
-                    for ivpa in 1:vpa.n
-                        pdf.charged.unnorm[ivpa,ivperp,iz,ir,is] = dfni_func(vpa.grid[ivpa],vperp.grid[ivperp],z.grid[iz],r.grid[ir],0.0)
-                        pdf.charged.norm[ivpa,ivperp,iz,ir,is] = pdf.charged.unnorm[ivpa,ivperp,iz,ir,is]
-
-                    end
-                end
-            end
+    begin_s_r_z_region()
+    @loop_s_r_z is ir iz begin
+        moments.charged.dens[iz,ir,is] = densi_func(z.grid[iz],r.grid[ir],0.0)
+        @loop_vperp_vpa ivperp ivpa begin
+            pdf.charged.unnorm[ivpa,ivperp,iz,ir,is] = dfni_func(vpa.grid[ivpa],vperp.grid[ivperp],z.grid[iz],r.grid[ir],0.0)
+            pdf.charged.norm[ivpa,ivperp,iz,ir,is] = pdf.charged.unnorm[ivpa,ivperp,iz,ir,is]
         end
     end
     # update upar, ppar, qpar, vth consistent with manufactured solns
@@ -210,27 +204,20 @@ function init_pdf_moments_manufactured_solns!(pdf, moments, vz, vr, vzeta, vpa,
     # get particle flux
     update_upar!(moments.charged.upar, pdf.charged.unnorm, vpa, vperp, z, r, composition)
     # convert from particle particle flux to parallel flow
+    begin_s_r_z_region()
     @loop_s_r_z is ir iz begin
         moments.charged.upar[iz,ir,is] /= moments.charged.dens[iz,ir,is]
     # update the thermal speed
         moments.charged.vth[iz,ir,is] = sqrt(2.0*moments.charged.ppar[iz,ir,is]/moments.charged.dens[iz,ir,is])
     end
 
     if n_neutral_species > 0
-        for isn in 1:n_neutral_species
-            for ir in 1:r.n
-                for iz in 1:z.n
-                    moments.neutral.dens[iz,ir,isn] = densn_func(z.grid[iz],r.grid[ir],0.0)
-                    for ivzeta in 1:vzeta.n
-                        for ivr in 1:vr.n
-                            for ivz in 1:vz.n
-                                pdf.neutral.unnorm[ivz,ivr,ivzeta,iz,ir,isn] = dfnn_func(vz.grid[ivz],vr.grid[ivr],vzeta.grid[ivzeta],z.grid[iz],r.grid[ir],0.0)
-                                pdf.neutral.norm[ivz,ivr,ivzeta,iz,ir,isn] = pdf.neutral.unnorm[ivz,ivr,ivzeta,iz,ir,isn]
-                            end
-                        end
-                    end
-
-                end
+        begin_sn_r_z_region()
+        @loop_sn_r_z isn ir iz begin
+            moments.neutral.dens[iz,ir,isn] = densn_func(z.grid[iz],r.grid[ir],0.0)
+            @loop_vzeta_vr_vz ivzeta ivr ivz begin
+                pdf.neutral.unnorm[ivz,ivr,ivzeta,iz,ir,isn] = dfnn_func(vz.grid[ivz],vr.grid[ivr],vzeta.grid[ivzeta],z.grid[iz],r.grid[ir],0.0)
+                pdf.neutral.norm[ivz,ivr,ivzeta,iz,ir,isn] = pdf.neutral.unnorm[ivz,ivr,ivzeta,iz,ir,isn]
             end
         end
         # get consistent moments with manufactured solutions 
@@ -241,8 +228,9 @@ function init_pdf_moments_manufactured_solns!(pdf, moments, vz, vr, vzeta, vpa,
         update_neutral_pzeta!(moments.neutral.pzeta, pdf.neutral.unnorm, vz, vr, vzeta, z, r, composition)
         #update ptot (isotropic pressure)
         if r.n > 1 #if 2D geometry
+            begin_sn_r_z_region()
             @loop_sn_r_z isn ir iz begin            
-            moments.neutral.ptot[iz,ir,isn] = (moments.neutral.pz[iz,ir,isn] + moments.neutral.pr[iz,ir,isn] + moments.neutral.pzeta[iz,ir,isn])/3.0
+                moments.neutral.ptot[iz,ir,isn] = (moments.neutral.pz[iz,ir,isn] + moments.neutral.pr[iz,ir,isn] + moments.neutral.pzeta[iz,ir,isn])/3.0
             end
         else #1D model
             moments.neutral.ptot .= moments.neutral.pz
@@ -252,6 +240,7 @@ function init_pdf_moments_manufactured_solns!(pdf, moments, vz, vr, vzeta, vpa,
         update_neutral_ur!(moments.neutral.ur, pdf.neutral.unnorm, vz, vr, vzeta, z, r, composition)
         update_neutral_uzeta!(moments.neutral.uzeta, pdf.neutral.unnorm, vz, vr, vzeta, z, r, composition)
         # now convert from particle particle flux to parallel flow
+        begin_sn_r_z_region()
         @loop_sn_r_z isn ir iz begin
             moments.neutral.uz[iz,ir,isn] /= moments.neutral.dens[iz,ir,isn]
             moments.neutral.ur[iz,ir,isn] /= moments.neutral.dens[iz,ir,isn]

src/input_structs.jl

@@ -38,6 +38,26 @@ struct time_input
     use_manufactured_solns::Bool
 end
 
+"""
+"""
+mutable struct advance_info
+    vpa_advection::Bool
+    z_advection::Bool
+    r_advection::Bool
+    neutral_z_advection::Bool
+    neutral_r_advection::Bool
+    cx_collisions::Bool
+    cx_collisions_1V::Bool
+    ionization_collisions::Bool
+    ionization_collisions_1V::Bool
+    source_terms::Bool
+    continuity::Bool
+    force_balance::Bool
+    energy::Bool
+    rk_coefs::Array{mk_float,2}
+    manufactured_solns_test::Bool
+end
+
 """
 """
 mutable struct advection_input_mutable
@@ -92,7 +112,7 @@ end
 
 """
 """
-struct grid_input
+struct grid_input{Tadvection<:Union{advection_input,Nothing}}
     # name of the variable associated with this coordinate
     name::String
     # number of grid points per element
@@ -108,7 +128,7 @@ struct grid_input
     # boundary option
     bc::String
     # struct containing advection speed options
-    advection::advection_input
+    advection::Tadvection
 end
 
 """