Re-implement fast iteration for ScaledInterpolation

Author: timholy

src/scaling/scaling.jl

@@ -31,30 +31,38 @@ check_ranges(::Any, ::Tuple{}, ::Tuple{}) = nothing
 check_range(::NoInterp, ax, r) = ax == r || throw(ArgumentError("The range $r did not equal the corresponding axis of the interpolation object $ax"))
 check_range(::Any, ax, r) = length(ax) == length(r) || throw(ArgumentError("The range $r is incommensurate with the corresponding axis $ax"))
 
+# With regards to size and [], ScaledInterpolation behaves like the underlying interpolation object
 size(sitp::ScaledInterpolation) = size(sitp.itp)
 axes(sitp::ScaledInterpolation) = axes(sitp.itp)
 
+@propagate_inbounds function Base.getindex(sitp::ScaledInterpolation{T,N}, i::Vararg{Int,N}) where {T,N}
+    sitp.itp[i...]
+end
+
 lbounds(sitp::ScaledInterpolation) = _lbounds(sitp.ranges, itpflag(sitp.itp))
 ubounds(sitp::ScaledInterpolation) = _ubounds(sitp.ranges, itpflag(sitp.itp))
 
 boundstep(r::StepRange) = r.step / 2
 boundstep(r::UnitRange) = 1//2
-
-lbound(ax::AbstractRange, ::DegreeBC, ::OnCell) = first(ax) - boundstep(ax)
-ubound(ax::AbstractRange, ::DegreeBC, ::OnCell) = last(ax) + boundstep(ax)
-lbound(ax::AbstractRange, ::DegreeBC, ::OnGrid) = first(ax)
-ubound(ax::AbstractRange, ::DegreeBC, ::OnGrid) = last(ax)
-
 """
 Returns *half* the width of one step of the range.
 
 This function is used to calculate the upper and lower bounds of `OnCell` interpolation objects.
 """ boundstep
 
+lbound(ax::AbstractRange, ::DegreeBC, ::OnCell) = first(ax) - boundstep(ax)
+ubound(ax::AbstractRange, ::DegreeBC, ::OnCell) = last(ax) + boundstep(ax)
+lbound(ax::AbstractRange, ::DegreeBC, ::OnGrid) = first(ax)
+ubound(ax::AbstractRange, ::DegreeBC, ::OnGrid) = last(ax)
+
+# For (), we scale the evaluation point
 function (sitp::ScaledInterpolation{T,N})(xs::Vararg{Number,N}) where {T,N}
     xl = coordslookup(itpflag(sitp.itp), sitp.ranges, xs)
     sitp.itp(xl...)
 end
+@inline function (sitp::ScaledInterpolation)(x::Vararg{UnexpandedIndexTypes})
+    sitp(to_indices(sitp, x)...)
+end
 
 (sitp::ScaledInterpolation{T,1}, x::Number, y::Int) where {T} = y == 1 ? sitp(x) : Base.throw_boundserror(sitp, (x, y))
 
