Re-implement GriddedInterpolation

Author: timholy

src/Interpolations.jl

@@ -33,7 +33,7 @@ using LinearAlgebra, SparseArrays
 using StaticArrays, WoodburyMatrices, Ratios, AxisAlgorithms, OffsetArrays
 
 using Base: @propagate_inbounds
-import Base: convert, size, axes, promote_rule, ndims, eltype, checkbounds
+import Base: convert, size, axes, promote_rule, ndims, eltype, checkbounds, axes1
 
 abstract type Flag end
 abstract type InterpolationType <: Flag end
@@ -300,7 +300,7 @@ end
 
 include("nointerp/nointerp.jl")
 include("b-splines/b-splines.jl")
-# include("gridded/gridded.jl")
+include("gridded/gridded.jl")
 include("extrapolation/extrapolation.jl")
 include("scaling/scaling.jl")
 include("utils.jl")

src/b-splines/indexing.jl

@@ -88,8 +88,8 @@ end
 end
 
 
-function weightedindexes(fs::F, itpflags::NTuple{N,Flag}, axs::NTuple{N,AbstractUnitRange}, xs::NTuple{N,Number}) where {F,N}
-    parts = map((flag, ax, x)->weightedindex_parts(fs, flag, ax, x), itpflags, axs, xs)
+function weightedindexes(fs::F, itpflags::NTuple{N,Flag}, knots::NTuple{N,AbstractVector}, xs::NTuple{N,Number}) where {F,N}
+    parts = map((flag, knotvec, x)->weightedindex_parts(fs, flag, knotvec, x), itpflags, knots, xs)
     weightedindexes(parts...)
 end
 
@@ -179,7 +179,7 @@ slot_substitute(kind1::Tuple{}, kind2::Tuple{}, p, v, g, h) = ()
 weightedindex_parts(fs::F, itpflag::BSpline, ax, x) where F =
     weightedindex_parts(fs, degree(itpflag), ax, x)
 
-function weightedindex_parts(fs::F, deg::Degree, ax, x) where F
+function weightedindex_parts(fs::F, deg::Degree, ax::AbstractUnitRange{<:Integer}, x) where F
     pos, δx = positions(deg, ax,  x)
     (position=pos, coefs=fmap(fs, deg, δx))
 end
