Skip to content

Commit

Permalink
Add GriddedInterpolation for non-uniform grids
Browse files Browse the repository at this point in the history
  • Loading branch information
timholy committed Aug 7, 2015
1 parent 28a08ca commit fad451a
Show file tree
Hide file tree
Showing 8 changed files with 226 additions and 15 deletions.
1 change: 1 addition & 0 deletions src/Interpolations.jl
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@ 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...)

include("b-splines/b-splines.jl")
include("gridded/gridded.jl")
include("extrapolation/extrapolation.jl")

end
25 changes: 16 additions & 9 deletions src/b-splines/b-splines.jl
Original file line number Diff line number Diff line change
Expand Up @@ -46,18 +46,25 @@ function interpolate{TWeights,TCoefs,IT<:DimSpec{BSpline},GT<:DimSpec{GridType}}
Apad, Pad = prefilter(TWeights, TCoefs, A, IT, GT)
BSplineInterpolation(TWeights, Apad, IT, GT, Pad)
end
interpolate{IT<:DimSpec{BSpline},GT<:DimSpec{GridType}}(A::AbstractArray, ::Type{IT}, ::Type{GT}) = interpolate(Float64, eltype(A), A, IT, GT)
interpolate{IT<:DimSpec{BSpline},GT<:DimSpec{GridType}}(A::AbstractArray{Float32}, ::Type{IT}, ::Type{GT}) = interpolate(Float32, Float32, A, IT, GT)
interpolate{IT<:DimSpec{BSpline},GT<:DimSpec{GridType}}(A::AbstractArray{Rational{Int}}, ::Type{IT}, ::Type{GT}) = interpolate(Rational{Int}, Rational{Int}, A, IT, GT)
function interpolate{T<:Integer,IT<:DimSpec{BSpline},GT<:DimSpec{GridType}}(A::AbstractArray{T}, ::Type{IT}, ::Type{GT})
TFloat = typeof(float(zero(T)))
interpolate(TFloat, TFloat, A, IT, GT)
function interpolate{IT<:DimSpec{BSpline},GT<:DimSpec{GridType}}(A::AbstractArray, ::Type{IT}, ::Type{GT})
interpolate(tweight(A), tcoef(A), A, IT, GT)
end

# We can't just return a tuple-of-types due to julia #12500
tweight(A::AbstractArray) = Float64
tweight(A::AbstractArray{Float32}) = Float32
tweight(A::AbstractArray{Rational{Int}}) = Rational{Int}
tweight{T<:Integer}(A::AbstractArray{T}) = typeof(float(zero(T)))

tcoef(A::AbstractArray) = eltype(A)
tcoef(A::AbstractArray{Float32}) = Float32
tcoef(A::AbstractArray{Rational{Int}}) = Rational{Int}
tcoef{T<:Integer}(A::AbstractArray{T}) = typeof(float(zero(T)))

interpolate!{TWeights,IT<:DimSpec{BSpline},GT<:DimSpec{GridType}}(::Type{TWeights}, A, ::Type{IT}, ::Type{GT}) = BSplineInterpolation(TWeights, prefilter!(TWeights, A, IT, GT), IT, GT, Val{0}())
interpolate!{IT<:DimSpec{BSpline},GT<:DimSpec{GridType}}(A::AbstractArray, ::Type{IT}, ::Type{GT}) = interpolate!(Float64, A, IT, GT)
interpolate!{IT<:DimSpec{BSpline},GT<:DimSpec{GridType}}(A::AbstractArray{Float32}, ::Type{IT}, ::Type{GT}) = interpolate!(Float32, A, IT, GT)
interpolate!{IT<:DimSpec{BSpline},GT<:DimSpec{GridType}}(A::AbstractArray{Rational{Int}}, ::Type{IT}, ::Type{GT}) = interpolate!(Rational{Int}, A, IT, GT)
function interpolate!{IT<:DimSpec{BSpline},GT<:DimSpec{GridType}}(A::AbstractArray, ::Type{IT}, ::Type{GT})
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]...)