@@ -134,167 +142,85 @@ rescale_gradient(r::UnitRange, g) = g
 Implements the chain rule dy/dx = dy/du * du/dx for use when calculating gradients with scaled interpolation objects.
 """ rescale_gradient
 
+### Iteration
 
-# ### Iteration
-# mutable struct ScaledIterator{CR<:CartesianIndices,SITPT,X1,Deg,T}
-#     rng::CR
-#     sitp::SITPT
-#     dx_1::X1
-#     nremaining::Int
-#     fx_1::X1
-#     itp_tail::NTuple{Deg,T}
-# end
-
-# nelements(::Union{Type{NoInterp},Type{Constant}}) = 1
-# nelements(::Type{Linear}) = 2
-# nelements(::Type{Q}) where {Q<:Quadratic} = 3
-
-# eachvalue_zero(::Type{R}, ::Type{BT}) where {R,BT<:Union{Type{NoInterp},Type{Constant}}} =
-#     (zero(R),)
-# eachvalue_zero(::Type{R}, ::Type{Linear}) where {R} = (zero(R),zero(R))
-# eachvalue_zero(::Type{R}, ::Type{Q}) where {R,Q<:Quadratic} = (zero(R),zero(R),zero(R))
-
-# """
-# `eachvalue(sitp)` constructs an iterator for efficiently visiting each
-# grid point of a ScaledInterpolation object in which a small grid is
-# being "scaled up" to a larger one.  For example, suppose you have a
-# core `BSpline` object defined on a 5x7x4 grid, and you are scaling it
-# to a 100x120x20 grid (via `linspace(1,5,100), linspace(1,7,120),
-# linspace(1,4,20)`).  You can perform interpolation at each of these
-# grid points via
-
-# ```
-#     function foo!(dest, sitp)
-#         i = 0
-#         for s in eachvalue(sitp)
-#             dest[i+=1] = s
-#         end
-#         dest
-#     end
-# ```
-
-# which should be more efficient than
-
-# ```
-#     function bar!(dest, sitp)
-#         for I in CartesianIndices(size(dest))
-#             dest[I] = sitp[I]
-#         end
-#         dest
-#     end
-# ```
-# """
-# function eachvalue(sitp::ScaledInterpolation{T,N}) where {T,N}
-#     ITPT = basetype(sitp)
-#     IT = itptype(ITPT)
-#     R = getindex_return_type(ITPT, Int)
-#     BT = bsplinetype(iextract(IT, 1))
-#     itp_tail = eachvalue_zero(R, BT)
-#     dx_1 = coordlookup(sitp.ranges[1], 2) - coordlookup(sitp.ranges[1], 1)
-#     ScaledIterator(CartesianIndices(ssize(sitp)), sitp, dx_1, 0, zero(dx_1), itp_tail)
-# end
+struct ScaledIterator{SITPT,CI,WIS}
+    sitp::SITPT   # ScaledInterpolation object
+    ci::CI        # the CartesianIndices object
+    wis::WIS      # WeightedIndex vectors
+    breaks1::Vector{Int}   # breaks along dimension 1 where new evaluations must occur
+end
 
-# function index_gen1(::Union{Type{NoInterp}, Type{BSpline{Constant}}})
-#     quote
-#         value = iter.itp_tail[1]
-#     end
-# end
+Base.IteratorSize(::Type{ScaledIterator{SITPT,CI,WIS}}) where {SITPT,CI<:CartesianIndices{N},WIS} where N = Base.HasShape{N}()
+Base.axes(iter::ScaledIterator) = axes(iter.ci)
+Base.size(iter::ScaledIterator) = size(iter.ci)
 
-# function index_gen1(::Type{BSpline{Linear}})
-#     quote
-#         p = iter.itp_tail
-#         value = c_1*p[1] + cp_1*p[2]
-#     end
-# end
+struct ScaledIterState{N,V}
+    cistate::CartesianIndex{N}
+    ibreak::Int
+    cached_evaluations::NTuple{N,V}
+end
 
-# function index_gen1(::Type{BSpline{Q}}) where Q<:Quadratic
-#     quote
-#         p = iter.itp_tail
-#         value = cm_1*p[1] + c_1*p[2] + cp_1*p[3]
-#     end
-# end
-# function index_gen_tail(B::Union{Type{NoInterp}, Type{BSpline{Constant}}}, ::Type{IT}, N) where IT
-#     [index_gen(B, IT, N, 0)]
-# end
+function eachvalue(sitp::ScaledInterpolation{T,N}) where {T,N}
+    itps = tcollect(itpflag, sitp.itp)
+    newaxes = map(r->Base.Slice(ceil(Int, first(r)):floor(Int, last(r))), sitp.ranges)
+    wis = dimension_wis(value_weights, itps, axes(sitp.itp), newaxes, sitp.ranges)
+    wis1 = wis[1]
+    i1 = first(axes(wis1, 1))
+    breaks1 = [i1]
+    for i in Iterators.drop(axes(wis1, 1), 1)
+        if indexes(wis1[i]) != indexes(wis1[i-1])
+            push!(breaks1, i)
+        end
+    end
+    push!(breaks1, last(axes(wis1, 1))+1)
+    ScaledIterator(sitp, CartesianIndices(newaxes), wis, breaks1)
+end
 
