Weights and indices for linear interpolation/extrapolation

fredrikpaues · March 8, 2022, 12:10pm

A performance-critical function of a bigger program that I’m writing consists of locating points on a grid. It amounts to finding the indices and weights needed for a linear interpolation/extrapolation—without doing the actual interpolation/extrapolation.

What I have written as of yet is below. I’ve been coding Julia for about a week and would very much appreciate some pointers on how I might improve what I’ve written.

Perhaps this functionality can be found in the Interpolations.jl package, but I haven’t been able to figure out how. The gridpoints vectors will typically be irregularly spaced.

function locate(
    value::AbstractFloat,
    gridpoints::AbstractVector{<:AbstractFloat};
    locate_below::Bool=false,
    locate_above::Bool=false,
)::Tuple{Integer, AbstractVector{<:AbstractFloat}}
    if value <= gridpoints[1]
        index = 1

        if !locate_below
            weights = [1.0, 0.0]
            return index, weights
        end
    elseif value >= gridpoints[end]
        index = length(gridpoints) - 1

        if !locate_above
            weights = [0.0, 1.0]
            return index, weights
        end
    else
        index = findlast(gridpoints .<= value)
    end

    weights = (
        [
            gridpoints[index + 1] - value,
            value - gridpoints[index]
        ] ./ (
            gridpoints[index + 1] - gridpoints[index]
        )
    )

    return index, weights
end
function locate(
    values::AbstractArray{<:AbstractFloat},
    gridpoints::AbstractVector{<:AbstractFloat};
    locate_below::Bool=false,
    locate_above::Bool=false,
)::Tuple{AbstractArray{Integer}, AbstractArray{<:AbstractFloat}}
    indices = ones(Integer, size(values))
    weights = zeros(2, size(values)...)

    for ix in CartesianIndices(values)
        indices[ix], weights[:, ix] = locate(
            values[ix],
            gridpoints;
            locate_below=locate_below,
            locate_above=locate_above
        )
    end

    return indices, weights
end

stevengj · March 8, 2022, 1:31pm

Have you tried one of the existing interpolations packages, like LinearInterpolations.jl or the other interpolations packages referenced on that page?

Some general tips:

Using things like [1.0, 0.0] as a return value, or anywhere in your code, allocates a small array on the heap. This is not a good idea in performance-critical code. e.g. use a tuple instead for a small fixed-length container (or use StaticArrays if you need other array operations). Doing ./ on arrays of length 2 is far slower than just doing two scalar divisions (unless you use StaticArrays) because of the heap allocations among other things, but it is fine on a tuple.
Your type annotations accomplish nothing for performance. (I would especially omit the return-type annotation.) value is probably overly restrictively typed; why not accept any Real value? See Functions · The Julia Language.
findlast(gridpoints .<= value) allocates large new array to hold gridpoints .<= value — not a good idea to repeat many times in an inner loop. You could do findlast(<=(value), gridpoints) to avoid this.
findlast is a linear search. Far better to require that the grid points be sorted and use searchsortedfirst to have a binary search. (If you have a regular grid, of course, you can avoid searching at all and just compute (value - start) / Δx.)

DNF · March 8, 2022, 4:10pm

fredrikpaues:

function locate(
    values::AbstractArray{<:AbstractFloat},
    gridpoints::AbstractVector{<:AbstractFloat};
    locate_below::Bool=false,
    locate_above::Bool=false,
)::Tuple{AbstractArray{Integer}, AbstractArray{<:AbstractFloat}}

It is a really unfortunate habit you’ve picked up with the output annotations. You should absolutely stop doing it. In this case, it doesn’t just ‘do nothing’, it can actually significantly harm performance, since you are forcing your function to return an array with an abstract element type, AbstractArray{Integer}.

Here’s an example:

foo(x)::AbstractArray{Integer} = reverse(x)
bar(x) = reverse(x)  # same but no annotation

1.8.0-beta1> x = rand(0:9, 1000);

