Initial functionality for partially-implicit 'IMEX' timestepping

moment_kinetics/Project.toml

@@ -17,6 +17,7 @@ LegendrePolynomials = "3db4a2ba-fc88-11e8-3e01-49c72059a882"
 LibGit2 = "76f85450-5226-5b5a-8eaa-529ad045b433"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LsqFit = "2fda8390-95c7-5789-9bda-21331edee243"
+MINPACK = "4854310b-de5a-5eb6-a2a5-c1dee2bd17f9"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267"
 Measures = "442fdcdd-2543-5da2-b0f3-8c86c306513e"

moment_kinetics/src/calculus.jl

@@ -362,7 +362,10 @@ function assign_endpoint!(df1d::AbstractArray{mk_float,Ndims},
     # test against coord name -- make sure to use exact string delimiters e.g. "x" not 'x'
     # test against Ndims (autodetermined) to choose which array slices to use in assigning endpoints
     #println("DEBUG MESSAGE: coord.name: ",coord.name," Ndims: ",Ndims," key: ",key)
-    if coord.name == "z" && Ndims==2
+    if coord.name == "z" && Ndims==1
+        df1d[j] = receive_buffer[]
+        #println("ASSIGNING DATA")
+    elseif coord.name == "z" && Ndims==2
         df1d[j,:] .= receive_buffer[:]
         #println("ASSIGNING DATA")
     elseif coord.name == "z" && Ndims==3
@@ -374,6 +377,9 @@ function assign_endpoint!(df1d::AbstractArray{mk_float,Ndims},
     elseif coord.name == "z" && Ndims==6
         df1d[:,:,:,j,:,:] .= receive_buffer[:,:,:,:,:]
         #println("ASSIGNING DATA")
+    elseif coord.name == "r" && Ndims==1
+        df1d[j] = receive_buffer[]
+        #println("ASSIGNING DATA")
     elseif coord.name == "r" && Ndims==2
         df1d[:,j] .= receive_buffer[:]
         #println("ASSIGNING DATA")

moment_kinetics/src/charge_exchange.jl

@@ -9,12 +9,12 @@ using ..looping
 using ..interpolation: interpolate_to_grid_vpa!
 
 """
-update the evolved pdf for each ion and electron species to account for
-charge exchange collisions between ions and neutrals
+update the evolved pdf for each ion species to account for charge exchange collisions
+between ions and neutrals
 """
-function charge_exchange_collisions_1V!(f_out, f_neutral_out, fvec_in, moments,
-                                        composition, vpa, vz, charge_exchange_frequency,
-                                        vpa_spectral, vz_spectral, dt)
+function ion_charge_exchange_collisions_1V!(f_out, fvec_in, moments, composition, vpa, vz,
+                                            charge_exchange_frequency, vpa_spectral,
+                                            vz_spectral, dt)
     # This routine assumes a 1D model with:
     # nvz = nvpa and identical vz and vpa grids 
 
@@ -32,19 +32,6 @@ function charge_exchange_collisions_1V!(f_out, f_neutral_out, fvec_in, moments,
                 moments.neutral.vth[:,:,is], moments, vpa, vz, charge_exchange_frequency,
                 vz_spectral, dt)
         end
-
-        begin_sn_r_z_region(no_synchronize=true)
-        @loop_sn isn begin
-            # apply CX collisions to all neutral species
-            # for each neutral species, obtain affect of charge exchange collisions
-            # with the corresponding ion species
-            @views charge_exchange_collisions_single_species!(
-                f_neutral_out[:,1,1,:,:,isn], fvec_in.pdf_neutral[:,1,1,:,:,isn],
-                fvec_in.pdf[:,1,:,:,isn], fvec_in.density[:,:,isn],
-                fvec_in.uz_neutral[:,:,isn], fvec_in.upar[:,:,isn],
-                moments.neutral.vth[:,:,isn], moments.ion.vth[:,:,isn], moments,
-                vz, vpa, charge_exchange_frequency, vpa_spectral, dt)
-        end
     else
         begin_s_r_z_region()
         @loop_s is begin
@@ -58,8 +45,35 @@ function charge_exchange_collisions_1V!(f_out, f_neutral_out, fvec_in, moments,
                         - fvec_in.pdf[ivpa,1,iz,ir,is]*fvec_in.density_neutral[iz,ir,is])
             end
         end
+    end
+end
+
+"""
+update the evolved pdf for each neutral species to account for charge exchange collisions
+between ions and neutrals
+"""
+function neutral_charge_exchange_collisions_1V!(f_neutral_out, fvec_in, moments,
+                                                composition, vpa, vz,
+                                                charge_exchange_frequency, vpa_spectral,
+                                                vz_spectral, dt)
+    # This routine assumes a 1D model with:
+    # nvz = nvpa and identical vz and vpa grids
 
-        begin_sn_r_z_region(no_synchronize=true)
+    if moments.evolve_density
+        begin_sn_r_z_region()
+        @loop_sn isn begin
+            # apply CX collisions to all neutral species
+            # for each neutral species, obtain affect of charge exchange collisions
+            # with the corresponding ion species
+            @views charge_exchange_collisions_single_species!(
+                f_neutral_out[:,1,1,:,:,isn], fvec_in.pdf_neutral[:,1,1,:,:,isn],
+                fvec_in.pdf[:,1,:,:,isn], fvec_in.density[:,:,isn],
+                fvec_in.uz_neutral[:,:,isn], fvec_in.upar[:,:,isn],
+                moments.neutral.vth[:,:,isn], moments.ion.vth[:,:,isn], moments,
+                vz, vpa, charge_exchange_frequency, vpa_spectral, dt)
+        end
+    else
+        begin_sn_r_z_region()
         @loop_sn isn begin
             # apply CX collisions to all neutral species
             # for each neutral species, obtain affect of charge exchange collisions
@@ -135,21 +149,10 @@ function charge_exchange_collisions_single_species!(f_out, pdf_in, pdf_other,
     end
 end
 
-function charge_exchange_collisions_3V!(f_out, f_neutral_out, f_neutral_gav_in, f_ion_vrvzvzeta_in, fvec_in, composition, vz, vr, vzeta, vpa, vperp, z, r,
-                                     charge_exchange_frequency, dt)
+function ion_charge_exchange_collisions_3V!(f_out, f_neutral_gav_in, fvec_in, composition,
+                                            vz, vr, vzeta, vpa, vperp, z, r,
+                                            charge_exchange_frequency, dt)
     # This routine assumes a 3V model with:
-    @boundscheck vz.n == size(f_neutral_out,1) || throw(BoundsError(f_neutral_out))
-    @boundscheck vr.n == size(f_neutral_out,2) || throw(BoundsError(f_neutral_out))
-    @boundscheck vzeta.n == size(f_neutral_out,3) || throw(BoundsError(f_neutral_out))
-    @boundscheck z.n == size(f_neutral_out,4) || throw(BoundsError(f_neutral_out))
-    @boundscheck r.n == size(f_neutral_out,5) || throw(BoundsError(f_neutral_out))
-    @boundscheck composition.n_neutral_species == size(f_neutral_out,6) || throw(BoundsError(f_neutral_out))
-    @boundscheck vz.n == size(f_ion_vrvzvzeta_in,1) || throw(BoundsError(f_ion_vrvzvzeta_in))
-    @boundscheck vr.n == size(f_ion_vrvzvzeta_in,2) || throw(BoundsError(f_ion_vrvzvzeta_in))
-    @boundscheck vzeta.n == size(f_ion_vrvzvzeta_in,3) || throw(BoundsError(f_ion_vrvzvzeta_in))
-    @boundscheck z.n == size(f_ion_vrvzvzeta_in,4) || throw(BoundsError(f_ion_vrvzvzeta_in))
-    @boundscheck r.n == size(f_ion_vrvzvzeta_in,5) || throw(BoundsError(f_ion_vrvzvzeta_in))
-    @boundscheck composition.n_neutral_species == size(f_ion_vrvzvzeta_in,6) || throw(BoundsError(f_ion_vrvzvzeta_in))
     @boundscheck vpa.n == size(f_out,1) || throw(BoundsError(f_out))
     @boundscheck vperp.n == size(f_out,2) || throw(BoundsError(f_out))
     @boundscheck z.n == size(f_out,3) || throw(BoundsError(f_out))
@@ -173,6 +176,26 @@ function charge_exchange_collisions_3V!(f_out, f_neutral_out, f_neutral_gav_in,
                     - fvec_in.pdf[ivpa,ivperp,iz,ir,is]*fvec_in.density_neutral[iz,ir,isn])
         end
     end
+end
+
+function neutral_charge_exchange_collisions_3V!(f_neutral_out, f_ion_vrvzvzeta_in,
+                                                fvec_in, composition, vz, vr, vzeta, vpa,
+                                                vperp, z, r, charge_exchange_frequency,
+                                                dt)
+    # This routine assumes a 3V model with:
+    @boundscheck vz.n == size(f_neutral_out,1) || throw(BoundsError(f_neutral_out))
+    @boundscheck vr.n == size(f_neutral_out,2) || throw(BoundsError(f_neutral_out))
+    @boundscheck vzeta.n == size(f_neutral_out,3) || throw(BoundsError(f_neutral_out))
+    @boundscheck z.n == size(f_neutral_out,4) || throw(BoundsError(f_neutral_out))
+    @boundscheck r.n == size(f_neutral_out,5) || throw(BoundsError(f_neutral_out))
+    @boundscheck composition.n_neutral_species == size(f_neutral_out,6) || throw(BoundsError(f_neutral_out))
+    @boundscheck vz.n == size(f_ion_vrvzvzeta_in,1) || throw(BoundsError(f_ion_vrvzvzeta_in))
+    @boundscheck vr.n == size(f_ion_vrvzvzeta_in,2) || throw(BoundsError(f_ion_vrvzvzeta_in))
+    @boundscheck vzeta.n == size(f_ion_vrvzvzeta_in,3) || throw(BoundsError(f_ion_vrvzvzeta_in))
+    @boundscheck z.n == size(f_ion_vrvzvzeta_in,4) || throw(BoundsError(f_ion_vrvzvzeta_in))
+    @boundscheck r.n == size(f_ion_vrvzvzeta_in,5) || throw(BoundsError(f_ion_vrvzvzeta_in))
+    @boundscheck composition.n_neutral_species == size(f_ion_vrvzvzeta_in,6) || throw(BoundsError(f_ion_vrvzvzeta_in))
+
     begin_sn_r_z_vzeta_vr_vz_region()
     @loop_sn_r_z_vzeta_vr_vz isn ir iz ivzeta ivr ivz begin
         # apply CX collisions to all neutral species

moment_kinetics/src/coordinates.jl

@@ -77,6 +77,14 @@ struct coordinate{T <: AbstractVector{mk_float}}
     scratch2::Array{mk_float,1}
     # scratch3 is an array used for intermediate calculations requiring n entries
     scratch3::Array{mk_float,1}
+    # scratch4 is an array used for intermediate calculations requiring n entries
+    scratch4::Array{mk_float,1}
+    # scratch5 is an array used for intermediate calculations requiring n entries
+    scratch5::Array{mk_float,1}
+    # scratch6 is an array used for intermediate calculations requiring n entries
+    scratch6::Array{mk_float,1}
+    # scratch7 is an array used for intermediate calculations requiring n entries
+    scratch7::Array{mk_float,1}
     # scratch_shared is a shared-memory array used for intermediate calculations requiring
     # n entries
     scratch_shared::T
@@ -221,10 +229,12 @@ function define_coordinate(input, parallel_io::Bool=false; run_directory=nothing
     coord = coordinate(input.name, n_global, n_local, input.ngrid,
         input.nelement_global, input.nelement_local, input.nrank, input.irank, input.L, grid,
         cell_width, igrid, ielement, imin, imax, igrid_full, input.discretization, input.fd_option, input.cheb_option,
-        input.bc, wgts, uniform_grid, duniform_dgrid, scratch, copy(scratch), copy(scratch), scratch_shared, scratch_shared2,
-        scratch_2d, copy(scratch_2d), advection, send_buffer, receive_buffer, input.comm,
-        local_io_range, global_io_range, element_scale, element_shift, input.element_spacing_option,
-        element_boundaries, radau_first_element, other_nodes, one_over_denominator)
+        input.bc, wgts, uniform_grid, duniform_dgrid, scratch, copy(scratch),
+        copy(scratch), copy(scratch), copy(scratch), copy(scratch), copy(scratch),
+        scratch_shared, scratch_shared2, scratch_2d, copy(scratch_2d), advection,
+        send_buffer, receive_buffer, input.comm, local_io_range, global_io_range,
+        element_scale, element_shift, input.element_spacing_option, element_boundaries,
+        radau_first_element, other_nodes, one_over_denominator)
 
     if coord.n == 1 && occursin("v", coord.name)
         spectral = null_velocity_dimension_info()
@@ -242,7 +252,8 @@ function define_coordinate(input, parallel_io::Bool=false; run_directory=nothing
     elseif input.discretization == "gausslegendre_pseudospectral"
         # create arrays needed for explicit GaussLegendre pseudospectral treatment in this
         # coordinate and create the matrices for differentiation
-        spectral = setup_gausslegendre_pseudospectral(coord, collision_operator_dim=collision_operator_dim)
+        spectral = setup_gausslegendre_pseudospectral(coord, collision_operator_dim=collision_operator_dim,
+                                                      dirichlet_bc=occursin("zero", coord.bc))
         # obtain the local derivatives of the uniform grid with respect to the used grid
         derivative!(coord.duniform_dgrid, coord.uniform_grid, coord, spectral)
     else

moment_kinetics/src/derivatives.jl

@@ -149,6 +149,33 @@ dfns (ion) -> [vpa,vperp,z,r,s]
 dfns (neutrals) -> [vz,vr,vzeta,z,r,sn]
 """
 
+#df/dz
+#1D version for f[z], used by implicit solvers
+function derivative_z!(dfdz::AbstractArray{mk_float,1}, f::AbstractArray{mk_float,1},
+        dfdz_lower_endpoints::AbstractArray{mk_float,0},
+        dfdz_upper_endpoints::AbstractArray{mk_float,0},
+        z_send_buffer::AbstractArray{mk_float,0},
+        z_receive_buffer::AbstractArray{mk_float,0}, z_spectral, z)
+
+    begin_serial_region()
+
+    @serial_region begin
+        # differentiate f w.r.t z
+        derivative!(dfdz, f, z, z_spectral)
+        # get external endpoints to reconcile via MPI
+        dfdz_lower_endpoints[] = z.scratch_2d[1,1]
+        dfdz_upper_endpoints[] = z.scratch_2d[end,end]
+    end
+
+    # now reconcile element boundaries across
+    # processes with large message involving all y
+    if z.nelement_local < z.nelement_global
+        reconcile_element_boundaries_MPI!(
+            dfdz, dfdz_lower_endpoints, dfdz_upper_endpoints, z_send_buffer,
+            z_receive_buffer, z)
+    end
+end
+
 #df/dz
 #2D version for f[z,r] -> Er, Ez, phi
 function derivative_z!(dfdz::AbstractArray{mk_float,2}, f::AbstractArray{mk_float,2},