-# function index_gen_tail(::Type{BSpline{Linear}}, ::Type{IT}, N) where IT
-#     [index_gen(BS1, IT, N, i) for i = 0:1]
-# end
+function dimension_wis(f::F, itps, axs, newaxes, ranges) where F
+    itpflag, ax, nax, r = itps[1], axs[1], newaxes[1], ranges[1]
+    function makewi(x)
+        pos, coefs = weightedindex_parts((f,), itpflag, ax, coordlookup(r, x))
+        maybe_weightedindex(pos, coefs[1])
+    end
+    (makewi.(nax), dimension_wis(f, Base.tail(itps), Base.tail(axs), Base.tail(newaxes), Base.tail(ranges))...)
+end
+dimension_wis(f, ::Tuple{}, ::Tuple{}, ::Tuple{}, ::Tuple{}) = ()
+
+function Base.iterate(iter::ScaledIterator)
+    ret = iterate(iter.ci)
+    ret === nothing && return nothing
+    item, cistate = ret
+    wis = getindex.(iter.wis, Tuple(item))
+    ces = cache_evaluations(iter.sitp.itp.coefs, indexes(wis[1]), weights(wis[1]), Base.tail(wis))
+    return _reduce(+, weights(wis[1]).*ces), ScaledIterState(cistate, first(iter.breaks1), ces)
+end
 
-# function index_gen_tail(::Type{BSpline{Q}}, ::Type{IT}, N) where {IT,Q<:Quadratic}
-#     [index_gen(BSpline{Q}, IT, N, i) for i = -1:1]
-# end
-# function nremaining_gen(::Union{Type{BSpline{Constant}}, Type{BSpline{Q}}}) where Q<:Quadratic
-#     quote
-#         EPS = 0.001*iter.dx_1
-#         floor(Int, iter.dx_1 >= 0 ?
-#               (min(length(range1)+EPS, round(Int,x_1) + 0.5) - x_1)/iter.dx_1 :
-#               (max(1-EPS, round(Int,x_1) - 0.5) - x_1)/iter.dx_1)
-#     end
-# end
+function Base.iterate(iter::ScaledIterator, state)
+    ret = iterate(iter.ci, state.cistate)
+    ret === nothing && return nothing
+    item, cistate = ret
+    i1 = item[1]
+    isnext1 = i1 == state.cistate[1]+1
+    if isnext1 && i1 < iter.breaks1[state.ibreak+1]
+        # We can use the previously cached values
+        wis1 = iter.wis[1][i1]
+        return _reduce(+, weights(wis1).*state.cached_evaluations), ScaledIterState(cistate, state.ibreak, state.cached_evaluations)
+    end
+    # Re-evaluate. We're being a bit lazy here: in some cases, some of the cached values could be reused
+    wis = getindex.(iter.wis, Tuple(item))
+    ces = cache_evaluations(iter.sitp.itp.coefs, indexes(wis[1]), weights(wis[1]), Base.tail(wis))
+    return _reduce(+, weights(wis[1]).*ces), ScaledIterState(cistate, isnext1 ? state.ibreak+1 : first(iter.breaks1), ces)
+end
 
