Inference and speeding in an interpolation

inference

#1

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.


#2

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...]

#3

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

#4

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

#5

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.


#6

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.


#7

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.


#8

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


#9

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