From bc81b9e66286854df933ba232a443f6074035b3d Mon Sep 17 00:00:00 2001 From: John Omotani Date: Sun, 26 Jan 2025 14:10:51 +0000 Subject: [PATCH] Allow for MKArray to not be an AbstractArray This is a bit of a faff, because arrays in post-processing are `Array`, not `MKArray`, so any function that we might want to use in post-processing has to allow for either type. --- moment_kinetics/Project.toml | 1 + moment_kinetics/ext/file_io_netcdf.jl | 5 +- moment_kinetics/src/analysis.jl | 12 +- moment_kinetics/src/boundary_conditions.jl | 36 +- moment_kinetics/src/calculus.jl | 60 ++-- moment_kinetics/src/chebyshev.jl | 42 ++- moment_kinetics/src/clenshaw_curtis.jl | 2 +- moment_kinetics/src/coordinates.jl | 14 +- moment_kinetics/src/derivatives.jl | 332 +++++++++--------- .../src/electron_fluid_equations.jl | 12 +- .../src/electron_kinetic_equation.jl | 2 +- moment_kinetics/src/em_fields.jl | 10 +- moment_kinetics/src/file_io.jl | 4 +- moment_kinetics/src/file_io_hdf5.jl | 4 +- moment_kinetics/src/finite_differences.jl | 10 +- moment_kinetics/src/fokker_planck.jl | 8 +- moment_kinetics/src/fokker_planck_calculus.jl | 16 +- moment_kinetics/src/gauss_legendre.jl | 18 +- moment_kinetics/src/interpolation.jl | 24 +- moment_kinetics/src/load_data.jl | 27 +- moment_kinetics/src/moment_constraints.jl | 12 +- moment_kinetics/src/moment_kinetics.jl | 5 + moment_kinetics/src/nonlinear_solvers.jl | 72 ++-- moment_kinetics/src/runge_kutta.jl | 48 +-- moment_kinetics/src/type_definitions.jl | 30 ++ moment_kinetics/src/utils.jl | 15 +- moment_kinetics/test/calculus_tests.jl | 21 +- moment_kinetics/test/jacobian_matrix_tests.jl | 4 +- .../test/nonlinear_solver_tests.jl | 4 +- 29 files changed, 454 insertions(+), 396 deletions(-) diff --git a/moment_kinetics/Project.toml b/moment_kinetics/Project.toml index 9828a3c7f..ac23858ff 100644 --- a/moment_kinetics/Project.toml +++ b/moment_kinetics/Project.toml @@ -52,6 +52,7 @@ manufactured_solns_ext = ["Symbolics", "IfElse"] [compat] julia = "1.9.0" +InboundsArrays = "0.2.1" [extras] Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" diff --git a/moment_kinetics/ext/file_io_netcdf.jl b/moment_kinetics/ext/file_io_netcdf.jl index 9b61197b4..a7f461eb1 100644 --- a/moment_kinetics/ext/file_io_netcdf.jl +++ b/moment_kinetics/ext/file_io_netcdf.jl @@ -14,6 +14,7 @@ import moment_kinetics.file_io: io_has_parallel, open_output_file_implementation import moment_kinetics.load_data: open_file_to_read, get_attribute, load_variable, load_slice using moment_kinetics.coordinates: coordinate using moment_kinetics.input_structs: netcdf +using moment_kinetics.type_definitions using NCDatasets @@ -105,7 +106,7 @@ function maybe_create_netcdf_dim(file_or_group::NCDataset, coord) end function write_single_value!(file_or_group::NCDataset, name, - value::Union{Number, AbstractString, AbstractArray{T,N}}, + value::Union{Number, AbstractString, AbstractMKArray{T,N}}, coords::Union{coordinate,NamedTuple}...; parallel_io, description=nothing, units=nothing, overwrite=false) where {T,N} @@ -209,7 +210,7 @@ function create_dynamic_variable!(file_or_group::NCDataset, name, type, end function append_to_dynamic_var(io_var::NCDatasets.CFVariable, - data::Union{Nothing,Number,AbstractArray{T,N}}, t_idx, + data::Union{Nothing,Number,AbstractMKArray{T,N}}, t_idx, parallel_io::Bool, coords...; only_root=false, write_from_this_rank=nothing) where {T,N} diff --git a/moment_kinetics/src/analysis.jl b/moment_kinetics/src/analysis.jl index 7f1ebb090..fc7722975 100644 --- a/moment_kinetics/src/analysis.jl +++ b/moment_kinetics/src/analysis.jl @@ -15,7 +15,7 @@ using ..interpolation: interpolate_to_grid_1d using ..load_data: open_readonly_output_file, get_nranks, load_pdf_data, load_rank_data using ..load_data: load_distributed_ion_pdf_slice using ..looping -using ..type_definitions: mk_int, mk_float +using ..type_definitions using ..velocity_moments: integrate_over_vspace using FFTW @@ -690,7 +690,7 @@ function steady_state_square_residuals(variable, variable_at_previous_time, dt, end variable_max = MPI.Allreduce(local_max, max, comm_world) end - if isa(dt, AbstractVector) + if isa(dt, MKOrAbstractVector) reshaped_dt = reshape(dt, tuple((1 for _ ∈ 1:t_dim-1)..., size(dt)...)) else reshaped_dt = dt @@ -699,9 +699,9 @@ function steady_state_square_residuals(variable, variable_at_previous_time, dt, if size(dt) == () # local_max_absolute should always be at least a 1d array of size 1, not # a 0d array, so that the MPI.Gather() below works correctly. - local_max_absolute = zeros(1) + local_max_absolute = mk_zeros(1) else - local_max_absolute = zeros(size(dt)) + local_max_absolute = mk_zeros(size(dt)) end else if size(dt) == () @@ -732,7 +732,7 @@ function steady_state_square_residuals(variable, variable_at_previous_time, dt, # have size 1. this_dims = tuple((1:t_dim-3)...) if this_dims === () - local_max_absolute = max.(local_max_absolute, [absolute_residual]) + local_max_absolute[1] = max(local_max_absolute[1], absolute_residual[]) else local_max_absolute = max.(local_max_absolute, vec(maximum(absolute_residual, @@ -996,7 +996,7 @@ phi_fit_result struct whose fields are: """ function fit_delta_phi_mode(t, z, delta_phi) # First fit a cosine to each time slice - results = allocate_float(3, size(delta_phi)[2]) + results = zeros(mk_float, 3, size(delta_phi)[2]) amplitude_guess = 1.0 offset_guess = 0.0 for (i, phi_z) in enumerate(eachcol(delta_phi)) diff --git a/moment_kinetics/src/boundary_conditions.jl b/moment_kinetics/src/boundary_conditions.jl index a6fcbdc89..60e686e51 100644 --- a/moment_kinetics/src/boundary_conditions.jl +++ b/moment_kinetics/src/boundary_conditions.jl @@ -14,7 +14,7 @@ using ..interpolation: interpolate_to_grid_1d! using ..looping using ..timer_utils using ..moment_kinetics_structs: scratch_pdf, em_fields_struct -using ..type_definitions: mk_float, mk_int +using ..type_definitions using ..velocity_moments: integrate_over_vspace, integrate_over_neutral_vspace, integrate_over_positive_vz, integrate_over_negative_vz @@ -75,10 +75,10 @@ end """ enforce boundary conditions on f in r """ -function enforce_r_boundary_condition!(f::AbstractArray{mk_float,5}, f_r_bc, bc::String, - adv, vpa, vperp, z, r, composition, end1::AbstractArray{mk_float,4}, - end2::AbstractArray{mk_float,4}, buffer1::AbstractArray{mk_float,4}, - buffer2::AbstractArray{mk_float,4}, r_diffusion::Bool) +function enforce_r_boundary_condition!(f::MKOrAbstractArray{mk_float,5}, f_r_bc, bc::String, + adv, vpa, vperp, z, r, composition, end1::MKOrAbstractArray{mk_float,4}, + end2::MKOrAbstractArray{mk_float,4}, buffer1::MKOrAbstractArray{mk_float,4}, + buffer2::MKOrAbstractArray{mk_float,4}, r_diffusion::Bool) nr = r.n @@ -125,9 +125,9 @@ end enforce boundary conditions on ion particle f in z """ function enforce_z_boundary_condition!(pdf, density, upar, ppar, phi, moments, bc::String, adv, - z, vperp, vpa, composition, end1::AbstractArray{mk_float,4}, - end2::AbstractArray{mk_float,4}, buffer1::AbstractArray{mk_float,4}, - buffer2::AbstractArray{mk_float,4}) + z, vperp, vpa, composition, end1::MKOrAbstractArray{mk_float,4}, + end2::MKOrAbstractArray{mk_float,4}, buffer1::MKOrAbstractArray{mk_float,4}, + buffer2::MKOrAbstractArray{mk_float,4}) # this block ensures periodic BC can be supported with distributed memory MPI if z.nelement_global > z.nelement_local # reconcile internal element boundaries across processes @@ -281,10 +281,10 @@ enforce boundary conditions on neutral particle distribution function end end -function enforce_neutral_r_boundary_condition!(f::AbstractArray{mk_float,6}, - f_r_bc::AbstractArray{mk_float,6}, adv, vz, vr, vzeta, z, r, composition, - end1::AbstractArray{mk_float,5}, end2::AbstractArray{mk_float,5}, - buffer1::AbstractArray{mk_float,5}, buffer2::AbstractArray{mk_float,5}, +function enforce_neutral_r_boundary_condition!(f::MKOrAbstractArray{mk_float,6}, + f_r_bc::MKOrAbstractArray{mk_float,6}, adv, vz, vr, vzeta, z, r, composition, + end1::MKOrAbstractArray{mk_float,5}, end2::MKOrAbstractArray{mk_float,5}, + buffer1::MKOrAbstractArray{mk_float,5}, buffer2::MKOrAbstractArray{mk_float,5}, r_diffusion) #f_initial, bc = r.bc @@ -334,8 +334,8 @@ enforce boundary conditions on neutral particle f in z function enforce_neutral_z_boundary_condition!(pdf, density, uz, pz, moments, density_ion, upar_ion, Er, boundary_distributions, adv, z, vzeta, vr, vz, composition, geometry, - end1::AbstractArray{mk_float,5}, end2::AbstractArray{mk_float,5}, - buffer1::AbstractArray{mk_float,5}, buffer2::AbstractArray{mk_float,5}) + end1::MKOrAbstractArray{mk_float,5}, end2::MKOrAbstractArray{mk_float,5}, + buffer1::MKOrAbstractArray{mk_float,5}, buffer2::MKOrAbstractArray{mk_float,5}) if z.nelement_global > z.nelement_local @@ -794,7 +794,7 @@ function enforce_neutral_wall_bc!(pdf, z, vzeta, vr, vz, pz, uz, density, wall_f zero_vz_ind = 0 for ivz ∈ 1:vz.n if vz.scratch2[ivz] <= -zero - pdf[ivz,:,:,1] .= N_in*pdf[ivz,:,:,1] + pdf[ivz,:,:,1] .= N_in .* pdf[ivz,:,:,1] else zero_vz_ind = ivz if abs(vz.scratch2[ivz]) < zero @@ -1046,14 +1046,14 @@ enforce zero boundary condition at vperp -> infinity """ function enforce_vperp_boundary_condition! end -function enforce_vperp_boundary_condition!(f::AbstractArray{mk_float,5}, bc, vperp, vperp_spectral, vperp_advect, diffusion) +function enforce_vperp_boundary_condition!(f::MKOrAbstractArray{mk_float,5}, bc, vperp, vperp_spectral, vperp_advect, diffusion) @loop_s is begin @views enforce_vperp_boundary_condition!(f[:,:,:,:,is], bc, vperp, vperp_spectral, vperp_advect[is], diffusion) end return nothing end -function enforce_vperp_boundary_condition!(f::AbstractArray{mk_float,4}, bc, vperp, vperp_spectral, vperp_advect, diffusion) +function enforce_vperp_boundary_condition!(f::MKOrAbstractArray{mk_float,4}, bc, vperp, vperp_spectral, vperp_advect, diffusion) @loop_r ir begin @views enforce_vperp_boundary_condition!(f[:,:,:,ir], bc, vperp, vperp_spectral, vperp_advect, diffusion, ir) @@ -1061,7 +1061,7 @@ function enforce_vperp_boundary_condition!(f::AbstractArray{mk_float,4}, bc, vpe return nothing end -function enforce_vperp_boundary_condition!(f::AbstractArray{mk_float,3}, bc, vperp, +function enforce_vperp_boundary_condition!(f::MKOrAbstractArray{mk_float,3}, bc, vperp, vperp_spectral, vperp_advect, diffusion, ir) if bc == "zero" || bc == "zero-impose-regularity" nvperp = vperp.n diff --git a/moment_kinetics/src/calculus.jl b/moment_kinetics/src/calculus.jl index adff79e30..a55d3ec0b 100644 --- a/moment_kinetics/src/calculus.jl +++ b/moment_kinetics/src/calculus.jl @@ -9,7 +9,7 @@ export integral using ..moment_kinetics_structs: discretization_info, null_spatial_dimension_info, null_velocity_dimension_info, weak_discretization_info using ..timer_utils -using ..type_definitions: mk_float, mk_int +using ..type_definitions using MPI using ..communication: block_rank using ..communication: _block_synchronize @@ -272,7 +272,7 @@ end """ """ -function derivative_elements_to_full_grid!(df1d, df2d, coord, adv_fac::AbstractArray{mk_float,1}) +function derivative_elements_to_full_grid!(df1d, df2d, coord, adv_fac::MKOrAbstractVector{mk_float}) # no changes need to be made for the derivative at points away from element boundaries elements_to_full_grid_interior_pts!(df1d, df2d, coord) # resolve the multi-valued nature of the derivative at element boundaries @@ -321,7 +321,7 @@ df is multi-valued at the overlapping point at the boundary between neighboring elements. here we choose to use the value of df from the upwind element. """ -function reconcile_element_boundaries_upwind!(df1d, df2d, coord, adv_fac::AbstractArray{mk_float,1}) +function reconcile_element_boundaries_upwind!(df1d, df2d, coord, adv_fac::MKOrAbstractVector{mk_float}) # note that the first ngrid points are classified as belonging to the first element # and the next ngrid-1 points belonging to second element, etc. @@ -426,8 +426,8 @@ updated to include each physical dimension required in the main code """ -function assign_endpoint!(df1d::AbstractArray{mk_float,Ndims}, - receive_buffer::AbstractArray{mk_float,Mdims},key::String,coord) where {Ndims,Mdims} +function assign_endpoint!(df1d::MKOrAbstractArray{mk_float,Ndims}, + receive_buffer::MKOrAbstractArray{mk_float,Mdims},key::String,coord) where {Ndims,Mdims} if key == "lower" j = 1 elseif key == "upper" @@ -480,11 +480,11 @@ function assign_endpoint!(df1d::AbstractArray{mk_float,Ndims}, end @timeit_debug global_timer reconcile_element_boundaries_MPI!( - df1d::AbstractArray{mk_float,Ndims}, - dfdx_lower_endpoints::AbstractArray{mk_float,Mdims}, - dfdx_upper_endpoints::AbstractArray{mk_float,Mdims}, - receive_buffer1::AbstractArray{mk_float,Mdims}, - receive_buffer2::AbstractArray{mk_float,Mdims}, + df1d::MKOrAbstractArray{mk_float,Ndims}, + dfdx_lower_endpoints::MKOrAbstractArray{mk_float,Mdims}, + dfdx_upper_endpoints::MKOrAbstractArray{mk_float,Mdims}, + receive_buffer1::MKOrAbstractArray{mk_float,Mdims}, + receive_buffer2::MKOrAbstractArray{mk_float,Mdims}, coord) where {Ndims,Mdims} = begin # synchronize buffers @@ -556,7 +556,7 @@ end _block_synchronize() end -function apply_adv_fac!(buffer::AbstractArray{mk_float,Ndims},adv_fac::AbstractArray{mk_float,Ndims},endpoints::AbstractArray{mk_float,Ndims},sgn::mk_int) where Ndims +function apply_adv_fac!(buffer::MKOrAbstractArray{mk_float,Ndims},adv_fac::MKOrAbstractArray{mk_float,Ndims},endpoints::MKOrAbstractArray{mk_float,Ndims},sgn::mk_int) where Ndims #buffer contains off-process endpoint #adv_fac < 0 is positive advection speed #adv_fac > 0 is negative advection speed @@ -578,13 +578,13 @@ function apply_adv_fac!(buffer::AbstractArray{mk_float,Ndims},adv_fac::AbstractA end @timeit_debug global_timer reconcile_element_boundaries_MPI!( - df1d::AbstractArray{mk_float,Ndims}, - adv_fac_lower_endpoints::AbstractArray{mk_float,Mdims}, - adv_fac_upper_endpoints::AbstractArray{mk_float,Mdims}, - dfdx_lower_endpoints::AbstractArray{mk_float,Mdims}, - dfdx_upper_endpoints::AbstractArray{mk_float,Mdims}, - receive_buffer1::AbstractArray{mk_float,Mdims}, - receive_buffer2::AbstractArray{mk_float,Mdims}, + df1d::MKOrAbstractArray{mk_float,Ndims}, + adv_fac_lower_endpoints::MKOrAbstractArray{mk_float,Mdims}, + adv_fac_upper_endpoints::MKOrAbstractArray{mk_float,Mdims}, + dfdx_lower_endpoints::MKOrAbstractArray{mk_float,Mdims}, + dfdx_upper_endpoints::MKOrAbstractArray{mk_float,Mdims}, + receive_buffer1::MKOrAbstractArray{mk_float,Mdims}, + receive_buffer2::MKOrAbstractArray{mk_float,Mdims}, coord) where {Ndims,Mdims} = begin # synchronize buffers @@ -660,11 +660,11 @@ end # Special version for pdf_electron with no r-dimension, which has the same number of # dimensions as an ion/neutral moment variable, but different dimensions. @timeit_debug global_timer reconcile_element_boundaries_MPI_z_pdf_vpavperpz!( - df1d::AbstractArray{mk_float,3}, - dfdx_lower_endpoints::AbstractArray{mk_float,2}, - dfdx_upper_endpoints::AbstractArray{mk_float,2}, - receive_buffer1::AbstractArray{mk_float,2}, - receive_buffer2::AbstractArray{mk_float,2}, coord) = begin + df1d::MKOrAbstractArray{mk_float,3}, + dfdx_lower_endpoints::MKOrAbstractArray{mk_float,2}, + dfdx_upper_endpoints::MKOrAbstractArray{mk_float,2}, + receive_buffer1::MKOrAbstractArray{mk_float,2}, + receive_buffer2::MKOrAbstractArray{mk_float,2}, coord) = begin # synchronize buffers # -- this all-to-all block communicate here requires that this function is NOT called from within a parallelised loop @@ -738,13 +738,13 @@ end # Special version for pdf_electron with no r-dimension, which has the same number of # dimensions as an ion/neutral moment variable, but different dimensions. @timeit_debug global_timer reconcile_element_boundaries_MPI_z_pdf_vpavperpz!( - df1d::AbstractArray{mk_float,3}, - adv_fac_lower_endpoints::AbstractArray{mk_float,2}, - adv_fac_upper_endpoints::AbstractArray{mk_float,2}, - dfdx_lower_endpoints::AbstractArray{mk_float,2}, - dfdx_upper_endpoints::AbstractArray{mk_float,2}, - receive_buffer1::AbstractArray{mk_float,2}, - receive_buffer2::AbstractArray{mk_float,2}, coord) = begin + df1d::MKOrAbstractArray{mk_float,3}, + adv_fac_lower_endpoints::MKOrAbstractArray{mk_float,2}, + adv_fac_upper_endpoints::MKOrAbstractArray{mk_float,2}, + dfdx_lower_endpoints::MKOrAbstractArray{mk_float,2}, + dfdx_upper_endpoints::MKOrAbstractArray{mk_float,2}, + receive_buffer1::MKOrAbstractArray{mk_float,2}, + receive_buffer2::MKOrAbstractArray{mk_float,2}, coord) = begin # synchronize buffers # -- this all-to-all block communicate here requires that this function is NOT called from within a parallelised loop diff --git a/moment_kinetics/src/chebyshev.jl b/moment_kinetics/src/chebyshev.jl index 56162ee28..ea159c751 100644 --- a/moment_kinetics/src/chebyshev.jl +++ b/moment_kinetics/src/chebyshev.jl @@ -24,25 +24,29 @@ using ..moment_kinetics_structs: discretization_info """ Chebyshev pseudospectral discretization """ -struct chebyshev_base_info{TForward <: FFTW.cFFTWPlan, TBackward <: AbstractFFTs.ScaledPlan} +struct chebyshev_base_info{TForward <: FFTW.cFFTWPlan, + TBackward <: AbstractFFTs.ScaledPlan, + Tc <: MKOrAbstractVector{Complex{mk_float}}, + Tm <: MKOrAbstractMatrix{mk_float}, + T <: MKOrAbstractVector{mk_float}} # 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::MKVector{Complex{mk_float}} + fext::Tc # Chebyshev spectral coefficients of distribution function f # first dimension contains location within element # second dimension indicates the element - f::MKMatrix{mk_float} + f::Tm # Chebyshev spectral coefficients of derivative of f - df::MKVector{mk_float} + df::T # 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 # elementwise differentiation matrix (ngrid*ngrid) - Dmat::MKMatrix{mk_float} + Dmat::Tm # elementwise differentiation vector (ngrid) for the point x = -1 - D0::MKVector{mk_float} + D0::T end struct chebyshev_info{TForward <: FFTW.cFFTWPlan, TBackward <: AbstractFFTs.ScaledPlan} <: discretization_info @@ -234,7 +238,7 @@ function scaled_chebyshev_radau_grid(ngrid, nelement_local, n, if irank == 0 # use a Chebyshev-Gauss-Radau element for the lowest element on rank 0 scale_factor = element_scale[1] shift = element_shift[1] - grid[imin[1]:imax[1]] .= (chebyshev_radau_grid[1:ngrid] * scale_factor) .+ shift + grid[imin[1]:imax[1]] .= (chebyshev_radau_grid[1:ngrid] .* scale_factor) .+ shift # account for the fact that the minimum index needed for the chebyshev_grid # within each element changes from 1 to 2 in going from the first element # to the remaining elements @@ -244,7 +248,7 @@ function scaled_chebyshev_radau_grid(ngrid, nelement_local, n, shift = element_shift[j] # reverse the order of the original chebyshev_grid (ran from [1,-1]) # and apply the scale factor and shift - grid[imin[j]:imax[j]] .= (reverse(chebyshev_grid)[k:ngrid] * scale_factor) .+ shift + grid[imin[j]:imax[j]] .= (reverse(chebyshev_grid)[k:ngrid] .* scale_factor) .+ shift end wgts = clenshaw_curtis_radau_weights(ngrid, nelement_local, n, imin, imax, element_scale) else @@ -257,7 +261,7 @@ function scaled_chebyshev_radau_grid(ngrid, nelement_local, n, shift = element_shift[j] # reverse the order of the original chebyshev_grid (ran from [1,-1]) # and apply the scale factor and shift - grid[imin[j]:imax[j]] .= (reverse(chebyshev_grid)[k:ngrid] * scale_factor) .+ shift + grid[imin[j]:imax[j]] .= (reverse(chebyshev_grid)[k:ngrid] .* scale_factor) .+ shift # after first element, increase minimum index for chebyshev_grid to 2 # to avoid double-counting boundary element k = 2 @@ -574,10 +578,10 @@ function clenshaw_curtis_weights(ngrid, nelement_local, n, imin, imax, element_s @inbounds begin # calculate the weights within a single element and # scale to account for modified domain (not [-1,1]) - wgts[1:ngrid] = w*element_scale[1] + wgts[1:ngrid] .= w .* element_scale[1] if nelement_local > 1 for j ∈ 2:nelement_local - wgts[imin[j]-1:imax[j]] .+= w*element_scale[j] + wgts[imin[j]-1:imax[j]] .+= w .* element_scale[j] end end end @@ -594,13 +598,13 @@ function clenshaw_curtis_radau_weights(ngrid, nelement_local, n, imin, imax, ele @inbounds begin # calculate the weights within a single element and # scale to account for modified domain (not [-1,1]) - wgts[1:ngrid] .= wgts_radau[1:ngrid]*element_scale[1] + wgts[1:ngrid] .= wgts_radau[1:ngrid] .* element_scale[1] if nelement_local > 1 for j ∈ 2:nelement_local # account for double-counting of points at inner element boundaries - wgts[imin[j]-1] += wgts_lobatto[1]*element_scale[j] + wgts[imin[j]-1] += wgts_lobatto[1] .* element_scale[j] # assign weights for interior of elements and one boundary point - wgts[imin[j]:imax[j]] .= wgts_lobatto[2:ngrid]*element_scale[j] + wgts[imin[j]:imax[j]] .= wgts_lobatto[2:ngrid] .* element_scale[j] end end end @@ -612,7 +616,7 @@ compute and return modified Chebyshev moments of the first kind: ∫dx Tᵢ(x) over range [-1,1] """ function chebyshevmoments(N) - μ = zeros(N) + μ = mk_zeros(N) @inbounds for i = 0:2:N-1 μ[i+1] = 2/(1-i^2) end @@ -646,7 +650,7 @@ function chebyshev_radau_points(n) return grid end -function chebyshev_radau_weights(moments::AbstractArray{mk_float,1}, n) +function chebyshev_radau_weights(moments::MKVector{mk_float}, n) # input should have values moments[j] = (cos(pi j) + 1)/(1-j^2) for j >= 0 nfft = 2*n - 1 # create array for moments on extended [0,2π] domain in theta = ArcCos[z] @@ -935,7 +939,7 @@ https://people.maths.ox.ac.uk/trefethen/pdetext.html D[j,j] = -sum(D[j,:]) end end - function Djk(x::AbstractArray{Float64,1},j::Int64,k::Int64,c_j::Float64,c_k::Float64) + function Djk(x::MKOrAbstractVector{mk_float},j::mk_int,k::mk_int,c_j::mk_float,c_k::mk_float) return (c_j/c_k)*((-1)^(k+j))/(x[j] - x[k]) end """ @@ -943,7 +947,7 @@ https://people.maths.ox.ac.uk/trefethen/pdetext.html Note that a similar function could be constructed for the Chebyshev-Lobatto grid, if desired. """ - function cheb_derivative_matrix_elementwise_radau_by_FFT!(D::AbstractArray{Float64,2}, coord, f, df, fext, forward) + function cheb_derivative_matrix_elementwise_radau_by_FFT!(D::MKOrAbstractMatrix{mk_float}, coord, f, df, fext, forward) ff_buffer = allocate_float(coord.ngrid) df_buffer = allocate_float(coord.ngrid) # use response matrix approach to calculate derivative matrix D @@ -962,7 +966,7 @@ https://people.maths.ox.ac.uk/trefethen/pdetext.html end end - function cheb_lower_endpoint_derivative_vector_elementwise_radau_by_FFT!(D::AbstractArray{Float64,1}, coord, f, df, fext, forward) + function cheb_lower_endpoint_derivative_vector_elementwise_radau_by_FFT!(D::MKOrAbstractVector{mk_float}, coord, f, df, fext, forward) ff_buffer = allocate_float(coord.ngrid) df_buffer = allocate_float(coord.ngrid) # use response matrix approach to calculate derivative vector D diff --git a/moment_kinetics/src/clenshaw_curtis.jl b/moment_kinetics/src/clenshaw_curtis.jl index df20dbb7d..8c97dd415 100644 --- a/moment_kinetics/src/clenshaw_curtis.jl +++ b/moment_kinetics/src/clenshaw_curtis.jl @@ -24,7 +24,7 @@ clenshawcurtisweights(μ) = clenshawcurtisweights!(copy(μ)) clenshawcurtisweights!(μ) = clenshawcurtisweights!(μ, plan_clenshawcurtis(μ)) function clenshawcurtisweights!(μ, plan) N = length(μ) - rmul!(μ, inv(N-1)) + μ .*= inv(N-1) plan*μ μ[1] *= 0.5; μ[N] *= 0.5 return μ diff --git a/moment_kinetics/src/coordinates.jl b/moment_kinetics/src/coordinates.jl index 057fec44f..5b307d232 100644 --- a/moment_kinetics/src/coordinates.jl +++ b/moment_kinetics/src/coordinates.jl @@ -26,13 +26,13 @@ using OrderedCollections: OrderedDict """ structure containing basic information related to coordinates """ -struct coordinate{Tshared <: Union{AbstractVector{mk_float}, AbstractInboundsArray{mk_float, 1}}, - Tsharedi <: Union{AbstractVector{mk_int}, AbstractInboundsArray{mk_int, 1}}, - Tunshared <: Union{AbstractVector{mk_float}, AbstractInboundsArray{mk_float, 1}}, - Tunsharedi <: Union{AbstractVector{mk_int}, AbstractInboundsArray{mk_int, 1}}, - Tunsharedmat <: Union{AbstractMatrix{mk_float}, AbstractInboundsArray{mk_float, 2}}, - Tunsharedmati <: Union{AbstractMatrix{mk_int}, AbstractInboundsArray{mk_int, 2}}, - Tunshared3 <: Union{AbstractArray{mk_float,3}, AbstractInboundsArray{mk_float, 3}}, +struct coordinate{Tshared <: MKOrAbstractVector{mk_float}, + Tsharedi <: MKOrAbstractVector{mk_int}, + Tunshared <: MKOrAbstractVector{mk_float}, + Tunsharedi <: MKOrAbstractVector{mk_int}, + Tunsharedmat <: MKOrAbstractMatrix{mk_float}, + Tunsharedmati <: MKOrAbstractMatrix{mk_int}, + Tunshared3 <: MKOrAbstractArray{mk_float,3}, Tbparams} # name is the name of the variable associated with this coordiante name::String diff --git a/moment_kinetics/src/derivatives.jl b/moment_kinetics/src/derivatives.jl index cd8cf84c0..2df4f6a37 100644 --- a/moment_kinetics/src/derivatives.jl +++ b/moment_kinetics/src/derivatives.jl @@ -14,7 +14,7 @@ export derivative_z!, derivative_z_chrg!, derivative_z_ntrl! using ..calculus: derivative!, second_derivative!, reconcile_element_boundaries_MPI!, reconcile_element_boundaries_MPI_z_pdf_vpavperpz!, apply_adv_fac! using ..communication -using ..type_definitions: mk_float +using ..type_definitions using ..looping using MPI @@ -29,11 +29,11 @@ dfns (neutrals) -> [vz,vr,vzeta,z,r,sn] #df/dr #2D version for f[z,r] -> Er, Ez, phi -function derivative_r!(dfdr::AbstractArray{mk_float,2}, f::AbstractArray{mk_float,2}, - dfdr_lower_endpoints::AbstractArray{mk_float,1}, - dfdr_upper_endpoints::AbstractArray{mk_float,1}, - r_receive_buffer1::AbstractArray{mk_float,1}, - r_receive_buffer2::AbstractArray{mk_float,1}, r_spectral, r) +function derivative_r!(dfdr::MKOrAbstractArray{mk_float,2}, f::MKOrAbstractArray{mk_float,2}, + dfdr_lower_endpoints::MKOrAbstractArray{mk_float,1}, + dfdr_upper_endpoints::MKOrAbstractArray{mk_float,1}, + r_receive_buffer1::MKOrAbstractArray{mk_float,1}, + r_receive_buffer2::MKOrAbstractArray{mk_float,1}, r_spectral, r) begin_z_region() @@ -56,11 +56,11 @@ end #df/dr #3D version for f[s,z,r] -> moments n, u, T etc -function derivative_r!(dfdr::AbstractArray{mk_float,3}, f::AbstractArray{mk_float,3}, - dfdr_lower_endpoints::AbstractArray{mk_float,2}, - dfdr_upper_endpoints::AbstractArray{mk_float,2}, - r_receive_buffer1::AbstractArray{mk_float,2}, - r_receive_buffer2::AbstractArray{mk_float,2}, r_spectral, r; neutrals=false) +function derivative_r!(dfdr::MKOrAbstractArray{mk_float,3}, f::MKOrAbstractArray{mk_float,3}, + dfdr_lower_endpoints::MKOrAbstractArray{mk_float,2}, + dfdr_upper_endpoints::MKOrAbstractArray{mk_float,2}, + r_receive_buffer1::MKOrAbstractArray{mk_float,2}, + r_receive_buffer2::MKOrAbstractArray{mk_float,2}, r_spectral, r; neutrals=false) # differentiate f w.r.t r if neutrals @@ -94,11 +94,11 @@ end #df/dr #5D version for f[vpa,vperp,z,r,s] -> ion particle dfn -function derivative_r!(dfdr::AbstractArray{mk_float,5}, f::AbstractArray{mk_float,5}, - dfdr_lower_endpoints::AbstractArray{mk_float,4}, - dfdr_upper_endpoints::AbstractArray{mk_float,4}, - r_receive_buffer1::AbstractArray{mk_float,4}, - r_receive_buffer2::AbstractArray{mk_float,4}, r_spectral, r) +function derivative_r!(dfdr::MKOrAbstractArray{mk_float,5}, f::MKOrAbstractArray{mk_float,5}, + dfdr_lower_endpoints::MKOrAbstractArray{mk_float,4}, + dfdr_upper_endpoints::MKOrAbstractArray{mk_float,4}, + r_receive_buffer1::MKOrAbstractArray{mk_float,4}, + r_receive_buffer2::MKOrAbstractArray{mk_float,4}, r_spectral, r) begin_s_z_vperp_vpa_region() @@ -120,11 +120,11 @@ function derivative_r!(dfdr::AbstractArray{mk_float,5}, f::AbstractArray{mk_floa end #6D version for f[vz,vz,vzeta,z,r,sn] -> neutral particle dfn (species indexing taken outside this loop) -function derivative_r!(dfdr::AbstractArray{mk_float,6}, f::AbstractArray{mk_float,6}, - dfdr_lower_endpoints::AbstractArray{mk_float,5}, - dfdr_upper_endpoints::AbstractArray{mk_float,5}, - r_receive_buffer1::AbstractArray{mk_float,5}, - r_receive_buffer2::AbstractArray{mk_float,5}, r_spectral, r) +function derivative_r!(dfdr::MKOrAbstractArray{mk_float,6}, f::MKOrAbstractArray{mk_float,6}, + dfdr_lower_endpoints::MKOrAbstractArray{mk_float,5}, + dfdr_upper_endpoints::MKOrAbstractArray{mk_float,5}, + r_receive_buffer1::MKOrAbstractArray{mk_float,5}, + r_receive_buffer2::MKOrAbstractArray{mk_float,5}, r_spectral, r) begin_sn_z_vzeta_vr_vz_region() @@ -155,11 +155,11 @@ 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) +function derivative_z!(dfdz::MKOrAbstractArray{mk_float,1}, f::MKOrAbstractArray{mk_float,1}, + dfdz_lower_endpoints::MKOrAbstractArray{mk_float,0}, + dfdz_upper_endpoints::MKOrAbstractArray{mk_float,0}, + z_send_buffer::MKOrAbstractArray{mk_float,0}, + z_receive_buffer::MKOrAbstractArray{mk_float,0}, z_spectral, z) begin_serial_region() @@ -182,11 +182,11 @@ 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}, - dfdz_lower_endpoints::AbstractArray{mk_float,1}, - dfdz_upper_endpoints::AbstractArray{mk_float,1}, - z_send_buffer::AbstractArray{mk_float,1}, - z_receive_buffer::AbstractArray{mk_float,1}, z_spectral, z) +function derivative_z!(dfdz::MKOrAbstractArray{mk_float,2}, f::MKOrAbstractArray{mk_float,2}, + dfdz_lower_endpoints::MKOrAbstractArray{mk_float,1}, + dfdz_upper_endpoints::MKOrAbstractArray{mk_float,1}, + z_send_buffer::MKOrAbstractArray{mk_float,1}, + z_receive_buffer::MKOrAbstractArray{mk_float,1}, z_spectral, z) begin_r_region() @@ -209,11 +209,11 @@ end #df/dz #3D version for f[z,r] -> moments n, u, T etc -function derivative_z!(dfdz::AbstractArray{mk_float,3}, f::AbstractArray{mk_float,3}, - dfdz_lower_endpoints::AbstractArray{mk_float,2}, - dfdz_upper_endpoints::AbstractArray{mk_float,2}, - z_send_buffer::AbstractArray{mk_float,2}, - z_receive_buffer::AbstractArray{mk_float,2}, z_spectral, z; neutrals=false) +function derivative_z!(dfdz::MKOrAbstractArray{mk_float,3}, f::MKOrAbstractArray{mk_float,3}, + dfdz_lower_endpoints::MKOrAbstractArray{mk_float,2}, + dfdz_upper_endpoints::MKOrAbstractArray{mk_float,2}, + z_send_buffer::MKOrAbstractArray{mk_float,2}, + z_receive_buffer::MKOrAbstractArray{mk_float,2}, z_spectral, z; neutrals=false) # differentiate f w.r.t z if neutrals @@ -248,11 +248,11 @@ end # df/dz # 3D version for f[vpa,vperp,z]. Uses modified function name to avoid clash with 'standard' # 3D version for ion/neutral moments. -function derivative_z_pdf_vpavperpz!(dfdz::AbstractArray{mk_float,3}, f::AbstractArray{mk_float,3}, - dfdz_lower_endpoints::AbstractArray{mk_float,2}, - dfdz_upper_endpoints::AbstractArray{mk_float,2}, - z_receive_buffer1::AbstractArray{mk_float,2}, - z_receive_buffer2::AbstractArray{mk_float,2}, z_spectral, z) +function derivative_z_pdf_vpavperpz!(dfdz::MKOrAbstractArray{mk_float,3}, f::MKOrAbstractArray{mk_float,3}, + dfdz_lower_endpoints::MKOrAbstractArray{mk_float,2}, + dfdz_upper_endpoints::MKOrAbstractArray{mk_float,2}, + z_receive_buffer1::MKOrAbstractArray{mk_float,2}, + z_receive_buffer2::MKOrAbstractArray{mk_float,2}, z_spectral, z) # differentiate f w.r.t z @loop_vperp_vpa ivperp ivpa begin @@ -272,11 +272,11 @@ function derivative_z_pdf_vpavperpz!(dfdz::AbstractArray{mk_float,3}, f::Abstrac end #5D version for f[vpa,vperp,z,r,s] -> dfn ions -function derivative_z!(dfdz::AbstractArray{mk_float,5}, f::AbstractArray{mk_float,5}, - dfdz_lower_endpoints::AbstractArray{mk_float,4}, - dfdz_upper_endpoints::AbstractArray{mk_float,4}, - z_send_buffer::AbstractArray{mk_float,4}, - z_receive_buffer::AbstractArray{mk_float,4}, z_spectral, z) +function derivative_z!(dfdz::MKOrAbstractArray{mk_float,5}, f::MKOrAbstractArray{mk_float,5}, + dfdz_lower_endpoints::MKOrAbstractArray{mk_float,4}, + dfdz_upper_endpoints::MKOrAbstractArray{mk_float,4}, + z_send_buffer::MKOrAbstractArray{mk_float,4}, + z_receive_buffer::MKOrAbstractArray{mk_float,4}, z_spectral, z) begin_s_r_vperp_vpa_region() @@ -298,11 +298,11 @@ function derivative_z!(dfdz::AbstractArray{mk_float,5}, f::AbstractArray{mk_floa end #4D version for f[vpa,vperp,z,r] -> dfn electron particles -function derivative_z!(dfdz::AbstractArray{mk_float,4}, f::AbstractArray{mk_float,4}, - dfdz_lower_endpoints::AbstractArray{mk_float,3}, - dfdz_upper_endpoints::AbstractArray{mk_float,3}, - z_send_buffer::AbstractArray{mk_float,3}, - z_receive_buffer::AbstractArray{mk_float,3}, z_spectral, z) +function derivative_z!(dfdz::MKOrAbstractArray{mk_float,4}, f::MKOrAbstractArray{mk_float,4}, + dfdz_lower_endpoints::MKOrAbstractArray{mk_float,3}, + dfdz_upper_endpoints::MKOrAbstractArray{mk_float,3}, + z_send_buffer::MKOrAbstractArray{mk_float,3}, + z_receive_buffer::MKOrAbstractArray{mk_float,3}, z_spectral, z) begin_r_vperp_vpa_region() @@ -324,11 +324,11 @@ function derivative_z!(dfdz::AbstractArray{mk_float,4}, f::AbstractArray{mk_floa end #6D version for f[vz,vr,vzeta,z,r] -> dfn neutral particles -function derivative_z!(dfdz::AbstractArray{mk_float,6}, f::AbstractArray{mk_float,6}, - dfdz_lower_endpoints::AbstractArray{mk_float,5}, - dfdz_upper_endpoints::AbstractArray{mk_float,5}, - z_send_buffer::AbstractArray{mk_float,5}, - z_receive_buffer::AbstractArray{mk_float,5}, z_spectral, z) +function derivative_z!(dfdz::MKOrAbstractArray{mk_float,6}, f::MKOrAbstractArray{mk_float,6}, + dfdz_lower_endpoints::MKOrAbstractArray{mk_float,5}, + dfdz_upper_endpoints::MKOrAbstractArray{mk_float,5}, + z_send_buffer::MKOrAbstractArray{mk_float,5}, + z_receive_buffer::MKOrAbstractArray{mk_float,5}, z_spectral, z) begin_sn_r_vzeta_vr_vz_region() @@ -359,11 +359,11 @@ dfns (neutrals) -> [vz,vr,vzeta,z,r,sn] #d2f/dr2 #2D version for f[z,r] -> Er, Ez, phi -function second_derivative_r!(d2fdr2::AbstractArray{mk_float,2}, f::AbstractArray{mk_float,2}, - d2fdr2_lower_endpoints::AbstractArray{mk_float,1}, - d2fdr2_upper_endpoints::AbstractArray{mk_float,1}, - r_receive_buffer1::AbstractArray{mk_float,1}, - r_receive_buffer2::AbstractArray{mk_float,1}, r_spectral, r) +function second_derivative_r!(d2fdr2::MKOrAbstractArray{mk_float,2}, f::MKOrAbstractArray{mk_float,2}, + d2fdr2_lower_endpoints::MKOrAbstractArray{mk_float,1}, + d2fdr2_upper_endpoints::MKOrAbstractArray{mk_float,1}, + r_receive_buffer1::MKOrAbstractArray{mk_float,1}, + r_receive_buffer2::MKOrAbstractArray{mk_float,1}, r_spectral, r) begin_z_region() @@ -385,11 +385,11 @@ end #d2f/dr2 #3D version for f[s,z,r] -> moments n, u, T etc -function second_derivative_r!(d2fdr2::AbstractArray{mk_float,3}, f::AbstractArray{mk_float,3}, - d2fdr2_lower_endpoints::AbstractArray{mk_float,2}, - d2fdr2_upper_endpoints::AbstractArray{mk_float,2}, - r_receive_buffer1::AbstractArray{mk_float,2}, - r_receive_buffer2::AbstractArray{mk_float,2}, r_spectral, r; neutrals=false) +function second_derivative_r!(d2fdr2::MKOrAbstractArray{mk_float,3}, f::MKOrAbstractArray{mk_float,3}, + d2fdr2_lower_endpoints::MKOrAbstractArray{mk_float,2}, + d2fdr2_upper_endpoints::MKOrAbstractArray{mk_float,2}, + r_receive_buffer1::MKOrAbstractArray{mk_float,2}, + r_receive_buffer2::MKOrAbstractArray{mk_float,2}, r_spectral, r; neutrals=false) # differentiate f w.r.t r if neutrals @@ -423,11 +423,11 @@ end #d2f/dr2 #5D version for f[vpa,vperp,z,r,s] -> ion particle dfn -function second_derivative_r!(d2fdr2::AbstractArray{mk_float,5}, f::AbstractArray{mk_float,5}, - d2fdr2_lower_endpoints::AbstractArray{mk_float,4}, - d2fdr2_upper_endpoints::AbstractArray{mk_float,4}, - r_receive_buffer1::AbstractArray{mk_float,4}, - r_receive_buffer2::AbstractArray{mk_float,4}, r_spectral, r) +function second_derivative_r!(d2fdr2::MKOrAbstractArray{mk_float,5}, f::MKOrAbstractArray{mk_float,5}, + d2fdr2_lower_endpoints::MKOrAbstractArray{mk_float,4}, + d2fdr2_upper_endpoints::MKOrAbstractArray{mk_float,4}, + r_receive_buffer1::MKOrAbstractArray{mk_float,4}, + r_receive_buffer2::MKOrAbstractArray{mk_float,4}, r_spectral, r) begin_s_z_vperp_vpa_region() @@ -448,11 +448,11 @@ function second_derivative_r!(d2fdr2::AbstractArray{mk_float,5}, f::AbstractArra end #6D version for f[vz,vz,vzeta,z,r,sn] -> neutral particle dfn (species indexing taken outside this loop) -function second_derivative_r!(d2fdr2::AbstractArray{mk_float,6}, f::AbstractArray{mk_float,6}, - d2fdr2_lower_endpoints::AbstractArray{mk_float,5}, - d2fdr2_upper_endpoints::AbstractArray{mk_float,5}, - r_receive_buffer1::AbstractArray{mk_float,5}, - r_receive_buffer2::AbstractArray{mk_float,5}, r_spectral, r) +function second_derivative_r!(d2fdr2::MKOrAbstractArray{mk_float,6}, f::MKOrAbstractArray{mk_float,6}, + d2fdr2_lower_endpoints::MKOrAbstractArray{mk_float,5}, + d2fdr2_upper_endpoints::MKOrAbstractArray{mk_float,5}, + r_receive_buffer1::MKOrAbstractArray{mk_float,5}, + r_receive_buffer2::MKOrAbstractArray{mk_float,5}, r_spectral, r) begin_sn_z_vzeta_vr_vz_region() @@ -482,11 +482,11 @@ dfns (neutrals) -> [vz,vr,vzeta,z,r,sn] #d2f/dz2 #2D version for f[z,r] -> Er, Ez, phi -function second_derivative_z!(d2fdz2::AbstractArray{mk_float,2}, f::AbstractArray{mk_float,2}, - d2fdz2_lower_endpoints::AbstractArray{mk_float,1}, - d2fdz2_upper_endpoints::AbstractArray{mk_float,1}, - z_send_buffer::AbstractArray{mk_float,1}, - z_receive_buffer::AbstractArray{mk_float,1}, z_spectral, z) +function second_derivative_z!(d2fdz2::MKOrAbstractArray{mk_float,2}, f::MKOrAbstractArray{mk_float,2}, + d2fdz2_lower_endpoints::MKOrAbstractArray{mk_float,1}, + d2fdz2_upper_endpoints::MKOrAbstractArray{mk_float,1}, + z_send_buffer::MKOrAbstractArray{mk_float,1}, + z_receive_buffer::MKOrAbstractArray{mk_float,1}, z_spectral, z) begin_r_region() @@ -508,11 +508,11 @@ end #d2f/dz2 #3D version for f[z,r] -> moments n, u, T etc -function second_derivative_z!(d2fdz2::AbstractArray{mk_float,3}, f::AbstractArray{mk_float,3}, - d2fdz2_lower_endpoints::AbstractArray{mk_float,2}, - d2fdz2_upper_endpoints::AbstractArray{mk_float,2}, - z_send_buffer::AbstractArray{mk_float,2}, - z_receive_buffer::AbstractArray{mk_float,2}, z_spectral, z; neutrals=false) +function second_derivative_z!(d2fdz2::MKOrAbstractArray{mk_float,3}, f::MKOrAbstractArray{mk_float,3}, + d2fdz2_lower_endpoints::MKOrAbstractArray{mk_float,2}, + d2fdz2_upper_endpoints::MKOrAbstractArray{mk_float,2}, + z_send_buffer::MKOrAbstractArray{mk_float,2}, + z_receive_buffer::MKOrAbstractArray{mk_float,2}, z_spectral, z; neutrals=false) # differentiate f w.r.t z if neutrals @@ -545,11 +545,11 @@ function second_derivative_z!(d2fdz2::AbstractArray{mk_float,3}, f::AbstractArra end #5D version for f[vpa,vperp,z,r,s] -> dfn ions -function second_derivative_z!(d2fdz2::AbstractArray{mk_float,5}, f::AbstractArray{mk_float,5}, - d2fdz2_lower_endpoints::AbstractArray{mk_float,4}, - d2fdz2_upper_endpoints::AbstractArray{mk_float,4}, - z_send_buffer::AbstractArray{mk_float,4}, - z_receive_buffer::AbstractArray{mk_float,4}, z_spectral, z) +function second_derivative_z!(d2fdz2::MKOrAbstractArray{mk_float,5}, f::MKOrAbstractArray{mk_float,5}, + d2fdz2_lower_endpoints::MKOrAbstractArray{mk_float,4}, + d2fdz2_upper_endpoints::MKOrAbstractArray{mk_float,4}, + z_send_buffer::MKOrAbstractArray{mk_float,4}, + z_receive_buffer::MKOrAbstractArray{mk_float,4}, z_spectral, z) begin_s_r_vperp_vpa_region() @@ -570,11 +570,11 @@ function second_derivative_z!(d2fdz2::AbstractArray{mk_float,5}, f::AbstractArra end #6D version for f[vz,vr,vzeta,z,r] -> dfn neutral particles -function second_derivative_z!(d2fdz2::AbstractArray{mk_float,6}, f::AbstractArray{mk_float,6}, - d2fdz2_lower_endpoints::AbstractArray{mk_float,5}, - d2fdz2_upper_endpoints::AbstractArray{mk_float,5}, - z_send_buffer::AbstractArray{mk_float,5}, - z_receive_buffer::AbstractArray{mk_float,5}, z_spectral, z) +function second_derivative_z!(d2fdz2::MKOrAbstractArray{mk_float,6}, f::MKOrAbstractArray{mk_float,6}, + d2fdz2_lower_endpoints::MKOrAbstractArray{mk_float,5}, + d2fdz2_upper_endpoints::MKOrAbstractArray{mk_float,5}, + z_send_buffer::MKOrAbstractArray{mk_float,5}, + z_receive_buffer::MKOrAbstractArray{mk_float,5}, z_spectral, z) begin_sn_r_vzeta_vr_vz_region() @@ -604,13 +604,13 @@ dfns (neutrals) -> [vz,vr,vzeta,z,r,sn] #df/dr #2D version for f[z,r] -> Er, Ez, phi -function derivative_r!(dfdr::AbstractArray{mk_float,2}, f::AbstractArray{mk_float,2}, - adv_fac, adv_fac_lower_buffer::AbstractArray{mk_float,1}, - adv_fac_upper_buffer::AbstractArray{mk_float,1}, - dfdr_lower_endpoints::AbstractArray{mk_float,1}, - dfdr_upper_endpoints::AbstractArray{mk_float,1}, - r_receive_buffer1::AbstractArray{mk_float,1}, - r_receive_buffer2::AbstractArray{mk_float,1}, r_spectral, r) +function derivative_r!(dfdr::MKOrAbstractArray{mk_float,2}, f::MKOrAbstractArray{mk_float,2}, + adv_fac, adv_fac_lower_buffer::MKOrAbstractArray{mk_float,1}, + adv_fac_upper_buffer::MKOrAbstractArray{mk_float,1}, + dfdr_lower_endpoints::MKOrAbstractArray{mk_float,1}, + dfdr_upper_endpoints::MKOrAbstractArray{mk_float,1}, + r_receive_buffer1::MKOrAbstractArray{mk_float,1}, + r_receive_buffer2::MKOrAbstractArray{mk_float,1}, r_spectral, r) begin_z_region() @@ -634,13 +634,13 @@ end #df/dr #3D version for f[z,r,s] -> moments n, u, T etc -function derivative_r!(dfdr::AbstractArray{mk_float,3}, f::AbstractArray{mk_float,3}, - adv_fac, adv_fac_lower_buffer::AbstractArray{mk_float,2}, - adv_fac_upper_buffer::AbstractArray{mk_float,2}, - dfdr_lower_endpoints::AbstractArray{mk_float,2}, - dfdr_upper_endpoints::AbstractArray{mk_float,2}, - r_receive_buffer1::AbstractArray{mk_float,2}, - r_receive_buffer2::AbstractArray{mk_float,2}, r_spectral, r; neutrals=false) +function derivative_r!(dfdr::MKOrAbstractArray{mk_float,3}, f::MKOrAbstractArray{mk_float,3}, + adv_fac, adv_fac_lower_buffer::MKOrAbstractArray{mk_float,2}, + adv_fac_upper_buffer::MKOrAbstractArray{mk_float,2}, + dfdr_lower_endpoints::MKOrAbstractArray{mk_float,2}, + dfdr_upper_endpoints::MKOrAbstractArray{mk_float,2}, + r_receive_buffer1::MKOrAbstractArray{mk_float,2}, + r_receive_buffer2::MKOrAbstractArray{mk_float,2}, r_spectral, r; neutrals=false) # differentiate f w.r.t r if neutrals @@ -678,13 +678,13 @@ end #df/dr #5D version for f[vpa,vperp,z,r,s] -> ion particle dfn -function derivative_r!(dfdr::AbstractArray{mk_float,5}, f::AbstractArray{mk_float,5}, - advect, adv_fac_lower_buffer::AbstractArray{mk_float,4}, - adv_fac_upper_buffer::AbstractArray{mk_float,4}, - dfdr_lower_endpoints::AbstractArray{mk_float,4}, - dfdr_upper_endpoints::AbstractArray{mk_float,4}, - r_receive_buffer1::AbstractArray{mk_float,4}, - r_receive_buffer2::AbstractArray{mk_float,4}, r_spectral, r) +function derivative_r!(dfdr::MKOrAbstractArray{mk_float,5}, f::MKOrAbstractArray{mk_float,5}, + advect, adv_fac_lower_buffer::MKOrAbstractArray{mk_float,4}, + adv_fac_upper_buffer::MKOrAbstractArray{mk_float,4}, + dfdr_lower_endpoints::MKOrAbstractArray{mk_float,4}, + dfdr_upper_endpoints::MKOrAbstractArray{mk_float,4}, + r_receive_buffer1::MKOrAbstractArray{mk_float,4}, + r_receive_buffer2::MKOrAbstractArray{mk_float,4}, r_spectral, r) begin_s_z_vperp_vpa_region() @@ -709,13 +709,13 @@ function derivative_r!(dfdr::AbstractArray{mk_float,5}, f::AbstractArray{mk_floa end #6D version for f[vz,vz,vzeta,z,r,sn] -> neutral particle dfn (species indexing taken outside this loop) -function derivative_r!(dfdr::AbstractArray{mk_float,6}, f::AbstractArray{mk_float,6}, - advect, adv_fac_lower_buffer::AbstractArray{mk_float,5}, - adv_fac_upper_buffer::AbstractArray{mk_float,5}, - dfdr_lower_endpoints::AbstractArray{mk_float,5}, - dfdr_upper_endpoints::AbstractArray{mk_float,5}, - r_receive_buffer1::AbstractArray{mk_float,5}, - r_receive_buffer2::AbstractArray{mk_float,5}, r_spectral, r) +function derivative_r!(dfdr::MKOrAbstractArray{mk_float,6}, f::MKOrAbstractArray{mk_float,6}, + advect, adv_fac_lower_buffer::MKOrAbstractArray{mk_float,5}, + adv_fac_upper_buffer::MKOrAbstractArray{mk_float,5}, + dfdr_lower_endpoints::MKOrAbstractArray{mk_float,5}, + dfdr_upper_endpoints::MKOrAbstractArray{mk_float,5}, + r_receive_buffer1::MKOrAbstractArray{mk_float,5}, + r_receive_buffer2::MKOrAbstractArray{mk_float,5}, r_spectral, r) begin_sn_z_vzeta_vr_vz_region() @@ -749,13 +749,13 @@ dfns (neutrals) -> [vz,vr,vzeta,z,r,sn] """ #2D version for f[z,r] -> Er, Ez, phi -function derivative_z!(dfdz::AbstractArray{mk_float,2}, f::AbstractArray{mk_float,2}, - adv_fac, adv_fac_lower_buffer::AbstractArray{mk_float,1}, - adv_fac_upper_buffer::AbstractArray{mk_float,1}, - dfdz_lower_endpoints::AbstractArray{mk_float,1}, - dfdz_upper_endpoints::AbstractArray{mk_float,1}, - z_send_buffer::AbstractArray{mk_float,1}, - z_receive_buffer::AbstractArray{mk_float,1}, z_spectral, z) +function derivative_z!(dfdz::MKOrAbstractArray{mk_float,2}, f::MKOrAbstractArray{mk_float,2}, + adv_fac, adv_fac_lower_buffer::MKOrAbstractArray{mk_float,1}, + adv_fac_upper_buffer::MKOrAbstractArray{mk_float,1}, + dfdz_lower_endpoints::MKOrAbstractArray{mk_float,1}, + dfdz_upper_endpoints::MKOrAbstractArray{mk_float,1}, + z_send_buffer::MKOrAbstractArray{mk_float,1}, + z_receive_buffer::MKOrAbstractArray{mk_float,1}, z_spectral, z) begin_r_region() @@ -778,13 +778,13 @@ function derivative_z!(dfdz::AbstractArray{mk_float,2}, f::AbstractArray{mk_floa end #3D version for f[z,r] -> moments n, u, T etc -function derivative_z!(dfdz::AbstractArray{mk_float,3}, f::AbstractArray{mk_float,3}, - adv_fac, adv_fac_lower_buffer::AbstractArray{mk_float,2}, - adv_fac_upper_buffer::AbstractArray{mk_float,2}, - dfdz_lower_endpoints::AbstractArray{mk_float,2}, - dfdz_upper_endpoints::AbstractArray{mk_float,2}, - z_send_buffer::AbstractArray{mk_float,2}, - z_receive_buffer::AbstractArray{mk_float,2}, z_spectral, z; neutrals=false) +function derivative_z!(dfdz::MKOrAbstractArray{mk_float,3}, f::MKOrAbstractArray{mk_float,3}, + adv_fac, adv_fac_lower_buffer::MKOrAbstractArray{mk_float,2}, + adv_fac_upper_buffer::MKOrAbstractArray{mk_float,2}, + dfdz_lower_endpoints::MKOrAbstractArray{mk_float,2}, + dfdz_upper_endpoints::MKOrAbstractArray{mk_float,2}, + z_send_buffer::MKOrAbstractArray{mk_float,2}, + z_receive_buffer::MKOrAbstractArray{mk_float,2}, z_spectral, z; neutrals=false) # differentiate f w.r.t z if neutrals @@ -823,13 +823,13 @@ end # df/dz # 3D version for f[vpa,vperp,z]. Uses modified function name to avoid clash with 'standard' # 3D version for ion/neutral moments. -function derivative_z_pdf_vpavperpz!(dfdz::AbstractArray{mk_float,3}, f::AbstractArray{mk_float,3}, - adv_fac, adv_fac_lower_buffer::AbstractArray{mk_float,2}, - adv_fac_upper_buffer::AbstractArray{mk_float,2}, - dfdz_lower_endpoints::AbstractArray{mk_float,2}, - dfdz_upper_endpoints::AbstractArray{mk_float,2}, - z_receive_buffer1::AbstractArray{mk_float,2}, - z_receive_buffer2::AbstractArray{mk_float,2}, z_spectral, z) +function derivative_z_pdf_vpavperpz!(dfdz::MKOrAbstractArray{mk_float,3}, f::MKOrAbstractArray{mk_float,3}, + adv_fac, adv_fac_lower_buffer::MKOrAbstractArray{mk_float,2}, + adv_fac_upper_buffer::MKOrAbstractArray{mk_float,2}, + dfdz_lower_endpoints::MKOrAbstractArray{mk_float,2}, + dfdz_upper_endpoints::MKOrAbstractArray{mk_float,2}, + z_receive_buffer1::MKOrAbstractArray{mk_float,2}, + z_receive_buffer2::MKOrAbstractArray{mk_float,2}, z_spectral, z) # differentiate f w.r.t z @loop_vperp_vpa ivperp ivpa begin @@ -851,13 +851,13 @@ function derivative_z_pdf_vpavperpz!(dfdz::AbstractArray{mk_float,3}, f::Abstrac end #5D version for f[vpa,vperp,z,r,s] -> dfn ion particles -function derivative_z!(dfdz::AbstractArray{mk_float,5}, f::AbstractArray{mk_float,5}, - advect, adv_fac_lower_buffer::AbstractArray{mk_float,4}, - adv_fac_upper_buffer::AbstractArray{mk_float,4}, - dfdz_lower_endpoints::AbstractArray{mk_float,4}, - dfdz_upper_endpoints::AbstractArray{mk_float,4}, - z_send_buffer::AbstractArray{mk_float,4}, - z_receive_buffer::AbstractArray{mk_float,4}, z_spectral, z) +function derivative_z!(dfdz::MKOrAbstractArray{mk_float,5}, f::MKOrAbstractArray{mk_float,5}, + advect, adv_fac_lower_buffer::MKOrAbstractArray{mk_float,4}, + adv_fac_upper_buffer::MKOrAbstractArray{mk_float,4}, + dfdz_lower_endpoints::MKOrAbstractArray{mk_float,4}, + dfdz_upper_endpoints::MKOrAbstractArray{mk_float,4}, + z_send_buffer::MKOrAbstractArray{mk_float,4}, + z_receive_buffer::MKOrAbstractArray{mk_float,4}, z_spectral, z) begin_s_r_vperp_vpa_region() @@ -882,13 +882,13 @@ function derivative_z!(dfdz::AbstractArray{mk_float,5}, f::AbstractArray{mk_floa end #4D version for f[vpa,vperp,z,r] -> dfn electron particles -function derivative_z!(dfdz::AbstractArray{mk_float,4}, f::AbstractArray{mk_float,4}, - advect, adv_fac_lower_buffer::AbstractArray{mk_float,3}, - adv_fac_upper_buffer::AbstractArray{mk_float,3}, - dfdz_lower_endpoints::AbstractArray{mk_float,3}, - dfdz_upper_endpoints::AbstractArray{mk_float,3}, - z_send_buffer::AbstractArray{mk_float,3}, - z_receive_buffer::AbstractArray{mk_float,3}, z_spectral, z) +function derivative_z!(dfdz::MKOrAbstractArray{mk_float,4}, f::MKOrAbstractArray{mk_float,4}, + advect, adv_fac_lower_buffer::MKOrAbstractArray{mk_float,3}, + adv_fac_upper_buffer::MKOrAbstractArray{mk_float,3}, + dfdz_lower_endpoints::MKOrAbstractArray{mk_float,3}, + dfdz_upper_endpoints::MKOrAbstractArray{mk_float,3}, + z_send_buffer::MKOrAbstractArray{mk_float,3}, + z_receive_buffer::MKOrAbstractArray{mk_float,3}, z_spectral, z) begin_r_vperp_vpa_region() @@ -913,13 +913,13 @@ function derivative_z!(dfdz::AbstractArray{mk_float,4}, f::AbstractArray{mk_floa end #6D version for f[vz,vr,vzeta,z,r,sn] -> dfn neutral particles -function derivative_z!(dfdz::AbstractArray{mk_float,6}, f::AbstractArray{mk_float,6}, - advect, adv_fac_lower_buffer::AbstractArray{mk_float,5}, - adv_fac_upper_buffer::AbstractArray{mk_float,5}, - dfdz_lower_endpoints::AbstractArray{mk_float,5}, - dfdz_upper_endpoints::AbstractArray{mk_float,5}, - z_send_buffer::AbstractArray{mk_float,5}, - z_receive_buffer::AbstractArray{mk_float,5}, z_spectral, z) +function derivative_z!(dfdz::MKOrAbstractArray{mk_float,6}, f::MKOrAbstractArray{mk_float,6}, + advect, adv_fac_lower_buffer::MKOrAbstractArray{mk_float,5}, + adv_fac_upper_buffer::MKOrAbstractArray{mk_float,5}, + dfdz_lower_endpoints::MKOrAbstractArray{mk_float,5}, + dfdz_upper_endpoints::MKOrAbstractArray{mk_float,5}, + z_send_buffer::MKOrAbstractArray{mk_float,5}, + z_receive_buffer::MKOrAbstractArray{mk_float,5}, z_spectral, z) begin_sn_r_vzeta_vr_vz_region() diff --git a/moment_kinetics/src/electron_fluid_equations.jl b/moment_kinetics/src/electron_fluid_equations.jl index c77d8c1ac..4bef03c1b 100644 --- a/moment_kinetics/src/electron_fluid_equations.jl +++ b/moment_kinetics/src/electron_fluid_equations.jl @@ -18,7 +18,7 @@ using ..input_structs using ..timer_utils using ..moment_kinetics_structs: electron_pdf_substruct, moments_electron_substruct using ..nonlinear_solvers -using ..type_definitions: mk_float +using ..type_definitions using ..velocity_moments: integrate_over_vspace using MPI @@ -640,11 +640,11 @@ Add just the braginskii conduction contribution to the electron pressure, and as we have to calculate qpar and dqpar_dz from ppar within this function (they are not pre-calculated). """ -function electron_braginskii_conduction!(ppar_out::AbstractVector{mk_float}, - ppar_in::AbstractVector{mk_float}, - dens::AbstractVector{mk_float}, - upar_e::AbstractVector{mk_float}, - upar_i::AbstractVector{mk_float}, +function electron_braginskii_conduction!(ppar_out::MKOrAbstractVector{mk_float}, + ppar_in::MKOrAbstractVector{mk_float}, + dens::MKOrAbstractVector{mk_float}, + upar_e::MKOrAbstractVector{mk_float}, + upar_i::MKOrAbstractVector{mk_float}, electron_moments, collisions, composition, z, z_spectral, scratch_dummy, dt, ir) diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl index d023dc29a..594861547 100644 --- a/moment_kinetics/src/electron_kinetic_equation.jl +++ b/moment_kinetics/src/electron_kinetic_equation.jl @@ -4888,7 +4888,7 @@ in the time derivative term as it is for the non-boundary points.] begin_vperp_vpa_region() update_electron_speed_z!(z_advect[1], upar, vth, vpa.grid, ir) @loop_vperp_vpa ivperp ivpa begin - @views z_advect[1].adv_fac[:,ivpa,ivperp,ir] = -z_speed[:,ivpa,ivperp] + @views @. z_advect[1].adv_fac[:,ivpa,ivperp,ir] = -z_speed[:,ivpa,ivperp] end #calculate the upwind derivative @views derivative_z_pdf_vpavperpz!(dpdf_dz, f, z_advect[1].adv_fac[:,:,:,ir], diff --git a/moment_kinetics/src/em_fields.jl b/moment_kinetics/src/em_fields.jl index 7b3544f08..08db47b41 100644 --- a/moment_kinetics/src/em_fields.jl +++ b/moment_kinetics/src/em_fields.jl @@ -129,10 +129,13 @@ function update_phi!(fields, fvec, vperp, z, r, composition, collisions, moments ## calculate the electric fields after obtaining phi #Er = - d phi / dr if r.n > 1 - derivative_r!(fields.Er,-fields.phi, + derivative_r!(fields.Er,fields.phi, scratch_dummy.buffer_z_1, scratch_dummy.buffer_z_2, scratch_dummy.buffer_z_3, scratch_dummy.buffer_z_4, r_spectral,r) + @loop_r_z ir iz begin + fields.Er[iz,ir] *= -1.0 + end if z.irank == 0 && fields.force_Er_zero_at_wall fields.Er[1,:] .= 0.0 end @@ -150,10 +153,13 @@ function update_phi!(fields, fvec, vperp, z, r, composition, collisions, moments kinetic_electrons_with_temperature_equation) if z.n > 1 # Ez = - d phi / dz - @views derivative_z!(fields.Ez,-fields.phi, + @views derivative_z!(fields.Ez, fields.phi, scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1], scratch_dummy.buffer_rs_4[:,1], z_spectral,z) + @loop_r_z ir iz begin + fields.Ez[iz,ir] *= -1.0 + end end end diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl index 199e69b49..e9e6641cf 100644 --- a/moment_kinetics/src/file_io.jl +++ b/moment_kinetics/src/file_io.jl @@ -690,7 +690,7 @@ function get_variable_keys end """ write_single_value!(file_or_group, name, - data::Union{Number, AbstractString, AbstractArray{T,N}}, + data::Union{Number, AbstractString, MKArray{T,N}}, coords::Union{coordinate,mk_int,NamedTuple}...; parallel_io, description=nothing, units=nothing, overwrite=false) where {T,N} @@ -2710,7 +2710,7 @@ function write_timing_data(io_moments, t_idx, dfns=false) # Gather the strings gathered_char_vector = MKVector{Char}(undef, sum(string_sizes)) - local_start_index = sum(string_sizes[1:comm_rank]) + 1 + local_start_index = sum(@view string_sizes[1:comm_rank]) + 1 local_end_index = local_start_index - 1 + string_sizes[comm_rank+1] gathered_char_vector[local_start_index:local_end_index] .= [new_names_string...] gathered_buffer = MPI.VBuffer(gathered_char_vector, string_sizes) diff --git a/moment_kinetics/src/file_io_hdf5.jl b/moment_kinetics/src/file_io_hdf5.jl index d647ab554..f717896cd 100644 --- a/moment_kinetics/src/file_io_hdf5.jl +++ b/moment_kinetics/src/file_io_hdf5.jl @@ -87,7 +87,7 @@ end # HDF5.H5DataStore is the supertype for HDF5.File and HDF5.Group function write_single_value!(file_or_group::HDF5.H5DataStore, name, - data::Union{Number, AbstractString, AbstractArray{T,N}}, + data::Union{Number, AbstractString, MKOrAbstractArray{T,N}}, coords::Union{coordinate,mk_int,NamedTuple}...; parallel_io, description=nothing, units=nothing, overwrite=false) where {T,N} @@ -259,7 +259,7 @@ function extend_time_index!(h5, t_idx) end function append_to_dynamic_var(io_var::HDF5.Dataset, - data::Union{Nothing,Number,AbstractArray{T,N}}, t_idx, + data::Union{Nothing,Number,MKOrAbstractArray{T,N}}, t_idx, parallel_io::Bool, coords::Union{coordinate,NamedTuple,Integer}...; only_root=false, write_from_this_rank=nothing) where {T,N} diff --git a/moment_kinetics/src/finite_differences.jl b/moment_kinetics/src/finite_differences.jl index 4cc4c9e6a..8d6e8f81c 100644 --- a/moment_kinetics/src/finite_differences.jl +++ b/moment_kinetics/src/finite_differences.jl @@ -2,7 +2,7 @@ """ module finite_differences -using ..type_definitions: mk_float +using ..type_definitions import ..calculus: elementwise_derivative!, second_derivative!, derivative_elements_to_full_grid! using ..moment_kinetics_structs: discretization_info @@ -371,7 +371,7 @@ take the derivative of input function f and return as df using second-order, centered differences. input/output array df is 2D array of size ngrid x nelement """ -function centered_second_order!(df::AbstractArray{mk_float,2}, f, del, bc, igrid, ielement) +function centered_second_order!(df::MKOrAbstractArray{mk_float,2}, f, del, bc, igrid, ielement) n = length(f) # get derivative at internal points for i ∈ 2:n-1 @@ -411,7 +411,7 @@ take the derivative of input function f and return as df using second-order, centered differences. input/output df is 1D array of size n (full grid) """ -function centered_second_order!(df::AbstractArray{mk_float,1}, f, del, bc, igrid, ielement) +function centered_second_order!(df::MKOrAbstractArray{mk_float,1}, f, del, bc, igrid, ielement) n = length(f) # get derivative at internal points for i ∈ 2:n-1 @@ -451,7 +451,7 @@ take the derivative of input function f and return as df using fourth-order, centered differences. input/output array df is 2D array of size ngrid x nelement """ -function centered_fourth_order!(df::AbstractArray{mk_float,2}, f, del, bc, igrid, ielement) +function centered_fourth_order!(df::MKOrAbstractArray{mk_float,2}, f, del, bc, igrid, ielement) n = length(f) # get derivative at internal points for i ∈ 3:n-2 @@ -505,7 +505,7 @@ Take the second derivative of input function f and return as df using second-ord centered differences. output array df is 2D array of size ngrid x nelement """ -function second_derivative_finite_difference!(df::AbstractArray{mk_float,2}, f, del, bc, igrid, ielement) +function second_derivative_finite_difference!(df::MKOrAbstractArray{mk_float,2}, f, del, bc, igrid, ielement) n = length(f) # get derivative at internal points for i ∈ 2:n-1 diff --git a/moment_kinetics/src/fokker_planck.jl b/moment_kinetics/src/fokker_planck.jl index f2f6aed37..6e78dab1e 100644 --- a/moment_kinetics/src/fokker_planck.jl +++ b/moment_kinetics/src/fokker_planck.jl @@ -46,7 +46,7 @@ using Dates using LinearAlgebra: lu, ldiv! using MPI using OrderedCollections: OrderedDict -using ..type_definitions: mk_float, mk_int +using ..type_definitions using ..array_allocation: allocate_float, allocate_shared_float using ..communication using ..velocity_moments: integrate_over_vspace @@ -519,9 +519,9 @@ are specified using analytical results. """ @timeit global_timer fokker_planck_collision_operator_weak_form_Maxwellian_Fsp!( ffs_in, nuref::mk_float, ms::mk_float, Zs::mk_float, - msp::AbstractArray{mk_float,1}, Zsp::AbstractArray{mk_float,1}, - densp::AbstractArray{mk_float,1}, uparsp::AbstractArray{mk_float,1}, - vthsp::AbstractArray{mk_float,1}, + msp::MKOrAbstractArray{mk_float,1}, Zsp::MKOrAbstractArray{mk_float,1}, + densp::MKOrAbstractArray{mk_float,1}, uparsp::MKOrAbstractArray{mk_float,1}, + vthsp::MKOrAbstractArray{mk_float,1}, fkpl_arrays::fokkerplanck_weakform_arrays_struct, vperp, vpa, vperp_spectral, vpa_spectral) = begin @boundscheck vpa.n == size(ffs_in,1) || throw(BoundsError(ffs_in)) diff --git a/moment_kinetics/src/fokker_planck_calculus.jl b/moment_kinetics/src/fokker_planck_calculus.jl index 09a941b5b..85e418525 100644 --- a/moment_kinetics/src/fokker_planck_calculus.jl +++ b/moment_kinetics/src/fokker_planck_calculus.jl @@ -1267,7 +1267,7 @@ function calculate_rosenbluth_potential_boundary_data!(rpbd::rosenbluth_potentia return nothing end -function multipole_H(vpa::mk_float,vperp::mk_float,Inn_vec::AbstractVector{mk_float}) +function multipole_H(vpa::mk_float,vperp::mk_float,Inn_vec::MKOrAbstractVector{mk_float}) (I00, I10, I20, I30, I40, I50, I60, I70, I80, I02, I12, I22, I32, I42, I52, I62, I04, I14, I24, I34, I44, @@ -1304,7 +1304,7 @@ function multipole_H(vpa::mk_float,vperp::mk_float,Inn_vec::AbstractVector{mk_fl return H_series end -function multipole_dHdvpa(vpa::mk_float,vperp::mk_float,Inn_vec::AbstractVector{mk_float}) +function multipole_dHdvpa(vpa::mk_float,vperp::mk_float,Inn_vec::MKOrAbstractVector{mk_float}) (I00, I10, I20, I30, I40, I50, I60, I70, I80, I02, I12, I22, I32, I42, I52, I62, I04, I14, I24, I34, I44, @@ -1341,7 +1341,7 @@ function multipole_dHdvpa(vpa::mk_float,vperp::mk_float,Inn_vec::AbstractVector{ return dHdvpa_series end -function multipole_dHdvperp(vpa::mk_float,vperp::mk_float,Inn_vec::AbstractVector{mk_float}) +function multipole_dHdvperp(vpa::mk_float,vperp::mk_float,Inn_vec::MKOrAbstractVector{mk_float}) (I00, I10, I20, I30, I40, I50, I60, I70, I80, I02, I12, I22, I32, I42, I52, I62, I04, I14, I24, I34, I44, @@ -1378,7 +1378,7 @@ function multipole_dHdvperp(vpa::mk_float,vperp::mk_float,Inn_vec::AbstractVecto return dHdvperp_series end -function multipole_G(vpa::mk_float,vperp::mk_float,Inn_vec::AbstractVector{mk_float}) +function multipole_G(vpa::mk_float,vperp::mk_float,Inn_vec::MKOrAbstractVector{mk_float}) (I00, I10, I20, I30, I40, I50, I60, I70, I80, I02, I12, I22, I32, I42, I52, I62, I04, I14, I24, I34, I44, @@ -1415,7 +1415,7 @@ function multipole_G(vpa::mk_float,vperp::mk_float,Inn_vec::AbstractVector{mk_fl return G_series end -function multipole_dGdvperp(vpa::mk_float,vperp::mk_float,Inn_vec::AbstractVector{mk_float}) +function multipole_dGdvperp(vpa::mk_float,vperp::mk_float,Inn_vec::MKOrAbstractVector{mk_float}) (I00, I10, I20, I30, I40, I50, I60, I70, I80, I02, I12, I22, I32, I42, I52, I62, I04, I14, I24, I34, I44, @@ -1452,7 +1452,7 @@ function multipole_dGdvperp(vpa::mk_float,vperp::mk_float,Inn_vec::AbstractVecto return dGdvperp_series end -function multipole_d2Gdvperp2(vpa::mk_float,vperp::mk_float,Inn_vec::AbstractVector{mk_float}) +function multipole_d2Gdvperp2(vpa::mk_float,vperp::mk_float,Inn_vec::MKOrAbstractVector{mk_float}) (I00, I10, I20, I30, I40, I50, I60, I70, I80, I02, I12, I22, I32, I42, I52, I62, I04, I14, I24, I34, I44, @@ -1489,7 +1489,7 @@ function multipole_d2Gdvperp2(vpa::mk_float,vperp::mk_float,Inn_vec::AbstractVec return d2Gdvperp2_series end -function multipole_d2Gdvperpdvpa(vpa::mk_float,vperp::mk_float,Inn_vec::AbstractVector{mk_float}) +function multipole_d2Gdvperpdvpa(vpa::mk_float,vperp::mk_float,Inn_vec::MKOrAbstractVector{mk_float}) (I00, I10, I20, I30, I40, I50, I60, I70, I80, I02, I12, I22, I32, I42, I52, I62, I04, I14, I24, I34, I44, @@ -1526,7 +1526,7 @@ function multipole_d2Gdvperpdvpa(vpa::mk_float,vperp::mk_float,Inn_vec::Abstract return d2Gdvperpdvpa_series end -function multipole_d2Gdvpa2(vpa::mk_float,vperp::mk_float,Inn_vec::AbstractVector{mk_float}) +function multipole_d2Gdvpa2(vpa::mk_float,vperp::mk_float,Inn_vec::MKOrAbstractVector{mk_float}) (I00, I10, I20, I30, I40, I50, I60, I70, I80, I02, I12, I22, I32, I42, I52, I62, I04, I14, I24, I34, I44, diff --git a/moment_kinetics/src/gauss_legendre.jl b/moment_kinetics/src/gauss_legendre.jl index 0f1e4af53..a6d023fbc 100644 --- a/moment_kinetics/src/gauss_legendre.jl +++ b/moment_kinetics/src/gauss_legendre.jl @@ -474,7 +474,7 @@ Or https://doc.nektar.info/tutorials/latest/fundamentals/differentiation/fundame Note that D has does not include a scaling factor """ -function gausslobattolegendre_differentiation_matrix!(D::AbstractArray{mk_float,2},x::AbstractArray{mk_float,1},ngrid::mk_int) +function gausslobattolegendre_differentiation_matrix!(D::MKOrAbstractArray{mk_float,2},x::MKOrAbstractArray{mk_float,1},ngrid::mk_int) D[:,:] .= 0.0 for ix in 1:ngrid for ixp in 1:ngrid @@ -506,7 +506,7 @@ https://doc.nektar.info/tutorials/latest/fundamentals/differentiation/fundamenta Note that D has does not include a scaling factor """ -function gaussradaulegendre_differentiation_matrix!(D::AbstractArray{mk_float,2},x::AbstractArray{mk_float,1},ngrid::Int64) +function gaussradaulegendre_differentiation_matrix!(D::MKOrAbstractArray{mk_float,2},x::MKOrAbstractArray{mk_float,1},ngrid::mk_int) D[:,:] .= 0.0 for ix in 1:ngrid for ixp in 1:ngrid @@ -616,7 +616,7 @@ polynomials with order \$k_{max} = 4N +1\$, with \$N=\$`ngrid` and the highest o that we integrate is \$P_{N-1}(x)P_{N-1}(x)x^2\$, which has order \$k=2N < k_{max}\$. """ -function GaussLegendre_weak_product_matrix!(QQ::AbstractArray{mk_float,2},ngrid,x,wgts,option;radau=false) +function GaussLegendre_weak_product_matrix!(QQ::MKOrAbstractArray{mk_float,2},ngrid,x,wgts,option;radau=false) # coefficient in expansion of # lagrange polys in terms of Legendre polys gamma = allocate_float(ngrid) @@ -754,7 +754,7 @@ matrix `Q` acts on two vectors `x1` and `x2` such that the quadratic form `y = x1 * Q * x2` is also a vector. See documentation of corresponding function for linear inner product matrices. """ -function GaussLegendre_weak_product_matrix!(QQ::AbstractArray{mk_float,3},ngrid,x,wgts,option;radau=false) +function GaussLegendre_weak_product_matrix!(QQ::MKOrAbstractArray{mk_float,3},ngrid,x,wgts,option;radau=false) # coefficient in expansion of # lagrange polys in terms of Legendre polys gamma = allocate_float(ngrid) @@ -1022,7 +1022,7 @@ in the function call, and create new matrices for this purpose in the `gausslegendre_info` struct. Currently the Laplacian matrix is supported with boundary conditions. """ -function setup_global_weak_form_matrix!(QQ_global::AbstractArray{mk_float,2}, +function setup_global_weak_form_matrix!(QQ_global::MKOrAbstractArray{mk_float,2}, lobatto::gausslegendre_base_info, radau::gausslegendre_base_info, coord,option; dirichlet_bc=false, periodic_bc=false) @@ -1100,7 +1100,7 @@ The shared points in the element assembly are averaged (instead of simply added) to be consistent with the `derivative_elements_to_full_grid!()` function in `calculus.jl`. """ -function setup_global_strong_form_matrix!(QQ_global::AbstractArray{mk_float,2}, +function setup_global_strong_form_matrix!(QQ_global::MKOrAbstractArray{mk_float,2}, lobatto::gausslegendre_base_info, radau::gausslegendre_base_info, coord,option; periodic_bc=false) @@ -1146,7 +1146,7 @@ function setup_global_strong_form_matrix!(QQ_global::AbstractArray{mk_float,2}, # upper boundary assembly on element if j == coord.nelement_local if periodic_bc && coord.nrank == 1 - QQ_global[imax[j],imin[j]-1:imax[j]] .+= QQ_j[ngrid,:] / 2.0 + QQ_global[imax[j],imin[j]-1:imax[j]] .+= QQ_j[ngrid,:] ./ 2.0 else QQ_global[imax[j],imin[j]-1:imax[j]] .+= QQ_j[ngrid,:] end @@ -1162,7 +1162,7 @@ end Construction function to provide the appropriate elemental matrix `Q` to the global matrix assembly functions. """ -function get_QQ_local!(QQ::AbstractArray{mk_float,2},ielement, +function get_QQ_local!(QQ::MKOrAbstractArray{mk_float,2},ielement, lobatto::gausslegendre_base_info, radau::gausslegendre_base_info, coord,option) @@ -1543,7 +1543,7 @@ end Construction function for nonlinear diffusion matrices, only used in the assembly of the collision operator """ -function get_QQ_local!(QQ::AbstractArray{mk_float,3}, +function get_QQ_local!(QQ::MKOrAbstractArray{mk_float,3}, ielement,lobatto::gausslegendre_base_info, radau::gausslegendre_base_info, coord,option) diff --git a/moment_kinetics/src/interpolation.jl b/moment_kinetics/src/interpolation.jl index 419c1de87..cc9c5993d 100644 --- a/moment_kinetics/src/interpolation.jl +++ b/moment_kinetics/src/interpolation.jl @@ -10,7 +10,7 @@ export interpolate_to_grid_z, interpolate_to_grid_1d!, interpolate_symmetric!, using ..array_allocation: allocate_float using ..moment_kinetics_structs: null_spatial_dimension_info, null_velocity_dimension_info -using ..type_definitions: mk_float, mk_int +using ..type_definitions """ single_element_interpolate!(result, newgrid, f, imin, imax, coord, spectral) @@ -294,7 +294,7 @@ end """ """ -function interpolate_to_grid_z!(result::AbstractArray{mk_float, 3}, newgrid, f::AbstractArray{mk_float, 3}, z, spectral) +function interpolate_to_grid_z!(result::MKOrAbstractArray{mk_float, 3}, newgrid, f::MKOrAbstractArray{mk_float, 3}, z, spectral) size_f = size(f) for is ∈ 1:size_f[3] for ivpa ∈ 1:size_f[1] @@ -307,7 +307,7 @@ end """ """ -function interpolate_to_grid_z(newgrid, f::AbstractArray{mk_float, 3}, z, spectral) +function interpolate_to_grid_z(newgrid, f::MKOrAbstractArray{mk_float, 3}, z, spectral) size_f = size(f) result = allocate_float(size_f[1], size(newgrid)[1], size_f[3]) @@ -318,7 +318,7 @@ end """ """ -function interpolate_to_grid_z!(result::AbstractArray{mk_float, 2}, newgrid, f::AbstractArray{mk_float, 2}, z, spectral) +function interpolate_to_grid_z!(result::MKOrAbstractArray{mk_float, 2}, newgrid, f::MKOrAbstractArray{mk_float, 2}, z, spectral) size_f = size(f) for is ∈ 1:size_f[2] @views interpolate_to_grid_1d!(result[:, is], newgrid, f[:, is], z, spectral) @@ -329,7 +329,7 @@ end """ """ -function interpolate_to_grid_z(newgrid, f::AbstractArray{mk_float, 2}, z, spectral) +function interpolate_to_grid_z(newgrid, f::MKOrAbstractArray{mk_float, 2}, z, spectral) size_f = size(f) result = allocate_float(size(newgrid)[1], size_f[2]) @@ -340,7 +340,7 @@ end """ """ -function interpolate_to_grid_z!(result::AbstractArray{mk_float, 1}, newgrid, f::AbstractArray{mk_float, 1}, z, spectral) +function interpolate_to_grid_z!(result::MKOrAbstractArray{mk_float, 1}, newgrid, f::MKOrAbstractArray{mk_float, 1}, z, spectral) interpolate_to_grid_1d!(result, newgrid, f, z, spectral) return nothing @@ -348,13 +348,13 @@ end """ """ -function interpolate_to_grid_z(newgrid, f::AbstractArray{mk_float, 1}, z, spectral) +function interpolate_to_grid_z(newgrid, f::MKOrAbstractArray{mk_float, 1}, z, spectral) return interpolate_to_grid_1d(newgrid, f, z, spectral) end """ """ -function interpolate_to_grid_vpa!(result::AbstractArray{mk_float, 3}, newgrid, f::AbstractArray{mk_float, 3}, vpa, spectral) +function interpolate_to_grid_vpa!(result::MKOrAbstractArray{mk_float, 3}, newgrid, f::MKOrAbstractArray{mk_float, 3}, vpa, spectral) size_f = size(f) for is ∈ 1:size_f[3] for iz ∈ 1:size_f[2] @@ -367,7 +367,7 @@ end """ """ -function interpolate_to_grid_vpa(newgrid, f::AbstractArray{mk_float, 3}, vpa, spectral) +function interpolate_to_grid_vpa(newgrid, f::MKOrAbstractArray{mk_float, 3}, vpa, spectral) size_f = size(f) result = allocate_float(size(newgrid)[1], size_f[2:3]...) @@ -378,8 +378,8 @@ end """ """ -function interpolate_to_grid_vpa!(result::AbstractVector{mk_float}, newgrid, - f::AbstractVector{mk_float}, vpa, spectral) +function interpolate_to_grid_vpa!(result::MKOrAbstractVector{mk_float}, newgrid, + f::MKOrAbstractVector{mk_float}, vpa, spectral) interpolate_to_grid_1d!(result, newgrid, f, vpa, spectral) @@ -388,7 +388,7 @@ end """ """ -function interpolate_to_grid_vpa(newgrid, f::AbstractVector{mk_float}, vpa, spectral) +function interpolate_to_grid_vpa(newgrid, f::MKOrAbstractVector{mk_float}, vpa, spectral) return interpolate_to_grid_1d(newgrid, f, vpa, spectral) end diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl index 0a54dbe41..b6ce7c862 100644 --- a/moment_kinetics/src/load_data.jl +++ b/moment_kinetics/src/load_data.jl @@ -3819,9 +3819,11 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing, !isa(it, mk_int) && push!(dims, nt) vartype = typeof(variable[1][1]) if vartype == mk_int - result = allocate_int(dims...) + #result = allocate_int(dims...) + result = zeros(mk_int, dims...) elseif vartype == mk_float - result = allocate_float(dims...) + #result = allocate_float(dims...) + result = zeros(mk_float, dims...) else error("Unsupported dtype for 1D variable $(variable.dtype)") end @@ -3835,9 +3837,11 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing, !isa(it, mk_int) && push!(dims, nt) vartype = typeof(variable[1][1,1]) if vartype == mk_int - result = allocate_int(dims...) + #result = allocate_int(dims...) + result = zeros(mk_int, dims...) elseif vartype == mk_float - result = allocate_float(dims...) + #result = allocate_float(dims...) + result = zeros(mk_float, dims...) else error("Unsupported dtype for 1D variable $(variable.dtype)") end @@ -3847,7 +3851,8 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing, !isa(iz, mk_int) && push!(dims, nz) !isa(ir, mk_int) && push!(dims, nr) !isa(it, mk_int) && push!(dims, nt) - result = allocate_float(dims...) + #result = allocate_float(dims...) + result = zeros(mk_float, dims...) elseif nd == 4 # moment variable with dimensions (z,r,s,t) # Get nspecies from the variable, not from run_info, because it might be @@ -3862,7 +3867,8 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing, push!(dims, nspecies) end !isa(it, mk_int) && push!(dims, nt) - result = allocate_float(dims...) + #result = allocate_float(dims...) + result = zeros(mk_float, dims...) elseif nd == 5 # electron distribution function variable with dimensions (vpa,vperp,z,r,t) dims = Vector{mk_int}() @@ -3871,7 +3877,8 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing, !isa(iz, mk_int) && push!(dims, nz) !isa(ir, mk_int) && push!(dims, nr) !isa(it, mk_int) && push!(dims, nt) - result = allocate_float(dims...) + #result = allocate_float(dims...) + result = zeros(mk_float, dims...) elseif nd == 6 # ion distribution function variable with dimensions (vpa,vperp,z,r,s,t) nspecies = size(variable[1], 5) @@ -3886,7 +3893,8 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing, push!(dims, nspecies) end !isa(it, mk_int) && push!(dims, nt) - result = allocate_float(dims...) + #result = allocate_float(dims...) + result = zeros(mk_float, dims...) elseif nd == 7 # neutral distribution function variable with dimensions (vz,vr,vzeta,z,r,s,t) nspecies = size(variable[1], 6) @@ -3902,7 +3910,8 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing, push!(dims, nspecies) end !isa(it, mk_int) && push!(dims, nt) - result = allocate_float(dims...) + #result = allocate_float(dims...) + result = zeros(mk_float, dims...) else error("Unsupported number of dimensions ($nd) for '$variable_name'.") end diff --git a/moment_kinetics/src/moment_constraints.jl b/moment_kinetics/src/moment_constraints.jl index 7f099b541..4fac549f7 100644 --- a/moment_kinetics/src/moment_constraints.jl +++ b/moment_kinetics/src/moment_constraints.jl @@ -9,7 +9,7 @@ using ..boundary_conditions: skip_f_electron_bc_points_in_Jacobian using ..communication: _block_synchronize using ..looping using ..timer_utils -using ..type_definitions: mk_float +using ..type_definitions using ..velocity_moments: integrate_over_vspace, update_ion_qpar! export hard_force_moment_constraints!, hard_force_moment_constraints_neutral!, @@ -91,7 +91,7 @@ function hard_force_moment_constraints!(f, moments, vpa) return A, B, C end @timeit global_timer hard_force_moment_constraints!( - f::AbstractArray{mk_float,4}, moments, vpa) = begin + f::MKOrAbstractArray{mk_float,4}, moments, vpa) = begin A = moments.electron.constraints_A_coefficient B = moments.electron.constraints_B_coefficient C = moments.electron.constraints_C_coefficient @@ -102,7 +102,7 @@ end end end @timeit global_timer hard_force_moment_constraints!( - f::AbstractArray{mk_float,5}, moments, vpa) = begin + f::MKOrAbstractArray{mk_float,5}, moments, vpa) = begin A = moments.ion.constraints_A_coefficient B = moments.ion.constraints_B_coefficient C = moments.ion.constraints_C_coefficient @@ -167,7 +167,7 @@ function hard_force_moment_constraints_neutral!(f, moments, vz) return A, B, C end @timeit global_timer hard_force_moment_constraints_neutral!( - f::AbstractArray{mk_float,6}, moments, vz) = begin + f::MKOrAbstractArray{mk_float,6}, moments, vz) = begin A = moments.neutral.constraints_A_coefficient B = moments.neutral.constraints_B_coefficient C = moments.neutral.constraints_C_coefficient @@ -193,8 +193,8 @@ r = \\hat{r} + (A + B w_{\\|} + C w_{\\|}^2) f Note this function assumes the input is given at a single spatial position. """ -function moment_constraints_on_residual!(residual::AbstractArray{T,N}, - f::AbstractArray{T,N}, moments, vpa) where {T,N} +function moment_constraints_on_residual!(residual::MKOrAbstractArray{T,N}, + f::MKOrAbstractArray{T,N}, moments, vpa) where {T,N} if N == 2 f = @view f[:,1] residual = @view residual[:,1] diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl index d0398a1d4..d2f13178e 100644 --- a/moment_kinetics/src/moment_kinetics.jl +++ b/moment_kinetics/src/moment_kinetics.jl @@ -4,8 +4,13 @@ module moment_kinetics export run_moment_kinetics +using HDF5 using MPI using StatsBase +try + using NCDatasets +catch +end # Include submodules from other source files # Note that order of includes matters - things used in one module must already diff --git a/moment_kinetics/src/nonlinear_solvers.jl b/moment_kinetics/src/nonlinear_solvers.jl index 15fa58d1a..6338d4300 100644 --- a/moment_kinetics/src/nonlinear_solvers.jl +++ b/moment_kinetics/src/nonlinear_solvers.jl @@ -564,7 +564,7 @@ old_precon_iterations = nl_solver_params.precon_iterations[] end @timeit_debug global_timer distributed_norm( - ::Val{:z}, residual::AbstractArray{mk_float, 1}, coords, + ::Val{:z}, residual::AbstractMKArray{mk_float, 1}, coords, rtol, atol, x) = begin z = coords.z @@ -600,7 +600,7 @@ end end @timeit_debug global_timer distributed_norm( - ::Val{:vpa}, residual::AbstractArray{mk_float, 1}, coords, + ::Val{:vpa}, residual::AbstractMKArray{mk_float, 1}, coords, rtol, atol, x) = begin # No parallelism needed when the implicit solve is over vpa - assume that this will be # called inside a parallelised s_r_z_vperp loop. @@ -616,7 +616,7 @@ end @timeit_debug global_timer distributed_norm( ::Val{:zvperpvpa}, - residual::Tuple{AbstractArray{mk_float, 1},AbstractArray{mk_float, 3}}, + residual::Tuple{AbstractMKArray{mk_float, 1},AbstractMKArray{mk_float, 3}}, coords, rtol, atol, x) = begin ppar_residual, pdf_residual = residual x_ppar, x_pdf = x @@ -679,7 +679,7 @@ end end @timeit_debug global_timer distributed_norm( - ::Val{:srzvperpvpa}, residual::AbstractArray{mk_float, 5}, coords, rtol, + ::Val{:srzvperpvpa}, residual::AbstractMKArray{mk_float, 5}, coords, rtol, atol, x) = begin n_ion_species = coords.s r = coords.r @@ -724,7 +724,7 @@ end end @timeit_debug global_timer distributed_dot( - ::Val{:z}, v::AbstractArray{mk_float, 1}, w::AbstractArray{mk_float, 1}, + ::Val{:z}, v::AbstractMKArray{mk_float, 1}, w::AbstractMKArray{mk_float, 1}, coords, rtol, atol, x) = begin z = coords.z @@ -761,7 +761,7 @@ end end @timeit_debug global_timer distributed_dot( - ::Val{:vpa}, v::AbstractArray{mk_float, 1}, w::AbstractArray{mk_float, 1}, coords, + ::Val{:vpa}, v::AbstractMKArray{mk_float, 1}, w::AbstractMKArray{mk_float, 1}, coords, rtol, atol, x) = begin # No parallelism needed when the implicit solve is over vpa - assume that this will be # called inside a parallelised s_r_z_vperp loop. @@ -774,8 +774,8 @@ end end @timeit_debug global_timer distributed_dot( - ::Val{:zvperpvpa}, v::Tuple{AbstractArray{mk_float, 1},AbstractArray{mk_float, 3}}, - w::Tuple{AbstractArray{mk_float, 1},AbstractArray{mk_float, 3}}, coords, + ::Val{:zvperpvpa}, v::Tuple{AbstractMKArray{mk_float, 1},AbstractMKArray{mk_float, 3}}, + w::Tuple{AbstractMKArray{mk_float, 1},AbstractMKArray{mk_float, 3}}, coords, rtol, atol, x) = begin v_ppar, v_pdf = v w_ppar, w_pdf = w @@ -835,8 +835,8 @@ end end @timeit_debug global_timer distributed_dot( - ::Val{:srzvperpvpa}, v::AbstractArray{mk_float, 5}, - w::AbstractArray{mk_float, 5}, coords, rtol, atol, x) = begin + ::Val{:srzvperpvpa}, v::AbstractMKArray{mk_float, 5}, + w::AbstractMKArray{mk_float, 5}, coords, rtol, atol, x) = begin n_ion_species = coords.s r = coords.r z = coords.z @@ -880,7 +880,7 @@ end # slow code @timeit_debug global_timer parallel_map( - ::Val{:z}, func, result::AbstractArray{mk_float, 1}) = begin + ::Val{:z}, func, result::AbstractMKArray{mk_float, 1}) = begin begin_z_region() @@ -891,7 +891,7 @@ end return nothing end @timeit_debug global_timer parallel_map( - ::Val{:z}, func, result::AbstractArray{mk_float, 1}, x1) = begin + ::Val{:z}, func, result::AbstractMKArray{mk_float, 1}, x1) = begin begin_z_region() @@ -902,11 +902,11 @@ end return nothing end @timeit_debug global_timer parallel_map( - ::Val{:z}, func, result::AbstractArray{mk_float, 1}, x1, x2) = begin + ::Val{:z}, func, result::AbstractMKArray{mk_float, 1}, x1, x2) = begin begin_z_region() - if isa(x2, AbstractArray) + if isa(x2, AbstractMKArray) @loop_z iz begin result[iz] = func(x1[iz], x2[iz]) end @@ -919,11 +919,11 @@ end return nothing end @timeit_debug global_timer parallel_map( - ::Val{:z}, func, result::AbstractArray{mk_float, 1}, x1, x2, x3) = begin + ::Val{:z}, func, result::AbstractMKArray{mk_float, 1}, x1, x2, x3) = begin begin_z_region() - if isa(x3, AbstractArray) + if isa(x3, AbstractMKArray) @loop_z iz begin result[iz] = func(x1[iz], x2[iz], x3[iz]) end @@ -937,7 +937,7 @@ end end @timeit_debug global_timer parallel_map( - ::Val{:vpa}, func, result::AbstractArray{mk_float, 1}) = begin + ::Val{:vpa}, func, result::AbstractMKArray{mk_float, 1}) = begin # No parallelism needed when the implicit solve is over vpa - assume that this will be # called inside a parallelised s_r_z_vperp loop. for i ∈ eachindex(result) @@ -946,7 +946,7 @@ end return nothing end @timeit_debug global_timer parallel_map( - ::Val{:vpa}, func, result::AbstractArray{mk_float, 1}, x1) = begin + ::Val{:vpa}, func, result::AbstractMKArray{mk_float, 1}, x1) = begin # No parallelism needed when the implicit solve is over vpa - assume that this will be # called inside a parallelised s_r_z_vperp loop. for i ∈ eachindex(result) @@ -955,10 +955,10 @@ end return nothing end @timeit_debug global_timer parallel_map( - ::Val{:vpa}, func, result::AbstractArray{mk_float, 1}, x1, x2) = begin + ::Val{:vpa}, func, result::AbstractMKArray{mk_float, 1}, x1, x2) = begin # No parallelism needed when the implicit solve is over vpa - assume that this will be # called inside a parallelised s_r_z_vperp loop. - if isa(x2, AbstractArray) + if isa(x2, AbstractMKArray) for i ∈ eachindex(result) result[i] = func(x1[i], x2[i]) end @@ -970,10 +970,10 @@ end return nothing end @timeit_debug global_timer parallel_map( - ::Val{:vpa}, func, result::AbstractArray{mk_float, 1}, x1, x2, x3) = begin + ::Val{:vpa}, func, result::AbstractMKArray{mk_float, 1}, x1, x2, x3) = begin # No parallelism needed when the implicit solve is over vpa - assume that this will be # called inside a parallelised s_r_z_vperp loop. - if isa(x3, AbstractArray) + if isa(x3, AbstractMKArray) for i ∈ eachindex(result) result[i] = func(x1[i], x2[i], x3[i]) end @@ -986,7 +986,7 @@ end end @timeit_debug global_timer parallel_map( - ::Val{:zvperpvpa}, func, result::Tuple{AbstractArray{mk_float, 1},AbstractArray{mk_float, 3}}) = begin + ::Val{:zvperpvpa}, func, result::Tuple{AbstractMKArray{mk_float, 1},AbstractMKArray{mk_float, 3}}) = begin result_ppar, result_pdf = result @@ -1005,7 +1005,7 @@ end return nothing end @timeit_debug global_timer parallel_map( - ::Val{:zvperpvpa}, func, result::Tuple{AbstractArray{mk_float, 1},AbstractArray{mk_float, 3}}, + ::Val{:zvperpvpa}, func, result::Tuple{AbstractMKArray{mk_float, 1},AbstractMKArray{mk_float, 3}}, x1) = begin result_ppar, result_pdf = result @@ -1026,7 +1026,7 @@ end return nothing end @timeit_debug global_timer parallel_map( - ::Val{:zvperpvpa}, func, result::Tuple{AbstractArray{mk_float, 1},AbstractArray{mk_float, 3}}, + ::Val{:zvperpvpa}, func, result::Tuple{AbstractMKArray{mk_float, 1},AbstractMKArray{mk_float, 3}}, x1, x2) = begin result_ppar, result_pdf = result @@ -1062,7 +1062,7 @@ end return nothing end @timeit_debug global_timer parallel_map( - ::Val{:zvperpvpa}, func, result::Tuple{AbstractArray{mk_float, 1},AbstractArray{mk_float, 3}}, + ::Val{:zvperpvpa}, func, result::Tuple{AbstractMKArray{mk_float, 1},AbstractMKArray{mk_float, 3}}, x1, x2, x3) = begin result_ppar, result_pdf = result @@ -1100,7 +1100,7 @@ end end @timeit_debug global_timer parallel_map( - ::Val{:srzvperpvpa}, func, result::AbstractArray{mk_float, 5}) = begin + ::Val{:srzvperpvpa}, func, result::AbstractMKArray{mk_float, 5}) = begin begin_s_r_z_vperp_vpa_region() @@ -1111,7 +1111,7 @@ end return nothing end @timeit_debug global_timer parallel_map( - ::Val{:srzvperpvpa}, func, result::AbstractArray{mk_float, 5}, x1) = begin + ::Val{:srzvperpvpa}, func, result::AbstractMKArray{mk_float, 5}, x1) = begin begin_s_r_z_vperp_vpa_region() @@ -1122,11 +1122,11 @@ end return nothing end @timeit_debug global_timer parallel_map( - ::Val{:srzvperpvpa}, func, result::AbstractArray{mk_float, 5}, x1, x2) = begin + ::Val{:srzvperpvpa}, func, result::AbstractMKArray{mk_float, 5}, x1, x2) = begin begin_s_r_z_vperp_vpa_region() - if isa(x2, AbstractArray) + if isa(x2, AbstractMKArray) @loop_s_r_z_vperp_vpa is ir iz ivperp ivpa begin result[ivpa,ivperp,iz,ir,is] = func(x1[ivpa,ivperp,iz,ir,is], x2[ivpa,ivperp,iz,ir,is]) end @@ -1139,12 +1139,12 @@ end return nothing end @timeit_debug global_timer parallel_map( - ::Val{:srzvperpvpa}, func, result::AbstractArray{mk_float, 5}, x1, x2, + ::Val{:srzvperpvpa}, func, result::AbstractMKArray{mk_float, 5}, x1, x2, x3) = begin begin_s_r_z_vperp_vpa_region() - if isa(x3, AbstractArray) + if isa(x3, AbstractMKArray) @loop_s_r_z_vperp_vpa is ir iz ivperp ivpa begin result[ivpa,ivperp,iz,ir,is] = func(x1[ivpa,ivperp,iz,ir,is], x2[ivpa,ivperp,iz,ir,is], x3[ivpa,ivperp,iz,ir,is]) end @@ -1158,7 +1158,7 @@ end end @timeit_debug global_timer parallel_delta_x_calc( - ::Val{:z}, delta_x::AbstractArray{mk_float, 1}, V, y) = begin + ::Val{:z}, delta_x::AbstractMKArray{mk_float, 1}, V, y) = begin begin_z_region() @@ -1173,7 +1173,7 @@ end end @timeit_debug global_timer parallel_delta_x_calc( - ::Val{:vpa}, delta_x::AbstractArray{mk_float, 1}, V, y) = begin + ::Val{:vpa}, delta_x::AbstractMKArray{mk_float, 1}, V, y) = begin # No parallelism needed when the implicit solve is over vpa - assume that this will be # called inside a parallelised s_r_z_vperp loop. ny = length(y) @@ -1186,7 +1186,7 @@ end end @timeit_debug global_timer parallel_delta_x_calc( - ::Val{:zvperpvpa}, delta_x::Tuple{AbstractArray{mk_float, 1},AbstractArray{mk_float, 3}}, V, + ::Val{:zvperpvpa}, delta_x::Tuple{AbstractMKArray{mk_float, 1},AbstractMKArray{mk_float, 3}}, V, y) = begin delta_x_ppar, delta_x_pdf = delta_x @@ -1214,7 +1214,7 @@ end end @timeit_debug global_timer parallel_delta_x_calc( - ::Val{:srzvperpvpa}, delta_x::AbstractArray{mk_float, 5}, V, y) = begin + ::Val{:srzvperpvpa}, delta_x::AbstractMKArray{mk_float, 5}, V, y) = begin begin_s_r_z_vperp_vpa_region() diff --git a/moment_kinetics/src/runge_kutta.jl b/moment_kinetics/src/runge_kutta.jl index 14e6f5255..d8fc34556 100644 --- a/moment_kinetics/src/runge_kutta.jl +++ b/moment_kinetics/src/runge_kutta.jl @@ -554,9 +554,9 @@ end # Ion distribution function function rk_update_loop_low_storage!(rk_coefs, rk_coefs_implicit, - new::AbstractArray{mk_float,5}, - old::AbstractArray{mk_float,5}, - first::AbstractArray{mk_float,5}, new_implicit, + new::MKArray{mk_float,5}, + old::MKArray{mk_float,5}, + first::MKArray{mk_float,5}, new_implicit, old_implicit, first_implicit; output=new) @boundscheck length(rk_coefs) == 3 @@ -580,7 +580,7 @@ function rk_update_loop_low_storage!(rk_coefs, rk_coefs_implicit, return nothing end function rk_update_loop!(rk_coefs, rk_coefs_implicit, - var_arrays::NTuple{N,AbstractArray{mk_float,5}}, + var_arrays::NTuple{N,MKArray{mk_float,5}}, var_arrays_implicit; output=var_arrays[N]) where N @boundscheck length(rk_coefs) ≥ N @@ -603,9 +603,9 @@ end # Ion moments function rk_update_loop_low_storage!(rk_coefs, rk_coefs_implicit, - new::AbstractArray{mk_float,3}, - old::AbstractArray{mk_float,3}, - first::AbstractArray{mk_float,3}, new_implicit, + new::MKArray{mk_float,3}, + old::MKArray{mk_float,3}, + first::MKArray{mk_float,3}, new_implicit, old_implicit, first_implicit; output=new) @boundscheck length(rk_coefs) == 3 @@ -629,7 +629,7 @@ function rk_update_loop_low_storage!(rk_coefs, rk_coefs_implicit, return nothing end function rk_update_loop!(rk_coefs, rk_coefs_implicit, - var_arrays::NTuple{N,AbstractArray{mk_float,3}}, + var_arrays::NTuple{N,MKArray{mk_float,3}}, var_arrays_implicit; output=var_arrays[N]) where N @boundscheck length(rk_coefs) ≥ N @@ -650,9 +650,9 @@ end # Electron distribution function function rk_update_loop_low_storage!(rk_coefs, rk_coefs_implicit, - new::AbstractArray{mk_float,4}, - old::AbstractArray{mk_float,4}, - first::AbstractArray{mk_float,4}, new_implicit, + new::MKArray{mk_float,4}, + old::MKArray{mk_float,4}, + first::MKArray{mk_float,4}, new_implicit, old_implicit, first_implicit; output=new) @boundscheck length(rk_coefs) == 3 @@ -676,7 +676,7 @@ function rk_update_loop_low_storage!(rk_coefs, rk_coefs_implicit, return nothing end function rk_update_loop!(rk_coefs, rk_coefs_implicit, - var_arrays::NTuple{N,AbstractArray{mk_float,4}}, + var_arrays::NTuple{N,MKArray{mk_float,4}}, var_arrays_implicit; output=var_arrays[N]) where N @boundscheck length(rk_coefs) ≥ N @@ -700,9 +700,9 @@ end # Electron moments function rk_update_loop_low_storage!(rk_coefs, rk_coefs_implicit, - new::AbstractArray{mk_float,2}, - old::AbstractArray{mk_float,2}, - first::AbstractArray{mk_float,2}, new_implicit, + new::MKArray{mk_float,2}, + old::MKArray{mk_float,2}, + first::MKArray{mk_float,2}, new_implicit, old_implicit, first_implicit; output=new) @boundscheck length(rk_coefs) == 3 @@ -726,7 +726,7 @@ function rk_update_loop_low_storage!(rk_coefs, rk_coefs_implicit, return nothing end function rk_update_loop!(rk_coefs, rk_coefs_implicit, - var_arrays::NTuple{N,AbstractArray{mk_float,2}}, + var_arrays::NTuple{N,MKArray{mk_float,2}}, var_arrays_implicit; output=var_arrays[N]) where N @boundscheck length(rk_coefs) ≥ N @@ -749,9 +749,9 @@ end # Neutral distribution function function rk_update_loop_neutrals_low_storage!(rk_coefs, rk_coefs_implicit, - new::AbstractArray{mk_float,6}, - old::AbstractArray{mk_float,6}, - first::AbstractArray{mk_float,6}, + new::MKArray{mk_float,6}, + old::MKArray{mk_float,6}, + first::MKArray{mk_float,6}, new_implicit, old_implicit, first_implicit; output=new) @boundscheck length(rk_coefs) == 3 @@ -776,7 +776,7 @@ function rk_update_loop_neutrals_low_storage!(rk_coefs, rk_coefs_implicit, return nothing end function rk_update_loop_neutrals!(rk_coefs, rk_coefs_implicit, - var_arrays::NTuple{N,AbstractArray{mk_float,6}}, + var_arrays::NTuple{N,MKArray{mk_float,6}}, var_arrays_implicit; output=var_arrays[N]) where N @boundscheck length(rk_coefs) ≥ N @@ -800,9 +800,9 @@ end # Neutral moments function rk_update_loop_neutrals_low_storage!(rk_coefs, rk_coefs_implicit, - new::AbstractArray{mk_float,3}, - old::AbstractArray{mk_float,3}, - first::AbstractArray{mk_float,3}, + new::MKArray{mk_float,3}, + old::MKArray{mk_float,3}, + first::MKArray{mk_float,3}, new_implicit, old_implicit, first_implicit; output=new) @boundscheck length(rk_coefs) == 3 @@ -827,7 +827,7 @@ function rk_update_loop_neutrals_low_storage!(rk_coefs, rk_coefs_implicit, return nothing end function rk_update_loop_neutrals!(rk_coefs, rk_coefs_implicit, - var_arrays::NTuple{N,AbstractArray{mk_float,3}}, + var_arrays::NTuple{N,MKArray{mk_float,3}}, var_arrays_implicit; output=var_arrays[N]) where N @boundscheck length(rk_coefs) ≥ N diff --git a/moment_kinetics/src/type_definitions.jl b/moment_kinetics/src/type_definitions.jl index 0f301bc0a..2ca0f5acf 100644 --- a/moment_kinetics/src/type_definitions.jl +++ b/moment_kinetics/src/type_definitions.jl @@ -8,6 +8,12 @@ export mk_zeros export MKArray export MKVector export MKMatrix +export AbstractMKArray +export AbstractMKVector +export AbstractMKMatrix +export MKOrAbstractArray +export MKOrAbstractVector +export MKOrAbstractMatrix export OptionsDict using InboundsArrays @@ -33,6 +39,30 @@ const MKVector{T} = MKArray{T,1} """ const MKMatrix{T} = MKArray{T,2} +""" +""" +const AbstractMKArray{T,N} = AbstractInboundsArray{T,N} + +""" +""" +const AbstractMKVector{T} = AbstractInboundsArray{T,1} + +""" +""" +const AbstractMKMatrix{T} = AbstractInboundsArray{T,2} + +""" +""" +const MKOrAbstractArray{T,N} = Union{AbstractInboundsArray{T,N}, AbstractArray{T,N}} where {T,N} + +""" +""" +const MKOrAbstractVector{T} = Union{AbstractInboundsArray{T, 1}, AbstractVector{T}} where T + +""" +""" +const MKOrAbstractMatrix{T} = Union{AbstractInboundsArray{T, 2}, AbstractMatrix{T}} where T + @inline function MKArray(args...) return InboundsArray(args...) end diff --git a/moment_kinetics/src/utils.jl b/moment_kinetics/src/utils.jl index 74b8a3bbb..398e1cacc 100644 --- a/moment_kinetics/src/utils.jl +++ b/moment_kinetics/src/utils.jl @@ -12,6 +12,7 @@ using ..input_structs using ..looping using ..moment_kinetics_input: mk_input using ..reference_parameters +using ..type_definitions # Import moment_kinetics so we can refer to it in docstrings import ..moment_kinetics @@ -356,7 +357,7 @@ the z direction. Reduces the result over the shared-memory block (handling distributed parallelism is left to the calling site). The result is only to be used on rank-0 of the shared-memory block. """ -function get_minimum_CFL_z(speed::AbstractArray{T,4} where T, z) +function get_minimum_CFL_z(speed::MKOrAbstractArray{T,4} where T, z) min_CFL = Inf dz = z.cell_width @@ -384,7 +385,7 @@ the vpa direction. Reduces the result over the shared-memory block (handling distributed parallelism is left to the calling site). The result is only to be used on rank-0 of the shared-memory block. """ -function get_minimum_CFL_vpa(speed::AbstractArray{T,4} where T, vpa) +function get_minimum_CFL_vpa(speed::MKOrAbstractArray{T,4} where T, vpa) min_CFL = Inf dvpa = vpa.cell_width @@ -412,7 +413,7 @@ neutrals in the z direction. Reduces the result over the shared-memory block (handling distributed parallelism is left to the calling site). The result is only to be used on rank-0 of the shared-memory block. """ -function get_minimum_CFL_neutral_z(speed::AbstractArray{T,5} where T, z) +function get_minimum_CFL_neutral_z(speed::MKOrAbstractArray{T,5} where T, z) min_CFL = Inf dz = z.cell_width @@ -440,7 +441,7 @@ neutrals in the vz direction. Reduces the result over the shared-memory block (handling distributed parallelism is left to the calling site). The result is only to be used on rank-0 of the shared-memory block. """ -function get_minimum_CFL_neutral_vz(speed::AbstractArray{T,5} where T, vz) +function get_minimum_CFL_neutral_vz(speed::MKOrAbstractArray{T,5} where T, vz) min_CFL = Inf dvz = vz.cell_width @@ -471,7 +472,7 @@ post-processing. """ function get_CFL end -function get_CFL!(CFL::AbstractArray{T,4}, speed::AbstractArray{T,4}, coord) where T +function get_CFL!(CFL::MKOrAbstractArray{T,4}, speed::MKOrAbstractArray{T,4}, coord) where T nmain, n2, n3, n4 = size(speed) @@ -482,7 +483,7 @@ function get_CFL!(CFL::AbstractArray{T,4}, speed::AbstractArray{T,4}, coord) whe return CFL end -function get_CFL!(CFL::AbstractArray{T,5}, speed::AbstractArray{T,5}, coord) where T +function get_CFL!(CFL::MKOrAbstractArray{T,5}, speed::MKOrAbstractArray{T,5}, coord) where T nmain, n2, n3, n4, n5 = size(speed) @@ -493,7 +494,7 @@ function get_CFL!(CFL::AbstractArray{T,5}, speed::AbstractArray{T,5}, coord) whe return CFL end -function get_CFL!(CFL::AbstractArray{T,6}, speed::AbstractArray{T,6}, coord) where T +function get_CFL!(CFL::MKOrAbstractArray{T,6}, speed::MKOrAbstractArray{T,6}, coord) where T nmain, n2, n3, n4, n5, n6 = size(speed) diff --git a/moment_kinetics/test/calculus_tests.jl b/moment_kinetics/test/calculus_tests.jl index 23a77e71a..4ea470b2f 100644 --- a/moment_kinetics/test/calculus_tests.jl +++ b/moment_kinetics/test/calculus_tests.jl @@ -2,6 +2,7 @@ module CalculusTests include("setup.jl") +using moment_kinetics.array_allocation: allocate_float using moment_kinetics.coordinates: define_test_coordinate using moment_kinetics.calculus: derivative!, second_derivative!, integral using moment_kinetics.calculus: laplacian_derivative! @@ -47,9 +48,9 @@ function runtests() element_spacing_option=element_spacing_option, collision_operator_dim=false) # create array for the function f(x) to be differentiated/integrated - f = Array{Float64,1}(undef, x.n) + f = allocate_float(x.n) # create array for the derivative df/dx - df = Array{Float64,1}(undef, x.n) + df = allocate_float(x.n) # initialize f for ix ∈ 1:x.n f[ix] = ( (cospi(2.0*x.grid[ix]/x.L)+sinpi(2.0*x.grid[ix]/x.L)) @@ -94,7 +95,7 @@ function runtests() collision_operator_dim=false) # create array for the derivative df/dx and the expected result - df = Array{Float64,1}(undef, x.n) + df = allocate_float(x.n) # initialize f and expected df offset = randn(rng) @@ -142,7 +143,7 @@ function runtests() collision_operator_dim=false) # create array for the derivative df/dx and the expected result - df = Array{Float64,1}(undef, x.n) + df = allocate_float(x.n) # initialize f and expected df offset = randn(rng) @@ -186,7 +187,7 @@ function runtests() collision_operator_dim=false) # create array for the derivative df/dx and the expected result - df = Array{Float64,1}(undef, x.n) + df = allocate_float(x.n) # initialize f and expected df offset = undef @@ -236,7 +237,7 @@ function runtests() collision_operator_dim=false) # create array for the derivative df/dx and the expected result - df = Array{Float64,1}(undef, x.n) + df = allocate_float(x.n) # initialize f and expected df offset = undef @@ -682,7 +683,7 @@ function runtests() # test polynomials up to order ngrid-1 for n ∈ 0:ngrid-1 # create array for the function f(x) to be differentiated/integrated - f = Array{Float64,1}(undef, x.n) + f = allocate_float(x.n) # create array for the derivative df/dx and the expected result df = similar(f) expected_df = similar(f) @@ -728,7 +729,7 @@ function runtests() # test polynomials up to order ngrid-1 for n ∈ 0:ngrid-1 # create array for the function f(x) to be differentiated/integrated - f = Array{Float64,1}(undef, x.n) + f = allocate_float(x.n) # create array for the derivative df/dx and the expected result df = similar(f) expected_df = similar(f) @@ -1009,7 +1010,7 @@ function runtests() # test polynomials up to order ngrid-1 for n ∈ 0:ngrid-1 # create array for the function f(x) to be differentiated/integrated - f = Array{Float64,1}(undef, x.n) + f = allocate_float(x.n) # create array for the derivative df/dx and the expected result df = similar(f) expected_df = similar(f) @@ -1054,7 +1055,7 @@ function runtests() # test polynomials up to order ngrid-1 for n ∈ 0:ngrid-1 # create array for the function f(x) to be differentiated/integrated - f = Array{Float64,1}(undef, x.n) + f = allocate_float(x.n) # create array for the derivative df/dx and the expected result df = similar(f) expected_df = similar(f) diff --git a/moment_kinetics/test/jacobian_matrix_tests.jl b/moment_kinetics/test/jacobian_matrix_tests.jl index 6721f7b30..370681afc 100644 --- a/moment_kinetics/test/jacobian_matrix_tests.jl +++ b/moment_kinetics/test/jacobian_matrix_tests.jl @@ -231,7 +231,7 @@ function cleanup_mk_state!(args...) return nothing end -function generate_norm_factor(perturbed_residual::AbstractArray{mk_float,3}) +function generate_norm_factor(perturbed_residual::MKArray{mk_float,3}) # half-width of the window for moving average w = 3 norm_factor_unsmoothed = mean(abs.(perturbed_residual); dims=3) @@ -249,7 +249,7 @@ function generate_norm_factor(perturbed_residual::AbstractArray{mk_float,3}) end return norm_factor end -function generate_norm_factor(perturbed_residual::AbstractArray{mk_float,1}) +function generate_norm_factor(perturbed_residual::MKArray{mk_float,1}) norm_factor_unsmoothed = mean(abs.(perturbed_residual); dims=1) end diff --git a/moment_kinetics/test/nonlinear_solver_tests.jl b/moment_kinetics/test/nonlinear_solver_tests.jl index 81d4c87c7..6f5d123f7 100644 --- a/moment_kinetics/test/nonlinear_solver_tests.jl +++ b/moment_kinetics/test/nonlinear_solver_tests.jl @@ -70,11 +70,11 @@ function linear_test() function rhs_func!(residual, x; krylov=false) if serial_solve - residual .= A * x - b + residual .= A * x .- b else begin_serial_region() @serial_region begin - residual .= A * x - b + residual .= A * x .- b end end return nothing