Inference and speeding in an interpolation

Out of convenience I’m wrapping a scaled interpolation from Interpolations.jl:

grid_x = 0:0.1:100.0
grid_y = 0:0.1:100.0

f(x, y) = (x + y)^2

farr = [f(x, y) for x in grid_x, y in grid_y]

using Interpolations

function interp1(grid_x, grid_y, farr)
    itp = interpolate(farr, BSpline(Cubic(Line())), OnGrid())
    sitp = scale(itp, grid_x, grid_y)
    return sitp
end

interp(grid_x, grid_y, xy) = interp1(grid_x, grid_y, farr)[xy...]

xy = [0.01, 0.02]

@code_warntype interp(grid_x, grid_y, xy)
Variables:
  #self# <optimized out>
  grid_x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}
  grid_y::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}
  xy::Array{Float64,1}
  itp::Interpolations.BSplineInterpolation{_,_,_,Interpolations.BSpline{Interpolations.Cubic{Interpolations.Line}},Interpolations.OnGrid,_} where _ where _ where _ where _
  sitp::Any

Body:
  begin 
      SSAValue(0) = Main.farr
      $(Expr(:inbounds, false))
      # meta: location In[4] interp1 3
      itp::Interpolations.BSplineInterpolation{_,_,_,Interpolations.BSpline{Interpolations.Cubic{Interpolations.Line}},Interpolations.OnGrid,_} where _ where _ where _ where _ = (Main.interpolate)(SSAValue(0), $(QuoteNode(Interpolations.BSpline{Interpolations.Cubic{Interpolations.Line}}())), $(QuoteNode(Interpolations.OnGrid())))::Interpolations.BSplineInterpolation{_,_,_,Interpolations.BSpline{Interpolations.Cubic{Interpolations.Line}},Interpolations.OnGrid,_} where _ where _ where _ where _ # line 4:
      sitp::Any = (Main.scale)(itp::Interpolations.BSplineInterpolation{_,_,_,Interpolations.BSpline{Interpolations.Cubic{Interpolations.Line}},Interpolations.OnGrid,_} where _ where _ where _ where _, grid_x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}, grid_y::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}})::Any
      # meta: pop location
      $(Expr(:inbounds, :pop))
      return (Core._apply)(Main.getindex, (Core.tuple)(sitp::Any)::Tuple{Any}, xy::Array{Float64,1})::Any
  end::Any

How can I improve inference to get rid of the Anys? I tried annotating the output and arguments of interp1 and interp but I still get some Anys in the body.

Inference fails because farr is global, and the compiler can’t ensure that you will not change it to something else (with a different type). Make it a parameter, eg

interp(grid_x, grid_y, farr, xy) = interp1(grid_x, grid_y, farr)[xy...]

I tried with making farr a const, and making it a parameter like you suggest, but I still get some Anys:

@code_warntype interp(grid_x, grid_y, farr, xy)
Variables:
  #self# <optimized out>
  grid_x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}
  grid_y::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}
  farr::Array{Float64,2}
  xy::Array{Float64,1}
  itp::Interpolations.BSplineInterpolation{Float64,2,Array{Float64,2},Interpolations.BSpline{Interpolations.Cubic{Interpolations.Line}},Interpolations.OnGrid,1}
  sitp::Interpolations.ScaledInterpolation{Float64,2,Interpolations.BSplineInterpolation{Float64,2,Array{Float64,2},Interpolations.BSpline{Interpolations.Cubic{Interpolations.Line}},Interpolations.OnGrid,1},Interpolations.BSpline{Interpolations.Cubic{Interpolations.Line}},Interpolations.OnGrid,Tuple{StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}

Body:
  begin 
      $(Expr(:inbounds, false))
      # meta: location In[7] interp1 3
      itp::Interpolations.BSplineInterpolation{Float64,2,Array{Float64,2},Interpolations.BSpline{Interpolations.Cubic{Interpolations.Line}},Interpolations.OnGrid,1} = $(Expr(:invoke, MethodInstance for interpolate(::Type{Float64}, ::Type{Float64}, ::Array{Float64,2}, ::Interpolations.BSpline{Interpolations.Cubic{Interpolations.Line}}, ::Interpolations.OnGrid), :(Interpolations.interpolate), Float64, Float64, :(farr), :($(QuoteNode(Interpolations.BSpline{Interpolations.Cubic{Interpolations.Line}}()))), :($(QuoteNode(Interpolations.OnGrid()))))) # line 4:
      sitp::Interpolations.ScaledInterpolation{Float64,2,Interpolations.BSplineInterpolation{Float64,2,Array{Float64,2},Interpolations.BSpline{Interpolations.Cubic{Interpolations.Line}},Interpolations.OnGrid,1},Interpolations.BSpline{Interpolations.Cubic{Interpolations.Line}},Interpolations.OnGrid,Tuple{StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}} = $(Expr(:invoke, MethodInstance for scale(::Interpolations.BSplineInterpolation{Float64,2,Array{Float64,2},Interpolations.BSpline{Interpolations.Cubic{Interpolations.Line}},Interpolations.OnGrid,1}, ::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}, ::Vararg{StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},N} where N), :(Main.scale), :(itp), :(grid_x), :(grid_y)))
      # meta: pop location
      $(Expr(:inbounds, :pop))
      return (Core._apply)(Main.getindex, (Core.tuple)(sitp::Interpolations.ScaledInterpolation{Float64,2,Interpolations.BSplineInterpolation{Float64,2,Array{Float64,2},Interpolations.BSpline{Interpolations.Cubic{Interpolations.Line}},Interpolations.OnGrid,1},Interpolations.BSpline{Interpolations.Cubic{Interpolations.Line}},Interpolations.OnGrid,Tuple{StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}})::Tuple{Interpolations.ScaledInterpolation{Float64,2,Interpolations.BSplineInterpolation{Float64,2,Array{Float64,2},Interpolations.BSpline{Interpolations.Cubic{Interpolations.Line}},Interpolations.OnGrid,1},Interpolations.BSpline{Interpolations.Cubic{Interpolations.Line}},Interpolations.OnGrid,Tuple{StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}}, xy::Array{Float64,1})::Any
  end::Any

Don’t pass xy as a Vector, as the type itself does not determine the number of elements. The following is inferred correctly:

using Interpolations
grid_x = 0:0.1:100.0
grid_y = 0:0.1:100.0
f(x, y) = (x + y)^2
farr = [f(x, y) for x in grid_x, y in grid_y]

function interp1(grid_x, grid_y, farr)
    itp = interpolate(farr, BSpline(Cubic(Line())), OnGrid())
    scale(itp, grid_x, grid_y)
end

interp(grid_x, grid_y, farr, x, y) = interp1(grid_x, grid_y, farr)[x, y]

@code_warntype interp1(grid_x, grid_y, farr) # OK
@code_warntype interp(grid_x, grid_y, farr, 0.01, 0.02) # OK

Also, I would suggest just saving the result of inter1p and working with that, not calling interpolate and scale each time you want an interpolation.

Yes, this works. So here the problem is that the compiler doesn’t know the number of elements? It would be really annoying to have to have to define the function for a large number of components of xy.

Yes, I can work like that, but eventually I will have to modify farr in the course of an algorithm. That was the whole idea of wrapping interpolate in a function.

Depending on how large the number of elements is, you can try something inferrable, eg

Thanks, I’ll play more with this, but simply calling interp on SVector(0.01, 0.02) still breaks inference.