@@ -198,16 +198,22 @@ function getindex_return_type(::Type{BSplineInterpolation{T,N,TCoefs,IT,Axs}}, :
 end
 
 # This handles round-towards-the-middle for points on half-integer edges
-roundbounds(x::Integer, bounds) = x
-function roundbounds(x, bounds)
+roundbounds(x::Integer, bounds::Tuple{Real,Real}) = x
+roundbounds(x::Integer, bounds::AbstractUnitRange) = x
+roundbounds(x::Number, bounds::Tuple{Real,Real}) = _roundbounds(x, bounds)
+roundbounds(x::Number, bounds::AbstractUnitRange) = _roundbounds(x, bounds)
+function _roundbounds(x::Number, bounds::Union{Tuple{Real,Real}, AbstractUnitRange})
     l, u = first(bounds), last(bounds)
     h = half(x)
     xh = x+h
     ifelse(x < u+half(u), floor(xh), ceil(xh)-1)
 end
 
-floorbounds(x::Integer, ax) = x
-function floorbounds(x, ax)
+floorbounds(x::Integer, ax::Tuple{Real,Real}) = x
+floorbounds(x::Integer, ax::AbstractUnitRange) = x
+floorbounds(x, ax::Tuple{Real,Real}) = _floorbounds(x, ax)
+floorbounds(x, ax::AbstractUnitRange) = _floorbounds(x, ax)
+function _floorbounds(x, ax::Union{Tuple{Real,Real}, AbstractUnitRange})
     l = first(ax)
     h = half(x)
     ifelse(x < l, floor(x+h), floor(x+zero(h)))

src/b-splines/linear.jl

@@ -21,7 +21,7 @@ a piecewise linear function connecting each pair of neighboring data points.
 """
 Linear
 
-function positions(::Linear, ax, x)
+function positions(::Linear, ax::AbstractUnitRange{<:Integer}, x)
     f = floor(x)
     # When x == last(ax) we want to use the x-1, x pair
     f = ifelse(x == last(ax), f - oneunit(f), f)

src/deprecations.jl

@@ -1,5 +1,6 @@
 # deprecate getindex for non-integer numeric indices
 @deprecate getindex(itp::AbstractInterpolation{T,N}, i::Vararg{Number,N}) where {T,N} itp(i...)
+@deprecate getindex(itp::AbstractInterpolation{T,N}, i::Vararg{ExpandedIndexTypes,N}) where {T,N} itp(i...)
 
 for T in (:Throw, :Flat, :Line, :Free, :Periodic, :Reflect, :InPlace, :InPlaceQ)
     @eval begin

src/gridded/constant.jl

@@ -1,3 +1,11 @@
+function base_rem(::Constant, knotv, ki, x)
+    l, u = knotv[ki], knotv[ki+1]
+
+    xm = roundbounds(x, bounds)
+    δx = x - xm
+    fast_trunc(Int, xm), δx
+end
+
 function define_indices_d(::Type{Gridded{Constant}}, d, pad)
     symix, symx = Symbol("ix_",d), Symbol("x_",d)
     symk, symkix = Symbol("k_",d), Symbol("kix_",d)

src/gridded/gridded.jl

@@ -1,33 +1,22 @@
 export Gridded
 
-struct Gridded{D<:Degree} <: InterpolationType end
-Gridded(::D) where {D<:Degree} = Gridded{D}()
-
-griddedtype(::Type{Gridded{D}}) where {D<:Degree} = D
+struct Gridded{D<:Degree} <: InterpolationType
+    degree::D
+end
 
 const GridIndex{T} = Union{AbstractVector{T}, Tuple}
 
-# Because Ranges check bounds on getindex, it's actually faster to convert the
-# knots to Vectors. It's also good to take a copy, so it doesn't get modified later.
-struct GriddedInterpolation{T,N,TCoefs,IT<:DimSpec{Gridded},K<:Tuple{Vararg{AbstractVector}}} <: AbstractInterpolation{T,N,IT,OnGrid}
+struct GriddedInterpolation{T,N,TCoefs,IT<:DimSpec{Gridded},K<:Tuple{Vararg{AbstractVector}}} <: AbstractInterpolation{T,N,IT}
     knots::K
     coefs::Array{TCoefs,N}
+    it::IT
 end
 function GriddedInterpolation(::Type{TWeights}, knots::NTuple{N,GridIndex}, A::AbstractArray{TCoefs,N}, it::IT) where {N,TCoefs,TWeights<:Real,IT<:DimSpec{Gridded},pad}
     isconcretetype(IT) || error("The b-spline type must be a leaf type (was $IT)")
     isconcretetype(TCoefs) || warn("For performance reasons, consider using an array of a concrete type (eltype(A) == $(eltype(A)))")
 
-    knts = mapcollect(knots...)
-    for (d,k) in enumerate(knts)
-        length(k) == size(A, d) || throw(DimensionMismatch("knot vectors must have the same number of elements as the corresponding dimension of the array"))
-        length(k) == 1 && error("dimensions of length 1 not yet supported")  # FIXME
-        issorted(k) || error("knot-vectors must be sorted in increasing order")
-        iextract(IT, d) != NoInterp || k == collect(1:size(A, d)) || error("knot-vector should be the range 1:$(size(A,d)) for the method Gridded{NoInterp}")
-    end
+    check_gridded(it, knots, axes(A))
     c = zero(TWeights)
-    for _ in 2:N
-        c *= c
-    end
     if isempty(A)
         T = Base.promote_op(*, typeof(c), eltype(A))
     else
@@ -36,39 +25,44 @@ function GriddedInterpolation(::Type{TWeights}, knots::NTuple{N,GridIndex}, A::A
     GriddedInterpolation{T,N,TCoefs,IT,typeof(knots)}(knots, A, it)
 end
 
-Base.parent(A::GriddedInterpolation) = A.coefs
+@inline function check_gridded(itpflag, knots, axs)
+    flag, ax1, k1 = getfirst(itpflag), axs[1], knots[1]
+    if flag isa NoInterp
+        k1 == ax1 || error("for NoInterp knot vector should be $ax1, got $k1")
+    else
+        axes(k1, 1) == ax1 || throw(DimensionMismatch("knot vectors must have the same axes as the corresponding dimension of the array"))
+    end
+    degree(flag) isa Union{NoInterp,Constant,Linear} || error("only Linear, Constant, and NoInterp supported, got $flag")
+    length(k1) == 1 && error("dimensions of length 1 not yet supported")  # FIXME
+    issorted(k1) || error("knot-vectors must be sorted in increasing order")
+    check_gridded(getrest(itpflag), Base.tail(knots), Base.tail(axs))
+end
+check_gridded(::Any, ::Tuple{}, ::Tuple{}) = nothing
+degree(flag::Gridded) = flag.degree
 
-# A type-stable version of map(collect, knots)
-mapcollect() = ()
-@inline mapcollect(k::AbstractVector) = (collect(k),)
-@inline mapcollect(k1::AbstractVector, k2::AbstractVector...) = (collect(k1), mapcollect(k2...)...)
+Base.parent(A::GriddedInterpolation) = A.coefs
+coefficients(A::GriddedInterpolation) = A.coefs
 
-# Utilities for working either with scalars or tuples/tuple-types
-iextract(::Type{T}, d) where {T<:Gridded} = T
-iextract(::Type{T}, d) where {T<:GridType} = T
+size(A::GriddedInterpolation) = size(A.coefs)
+axes(A::GriddedInterpolation) = axes(A.coefs)
 
-@generated function size(itp::GriddedInterpolation{T,N,TCoefs,IT,K,pad}, d) where {T,N,TCoefs,IT,K,pad}
-    quote
-        d <= $N ? size(itp.coefs, d) - 2*padextract($pad, d) : 1
-    end
-end
+itpflag(A::GriddedInterpolation) = A.it
 
 function interpolate(::Type{TWeights}, ::Type{TCoefs}, knots::NTuple{N,GridIndex}, A::AbstractArray{Tel,N}, it::IT) where {TWeights,TCoefs,Tel,N,IT<:DimSpec{Gridded}}
-    GriddedInterpolation(TWeights, knots, A, it, Val{0}())
+    GriddedInterpolation(TWeights, knots, A, it)
 end
 function interpolate(knots::NTuple{N,GridIndex}, A::AbstractArray{Tel,N}, it::IT) where {Tel,N,IT<:DimSpec{Gridded}}
     interpolate(tweight(A), tcoef(A), knots, A, it)
 end
 
-interpolate!(::Type{TWeights}, knots::NTuple{N,GridIndex}, A::AbstractArray{Tel,N}, it::IT) where {TWeights,Tel,N,IT<:DimSpec{Gridded}} = GriddedInterpolation(TWeights, knots, A, it, Val{0}())
+interpolate!(::Type{TWeights}, knots::NTuple{N,GridIndex}, A::AbstractArray{Tel,N}, it::IT) where {TWeights,Tel,N,IT<:DimSpec{Gridded}} =
+    GriddedInterpolation(TWeights, knots, A, it)
 function interpolate!(knots::NTuple{N,GridIndex}, A::AbstractArray{Tel,N}, it::IT) where {Tel,N,IT<:DimSpec{Gridded}}
     interpolate!(tweight(A), tcoef(A), knots, A, it)
 end
 
-lbound(itp::GriddedInterpolation, d) = itp.knots[d][1]
-ubound(itp::GriddedInterpolation, d) = itp.knots[d][end]
-lbound(itp::GriddedInterpolation, d, inds) = itp.knots[d][1]
-ubound(itp::GriddedInterpolation, d, inds) = itp.knots[d][end]
+lbounds(itp::GriddedInterpolation) = first.(itp.knots)
+ubounds(itp::GriddedInterpolation) = last.(itp.knots)
 
 include("constant.jl")
 include("linear.jl")