Skip to content

Commit

Permalink
Merge pull request #219 from mabarnes/imex
Browse files Browse the repository at this point in the history 
Initial functionality for partially-implicit 'IMEX' timestepping
  • Loading branch information
@johnomotani
johnomotani authored Jul 30, 2024
2 parents 24938f2 + 4f31b31 commit 352297f
Show file tree
Hide file tree
Showing 27 changed files with 5,236 additions and 1,223 deletions.
700 changes: 557 additions & 143 deletions makie_post_processing/makie_post_processing/src/makie_post_processing.jl
View file

Large diffs are not rendered by default.

1 change: 1 addition & 0 deletions moment_kinetics/Project.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
8 changes: 7 additions & 1 deletion moment_kinetics/src/calculus.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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")
Expand Down
89 changes: 56 additions & 33 deletions moment_kinetics/src/charge_exchange.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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))
Expand All @@ -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
Expand Down
21 changes: 16 additions & 5 deletions moment_kinetics/src/coordinates.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand All @@ -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
Expand Down
27 changes: 27 additions & 0 deletions moment_kinetics/src/derivatives.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand Down
Loading

0 comments on commit 352297f

Please sign in to comment.