Introduce and implement WeightedIndex

Author: timholy

src/Interpolations.jl

@@ -127,6 +127,80 @@ count_interp_dims(it::Type{IT}, n) where IT<:Tuple{Vararg{InterpolationType,N}}
     _count_interp_dims(c + count_interp_dims(IT1), args...)
 _count_interp_dims(c) = c
 
+
+"""
+    wi = WeightedIndex(indexes, weights)
+
+Construct a weighted index `wi`, which can be thought of as a generalization of an
+ordinary array index to the context of interpolation.
+For an ordinary vector `a`, `a[i]` extracts the element at index `i`.
+When interpolating, one is typically interested in a range of indexes and the output is
+some weighted combination of array values at these indexes.
+For example, for linear interpolation between `i` and `i+1` we have
+
+    ret = (1-f)*a[i] + f*a[i]
+
+This can be represented `a[wi]`, where
+
+    wi = WeightedIndex(i:i+1, (1-f, f))
+
+i.e.,
+
+    ret = sum(a[indexes] .* weights)
+
+Linear interpolation thus constructs weighted indices using a 2-tuple for `weights` and
+a length-2 `indexes` range.
+Higher-order interpolation would involve more positions and weights (e.g., 3-tuples for
+quadratic interpolation, 4-tuples for cubic).
+
+In multiple dimensions, separable interpolation schemes are implemented in terms
+of multiple weighted indices, accessing `A[wi1, wi2, ...]` where each `wi` is the
+`WeightedIndex` along the corresponding dimension.
+
+For value interpolation, `weights` will typically sum to 1.
+However, for gradient and Hessian computation this will not necessarily be true.
+For example, the gradient of one-dimensional linear interpolation can be represented as
+
+    gwi = WeightedIndex(i:i+1, (-1, 1))
+    g1 = a[gwi]
+
+For a three-dimensional array `A`, one might compute `∂A/∂x₂` (the second component
+of the gradient) as `A[wi1, gwi2, wi3]`, where `wi1` and `wi3` are "value" weights
+and `gwi2` "gradient" weights.
+
+`indexes` may be supplied as a range or as a tuple of the same length as `weights`.
+The latter is applicable, e.g., for periodic boundary conditions.
+"""
+abstract type WeightedIndex{L,W} end
+
+# Type to use when array locations are adjacent. This may offer more opportunities
+# for compiler optimizations (e.g., SIMD).
+struct WeightedAdjIndex{L,W} <: WeightedIndex{L,W}
+    istart::Int
+    weights::NTuple{L,W}
+end
+# Type to use with non-adjacent locations. E.g., periodic boundary conditions.
+struct WeightedArbIndex{L,W} <: WeightedIndex{L,W}
+    indexes::NTuple{L,Int}
+    weights::NTuple{L,W}
+end
+
+function WeightedIndex(indexes::AbstractUnitRange{<:Integer}, weights::NTuple{L,Any}) where L
+    @noinline mismatch(indexes, weights) = throw(ArgumentError("the length of indexes must match weights, got $indexes vs $weights"))
+    length(indexes) == L || mismatch(indexes, weights)
+    WeightedAdjIndex(first(indexes), promote(weights...))
+end
+WeightedIndex(istart::Integer, weights::NTuple{L,Any}) where L =
+    WeightedAdjIndex(istart, promote(weights...))
+WeightedIndex(indexes::NTuple{L,Integer}, weights::NTuple{L,Any}) where L =
+    WeightedArbIndex(indexes, promote(weights...))
+
+weights(wi::WeightedIndex) = wi.weights
+indexes(wi::WeightedAdjIndex) = wi.istart
+indexes(wi::WeightedArbIndex) = wi.indexes
+
+
+
 """
     w = value_weights(degree, δx)
 

src/b-splines/b-splines.jl

@@ -80,12 +80,12 @@ _ubounds(axs, itp) = (ubound(axs[1], getfirst(itp)), _ubounds(Base.tail(axs), ge
 _lbounds(::Tuple{}, itp) = ()
 _ubounds(::Tuple{}, itp) = ()
 
-lbound(ax::AbstractUnitRange, bs::BSpline)   = lbound(ax, degree(bs))
-lbound(ax::AbstractUnitRange, deg::Degree)   = first(ax)
-lbound(ax::AbstractUnitRange, deg::DegreeBC) = lbound(ax, deg, deg.bc.gt)
-ubound(ax::AbstractUnitRange, bs::BSpline)   = ubound(ax, degree(bs))
-ubound(ax::AbstractUnitRange, deg::Degree)   = last(ax)
-ubound(ax::AbstractUnitRange, deg::DegreeBC) = ubound(ax, deg, deg.bc.gt)
+lbound(ax::AbstractRange, bs::BSpline)   = lbound(ax, degree(bs))
+lbound(ax::AbstractRange, deg::Degree)   = first(ax)
+lbound(ax::AbstractRange, deg::DegreeBC) = lbound(ax, deg, deg.bc.gt)
+ubound(ax::AbstractRange, bs::BSpline)   = ubound(ax, degree(bs))
+ubound(ax::AbstractRange, deg::Degree)   = last(ax)
+ubound(ax::AbstractRange, deg::DegreeBC) = ubound(ax, deg, deg.bc.gt)
 
 lbound(ax::AbstractUnitRange, ::DegreeBC, ::OnCell) = first(ax) - 0.5
 ubound(ax::AbstractUnitRange, ::DegreeBC, ::OnCell) = last(ax) + 0.5

src/b-splines/constant.jl

@@ -2,20 +2,18 @@ struct Constant <: Degree{0} end
 
 """
 Constant b-splines are *nearest-neighbor* interpolations, and effectively