1.8.0-beta1> @btime foo($x);
  5.133 μs (2 allocations: 15.88 KiB)

1.8.0-beta1> @btime bar($x);
  1.120 μs (1 allocation: 7.94 KiB)

If you then use this vector later in your calculations, it is very bad for performance.

You should use output annotations only when there is a very good reason. Right now, they make the code less readable, confusing, restrictive, error prone, and slow.

(In fact, I suggest you start writing your code with zero type annotations, and then just add them when necessary. The exception being if you define new structs.)

Good luck!

markmbaum · March 8, 2022, 4:33pm

There is a simple bisection search function inside of BasicInterpolators.jl called findcell if that suits your particular need. You could copy it from here or just install the package and put using BasicInterpolators: findcell in your code.

The function has a docstring, but it simply finds the index of the grid cell containing some value along a sorted vector of coordinates. Once you have that index, you can do a linear interpolation pretty easily.

DNF · March 8, 2022, 4:44pm

Just for fun, here’s an even more catastrophic example of the effects of that annotation:

1.8.0-beta1> x = 1:1_000_000;

1.8.0-beta1> @btime bar($x);
  7.900 ns (0 allocations: 0 bytes)

1.8.0-beta1> @btime foo($x);
  7.857 ms (999491 allocations: 22.88 MiB)

Note nanosecond vs millisecond. The type annotated version is literally a million times slower, because bar does basically no work at all (independent of vector length), while foo scales linearly and allocates heavily.

fredrikpaues · March 9, 2022, 5:36am

Thank you!

I haven’t look into other interpolation packages except Interpolations.jl. Thanks for pointing out that there are alternatives.

In order to address this, it seems that I need to change how I use the function (as i, w .= locate(value, gridpoints) throws an error). I will be sure to look into it.
I mostly put in the type notation to discipline myself while coding and thought that they at worst didn’t affect performance at all. But after this and the replies by @DNF, I will be sure to stop with this unfortunate habit.
This improved performance of that line a lot!
Using the searchsortedlast function improved performance even more!

fredrikpaues · March 9, 2022, 5:41am

Thank you for pointing this out! As I wrote in my last reply I hadn’t considered that the type annotations could actually hurt performance. I will immediately stop with this.

But I remember having read that type annotations sometimes can help the compiler. How do I figure out if a particular variable would benefit from a type annotation?
If I hypothetically switched to concrete types, could that improve performance? Or would it at best not hurt it?
In the code in my original post, I dabbled in using multiple dispatch. To do that, I’ve gotten the impression that I need to use type annotations. Should I just use Number for the first/scalar case and AbstractArray for the second? (Could I somehow write my function so that I could call it with dot notation, i.e. locate.(values, gridpoints)?)

fredrikpaues · March 9, 2022, 5:47am

I will try this, but I’m a bit skeptical. As the findcell function only returns the index, it seems that I would have to check the boundaries another time to compute the weights. Perhaps I could improve performance by computing the index and the weights myself using the bisection method as findcell

DNF · March 9, 2022, 8:25am

I think that as a first approximation you can assume ‘never’.There may be some few cases, it used to be that things like

x += boolean_expression::Int

where x is an Int, but I don’t think that’s been needed for a long time. There’s also a long-standing issue with captured variables in closures, that I think is still ongoing.

But if you yourself can predict the output type based solely on the input types, then the compiler can too, virtually always, and better than a human. But if you have an actual type unstable function, and you have privileged knowledge about the computation from outside the type domain, then you can try. But first you should verify that there is an actual instability, using e.g. @code_warntype, otherwise, there is definitely no point.

In function signatures, no. In output signatures, concrete is better than abstract, apparently, but it’s better to leave it off, unless you discover a bona fide instability.

In struct definitions, yes, you should use concrete or parametric types.

Yes, you need type annotations to define different methods of the same function for different input types. This is actually the primary purpose of function type signatures.

