Merge pull request #50 from tlycken/more-extraps

RFC: More extrapolation behaviors

Author: Tomas Lycken

src/Interpolations.jl

@@ -67,8 +67,14 @@ lbound{T,N}(itp::AbstractInterpolation{T,N}, d) = convert(T, 1)
 ubound{T,N}(itp::AbstractInterpolation{T,N}, d) = convert(T, size(itp, d))
 itptype{T,N,IT,GT}(itp::AbstractInterpolation{T,N,IT,GT}) = IT
 gridtype{T,N,IT,GT}(itp::AbstractInterpolation{T,N,IT,GT}) = GT
-
-@inline gradient{T,N}(itp::AbstractInterpolation{T,N}, xs...) = gradient!(Array(T,N), itp, xs...)
+count_interp_dims{T<:AbstractInterpolation}(::Type{T}, N) = N
+
+@generated function gradient{T,N}(itp::AbstractInterpolation{T,N}, xs...)
+    n = count_interp_dims(itp, N)
+    Tg = promote_type(T, [x <: AbstractArray ? eltype(x) : x for x in xs]...)
+    xargs = [:(xs[$d]) for d in 1:length(xs)]
+    :(gradient!(Array($Tg,$n), itp, $(xargs...)))
+end
 
 include("nointerp/nointerp.jl")
 include("b-splines/b-splines.jl")

src/b-splines/b-splines.jl

@@ -41,6 +41,8 @@ ubound{T,N,TCoefs,IT}(itp::BSplineInterpolation{T,N,TCoefs,IT,OnGrid}, d) = conv
 lbound{T,N,TCoefs,IT}(itp::BSplineInterpolation{T,N,TCoefs,IT,OnCell}, d) = convert(T, .5)
 ubound{T,N,TCoefs,IT}(itp::BSplineInterpolation{T,N,TCoefs,IT,OnCell}, d) = convert(T, size(itp, d) + .5)
 
+count_interp_dims{T,N,TCoefs,IT<:DimSpec{InterpolationType},GT<:DimSpec{GridType},pad}(::Type{BSplineInterpolation{T,N,TCoefs,IT,GT,pad}}, n) = count_interp_dims(IT, n)
+
 @generated function size{T,N,TCoefs,IT,GT,pad}(itp::BSplineInterpolation{T,N,TCoefs,IT,GT,pad}, d)
     quote
         d <= $N ? size(itp.coefs, d) - 2*padextract($pad, d) : 1