-return `A[round(Int,x)]` when interpolating
+return `A[round(Int,x)]` when interpolating.
 """
 Constant
 
-function base_rem(::Constant, bounds, x)
-    xm = roundbounds(x, bounds)
+function positions(::Constant, ax, x)  # discontinuity occurs at half-integer locations
+    xm = roundbounds(x, ax)
     δx = x - xm
     fast_trunc(Int, xm), δx
 end
 
-expand_index(::Constant, xi::Number, ax::AbstractUnitRange, δx) = (xi,)
-
-value_weights(::Constant, δx) = (oneunit(δx),)
-gradient_weights(::Constant, δx) = (zero(δx),)
-hessian_weights(::Constant, δx) = (zero(δx),)
+value_weights(::Constant, δx) = (1,)
+gradient_weights(::Constant, δx) = (0,)
+hessian_weights(::Constant, δx) = (0,)
 
 padded_axis(ax::AbstractUnitRange, ::BSpline{Constant}) = ax

src/b-splines/cubic.jl

@@ -27,44 +27,34 @@ When we derive boundary conditions we will use derivatives `y_0'(x)` and
 """
 Cubic
 
-function base_rem(::Cubic, bounds, x)
-    xf = floorbounds(x, bounds)
-    xf -= ifelse(xf > last(bounds)-1, oneunit(xf), zero(xf))
+function positions(deg::Cubic, ax, x)
+    xf = floorbounds(x, ax)
+    xf -= ifelse(xf > last(ax)-1, oneunit(xf), zero(xf))
     δx = x - xf
-    fast_trunc(Int, xf), δx
+    expand_index(deg, fast_trunc(Int, xf), ax, δx), δx
 end
 
-expand_index(::Cubic{BC}, xi::Number, ax::AbstractUnitRange, δx) where BC = (xi-1, xi, xi+1, xi+2)
+expand_index(::Cubic{BC}, xi::Number, ax::AbstractUnitRange, δx) where BC = xi-1
 expand_index(::Cubic{Periodic{GT}}, xi::Number, ax::AbstractUnitRange, δx) where GT<:GridType =
     (modrange(xi-1, ax), modrange(xi, ax), modrange(xi+1, ax), modrange(xi+2, ax))
 
-# expand_coefs(::Type{BSpline{Cubic{BC}}}, δx) = cvcoefs(δx)
-# expand_coefs(::Type{BSpline{Cubic{BC}}}, dref, d, δx) = ifelse(d==dref, cgcoefs(δx), cvcoefs(δx))
-# function expand_coefs(::Type{BSpline{Cubic{BC}}}, dref1, dref2, d, δx)
-#     if dref1 == dref2
-#         d == dref1 ? chcoefs(δx) : cvcoefs(δx)
-#     else
-#         d == dref1 | d == dref2 ? cgcoefs(δx) : cvcoefs(δx)
-#     end
-# end
-
-function value_weights(::BSpline{<:Cubic}, δx)
+function value_weights(::Cubic, δx)
     x3, xcomp3 = cub(δx), cub(1-δx)
     (SimpleRatio(1,6) * xcomp3,
      SimpleRatio(2,3) - sqr(δx) + SimpleRatio(1,2)*x3,
      SimpleRatio(2,3) - sqr(1-δx) + SimpleRatio(1,2)*xcomp3,
      SimpleRatio(1,6) * x3)
 end
 
-function gradient_weights(::BSpline{<:Cubic}, δx)
+function gradient_weights(::Cubic, δx)
     x2, xcomp2 = sqr(δx), sqr(1-δx)
     (-SimpleRatio(1,2) * xcomp2,
      -2*δx + SimpleRatio(3,2)*x2,
      +2*(1-δx) - SimpleRatio(3,2)*xcomp2,
      SimpleRatio(1,2) * x2)
 end
 
-hessian_weights(::BSpline{<:Cubic}, δx) = (1-δx, 3*δx-2, 3*(1-δx)-2, δx)
+hessian_weights(::Cubic, δx) = (1-δx, 3*δx-2, 3*(1-δx)-2, δx)
 
 
 # ------------ #