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