Yes, from a very brief look, I think you need only one method. Just keep the first one, you could give it the signature

locate(::Real, ::AbstractVector, ::Bool, ::Bool)

or just drop the types completely. And then, when you have a vector of values, you call

locate.(values, Ref(gridpoints), locate_below, locate_above)

The Ref() protects gridpoints against getting broadcasted over. This will give you a somewhat different output type, something like Vector{Tuple{Int, Vector}} (don’t annotate your function with it, though ) I think someone already mentioned that weights ought to be a tuple instead of a vector, if you change that, the output would be, approximately, Tuple{Int, Tuple{Float64, Float64}}.

stevengj · March 9, 2022, 12:33pm

If you use a tuple for w instead of an array, you can just use i, w = locate(...).

markmbaum · March 11, 2022, 10:45pm

Here is an example script that wraps the findcell function to handle boundaries with extrapolation or not, plotting a test case with with the PyPlot package.

using PyPlot

pygui(true)

##

function findcell(xᵢ, X)
    #number of grid cells
    n = length(X)
    #handle boundaries
    @inbounds (xᵢ <= X[1]) && return(1)
    @inbounds (xᵢ >= X[n]) && return(n-1)
    #bisection search for the containing cell
    L = 0
    H = n
    while H - L > 1
        M = (H + L) ÷ 2
        @inbounds if X[M] > xᵢ
            H = M
        else
            L = M
        end
    end
    return L
end

function locate(value, gridpoints, locate_below=false, locate_above=true)
    #get the cell index or boundaries
    index = findcell(value, gridpoints)
    #width of target cell
    Δ = gridpoints[index+1] - gridpoints[index]
    #get interpolation weights, handling boundaries
    weights = if (value < gridpoints[1]) & !locate_below
        (1.0, 0.0)
    elseif (value > gridpoints[end]) & !locate_above
        (0.0, 1.0)
    else
        (
            (gridpoints[index+1] - value)/Δ,
            (value - gridpoints[index])/Δ
        )
    end
    return index, weights
end

##

#grid coordinates
gridpoints = [0.5, 1, 2, 2.5, 4]
#values to interpolate
gridvalues = 2*rand(5) .- 1

#coordinates to interpolate
x = LinRange(-1, 6, 1000)
#interpolation values
y = zeros(length(x))

figure()
#plot the interpolation points
plot(gridpoints, gridvalues, ".")
#interpolate and plot with different boundary behaviors
for extrapolate ∈ (true, false)
    for i ∈ 1:length(x)
        index, weights = locate(x[i], gridpoints, extrapolate, extrapolate)
        y[i] = sum(weights .* gridvalues[index:index+1])
    end
    plot(x, y, label="extrapolate=$extrapolate")
end
legend()

Hopefully this is helpful. There shouldn’t be any need to annotate any of the function inputs or outputs here.

As you did in the initial post, you can write another function to wrap the locate function if you need to store a bunch of indices and weights.

fredrikpaues · March 12, 2022, 10:07am

Thanks! I ended up doing something very similar, but that avoids the duplicated boundary handling (once in findcell and once in locate).

fredrikpaues · March 12, 2022, 10:21am

This is what I have right now.

Weights are tuples (as advised by @stevengj).
I have tried to avoid the type annotations (@stevengj and @DNF) except for multiple dispatch and the booleans. (Perhaps I should skip the Bool:s as well ).
I use binary search for the interior points (@stevengj and @markmbaum).
I found that dot notation (locate.()) performed worse than just implementing the values::AbstractArray{<:Real} case myself. That way, I could also (a) control the structure of the output to something that was a little more tractable and (b) ensure that gridpoints was a tuple, which turned out to improve performance.

Perhaps performance could be improved by using @inbounds when checking the boundaries (as is done in findcell suggested by @markmbaum), but I haven’t managed to figure out what exactly that does