-# function nremaining_gen(::Type{BSpline{Linear}})
-#     quote
-#         EPS = 0.001*iter.dx_1
-#         floor(Int, iter.dx_1 >= 0 ?
-#               (min(length(range1)+EPS, floor(Int,x_1) + 1) - x_1)/iter.dx_1 :
-#               (max(1-EPS, floor(Int,x_1)) - x_1)/iter.dx_1)
-#     end
-# end
-# function next_gen(::Type{ScaledIterator{CR,SITPT,X1,Deg,T}}) where {CR,SITPT,X1,Deg,T}
-#     N = ndims(CR)
-#     ITPT = basetype(SITPT)
-#     IT = itptype(ITPT)
-#     BS1 = iextract(IT, 1)
-#     BS1 == NoInterp && error("eachvalue is not implemented (and does not make sense) for NoInterp along the first dimension")
-#     pad = padding(ITPT)
-#     x_syms = [Symbol("x_", i) for i = 1:N]
-#     interp_index(IT, i) = iextract(IT, i) != NoInterp ?
-#         :($(x_syms[i]) = coordlookup(sitp.ranges[$i], state[$i])) :
-#         :($(x_syms[i]) = state[$i])
-#     # Calculations for the first dimension
-#     interp_index1 = interp_index(IT, 1)
-#     indices1 = define_indices_d(BS1, 1, padextract(pad, 1))
-#     coefexprs1 = coefficients(BS1, N, 1)
-#     nremaining_expr = nremaining_gen(BS1)
-#     # Calculations for the rest of the dimensions
-#     interp_indices_tail = map(i -> interp_index(IT, i), 2:N)
-#     indices_tail = [define_indices_d(iextract(IT, i), i, padextract(pad, i)) for i = 2:N]
-#     coefexprs_tail = [coefficients(iextract(IT, i), N, i) for i = 2:N]
-#     value_exprs_tail = index_gen_tail(BS1, IT, N)
-#     quote
-#         sitp = iter.sitp
-#         itp = sitp.itp
-#         inds_itp = axes(itp)
-#         if iter.nremaining > 0
-#             iter.nremaining -= 1
-#             iter.fx_1 += iter.dx_1
-#         else
-#             range1 = sitp.ranges[1]
-#             $interp_index1
-#             $indices1
-#             iter.nremaining = $nremaining_expr
-#             iter.fx_1 = fx_1
-#             $(interp_indices_tail...)
-#             $(indices_tail...)
-#             $(coefexprs_tail...)
-#             @inbounds iter.itp_tail = ($(value_exprs_tail...),)
-#         end
-#         fx_1 = iter.fx_1
-#         $coefexprs1
-#         $(index_gen1(BS1))
-#     end
-# end
+_reduce(op, list) = op(list[1], _reduce(op, Base.tail(list)))
+_reduce(op, list::Tuple{Number}) = list[1]
+_reduce(op, list::Tuple{}) = error("cannot reduce an empty list")
 
-# @generated function iterate(iter::ScaledIterator{CR,ITPT}, state::Union{Nothing,CartesianIndex{N}} = nothing) where {CR,ITPT,N}
-#     value_expr = next_gen(iter)
-#     quote
-#         rng_next = state ≡ nothing ? iterate(iter.rng) : iterate(iter.rng, state)
-#         rng_next ≡ nothing && return nothing
-#         state = rng_next[2]
-#         $value_expr
-#         (value, state)
-#     end
-# end
+# We use weights only as a ruler to determine when we are done
+cache_evaluations(coefs, i::Int, weights, rest) = (coefs[i, rest...], cache_evaluations(coefs, i+1, Base.tail(weights), rest)...)
+cache_evaluations(coefs, indexes, weights, rest) = (coefs[indexes[1], rest...], cache_evaluations(coefs, Base.tail(indexes), Base.tail(weights), rest)...)
+cache_evaluations(coefs, ::Int, ::Tuple{}, rest) = ()
+cache_evaluations(coefs, ::Any, ::Tuple{}, rest) = ()
 
-# ssize(sitp::ScaledInterpolation{T,N}) where {T,N} = map(r->round(Int, last(r)-first(r)+1), sitp.ranges)::NTuple{N,Int}
+ssize(sitp::ScaledInterpolation{T,N}) where {T,N} = map(r->round(Int, last(r)-first(r)+1), sitp.ranges)::NTuple{N,Int}

test/readme-examples.jl

@@ -18,7 +18,7 @@ using Interpolations, Test
     @test v ≈ (0.9*(0.8*A[3,4]+0.2*A[4,4]) + 0.1*(0.8*A[3,5]+0.2*A[4,5]))
 
     # Quadratic interpolation with reflecting boundary conditions
-    # Quadratic is the lowest order that has continuous gradien
+    # Quadratic is the lowest order that has continuous gradient
     itp = interpolate(A, BSpline(Quadratic(Reflect(OnCell()))))
 
     # Linear interpolation in the first dimension, and no interpolation (just lookup) in the second
@@ -27,23 +27,23 @@ using Interpolations, Test
     @test v ≈ (0.35*A[3,5] + 0.65*A[4,5])
 
 
-    # ## Scaled Bsplines
-    # A_x = 1.:2.:40.
-    # A = [log(x) for x in A_x]
-    # itp = interpolate(A, BSpline(Cubic(Line())), OnGrid())
-    # sitp = scale(itp, A_x)
-    # @test sitp(3.) ≈ log(3.) # exactly log(3.)
-    # @test sitp(3.5) ≈ log(3.5) atol=.1 # approximately log(3.5)
+    ## Scaled Bsplines
+    A_x = 1.:2.:40.
+    A = [log(x) for x in A_x]
+    itp = interpolate(A, BSpline(Cubic(Line())), OnGrid())
+    sitp = scale(itp, A_x)
+    @test sitp(3.) ≈ log(3.) # exactly log(3.)
+    @test sitp(3.5) ≈ log(3.5) atol=.1 # approximately log(3.5)
 
     # For multidimensional uniformly spaced grids
     A_x1 = 1:.1:10
     A_x2 = 1:.5:20
     f(x1, x2) = log(x1+x2)
     A = [f(x1,x2) for x1 in A_x1, x2 in A_x2]
     itp = interpolate(A, BSpline(Cubic(Line(OnGrid()))))