@@ -71,23 +73,6 @@ function interpolate!{IT<:DimSpec{BSpline},GT<:DimSpec{GridType}}(A::AbstractArr
     interpolate!(tweight(A), A, IT, GT)
 end
 
-define_indices{IT}(::Type{IT}, N, pad) = Expr(:block, Expr[define_indices_d(iextract(IT, d), d, padextract(pad, d)) for d = 1:N]...)
-
-coefficients{IT}(::Type{IT}, N) = Expr(:block, Expr[coefficients(iextract(IT, d), N, d) for d = 1:N]...)
-
-function gradient_coefficients{IT<:DimSpec{BSpline}}(::Type{IT}, N, dim)
-    exs = Expr[d==dim ? gradient_coefficients(iextract(IT, dim), d) :
-                        coefficients(iextract(IT, d), N, d) for d = 1:N]
-    Expr(:block, exs...)
-end
-
-index_gen{IT}(::Type{IT}, N::Integer, offsets...) = index_gen(iextract(IT, min(length(offsets)+1, N)), IT, N, offsets...)
-
-@generated function gradient{T,N,TCoefs,IT,GT,pad}(itp::BSplineInterpolation{T,N,TCoefs,IT,GT,pad}, xs...)
-    n = count_interp_dims(IT, N)
-    :(gradient!(Array(T,$n), itp, xs...))
-end
-
 include("constant.jl")
 include("linear.jl")
 include("quadratic.jl")

src/b-splines/indexing.jl

@@ -4,6 +4,18 @@ import Base.getindex
 
 Base.linearindexing{T<:AbstractInterpolation}(::Type{T}) = Base.LinearSlow()
 
+define_indices{IT}(::Type{IT}, N, pad) = Expr(:block, Expr[define_indices_d(iextract(IT, d), d, padextract(pad, d)) for d = 1:N]...)
+
+coefficients{IT}(::Type{IT}, N) = Expr(:block, Expr[coefficients(iextract(IT, d), N, d) for d = 1:N]...)
+
+function gradient_coefficients{IT<:DimSpec{BSpline}}(::Type{IT}, N, dim)
+    exs = Expr[d==dim ? gradient_coefficients(iextract(IT, dim), d) :
+                        coefficients(iextract(IT, d), N, d) for d = 1:N]
+    Expr(:block, exs...)
+end
+
+index_gen{IT}(::Type{IT}, N::Integer, offsets...) = index_gen(iextract(IT, min(length(offsets)+1, N)), IT, N, offsets...)
+
 function getindex_impl{T,N,TCoefs,IT<:DimSpec{BSpline},GT<:DimSpec{GridType},Pad}(itp::Type{BSplineInterpolation{T,N,TCoefs,IT,GT,Pad}})
     meta = Expr(:meta, :inline)
     quote
@@ -57,7 +69,6 @@ function gradient_impl{T,N,TCoefs,IT<:DimSpec{BSpline},GT<:DimSpec{GridType},Pad
     end
 end
 
-
 @generated function gradient!{T,N}(g::AbstractVector, itp::BSplineInterpolation{T,N}, xs::Number...)
     length(xs) == N || error("Can only be called with $N indexes")
     gradient_impl(itp)

src/extrapolation/constant.jl

@@ -0,0 +1,87 @@
+"""
+The `extrap_prep` function is used by `getindex_impl` to generate the body of
+the `getindex` function for extrapolation objects.
+
+The methods of `extrap_prep` work in "layers", iteratively working out the exact
+expression needed.
+
+The first layer takes a specification of the extrapolation scheme(s) to be used
+and a `Val` object that specifies the dimensionality of the extrapolation object:
+`extrap_prep{T,N}(::Type{T}, Val{N})`. These methods only dispatch to the second
+layer, and need not be extended for new schemes.
+
+The second layer also takes a `Val` object that specifies a single dimension on
+which to work: `extrap_prep{T,N,d}(::Type{T}, ::Val{N}, ::Val{d}). The methods
+with this signature in src/extrapolation/extrap_prep.jl simply expand into a
+block with sub-expressions for handling too-low and too-high values separately
+(the third layer), but specific interpolation schemes can provide more specific
+methods for this layer that handle both ends simultaneously. For example, the
+`Flat` scheme has a layer-2 method that uses `clamp` to restrict the coordinate
+when used in both directions, but uses `min` and `max` when handling each end
+separately.
+
+The third layer, to which the second dispatches if no scheme-specific method is
+found, adds a final `Val` object with a symbol `:lo` or `:hi`:
+`extrap_prep{T,N,d,l}(::Type{T}, ::Val{N}, ::Val{d}, ::Val{l})`. These methods
+must be specified for each extrapolation scheme. However, the general framework
+takes care of expanding all possible tuple combinations, so individual schemes
+need only care about e.g. `T==Flat`.
+
+In addition to these methods, there is a similar three-layer method hierarchy
+for gradient evaluation, in which a `Val{:gradient}` is prepended to the other
+arguments:
+`extrap_prep{T,N,d,l}(::Val{:gradient}`, ::Type{T}, ::Val{N}, ::Val{d}, ::Val{l})`
+If nothing else is specified for the individual schemes, these methods forward
+to the same methods without the `:gradient` argument, i.e. the same behavior as
+for value extrapolation. This works well with all schemes that are simple
+coordinate transformations, but for anything else methods for the low- and high-
+value cases need to be implemented for each scheme.
+""" extrap_prep
+
+extrap_prep{T}(::Type{T}, n::Val{1}) = extrap_prep(T, n, Val{1}())
+extrap_prep{T}(::Type{Tuple{T}}, n::Val{1}) = extrap_prep(T, n)
+extrap_prep{T}(::Type{Tuple{T,T}}, n::Val{1}) = extrap_prep(T, n)
+extrap_prep{T}(::Type{Tuple{Tuple{T,T}}}, n::Val{1}) = extrap_prep(T, n)
+function extrap_prep{S,T}(::Type{Tuple{S,T}}, n::Val{1})
+    quote
+        $(extrap_prep(S, n, Val{1}(), Val{:lo}()))
+        $(extrap_prep(T, n, Val{1}(), Val{:hi}()))
+    end
+end
+extrap_prep{S,T}(::Type{Tuple{Tuple{S,T}}}, n::Val{1}) = extrap_prep(Tuple{S,T}, n)
+
+# needed for ambiguity resolution
+extrap_prep{T<:Tuple}(::Type{T}, ::Val{1}) = :(throw(ArgumentError("The 1-dimensional extrap configuration $T is not supported")))
+
+function extrap_prep{T,N}(::Type{T}, n::Val{N})
+    exprs = Expr[]
+    for d in 1:N
+        push!(exprs, extrap_prep(T, n, Val{d}()))
+    end
+    return Expr(:block, exprs...)
+end
+function extrap_prep{N,T<:Tuple}(::Type{T}, n::Val{N})
+    length(T.parameters) == N || return :(throw(ArgumentError("The $N-dimensional extrap configuration $T is not supported - must be a tuple of length $N (was length $(lenght(T.parameters)))")))
+    exprs = Expr[]
+    for d in 1:N
+        Tdim = T.parameters[d]
+        if Tdim <: Tuple
+            length(Tdim.parameters) == 2 || return :(throw(ArgumentError("The extrap configuration $Tdim for dimension $d is not supported - must be a tuple of length 2 or a simple configuration type")))
+            if Tdim.parameters[1] != Tdim.parameters[2]
+                push!(exprs, extrap_prep(Tdim, n, Val{d}()))
+            else
+                push!(exprs, extrap_prep(Tdim.parameters[1], n, Val{d}()))
+            end
+        else
+            push!(exprs, extrap_prep(Tdim, n, Val{d}()))
+        end
+    end
+    return Expr(:block, exprs...)
+end
+extrap_prep{T,N,d}(::Type{T}, n::Val{N}, dim::Val{d}) = extrap_prep(Tuple{T,T}, n, dim)
+function extrap_prep{S,T,N,d}(::Type{Tuple{S,T}}, n::Val{N}, dim::Val{d})
+    quote
+        $(extrap_prep(S, n, dim, Val{:lo}()))
+        $(extrap_prep(T, n, dim, Val{:hi}()))
+    end
+end

src/extrapolation/error.jl

@@ -0,0 +1,50 @@
+# See ?extrap_prep for documentation for all these methods
+
+extrap_prep{T}(g::Val{:gradient}, ::Type{T}, n::Val{1}) = extrap_prep(g, T, n, Val{1}())
+extrap_prep{T}(g::Val{:gradient}, ::Type{Tuple{T}}, n::Val{1}) = extrap_prep(g, T, n)
+extrap_prep{T}(g::Val{:gradient}, ::Type{Tuple{T,T}}, n::Val{1}) = extrap_prep(g, T, n)
+extrap_prep{T}(g::Val{:gradient}, ::Type{Tuple{Tuple{T,T}}}, n::Val{1}) = extrap_prep(g, T, n)
+function extrap_prep{S,T}(g::Val{:gradient}, ::Type{Tuple{S,T}}, ::Val{1})
+    quote
+        $(extrap_prep(g, S, n, Val{1}(), Val{:lo}()))
+        $(extrap_prep(g, T, n, Val{1}(), Val{:hi}()))
+    end
+end
+extrap_prep{S,T}(g::Val{:gradient}, ::Type{Tuple{Tuple{S,T}}}, n::Val{1}) = extrap_prep(g, Tuple{S,T}, n)
+# needed for ambiguity resolution
+extrap_prep{T<:Tuple}(::Val{:gradient}, ::Type{T}, ::Val{1}) = :(throw(ArgumentError("The 1-dimensional extrap configuration $T is not supported")))
+
+
+function extrap_prep{T,N}(g::Val{:gradient}, ::Type{T}, n::Val{N})
+    Expr(:block, [extrap_prep(g, T, n, Val{d}()) for d in 1:N]...)
+end
+
+function extrap_prep{T<:Tuple,N}(g::Val{:gradient}, ::Type{T}, n::Val{N})
+    length(T.parameters) == N || return :(throw(ArgumentError("The $N-dimensional extrap configuration $T is not supported")))
+    exprs = Expr[]
+    for d in 1:N
+        Tdim = T.parameters[d]
+        if Tdim <: Tuple
+            length(Tdim.parameters) == 2 || return :(throw(ArgumentError("The extrap configuration $Tdim for dimension $d is not supported - must be a tuple of length 2 or a simple configuration type"))
+                )
+            if Tdim.parameters[1] != Tdim.parameters[2]
+                push!(exprs, extrap_prep(g, Tdim, n, Val{d}()))
+            else
+                push!(exprs, extrap_prep(g, Tdim.parameters[1], n, Val{d}()))
+            end
+        else
+            push!(exprs, extrap_prep(g, Tdim, n, Val{d}()))
+        end
+    end
+    return Expr(:block, exprs...)
+end
+
+function extrap_prep{S,T,N,d}(g::Val{:gradient}, ::Type{Tuple{S,T}}, n::Val{N}, dim::Val{d})
+    quote
+        $(extrap_prep(g, S, n, dim, Val{:lo}()))
+        $(extrap_prep(g, T, n, dim, Val{:hi}()))
+    end
+end
+
+extrap_prep{T,N,d}(g::Val{:gradient}, ::Type{T}, n::Val{N}, dim::Val{d}) = extrap_prep(g, Tuple{T,T}, n, dim)
+extrap_prep(g::Val{:gradient}, args...) = extrap_prep(args...)

src/extrapolation/extrap_prep.jl

@@ -1,10 +1,77 @@
 export Throw,
        FilledExtrapolation   # for direct control over typeof(fillvalue)
 
-include("error.jl")
+type Extrapolation{T,N,ITPT,IT,GT,ET} <: AbstractInterpolationWrapper{T,N,ITPT,IT,GT}
+    itp::ITPT
+end
+Extrapolation{T,ITPT,IT,GT,ET}(::Type{T}, N, itp::ITPT, ::Type{IT}, ::Type{GT}, ::Type{ET}) =
+    Extrapolation{T,N,ITPT,IT,GT,ET}(itp)
 
-include("constant.jl")
+"""
+`extrapolate(itp, scheme)` adds extrapolation behavior to an interpolation object, according to the provided scheme.
 
-include("indexing.jl")
+The scheme can take any of these values:
+
+* `Throw` - throws a BoundsError for out-of-bounds indices
+* `Flat` - for constant extrapolation, taking the closest in-bounds value
+* `Linear` - linear extrapolation (the wrapped interpolation object must support gradient)
+* `Reflect` - reflecting extrapolation
+* `Periodic` - periodic extrapolation
+
+You can also combine schemes in tuples. For example, the scheme `Tuple{Linear, Flat}` will use linear extrapolation in the first dimension, and constant in the second.
+
+Finally, you can specify different extrapolation behavior in different direction. `Tuple{Tuple{Linear,Flat}, Flat}` will extrapolate linearly in the first dimension if the index is too small, but use constant etrapolation if it is too large, and always use constant extrapolation in the second dimension.
+"""
+extrapolate{T,N,IT,GT,ET}(itp::AbstractInterpolation{T,N,IT,GT}, ::Type{ET}) =
+    Extrapolation(T,N,itp,IT,GT,ET)
+
+include("throw.jl")
+include("flat.jl")
+include("linear.jl")
+include("reflect.jl")
+include("periodic.jl")
+
+include("extrap_prep.jl")
+include("extrap_prep_gradient.jl")
+
+"""
+`getindex_impl(::Type{E<:Extrapolation}, xs...)`
+
+Generates an expression to be used
+as the function body of the getindex method for the given type of extrapolation
+and indices. The heavy lifting is done by the `extrap_prep` function; see
+`?extrap_prep` for details.
+"""
+function getindex_impl{T,N,ITPT,IT,GT,ET}(etp::Type{Extrapolation{T,N,ITPT,IT,GT,ET}}, xs...)
+    coords = [symbol("xs_",d) for d in 1:N]
+    quote
+        @nexprs $N d->(xs_d = xs[d])
+        $(extrap_prep(ET, Val{N}()))
+        etp.itp[$(coords...)]
+    end
+end
+
+@generated function getindex{T,N,ITPT,IT,GT,ET}(etp::Extrapolation{T,N,ITPT,IT,GT,ET}, xs...)
+    getindex_impl(etp, xs...)
+end
+
+
+function gradient!_impl{T,N,ITPT,IT,GT,ET}(g, etp::Type{Extrapolation{T,N,ITPT,IT,GT,ET}}, xs...)
+    coords = [symbol("xs_", d) for d in 1:N]
+    quote
+        @nexprs $N d->(xs_d = xs[d])
+        $(extrap_prep(Val{:gradient}(), ET, Val{N}()))
+        gradient!(g, etp.itp, $(coords...))
+    end
+end
+
+
+@generated function gradient!{T,N,ITPT,IT,GT,ET}(g::AbstractVector, etp::Extrapolation{T,N,ITPT,IT,GT,ET}, xs...)
+    gradient!_impl(g, etp, xs...)
+end
+
+lbound(etp::Extrapolation, d) = lbound(etp.itp, d)
+ubound(etp::Extrapolation, d) = ubound(etp.itp, d)
+size(etp::Extrapolation, d) = size(etp.itp, d)
 
 include("filled.jl")

src/extrapolation/extrap_prep_gradient.jl

@@ -0,0 +1,42 @@
+function extrap_prep{N,d}(::Type{Flat}, ::Val{N}, ::Val{d})
+    xs_d = symbol("xs_", d)
+    :($xs_d = clamp($xs_d, lbound(etp, $d), ubound(etp, $d)))
+end
+
+function extrap_prep{N,d}(::Type{Flat}, ::Val{N}, ::Val{d}, ::Val{:lo})
+    xs_d = symbol("xs_", d)
+    :($xs_d = max($xs_d, lbound(etp, $d)))
+end
+
+function extrap_prep{N,d}(::Type{Flat}, ::Val{N}, ::Val{d}, ::Val{:hi})
+    xs_d = symbol("xs_", d)
+    :($xs_d = min($xs_d, ubound(etp, $d)))
+end
+
+function extrap_prep{N,d}(::Val{:gradient}, ::Type{Flat}, ::Val{N}, ::Val{d}, ::Val{:lo})
+    coords = [symbol("xs_", k) for k in 1:N]
+    xs_d = coords[d]
+
+    quote
+        if $xs_d < lbound(etp, $d)
+            $xs_d = lbound(etp, $d)
+            gradient!(g, etp.itp, $(coords...))
+            g[$d] = 0
+            return g
+        end
+    end
+end
+
+function extrap_prep{N,d}(::Val{:gradient}, ::Type{Flat}, ::Val{N}, ::Val{d}, ::Val{:hi})
+    coords = [symbol("xs_", k) for k in 1:N]
+    xs_d = coords[d]
+
+    quote
+        if $xs_d > ubound(etp, $d)
+            $xs_d = ubound(etp, $d)
+            gradient!(g, etp.itp, $(coords...))
+            g[$d] = 0
+            return g
+        end
+    end
+end

src/extrapolation/extrapolation.jl

@@ -1,13 +0,0 @@
-@generated function getindex{T}(etp::AbstractExtrapolation{T,1}, x)
-    quote
-        $(extrap_prep(etp, x))
-        etp.itp[x]
-    end
-end
-
-@generated function getindex{T,N,ITP,GT}(etp::AbstractExtrapolation{T,N,ITP,GT}, xs...)
-    quote
-        $(extrap_prep(etp, xs...))
-        etp.itp[xs...]
-    end
-end