Fixes for stricter InboundsArray that is not an AbstractArray

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, MKOrAbtractVector)
+        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))

moment_kinetics/src/chebyshev.jl

@@ -238,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
@@ -248,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
@@ -261,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
@@ -578,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
@@ -598,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
@@ -616,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

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 μ

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
 

moment_kinetics/src/file_io.jl

@@ -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)

moment_kinetics/src/gauss_legendre.jl

@@ -1146,7 +1146,7 @@ function setup_global_strong_form_matrix!(QQ_global::MKOrAbstractArray{mk_float,
         # 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

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

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