Type stability and OffsetArray components of a struct

I found that @code_warntype indicated a type instability in a finite difference code. The same issue arises in the following simple example.

using OffsetArrays

abstract type AbstractGrid end

struct Grid <: AbstractGrid
    x::Vector{Float64}
end 

function Grid(L, N)
    x = range(0, L, N+1)
    return Grid(x)
end 

struct OGrid <: AbstractGrid
    x::OffsetVector{Float64}
end 

function OGrid(L, N)
    x_ = range(0, L, N+1)
    x = OffsetVector(x_, 0:N)
    return OGrid(x)
end 

function average(f::Function, grid::AbstractGrid)
    x = grid.x
    s = 0.0
    for k in eachindex(x)
        s+= f(x[k])
    end 
    return s / length(x)
end

If I define

grid = Grid(2.0, 10)
ogrid = OGrid(2.0, 10)
f(x::Float64)::Float64 = exp(-x) * cos(x)

then average(f, grid) and average(f, ogrid) both return 0.3123517502648151, as expected. However,

@code_warntype average(f, grid)

shows no type instabilities, whereas

@code_warntype average(f, ogrid)

flags x, s and k in red. Is there a way to avoid this problem when using an OffsetArray in a struct?

You need a parameter for the parent array.

struct OffsetArray{T,N,AA<:AbstractArray{T,N}} <: AbstractArray{T,N}
    parent::AA
    offsets::NTuple{N,Int}
...
end

const OffsetVector{T,AA<:AbstractVector{T}} = OffsetArray{T,1,AA}
2 Likes

Thanks for the tip. Yes, replacing

struct OGrid <: AbstractGrid
    x::OffsetVector{Float64}
end 

to

struct OGrid <: AbstractGrid
    x::OffsetVector{Float64, Vector{Float64}}
end 

fixed the problem.

1 Like

I don’t know how well your MWE approximates your actual use case, but I would just point out that by forcing Grid and OGrid to use Vector{Float64} you are converting ranges to fully allocated arrays, which seems a bit wasteful, given that Julia provides the very nice, lightweight range datatypes.

Unless you mostly use non-uniformly spaced grids, isn’t it better to use ranges?

Yes, for uniform grids it makes sense to save on storage using ranges.

struct OGrid <: AbstractGrid
    x::OVector{Float64, StepRangeLen{Float64, TwicePrecision{Float64},
               TwicePrecision{Float64}, Int64}}
end

if you need to put that into some other data structure and is annoyed by having to type out type parameter, you can try FHist.jl/src/polybinedges.jl at 26b3c1c9cdd162e4906581316e07e525f29ac328 · Moelf/FHist.jl · GitHub