-    # sitp = scale(itp, A_x1, A_x2)
-    # @test sitp(5., 10.) ≈ log(5 + 10) # exactly log(5 + 10)
-    # @test sitp(5.6, 7.1) ≈ log(5.6 + 7.1) atol=.1 # approximately log(5.6 + 7.1)
+    sitp = scale(itp, A_x1, A_x2)
+    @test sitp(5., 10.) ≈ log(5 + 10) # exactly log(5 + 10)
+    @test sitp(5.6, 7.1) ≈ log(5.6 + 7.1) atol=.1 # approximately log(5.6 + 7.1)
 
     # ## Gridded interpolation
     # A = rand(8,20)

test/scaling/scaling.jl

@@ -50,45 +50,29 @@ using Test, LinearAlgebra
     sitp32 = @inferred scale(interpolate(Float32[testfunction(x,y) for x in -5:.5:5, y in -4:.2:4], BSpline(Quadratic(Flat(OnGrid())))), -5f0:.5f0:5f0, -4f0:.2f0:4f0)
     @test typeof(@inferred(sitp32(-3.4f0, 1.2f0))) == Float32
 
-    # # Iteration
-    # itp = interpolate(rand(3,3,3), BSpline(Quadratic(Flat())), OnCell())
-    # knots = map(d->1:10:21, 1:3)
-    # sitp = @inferred scale(itp, knots...)
-
-    # iter = @inferred(eachvalue(sitp))
-
-    # iter_next = iterate(iter)
-    # @test iter_next isa Tuple
-    # @test iter_next[1] isa Float64
-    # state = iter_next[2]
-    # inferred_next = Base.return_types(iterate, (typeof(iter),))
-    # @test length(inferred_next) == 1
-    # @test inferred_next[1] == Union{Nothing,Tuple{Float64,typeof(state)}}
-    # iter_next = iterate(iter, state)
-    # @test iter_next isa Tuple
-    # @test iter_next[1] isa Float64
-    # inferred_next = Base.return_types(iterate, (typeof(iter),typeof(state)))
-    # state = iter_next[2]
-    # @test length(inferred_next) == 1
-    # @test inferred_next[1] == Union{Nothing,Tuple{Float64,typeof(state)}}
-
-    # function foo!(dest, sitp)
-    #     i = 0
-    #     for s in eachvalue(sitp)
-    #         dest[i+=1] = s
-    #     end
-    #     dest
-    # end
-    # function bar!(dest, sitp)
-    #     for I in CartesianIndices(size(dest))
-    #         dest[I] = sitp[I]
-    #     end
-    #     dest
-    # end
-    # rfoo = Array{Float64}(undef, Interpolations.ssize(sitp))
-    # rbar = similar(rfoo)
-    # foo!(rfoo, sitp)
-    # bar!(rbar, sitp)
-    # @test rfoo ≈ rbar
-
+    # Iteration
+    itp = interpolate(rand(3,3,3), BSpline(Quadratic(Flat(OnCell()))))
+    knots = map(d->1:10:21, 1:3)
+    sitp = @inferred scale(itp, knots...)
+
+    iter = @inferred(eachvalue(sitp))
+
+    function foo!(dest, sitp)
+        i = 0
+        for s in eachvalue(sitp)
+            dest[i+=1] = s
+        end
+        dest
+    end
+    function bar!(dest, sitp)
+        for I in CartesianIndices(size(dest))
+            dest[I] = sitp(I)
+        end
+        dest
+    end
+    rfoo = Array{Float64}(undef, Interpolations.ssize(sitp))
+    rbar = similar(rfoo)
+    foo!(rfoo, sitp)
+    bar!(rbar, sitp)
+    @test rfoo ≈ rbar
 end