From 81af1a98d7c6c79313170982084e81d62c871b27 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
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/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 +-
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 +-
27 files changed, 450 insertions(+), 393 deletions(-)
diff --git a/moment_kinetics/Project.toml b/moment_kinetics/Project.toml
index 9828a3c7fd..ac23858fff 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/src/analysis.jl b/moment_kinetics/src/analysis.jl
index 7f1ebb090a..fc77229750 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 a6fcbdc897..60e686e51f 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 adff79e306..a55d3ec0b0 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 56162ee287..ea159c751e 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 df20dbb7dc..8c97dd4154 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 057fec44f2..5b307d232c 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 cd8cf84c01..2df4f6a370 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 c77d8c1ac3..4bef03c1ba 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/em_fields.jl b/moment_kinetics/src/em_fields.jl
index 7b3544f089..08db47b41a 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 199e69b49f..e9e6641cfa 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 d647ab5541..f717896cd9 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 4cc4c9e6a0..8d6e8f81c6 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 f2f6aed37a..6e78dab1e1 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 09a941b5b8..85e4185251 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 0f1e4af53b..a6d023fbca 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 419c1de876..cc9c5993db 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 0a54dbe412..b6ce7c8629 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 7f099b541c..4fac549f76 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 d0398a1d42..d2f13178ec 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 15fa58d1ae..6338d43005 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 14e6f52550..d8fc345564 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 0f301bc0af..2ca0f5acf1 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 74b8a3bbb3..398e1cacc2 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 23a77e71a4..4ea470b2f2 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 6721f7b302..370681afc4 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 81d4c87c76..6f5d123f75 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