function locate(
    value::Real,
    gridpoints;
    locate_below::Bool=false,
    locate_above::Bool=false,
)
    if value <= gridpoints[1]
        index = 1

        if !locate_below
            weights = (1.0, 0.0)
            return index, weights
        end
    elseif value >= gridpoints[end]
        index = length(gridpoints) - 1

        if !locate_above
            weights = (0.0, 1.0)
            return index, weights
        end
    else
        # Implement binary search as
        #     `index = searchsortedlast(gridpoints, value)`
        # but with support for tuple gridpoints
        low = 0
        high = length(gridpoints) + 1
        @inbounds while low < high - 1
            mid = low + ((high - low) >>> 0x01)  # floor(Integer, (low + high) / 2)
            if value < gridpoints[mid]
                high = mid
            else
                low = mid
            end
        end

        index = low
    end

    weights = (
        (
            gridpoints[index + 1] - value,
            value - gridpoints[index]
        ) ./ (
            gridpoints[index + 1] - gridpoints[index]
        )
    )

    return index, weights
end
function locate(
    values::AbstractArray{<:Real},
    gridpoints;
    locate_below::Bool=false,
    locate_above::Bool=false,
)
    if !(typeof(gridpoints) <: Tuple{Vararg{<:Real}})
        gridpoints = Tuple{Vararg{<:Real}}(gridpoints)
    end

    indices = ones(Integer, size(values))
    weights = Array{Tuple, ndims(values)}(undef, size(values))

    for ix in eachindex(values)
        indices[ix], weights[ix] = locate(
            values[ix],
            gridpoints;
            locate_below=locate_below,
            locate_above=locate_above
        )
    end

    return indices, weights
end

stevengj · March 12, 2022, 1:46pm

I’m confused, why can’t you use searchsortedlast? You almost certainly don’t want to use a tuple to store a long array of grid points.

fredrikpaues · March 12, 2022, 2:16pm

I found it was faster when I tried it with @btime, which is why I change gridpoints to a tuple in the array case. Maybe I was unsuccessful in my thinking here I will check again in a few hours and post my tests here.

fredrikpaues · March 12, 2022, 7:29pm

So these are the tests that I ran.

First, I ran these but without the part in locate that converts gridpoints to a tuple.

julia> logrange(x1, x2, n) = (10^y for y in range(log10(x1), log10(x2); length=n));
julia> g_vector = collect(logrange(0.1, 0.9, 100));
julia> g_tuple = Tuple{Real}(g_vector);
julia> using StaticArrays;
julia> g_svector = SVector{length(g_vector), Real}(g_vector);
julia> values = rand(200);
julia> using BenchmarkTools;
julia> @btime locate(values, g_vector); @btime locate(values, g_tuple); @btime locate(values, g_svector);
3.612 μs (203 allocations: 9.81 KiB)
2.100 μs (203 allocations: 9.81 KiB)
187.400 μs (5183 allocations: 103.09 KiB)
julia> @btime locate(values, g_vector); @btime locate(values, g_tuple); @btime locate(values, g_svector);
3.600 μs (203 allocations: 9.81 KiB)
2.089 μs (203 allocations: 9.81 KiB)
188.200 μs (5183 allocations: 103.09 KiB)

And then I ran the same tests but did convert gridpoints to a tuple.

julia> @btime locate(values, g_vector); @btime locate(values, g_tuple); @btime locate(values, g_svector);
32.100 μs (1115 allocations: 32.20 KiB)
20.200 μs (1009 allocations: 28.78 KiB)
46.300 μs (1519 allocations: 42.36 KiB)
julia> @btime locate(values, g_vector); @btime locate(values, g_tuple); @btime locate(values, g_svector);
32.400 μs (1115 allocations: 32.20 KiB)
20.200 μs (1009 allocations: 28.78 KiB)
46.800 μs (1519 allocations: 42.36 KiB)

So, tuples perform the best. But apparently my decision to convert was rather stupid. (I think that I originally tested with fewer nodes in my grids, but these number are more in line with what I will use the function for.)