Expand Down
9 changes: 3 additions & 6 deletions src/b-splines/indexing.jl
Original file line number Diff line number Diff line change
Expand Up @@ -23,14 +23,10 @@ function getindex_impl{T,N,TCoefs,IT<:DimSpec{BSpline},GT<:DimSpec{GridType},Pad
end
end

@generated function getindex{T,N}(itp::BSplineInterpolation{T,N}, xs...)
@generated function getindex{T,N}(itp::BSplineInterpolation{T,N}, xs::Number...)
getindex_impl(itp)
end

@generated function getindex{T,N}(itp::BSplineInterpolation{T,N}, index::CartesianIndex{N})
:(getindex(itp, $(Base.IteratorsMD.cartindex_exprs((index,), (:index,))...)))
end

function gradient_impl{T,N,TCoefs,IT<:DimSpec{BSpline},GT<:DimSpec{GridType},Pad}(itp::Type{BSplineInterpolation{T,N,TCoefs,IT,GT,Pad}})
meta = Expr(:meta, :inline)
# For each component of the gradient, alternately calculate
Expand Down Expand Up @@ -68,7 +64,8 @@ end
end

@generated function gradient!{T,N}(g::AbstractVector, itp::BSplineInterpolation{T,N}, index::CartesianIndex{N})
:(gradient!(g, itp, $(Base.IteratorsMD.cartindex_exprs((index,), (:index,))...)))
args = [:(index[$d]) for d = 1:N]
:(gradient!(g, itp, $(args...)))
end

offsetsym(off, d) = off == -1 ? symbol(string("ixm_", d)) :
Expand Down
55 changes: 55 additions & 0 deletions src/gridded/gridded.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
export Gridded

immutable Gridded{D<:Degree} <: InterpolationType end
Gridded{D<:Degree}(::Type{D}) = Gridded{D}

griddedtype{D<:Degree}(::Type{Gridded{D}}) = D

typealias GridIndex{T} Union(AbstractVector{T}, Tuple)

immutable GriddedInterpolation{T,N,TCoefs,IT<:DimSpec{Gridded},K<:Tuple{Vararg{GridIndex}},pad} <: AbstractInterpolation{T,N,IT,OnGrid}
knots::K
coefs::Array{TCoefs,N}
end
function GriddedInterpolation{N,TCoefs,TWeights<:Real,IT<:DimSpec{Gridded},pad}(::Type{TWeights}, knots::NTuple{N,GridIndex}, A::AbstractArray{TCoefs,N}, ::Type{IT}, ::Val{pad})
isleaftype(IT) || error("The b-spline type must be a leaf type (was $IT)")
isleaftype(TCoefs) || warn("For performance reasons, consider using an array of a concrete type (eltype(A) == $(eltype(A)))")

for (d,k) in enumerate(knots)
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")
end
c = one(TWeights)
for _ in 2:N
c *= c
end
T = typeof(c * one(TCoefs))

GriddedInterpolation{T,N,TCoefs,IT,typeof(knots),pad}(knots, A)
end

# Utilities for working either with scalars or tuples/tuple-types
iextract{T<:Gridded}(::Type{T}, d) = T
iextract{T<:GridType}(::Type{T}, d) = T

@generated function size{T,N,TCoefs,IT,K,pad}(itp::GriddedInterpolation{T,N,TCoefs,IT,K,pad}, d)
quote
d <= $N ? size(itp.coefs, d) - 2*padextract($pad, d) : 1
end
end

function interpolate{TWeights,TCoefs,Tel,N,IT<:DimSpec{Gridded}}(::Type{TWeights}, ::Type{TCoefs}, knots::NTuple{N,GridIndex}, A::AbstractArray{Tel,N}, ::Type{IT})
GriddedInterpolation(TWeights, knots, A, IT, Val{0}())
end
function interpolate{Tel,N,IT<:DimSpec{Gridded}}(knots::NTuple{N,GridIndex}, A::AbstractArray{Tel,N}, ::Type{IT})
interpolate(tweight(A), tcoef(A), knots, A, IT)
end

interpolate!{TWeights,Tel,N,IT<:DimSpec{Gridded}}(::Type{TWeights}, knots::NTuple{N,GridIndex}, A::AbstractArray{Tel,N}, ::Type{IT}) = GriddedInterpolation(TWeights, knots, A, IT, Val{0}())
function interpolate!{Tel,N,IT<:DimSpec{Gridded}}(knots::NTuple{N,GridIndex}, A::AbstractArray{Tel,N}, ::Type{IT})
interpolate!(tweight(A), tcoef(A), knots, A, IT)
end

include("linear.jl")
include("indexing.jl")
86 changes: 86 additions & 0 deletions src/gridded/indexing.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,86 @@
using Base.Cartesian

import Base.getindex

# Indexing at a point
@generated function getindex{T,N,TCoefs,IT<:DimSpec{Gridded},K,P}(itp::GriddedInterpolation{T,N,TCoefs,IT,K,P}, x::Number...)
length(x) == N || error("Can only be called with $N indexes")
meta = Expr(:meta, :inline)
quote
$meta
@nexprs $N d->begin
x_d = x[d]
k_d = itp.knots[d]
ix_d = searchsortedfirst(k_d, x_d, 1, length(k_d), Base.Order.ForwardOrdering())
end
$(define_indices(IT, N, P))
$(coefficients(IT, N))
@inbounds ret = $(index_gen(IT, N))
ret
end
end

# Because of the "vectorized" definition below, we need a definition for CartesianIndex
@generated function getindex{T,N}(itp::GriddedInterpolation{T,N}, index::CartesianIndex{N})
args = [:(index[$d]) for d = 1:N]
:(getindex(itp, $(args...)))
end

# Indexing with vector inputs. Here, it pays to pre-process the input indexes,
# because N*n is much smaller than n^N.
# TODO: special-case N=1, because there is no reason to separately cache the indexes.
@generated function getindex!{T,N,TCoefs,IT<:DimSpec{Gridded},K,P}(dest, itp::GriddedInterpolation{T,N,TCoefs,IT,K,P}, xv...)
length(xv) == N || error("Can only be called with $N indexes")
indexes_exprs = Expr[define_indices_d(iextract(IT, d), d, P) for d = 1:N]
coefficient_exprs = Expr[coefficients(iextract(IT, d), N, d) for d = 1:N]
# A manual @nloops (the interaction of d with the two exprs above is tricky...)
ex = :(@inbounds @nref($N,dest,i) = $(index_gen(IT, N)))
for d = 1:N
isym, xsym, xvsym, ixsym, ixvsym = symbol("i_",d), symbol("x_",d), symbol("xv_",d), symbol("ix_",d), symbol("ixv_",d)
ex = quote
for $isym = 1:length($xvsym)
$xsym = $xvsym[$isym]
$ixsym = $ixvsym[$isym]
$(indexes_exprs[d])
$(coefficient_exprs[d])
$ex
end
end
end
quote
@nexprs $N d->begin
xv_d = xv[d]
k_d = itp.knots[d]
ixv_d = Array(Int, length(xv_d)) # ixv_d[i] is the smallest value such that k_d[ixv_d[i]] <= x_d[i]
# If x_d is sorted and has quite a few entries, it's better to match
# entries of x_d and k_d by iterating through them both in unison.
l_d = length(k_d) # FIXME: check l_d == 1 someday, see FIXME above
# estimate the time required for searchsortedfirst vs. linear traversal
den = 5*log(l_d) - 1 # 5 is arbitrary, for now (it's the coefficient of ssf compared to the while loop below)
ascending = den*length(xv_d) > l_d # if this is (or becomes) false, use searchsortedfirst
i = 2 # this clamps ixv_d .>= 1
knext = k_d[i]
xjold = xv_d[1]
for j = 1:length(xv_d)
xj = xv_d[j]
ascending = ascending & (xj >= xjold)
if ascending
while i < length(k_d) && knext < xj
knext = k_d[i+=1]
end
ixv_d[j] = i-1
xjold = xj
else
ixv_d[j] = searchsortedfirst(k_d, xj, 1, l_d, Base.Order.ForwardOrdering())
end
end
end
$ex
dest
end
end

function getindex{T,N,TCoefs,IT<:DimSpec{Gridded},K,P}(itp::GriddedInterpolation{T,N,TCoefs,IT,K,P}, x...)
dest = Array(T, map(length, x))::Array{T,N}
getindex!(dest, itp, x...)
end
33 changes: 33 additions & 0 deletions src/gridded/linear.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
function define_indices_d(::Type{Gridded{Linear}}, d, pad)
symix, symixp, symx = symbol("ix_",d), symbol("ixp_",d), symbol("x_",d)
quote
$symix = clamp($symix, 1, size(itp, $d)-1)
$symixp = $symix + 1
end
end

function coefficients(::Type{Gridded{Linear}}, N, d)
symix, symixp, symx = symbol("ix_",d), symbol("ixp_",d), symbol("x_",d)
sym, symp, symfx = symbol(string("c_",d)), symbol(string("cp_",d)), symbol(string("fx_",d))
symk, symkix = symbol("k_",d), symbol("kix_",d)
quote
$symkix = $symk[$symix]
$symfx = ($symx - $symkix)/($symk[$symixp] - $symkix)
$sym = 1 - $symfx
$symp = $symfx
end
end

# This assumes fractional values 0 <= fx_d <= 1, integral values ix_d and ixp_d (typically ixp_d = ix_d+1,
#except at boundaries), and an array itp.coefs
function index_gen{IT<:DimSpec{Gridded}}(::Type{Gridded{Linear}}, ::Type{IT}, N::Integer, offsets...)
if length(offsets) < N
d = length(offsets)+1
sym = symbol("c_"*string(d))
symp = symbol("cp_"*string(d))
return :($sym * $(index_gen(IT, N, offsets..., 0)) + $symp * $(index_gen(IT, N, offsets..., 1)))
else
indices = [offsetsym(offsets[d], d) for d = 1:N]
return :(itp.coefs[$(indices...)])
end
end
29 changes: 29 additions & 0 deletions test/gridded/linear.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
module LinearTests

using Interpolations, Base.Test

a = rand(5)
knots = (collect(linspace(1,length(a),length(a))),)
itp = @inferred(interpolate(knots, a, Gridded{Linear}))
@inferred(getindex(itp, 2))
@inferred(getindex(itp, CartesianIndex((2,))))
for i = 2:length(a)-1
@test_approx_eq itp[i] a[i]
@test_approx_eq itp[CartesianIndex((i,))] a[i]
end
@inferred(getindex(itp, knots...))
@test_approx_eq itp[knots...] a

A = rand(6,5)
knots = (collect(linspace(1,size(A,1),size(A,1))),collect(linspace(1,size(A,2),size(A,2))))
itp = @inferred(interpolate(knots, A, Gridded{Linear}))
@inferred(getindex(itp, 2, 2))
@inferred(getindex(itp, CartesianIndex((2,2))))
for j = 2:size(A,2)-1, i = 2:size(A,1)-1
@test_approx_eq itp[i,j] A[i,j]
@test_approx_eq itp[CartesianIndex((i,j))] A[i,j]
end
@test_approx_eq itp[knots...] A
@inferred(getindex(itp, knots...))

end
3 changes: 3 additions & 0 deletions test/runtests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,9 @@ include("extrapolation/runtests.jl")
# # test gradient evaluation
include("gradient.jl")

# gridded interpolation tests
include("gridded/runtests.jl")

# # test interpolation with specific types
# include("typing.jl")

Expand Down

0 comments on commit fad451a

Please sign in to comment.