diff --git a/Project.toml b/Project.toml index 9090b47..450cc13 100644 --- a/Project.toml +++ b/Project.toml @@ -18,6 +18,7 @@ DataInterpolationsNDSymbolicsExt = ["SymbolicUtils", "Symbolics"] [compat] Adapt = "4.3.0" +Aqua = "0.8" DataInterpolations = "8" EllipsisNotation = "1.8.0" ForwardDiff = "0" @@ -32,6 +33,7 @@ Test = "1" julia = "1" [extras] +Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" @@ -42,4 +44,4 @@ Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["DataInterpolations", "ForwardDiff", "Pkg", "Random", "SafeTestsets", "Symbolics", "Test", "SymbolicUtils"] +test = ["Aqua", "DataInterpolations", "ForwardDiff", "Pkg", "Random", "SafeTestsets", "Symbolics", "Test", "SymbolicUtils"] diff --git a/README.md b/README.md index b8f4b7e..638d816 100644 --- a/README.md +++ b/README.md @@ -14,7 +14,7 @@ DataInterpolationsND.jl is a library for interpolating arbitrarily high dimensio For one dimensional interpolation see also [DataInterpolations.jl](https://github.com/SciML/DataInterpolations.jl). -DataInterpolationsND.jl is similar in functionality to [Interpolations.jl](https://github.com/JuliaMath/Interpolations.jl), a well established interpolation package that is currently much more feature rich than DataInterpolationsND.jl. We hope to justify the existence of DataInterpolationsND.jl through its use of the KernelAbstractions and its planned integration with the SciML ecosystem. +DataInterpolationsND.jl is similar in functionality to [Interpolations.jl](https://github.com/JuliaMath/Interpolations.jl), a well established interpolation package that is currently much more feature rich than DataInterpolationsND.jl. We hope to justify the existence of DataInterpolationsND.jl through its use of KernelAbstractions and its planned integration with the SciML ecosystem. ## API @@ -28,7 +28,7 @@ t2 = cumsum(rand(7)) interpolation_dimensions = ( LinearInterpolationDimension(t1), - LinearInterpolationDimension(t2) + ConstantInterpolationDimension(t2) ) # The outputs will be vectors of length 2 @@ -51,7 +51,7 @@ If we provide `t_eval` for the interpolation dimensions, we can evaluate at thes ```julia interpolation_dimensions = ( LinearInterpolationDimension(t1; t_eval = range(first(t1), last(t1); length = 100)), - LinearInterpolationDimension(t2; t_eval = range(first(t2), last(t2); length = 100)) + ConstantInterpolationDimension(t2; t_eval = range(first(t2), last(t2); length = 100)) ) interp = NDInterpolation(u, interpolation_dimensions) diff --git a/docs/src/index.md b/docs/src/index.md index 98d93a0..ab3b014 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -16,7 +16,7 @@ t2 = cumsum(rand(7)) interpolation_dimensions = ( LinearInterpolationDimension(t1), - LinearInterpolationDimension(t2) + ConstantInterpolationDimension(t2) ) # The outputs will be vectors of length 2 @@ -39,7 +39,7 @@ If we provide `t_eval` for the interpolation dimensions, we can evaluate at thes ```julia interpolation_dimensions = ( LinearInterpolationDimension(t1; t_eval = range(first(t1), last(t1); length = 100)), - LinearInterpolationDimension(t2; t_eval = range(first(t2), last(t2); length = 100)) + ConstantInterpolationDimension(t2; t_eval = range(first(t2), last(t2); length = 100)) ) interp = NDInterpolation(u, interpolation_dimensions) diff --git a/docs/src/usage.md b/docs/src/usage.md index 47c6b13..428d627 100644 --- a/docs/src/usage.md +++ b/docs/src/usage.md @@ -9,7 +9,7 @@ Random.seed!(1) interp_dims = ( LinearInterpolationDimension(cumsum(0.5 .+ rand(5))), - LinearInterpolationDimension(cumsum(0.5 .+ rand(10))) + ConstantInterpolationDimension(cumsum(0.5 .+ rand(10))) ) t_eval_1 = range( @@ -85,7 +85,7 @@ using LinearAlgebra interp_dims = ( LinearInterpolationDimension(interp_dims[1].t; t_eval = t_eval_1), - LinearInterpolationDimension(interp_dims[2].t; t_eval = t_eval_2) + ConstantInterpolationDimension(interp_dims[2].t; t_eval = t_eval_2) ) itp = NDInterpolation(u, interp_dims) diff --git a/src/DataInterpolationsND.jl b/src/DataInterpolationsND.jl index 1d407ef..49cb7f0 100644 --- a/src/DataInterpolationsND.jl +++ b/src/DataInterpolationsND.jl @@ -13,13 +13,13 @@ struct EmptyCache <: AbstractInterpolationCache end """ NDInterpolation(u, interp_dims; cache = EmptyCache()) -The interpolation object containing the interpolation dimensions and the data to interpolate `u`. -Given the number of interpolation dimensions `N_in`, for first `N_in` dimensions of `u` -the size of `u` along that dimension must match the length of `t` of the corresponding interpolation dimension. +The interpolation object containing the interpolation dimensions and the data to be interpolated `u`. +For non-trivial interpolation dimensions (i.e. not `NoInterpolationDimension`) the length of its `t` must match +the size of `u` along the corresponding dimension. ## Arguments - - `interp_dims`: A tuple of identically typed interpolation dimensions. + - `interp_dims`: A tuple of interpolation dimensions of length `≤ ndims(u)`. - `u`: The array to be interpolated. ## Keyword Arguments @@ -27,34 +27,43 @@ the size of `u` along that dimension must match the length of `t` of the corresp - `cache`: Optional global cache """ struct NDInterpolation{ - N_in, N_out, - ID <: AbstractInterpolationDimension, + N, + N_in, + N_out, gType <: AbstractInterpolationCache, + D <: Tuple{Vararg{<:AbstractInterpolationDimension}}, uType <: AbstractArray, } u::uType - interp_dims::NTuple{N_in, ID} + interp_dims::D cache::gType - function NDInterpolation(u, interp_dims, cache) + function NDInterpolation(u::AbstractArray{<:Any, N}, interp_dims, cache) where {N} if interp_dims isa AbstractInterpolationDimension interp_dims = (interp_dims,) end - N_in = length(interp_dims) - N_out = ndims(u) - N_in - @assert N_out ≥ 0 "The number of dimensions of u must be at least the number of interpolation dimensions." + @assert length(interp_dims) <= N "The number of interpolation dimensions must be no more than the dimensions of u. got $(length(interp_dims)) and $N" + interp_dims = _add_trailing_no_interp_dims(interp_dims, Val{N}()) + N_in = length(remove(NoInterpolationDimension, interp_dims)) + N_out = length(keep(NoInterpolationDimension, interp_dims)) validate_size_u(interp_dims, u) validate_cache(cache, interp_dims, u) - return new{N_in, N_out, eltype(interp_dims), typeof(cache), typeof(u)}( + return new{N, N_in, N_out, typeof(cache), typeof(interp_dims), typeof(u)}( u, interp_dims, cache ) end end - # Constructor with optional global cache function NDInterpolation(u, interp_dims; cache = EmptyCache()) return NDInterpolation(u, interp_dims, cache) end +function _add_trailing_no_interp_dims(dim::AbstractInterpolationDimension, n) + return _add_trailing_interp_dims((dim,), n) +end +function _add_trailing_no_interp_dims(dims::Tuple, ::Val{N}) where {N} + return (dims..., ntuple(_ -> NoInterpolationDimension(), Val{N - length(dims)}())...) +end + @adapt_structure NDInterpolation include("interpolation_dimensions.jl") @@ -76,15 +85,22 @@ function (interp::NDInterpolation)( end # In place single input evaluation -function (interp::NDInterpolation{N_in})( - out::Union{Number, AbstractArray{<:Number}}, - t::Tuple{Vararg{Number, N_in}}; - derivative_orders::NTuple{N_in, <:Integer} = ntuple(_ -> 0, N_in) - ) where {N_in} - validate_derivative_orders(derivative_orders, interp) - idx = get_idx(interp.interp_dims, t) - @assert size(out) == size(interp.u)[(N_in + 1):end] "The size of out must match the size of the last N_out dimensions of u." - return _interpolate!(out, interp, t, idx, derivative_orders, nothing) +function (interp::NDInterpolation{N, N_in, N_out})( + out::Union{Number, AbstractArray{<:Number, N_out}}, + t::Tuple{Vararg{Any, N_in}}; + derivative_orders::Tuple{Vararg{Integer, N_in}} = ntuple(_ -> 0, N_in) + ) where {N, N_in, N_out} + (; interp_dims, u) = interp + # We need to add the NoInterpolationDimensions for all arguments to _interpolate. + # Currently interp() accepts only the interpolating dimensions. + # We could switch this so it needs indices for NoInterpolationDimension, but thats a breaking change. + d_o_all = insertat(NoInterpolationDimension, 0, derivative_orders, interp_dims) + validate_derivative_order(d_o_all, interp) + t_all = insertat(NoInterpolationDimension, Colon(), t, interp.interp_dims) + idx_all = get_idx(interp_dims, t_all) + @assert size(out) == keepat(NoInterpolationDimension, size(u), interp_dims) "The size of out must match the size of the NoInterpolationDimensions of u." + multi_point_index = nothing + return _interpolate!(out, interp, t_all, idx_all, d_o_all, multi_point_index) end # Out of place single input evaluation @@ -94,9 +110,7 @@ function (interp::NDInterpolation)(t::Tuple{Vararg{Number}}; kwargs...) end export NDInterpolation, LinearInterpolationDimension, ConstantInterpolationDimension, - BSplineInterpolationDimension, NURBSWeights, + BSplineInterpolationDimension, NURBSWeights, NoInterpolationDimension, eval_unstructured, eval_unstructured!, eval_grid, eval_grid! -include("precompilation.jl") - end # module DataInterpolationsND diff --git a/src/interpolation_dimensions.jl b/src/interpolation_dimensions.jl index 6260e2a..8050151 100644 --- a/src/interpolation_dimensions.jl +++ b/src/interpolation_dimensions.jl @@ -1,3 +1,10 @@ +""" + NoInterpolationDimension() + +A dimension that does not perform interpolation. +""" +struct NoInterpolationDimension <: AbstractInterpolationDimension end + """ LinearInterpolationDimension(t; t_eval = similar(t, 0)) @@ -161,15 +168,9 @@ function BSplineInterpolationDimension( synchronize(backend) idx_eval = similar(t_eval, Int) - basis_function_eval = similar( - t_eval, - typeof(inv(one(eltype(t))) * inv(one(eltype(t_eval)))), - ( - length(t_eval), - degree + 1, - max_derivative_order_eval + 1, - ) - ) + T = typeof(inv(one(eltype(t))) * inv(one(eltype(t_eval)))) + s = (length(t_eval), degree + 1, max_derivative_order_eval + 1) + basis_function_eval = similar(t_eval, T, s) itp_dim = BSplineInterpolationDimension( t, knots_all, t_eval, idx_eval, degree, max_derivative_order_eval, basis_function_eval, multiplicities diff --git a/src/interpolation_methods.jl b/src/interpolation_methods.jl index cff6b5f..9e975d2 100644 --- a/src/interpolation_methods.jl +++ b/src/interpolation_methods.jl @@ -1,141 +1,110 @@ -function _interpolate!( - out, - A::NDInterpolation{N_in, N_out, ID}, - t::Tuple{Vararg{Number, N_in}}, - idx::NTuple{N_in, <:Integer}, - derivative_orders::NTuple{N_in, <:Integer}, +Base.@propagate_inbounds function _interpolate!( + out::Union{Number, AbstractArray}, + A::NDInterpolation{N}, + ts::Tuple{Vararg{Any, N}}, + idx::Tuple{Vararg{Any, N}}, + derivative_orders::Tuple{Vararg{Any, N}}, multi_point_index - ) where {N_in, N_out, ID <: LinearInterpolationDimension} - out = make_zero!!(out) - any(>(1), derivative_orders) && return out + ) where {N} + (; interp_dims, cache, u) = A - tᵢ = ntuple(i -> A.interp_dims[i].t[idx[i]], N_in) - tᵢ₊₁ = ntuple(i -> A.interp_dims[i].t[idx[i] + 1], N_in) - - # Size of the (hyper)rectangle `t` is in - t_vol = one(eltype(tᵢ)) - for (t₁, t₂) in zip(tᵢ, tᵢ₊₁) - t_vol *= t₂ - t₁ + out, + valid_derivative_orders = check_derivative_order( + interp_dims, derivative_orders, ts, out + ) + valid_derivative_orders || return out # Array was zeroed out in this case + if isnothing(multi_point_index) + multi_point_index = map(_ -> nothing, interp_dims) end - # Loop over the corners of the (hyper)rectangle `t` is in - for I in Iterators.product(ntuple(i -> (false, true), N_in)...) - c = eltype(out)(inv(t_vol)) - for (t_, right_point, d, t₁, t₂) in zip(t, I, derivative_orders, tᵢ, tᵢ₊₁) - c *= if right_point - iszero(d) ? t_ - t₁ : one(t_) - else - iszero(d) ? t₂ - t_ : -one(t_) - end - end - J = (ntuple(i -> idx[i] + I[i], N_in)..., ..) - if iszero(N_out) - out += c * A.u[J...] - else - @. out += c * A.u[J...] + # Setup + out = make_zero!!(out) # TODO remove this + stencils = map(stencil, interp_dims) + coeffs = map(coefficients, interp_dims, derivative_orders, multi_point_index, ts, idx) + denom = zero(eltype(u)) + + # TODO this can be a single unrolled broadcast rather than a loop of .+= + for I in Iterators.product(stencils...) + J = map(index, interp_dims, ts, idx, I) + product = prod(map(getindex, coeffs, I)) + + if cache isa NURBSWeights + K = removeat(NoInterpolationDimension, J, interp_dims) + product *= cache.weights[K...] + denom += product end - end - return out -end -function _interpolate!( - out, - A::NDInterpolation{N_in, N_out, ID}, - t::Tuple{Vararg{Number, N_in}}, - idx::NTuple{N_in, <:Integer}, - derivative_orders::NTuple{N_in, <:Integer}, - multi_point_index - ) where {N_in, N_out, ID <: ConstantInterpolationDimension} - if any(>(0), derivative_orders) - return if any(i -> !isempty(searchsorted(A.interp_dims[i].t, t[i])), 1:N_in) - typed_nan(out) + if out isa AbstractArray + out .+= product .* u[J...] else - out + out += product * u[J...] end end - # Use a new variable to avoid shadowing/capturing idx in the ntuple closure - idx_adjusted = ntuple(N_in) do i - t[i] >= A.interp_dims[i].t[end] ? length(A.interp_dims[i].t) : idx[i] - end - if iszero(N_out) - out = A.u[idx_adjusted...] - else - out .= A.u[idx_adjusted...] - end - return out -end - -# BSpline evaluation -function _interpolate!( - out, - A::NDInterpolation{N_in, N_out, ID}, - t::Tuple{Vararg{Number, N_in}}, - idx::NTuple{N_in, <:Integer}, - derivative_orders::NTuple{N_in, <:Integer}, - multi_point_index - ) where {N_in, N_out, ID <: BSplineInterpolationDimension} - (; interp_dims) = A - - out = make_zero!!(out) - degrees = ntuple(dim_in -> interp_dims[dim_in].degree, N_in) - basis_function_vals = get_basis_function_values_all( - A, t, idx, derivative_orders, multi_point_index - ) - for I in CartesianIndices(ntuple(dim_in -> 1:(degrees[dim_in] + 1), N_in)) - B_product = prod(dim_in -> basis_function_vals[dim_in][I[dim_in]], 1:N_in) - cp_index = ntuple( - dim_in -> idx[dim_in] + I[dim_in] - degrees[dim_in] - 1, N_in - ) - if iszero(N_out) - out += B_product * A.u[cp_index...] + if cache isa NURBSWeights + if out isa AbstractArray + out ./= denom else - out .+= B_product * view(A.u, cp_index..., ..) + out /= denom end end return out end -# NURBS evaluation -function _interpolate!( - out, - A::NDInterpolation{N_in, N_out, ID, <:NURBSWeights}, - t::Tuple{Vararg{Number, N_in}}, - idx::NTuple{N_in, <:Integer}, - derivative_orders::NTuple{N_in, <:Integer}, - multi_point_index - ) where {N_in, N_out, ID <: BSplineInterpolationDimension} - (; interp_dims, cache) = A - - out = make_zero!!(out) - degrees = ntuple(dim_in -> interp_dims[dim_in].degree, N_in) - basis_function_vals = get_basis_function_values_all( - A, t, idx, derivative_orders, multi_point_index - ) - - denom = zero(eltype(t)) - - for I in CartesianIndices(ntuple(dim_in -> 1:(degrees[dim_in] + 1), N_in)) - B_product = prod(dim_in -> basis_function_vals[dim_in][I[dim_in]], 1:N_in) - cp_index = ntuple( - dim_in -> idx[dim_in] + I[dim_in] - degrees[dim_in] - 1, N_in - ) - weight = cache.weights[cp_index...] - product = weight * B_product - denom += product - if iszero(N_out) - out += product * A.u[cp_index...] +function check_derivative_order(dims::Tuple, derivative_orders::Tuple, ts::Tuple, out) + itr = map(tuple, dims, derivative_orders, ts) + # Fold over itr for all dims, combining out and valid + return foldl(itr; init = (out, true)) do (acc_out, acc_valid), (d, d_o, t) + dim_out, dim_valid = check_derivative_order(d, d_o, t, acc_out) + dim_out, dim_valid & acc_valid + end +end +check_derivative_order(::AbstractInterpolationDimension, d_o, t, out) = (out, true) +check_derivative_order(::LinearInterpolationDimension, d_o, t, out) = (out, d_o <= 1) +function check_derivative_order(d::ConstantInterpolationDimension, d_o, t, out) + if d_o > 0 + # Check if t is on the boundary between constant steps and if so return nans + return if isempty(searchsorted(d.t, t)) + (out, false) else - out .+= product * view(A.u, cp_index..., ..) + (typed_nan(out), false) end - end - - if iszero(N_out) - out /= denom else - out ./= denom + (out, true) end +end - return out +stencil(::LinearInterpolationDimension) = (1, 2) +stencil(::ConstantInterpolationDimension) = 1 +stencil(::NoInterpolationDimension) = 1 +stencil(d::BSplineInterpolationDimension) = 1:(d.degree + 1) + +# Precalculate coefficients +function coefficients( + d::LinearInterpolationDimension, derivative_order, multi_point_index, t, i + ) + t₁ = d.t[i] + t₂ = d.t[i + 1] + t_vol_inv = inv(t₂ - t₁) + a = (iszero(derivative_order) ? t₂ - t : -one(t)) * t_vol_inv + b = (iszero(derivative_order) ? t - t₁ : one(t)) * t_vol_inv + return (a, b) end +function coefficients( + ::ConstantInterpolationDimension, derivative_order, multi_point_index, t, i + ) + return true +end +coefficients(::NoInterpolationDimension, derivative_order, multi_point_index, t, i) = true +function coefficients( + d::BSplineInterpolationDimension, derivative_order, multi_point_index, t, i + ) + return get_basis_function_values(d, t, i, derivative_order, multi_point_index) +end + +index(::LinearInterpolationDimension, t, idx, i) = idx + i - 1 +# TODO: this should happen outside of the loop +index(d::ConstantInterpolationDimension, t, idx, i) = t >= d.t[end] ? length(d.t) : idx[i] +index(::NoInterpolationDimension, t, idx, i) = idx +index(d::BSplineInterpolationDimension, t, idx, i) = idx + i - d.degree - 1 diff --git a/src/interpolation_parallel.jl b/src/interpolation_parallel.jl index 8571fe1..b58726a 100644 --- a/src/interpolation_parallel.jl +++ b/src/interpolation_parallel.jl @@ -33,18 +33,24 @@ length for each interpolation dimension and the interpolation is evaluated at th """ function eval_unstructured!( out::AbstractArray, - interp::NDInterpolation{N_in}; + interp::NDInterpolation{N, N_in}; derivative_orders::NTuple{N_in, <:Integer} = ntuple(_ -> 0, N_in) - ) where {N_in} - validate_derivative_orders(derivative_orders, interp; multi_point = true) + ) where {N, N_in} + validate_derivative_order(derivative_orders, interp; multi_point = true) backend = get_backend(out) - @assert all(i -> length(interp.interp_dims[i].t_eval) == size(out, 1), N_in) "The t_eval of all interpolation dimensions must have the same length as the first dimension of out." + no_interp_inds = map(_ -> Colon(), keep(NoInterpolationDimension, interp.interp_dims)) + @assert all( + d -> length(d.t_eval) == size(out, 1), + remove(NoInterpolationDimension, interp.interp_dims) + ) "The t_eval of all interpolation dimensions must have the same length as the first dimension of out." @assert size(out)[2:end] == get_output_size(interp) "The size of the last N_out dimensions of out must be the same as the output size of the interpolation." + eval_grid = false eval_kernel(backend)( out, interp, derivative_orders, - false, + no_interp_inds, + eval_grid, ndrange = size(out, 1) ) synchronize(backend) @@ -63,9 +69,9 @@ out of place. - `derivative_orders`: The partial derivative order for each interpolation dimension. Defaults to `0` for each. """ -function eval_grid(interp::NDInterpolation{N_in}; kwargs...) where {N_in} - grid_size = map(itp_dim -> length(itp_dim.t_eval), interp.interp_dims) - out = similar(interp.u, (grid_size..., get_output_size(interp)...)) +function eval_grid(interp::NDInterpolation; kwargs...) + # TODO: do we need to promote the type here, e.g. for eltype(u) <: Integer ? + out = similar(interp.u, grid_size(interp)) return eval_grid!(out, interp; kwargs...) end @@ -84,51 +90,90 @@ in place. """ function eval_grid!( out::AbstractArray, - interp::NDInterpolation{N_in}; - derivative_orders::NTuple{N_in, <:Integer} = ntuple(_ -> 0, N_in) - ) where {N_in} - validate_derivative_orders(derivative_orders, interp; multi_point = true) + interp::NDInterpolation{N, N_in, N_out}; + derivative_orders::NTuple{N_in, <:Integer} = ntuple(_ -> 0, N_in), + no_interp_inds = map(_ -> Colon(), N_out) + ) where {N, N_in, N_out} + validate_t_eval_lengths(out, interp) + validate_output_size(out, interp) + validate_derivative_order(derivative_orders, interp; multi_point = true) backend = get_backend(out) - @assert all(i -> size(out, i) == length(interp.interp_dims[i].t_eval), N_in) "For the first N_in dimensions of out the length must match the t_eval of the corresponding interpolation dimension." - @assert size(out)[(N_in + 1):end] == get_output_size(interp) "The size of the last N_out dimensions of out must be the same as the output size of the interpolation." + eval_grid = true eval_kernel(backend)( out, interp, derivative_orders, - true, - ndrange = size(out)[1:N_in] + no_interp_inds, + eval_grid, + ndrange = get_ndrange(interp) ) synchronize(backend) return out end +function validate_t_eval_lengths(out, interp) + (; interp_dims) = interp + interp_sizes = removeat(NoInterpolationDimension, size(out), interp_dims) + interp_t_eval_lengths = map( + d -> length(d.t_eval), remove(NoInterpolationDimension, interp_dims) + ) + return all(map(==, interp_sizes, interp_t_eval_lengths)) || + throw(ArgumentError("The length must match the t_eval of the corresponding interpolation dimension.")) +end + +function validate_output_size(out, interp) + return keepat(NoInterpolationDimension, size(out), interp.interp_dims) == + get_output_size(interp) || + throw(ArumentError("The size of the NoInterpolationDimension dimensions of `out` must be the same as the output size of the interpolation.")) +end + @kernel function eval_kernel( out, - @Const(A), + A, # @Const(A), TODO: somehow this now hits a bug in KernelAbstractions where elsize is not defined for Const derivative_orders, + no_interp_inds, eval_grid ) - N_in = length(A.interp_dims) - N_out = ndims(A.u) - N_in - - k = @index(Global, NTuple) + (; interp_dims) = A + N_out = length(keep(NoInterpolationDimension, interp_dims)) + I = @index(Global, NTuple) + d_o = insertat(NoInterpolationDimension, 0, derivative_orders, interp_dims) + # TODO eval_grid should be a static bool or Val so this branch can be deleted if eval_grid - t_eval = ntuple(i -> A.interp_dims[i].t_eval[k[i]], N_in) - idx_eval = ntuple(i -> A.interp_dims[i].idx_eval[k[i]], N_in) - else - t_eval = ntuple(i -> A.interp_dims[i].t_eval[only(k)], N_in) - idx_eval = ntuple(i -> A.interp_dims[i].idx_eval[only(k)], N_in) + # Insert no_interp_inds in I + k = insertat(NoInterpolationDimension, no_interp_inds, I, interp_dims) + else # eval_unstructured + # The same index is used for all interp dimensions in eval_unstructured + I1 = map(_ -> only(I), remove(NoInterpolationDimension, interp_dims)) + # Insert no_interp_inds in I1 + k = insertat(NoInterpolationDimension, no_interp_inds, I1, interp_dims) end + t_eval = map(get_t_eval, interp_dims, k) + idx_eval = map(get_idx_eval, interp_dims, k) if iszero(N_out) - out[k...] = _interpolate!( - make_out(A, t_eval), A, t_eval, idx_eval, derivative_orders, k + @inbounds out[k...] = _interpolate!( + make_out(A, t_eval), A, t_eval, idx_eval, d_o, k ) else _interpolate!( - view(out, k..., ..), - A, t_eval, idx_eval, derivative_orders, k + view(out, k...), + A, t_eval, idx_eval, d_o, k ) end end + +get_t_eval(d::AbstractInterpolationDimension, i::Integer) = d.t_eval[i] +get_t_eval(d::AbstractInterpolationDimension, i::Colon) = d.t_eval +get_t_eval(d::AbstractInterpolationDimension, i::AbstractVector) = d.t_eval[i] +get_t_eval(::NoInterpolationDimension, i::Integer) = nothing +get_t_eval(::NoInterpolationDimension, i::Colon) = nothing +get_t_eval(::NoInterpolationDimension, i::AbstractVector) = nothing + +get_idx_eval(d::AbstractInterpolationDimension, i::Integer) = d.idx_eval[i] +get_idx_eval(d::AbstractInterpolationDimension, i::AbstractVector) = d.idx_eval[i] +get_idx_eval(d::AbstractInterpolationDimension, i::Colon) = d.idx_eval +get_idx_eval(::NoInterpolationDimension, i::Integer) = i +get_idx_eval(::NoInterpolationDimension, i::Colon) = i +get_idx_eval(::NoInterpolationDimension, i::AbstractVector) = i diff --git a/src/interpolation_utils.jl b/src/interpolation_utils.jl index b0c2db9..3c313a8 100644 --- a/src/interpolation_utils.jl +++ b/src/interpolation_utils.jl @@ -2,33 +2,45 @@ trivial_range(i::Integer) = i:i Base.length(itp_dim::AbstractInterpolationDimension) = length(itp_dim.t) -function validate_derivative_orders( - derivative_orders::NTuple{N_in, <:Integer}, - ::NDInterpolation{N_in}; - kwargs... - ) where {N_in} - return @assert all(≥(0), derivative_orders) "Derivative orders must me non-negative." -end - -function validate_derivative_orders( - derivative_orders::NTuple{N_in, <:Integer}, - A::NDInterpolation{N_in, N_out, <:BSplineInterpolationDimension}; +function validate_derivative_order( + derivative_orders::NTuple, + A::NDInterpolation; multi_point::Bool = false - ) where {N_in, N_out} - @assert all(≥(0), derivative_orders) "Derivative orders must me non-negative." - + ) + return map(derivative_orders, A.interp_dims) do d_o, d + validate_derivative_order(d_o, d; multi_point, cache = A.cache) + end +end +function validate_derivative_order( + derivative_order::Integer, + interp_dim::BSplineInterpolationDimension; + multi_point::Bool, + cache + ) if multi_point - @assert all( - i -> derivative_orders[i] ≤ A.interp_dims[i].max_derivative_order_eval, 1:N_in - ) "For BSpline interpolation, when using multi-point evaluation the derivative orders cannot be \ - larger than the `max_derivative_order_eval` eval of of the `BSplineInterpolationDimension`. If you want \ - to compute higher order multi-point derivatives, pass a larger `max_derivative_order_eval` to the \ - `BSplineInterpolationDimension` constructor(s)." + @assert derivative_order ≤ interp_dim.max_derivative_order_eval """ + For BSpline interpolation, when using multi-point evaluation the derivative orders cannot be + larger than the `max_derivative_order_eval` eval of of the `BSplineInterpolationDimension`. If you want + to compute higher order multi-point derivatives, pass a larger `max_derivative_order_eval` to the + `BSplineInterpolationDimension` constructor(s). + """ end + return validate_derivative_order_by_cache(cache, derivative_order) +end +function validate_derivative_order( + derivative_order::Integer, + interp_dim::AbstractInterpolationDimension; + multi_point::Bool, + cache + ) + return validate_derivative_order_by_cache(cache, derivative_order) +end - return if A.cache isa NURBSWeights - @assert all(==(0), derivative_orders) "Currently partial derivatives of NURBS are not supported." - end +function validate_derivative_order_by_cache(::NURBSWeights, derivative_order) + return @assert derivative_order == 0 "Currently partial derivatives of NURBS are not supported." +end +function validate_derivative_order_by_cache(::Any, derivative_order) + return @assert derivative_order >= 0 "Derivative orders must be non-negative." end function validate_t(t) @@ -36,75 +48,97 @@ function validate_t(t) return @assert all(>(0), diff(t)) "The elements of t must be sorted and unique." end -function validate_size_u( - interp_dims::NTuple{N_in, <:AbstractInterpolationDimension}, - u::AbstractArray - ) where {N_in} - return @assert ntuple(i -> length(interp_dims[i]), N_in) == size(u)[1:N_in] "For the first N_in dimensions of u the length must match the t of the corresponding interpolation dimension." +validate_size_u(interp::NDInterpolation, u) = validate_size_u(interp.interp_dims, u) +validate_size_u(interp_dims::Tuple, u::Number) = nothing +function validate_size_u(interp_dims::Tuple, u::AbstractArray) + return validate_size_u(interp_dims, axes(u)) +end +validate_size_u(interp_dims::Tuple, ax::Tuple) = map(validate_size_u, interp_dims, ax) +validate_size_u(interp_dim::NoInterpolationDimension, ax::AbstractRange) = nothing +function validate_size_u(interp_dim::AbstractInterpolationDimension, ax::AbstractRange) + return @assert length(interp_dim) == length(ax) "The size if `u` must match the t of the corresponding interpolation dimension. Got $interp_dim and $ax" end function validate_size_u( - interp_dims::NTuple{N_in, <:BSplineInterpolationDimension}, - u::AbstractArray - ) where {N_in} - expected_size = ntuple(dim_in -> get_n_basis_functions(interp_dims[dim_in]), N_in) - return @assert expected_size == size(u)[1:N_in] "Expected the size of the first N_in dimensions of u to be $expected_size based on the BSplineInterpolation properties." + interp_dim::BSplineInterpolationDimension, + ax::AbstractRange + ) + expected_size = get_n_basis_functions(interp_dim) + return @assert expected_size == length(ax) "Expected the size to be $expected_size based on the BSplineInterpolation properties, got $(length(ax))." end function validate_cache( - ::EmptyCache, ::NTuple{N_in, ID}, ::AbstractArray - ) where {N_in, ID} + cache::AbstractInterpolationCache, dims::Tuple, u::AbstractArray + ) + return ntuple(length(dims)) do n + validate_cache(cache, dims[n], u, n) + end +end +function validate_cache( + ::EmptyCache, ::AbstractInterpolationDimension, ::AbstractArray, ::Int + ) + return nothing +end +validate_cache(::EmptyCache, ::NoInterpolationDimension, ::AbstractArray, ::Int) = nothing +function validate_cache( + ::AbstractInterpolationCache, ::NoInterpolationDimension, ::AbstractArray, ::Int + ) return nothing end - function validate_cache( nurbs_weights::NURBSWeights, - ::NTuple{N_in, BSplineInterpolationDimension}, - u::AbstractArray - ) where {N_in} - size_expected = size(u)[1:N_in] - return @assert size(nurbs_weights.weights) == size_expected "The size of the weights array must match the length of the first N_in dimensions of u ($size_expected)." + ::BSplineInterpolationDimension, + u::AbstractArray, + n::Int + ) + size_expected = size(u, n) + return @assert size(nurbs_weights.weights, n) == size_expected "The size of the weights array must match the length of the first N_in dimensions of u ($size_expected), got $(size(nurbs_weights.weights, n))." end function validate_cache( - ::gType, ::NTuple{N_in, ID}, ::AbstractArray - ) where {gType, N_in, ID} + ::gType, ::ID, ::AbstractArray, ::Int + ) where { + gType, ID <: AbstractInterpolationDimension, + } return @error("Interpolation dimension type $ID is not compatible with global cache type $gType.") end -function get_ts( - interp_dims::NTuple{ - N_in, AbstractInterpolationDimension, - } - ) where {N_in} - return ntuple(i -> interp_dims[i].t, N_in) +function get_output_size(interp::NDInterpolation) + I = map(interp.interp_dims, axes(interp.u)) do d, ax + d isa NoInterpolationDimension ? length(ax) : nothing + end + return remove(Nothing, I) end -function get_output_size(interp::NDInterpolation{N_in}) where {N_in} - return size(interp.u)[(N_in + 1):end] +function grid_size(interp::NDInterpolation) + (; interp_dims) = interp + # Get the size of dims that are not NoInterpolationDimension + interp_size = map(d -> length(d.t_eval), remove(NoInterpolationDimension, interp_dims)) + # Get the size of NoInterpolationDimension dims + nointerp_size = get_output_size(interp) + # Insert the nointerp sizes back into the interp_size tuple + return insertat(NoInterpolationDimension, nointerp_size, interp_size, interp_dims) end make_zero!!(::T) where {T <: Number} = zero(T) -function make_zero!!(v::T) where {T <: AbstractArray} - v .= 0 +function make_zero!!(v::AbstractArray) + fill!(v, zero(eltype(v))) return v end function make_out( - interp::NDInterpolation{N_in, 0}, + interp::NDInterpolation{<:Any, N_in, 0}, t::NTuple{N_in, >:Number} ) where {N_in} - return zero(promote_type(eltype(interp.u), map(typeof, t)...)) + return zero(eltype(interp.u)) end - function make_out( - interp::NDInterpolation{N_in}, + interp::NDInterpolation{<:Any, N_in}, t::NTuple{N_in, >:Number} ) where {N_in} - return similar( - interp.u, promote_type(eltype(interp.u), map(eltype, t)...), get_output_size(interp) - ) + T = promote_type(eltype(interp.u), map(eltype, t)...) + return similar(interp.u, T, get_output_size(interp)) end get_left(::AbstractInterpolationDimension) = false @@ -138,12 +172,10 @@ function get_idx( clamp(searchsortedlast(t, t_eval) + idx_shift, lb, ub) end end - -function get_idx( - interp_dims::NTuple{N_in}, - t::Tuple{Vararg{Number, N_in}} - ) where {N_in} - return ntuple(dim_in -> get_idx(interp_dims[dim_in], t[dim_in]), N_in) +# t_eval must already be an index for NoInterpolationDimension +get_idx(::NoInterpolationDimension, t_eval::Union{Colon, Int, AbstractArray{Int}}) = t_eval +function get_idx(interp_dims::Tuple{Vararg{Any, N}}, t::Tuple{Vararg{Any, N}}) where {N} + return map(get_idx, interp_dims, t) end function set_eval_idx!( @@ -176,3 +208,85 @@ end typed_nan(::T) where {T <: Integer} = zero(T) typed_nan(::T) where {T <: AbstractFloat} = T(NaN) + +# Get the KernelAbstractions nd_range, over interpolated dimensions +function get_ndrange(interp::NDInterpolation) + I = map(interp.interp_dims) do d + d isa NoInterpolationDimension ? nothing : length(d.t_eval) + end + return remove(Nothing, I) +end + +# Some tuple handling primitives +# TODO: make sure these compile away completely + +# Remove objects of type `T` from `in` (reworked from DimensionalData.jl) +remove(::Type{T}, in) where {T} = _remove(T, in...) +_remove(::Type{T}, x, xs...) where {T} = (x, _remove(T, xs...)...) +_remove(::Type{T}, x::T, xs...) where {T} = _remove(T, xs...) +_remove(::Type) = () + +# Keep only objects of type `T` from `in` +keep(::Type{T}, in) where {T} = _keep(T, in...) +_keep(::Type{T}, x, xs...) where {T} = _keep(T, xs...) +_keep(::Type{T}, x::T, xs...) where {T} = (x, _keep(T, xs...)...) +_keep(::Type) = () + +# Remove values from `in` where `matches` are of type `T` +function removeat(::Type{T}, in::Tuple, matches::Tuple) where {T} + @assert length(in) == length(matches) + return _removeat(T, in, matches...) +end +# If `!(m isa T)` take from `in` +function _removeat(::Type{T}, in::Tuple{<:Any, Vararg}, m, ms...) where {T} + return (first(in), _removeat(T, Base.tail(in), ms...)...) +end +# If `m isa T` remove +function _removeat(::Type{T}, in::Tuple{<:Any, Vararg}, m::T, ms...) where {T} + return _removeat(T, Base.tail(in), ms...) +end +# `in` can be empty if there are no remaining `m` +_removeat(::Type, in::Tuple{}) = () + +# Keep only values from `in` where `matches` are of type `T` +function keepat(::Type{T}, in::Tuple, matches::Tuple) where {T} + @assert length(in) == length(matches) + return _keepat(T, in, matches...) +end +# If `!(m isa T)` disgaurd +function _keepat(::Type{T}, in::Tuple{<:Any, Vararg}, m, ms...) where {T} + return _keepat(T, Base.tail(in), ms...) +end +# If `m isa T` keep +function _keepat(::Type{T}, in::Tuple{<:Any, Vararg}, m::T, ms...) where {T} + return (first(in), _keepat(T, Base.tail(in), ms...)...) +end +# `in` can be empty if there are no remaining `m` +_keepat(::Type, in::Tuple{}) = () + +# Insert x into `in` where `matches` are of type `T`, +# otherwise output one of `in` for each `m`. +function insertat(::Type{T}, x, in::Tuple, matches::Tuple) where {T} + return _insertat(T, x, in, matches...) +end +# If `!(m isa T)` take from `in` +function _insertat(::Type{T}, x, in::Tuple{<:Any, Vararg}, m, ms...) where {T} + return (first(in), _insertat(T, x, Base.tail(in), ms...)...) +end +# If `m isa T` insert `x` +function _insertat(::Type{T}, x, in::Tuple{<:Any, Vararg}, m::T, ms...) where {T} + return (x, _insertat(T, x, in, ms...)...) +end +# For Tuple x we insert the first +function _insertat(::Type{T}, xs::Tuple, in::Tuple{<:Any, Vararg}, m::T, ms...) where {T} + return (first(xs), _insertat(T, Base.tail(xs), in, ms...)...) +end +# `in` can be empty if there are no remaining `m` +_insertat(::Type, x, in::Tuple{}) = () +# `in` can also be empty if all trailing ms are T +function _insertat(::Type{T}, x, in::Tuple{}, m::T, ms::T...) where {T} + return (x, _insertat(T, x, in, ms...)...) +end +function _insertat(::Type{T}, xs::Tuple, in::Tuple{}, m::T, ms::T...) where {T} + return (first(xs), _insertat(T, Base.tail(xs), in, ms...)...) +end diff --git a/src/spline_utils.jl b/src/spline_utils.jl index 248b958..ba0f98f 100644 --- a/src/spline_utils.jl +++ b/src/spline_utils.jl @@ -66,9 +66,8 @@ function get_basis_function_values( t::T_t, idx::Integer, derivative_order::Integer, - multi_point_index::Nothing, - dim_in::Integer - ) where {T_t <: Number} + multi_point_index::Nothing + ) where {T_t} (; degree, knots_all) = itp_dim T_eval = eltype(itp_dim.basis_function_eval) # Use a concrete zero value for type stability @@ -120,8 +119,7 @@ function get_basis_function_values( t::Number, idx::Integer, derivative_order::Integer, - multi_point_index::Number, - dim_in::Integer + multi_point_index::Integer ) return view( itp_dim.basis_function_eval, @@ -129,35 +127,17 @@ function get_basis_function_values( ) end -# Get the basis function values for one point in a -# grid evaluation (given by the tuple multi point index) -function get_basis_function_values( - itp_dim::BSplineInterpolationDimension, - t::Number, - idx::Integer, - derivative_order::Integer, - multi_point_index::NTuple{N_in, <:Integer}, - dim_in::Integer - ) where {N_in} - return view( - itp_dim.basis_function_eval, - multi_point_index[dim_in], :, derivative_order + 1 - ) -end - # Get all basis function values to evaluate a BSpline interpolation in t function get_basis_function_values_all( - A::NDInterpolation{N_in, N_out, <:BSplineInterpolationDimension}, - t::Tuple{Vararg{Number, N_in}}, - idx::NTuple{N_in, <:Integer}, - derivative_orders::NTuple{N_in, <:Integer}, + A::NDInterpolation, + ts::Tuple, + idx::Tuple, + derivative_orders::Tuple, multi_point_index - ) where {N_in, N_out} - return ntuple( - dim_in -> get_basis_function_values( - A.interp_dims[dim_in], t[dim_in], idx[dim_in], derivative_orders[dim_in], multi_point_index, dim_in - ), - N_in + ) + return map( + get_basis_function_values, A.interp_dims, ts, + idx, derivative_orders, multi_point_index ) end @@ -176,17 +156,12 @@ end ) i, derivative_order_plus_1 = @index(Global, NTuple) - itp_dim.basis_function_eval[ - i, - :, - derivative_order_plus_1, - ] .= get_basis_function_values( + itp_dim.basis_function_eval[i, :, derivative_order_plus_1] .= get_basis_function_values( itp_dim, itp_dim.t_eval[i], itp_dim.idx_eval[i], derivative_order_plus_1 - 1, - nothing, - 0 + nothing ) end diff --git a/test/test_interpolations.jl b/test/test_interpolations.jl index 1a3075e..7b8e519 100644 --- a/test/test_interpolations.jl +++ b/test/test_interpolations.jl @@ -1,10 +1,28 @@ -using DataInterpolationsND: AbstractInterpolationDimension, EmptyCache using DataInterpolationsND using Random +using DataInterpolationsND: AbstractInterpolationDimension, EmptyCache, + keepat, removeat, insertat, remove, keep + +@testset "tuple primitives" begin + # These help with handling NoInterpolationDimension and its consequences + @test insertat(Nothing, 1, (2,), (nothing, 4, nothing)) === (1, 2, 1) + @test insertat(Float64, missing, (1,), (8, 9.0)) === (1, missing) + @test keepat(Nothing, (1, 2, 3), (nothing, 4, nothing)) === (1, 3) + @test keepat(Float64, (1, 2), (8, 9.0)) === (2,) + @test removeat(Nothing, (1, 2, 3), (nothing, 4, nothing)) === (2,) + @test removeat(Float64, (1, 2), (8, 9.0)) === (1,) + @test remove(Nothing, (1, nothing, 3)) === (1, 3) + @test keep(Int, (1, nothing, 3)) === (1, 3) +end function test_globally_constant( - ID::Type{<:AbstractInterpolationDimension}; args1 = (), args2 = (), kwargs1 = (), - kwargs2 = (), cache = EmptyCache(), test_derivatives = true + ID::Type{<:AbstractInterpolationDimension}; + args1 = (), + args2 = (), + kwargs1 = (), + kwargs2 = (), + cache = EmptyCache(), + test_derivatives = true ) t1 = [-3.14, 1.0, 3.0, 7.6, 12.8] t2 = [-2.71, 1.41, 12.76, 50.2, 120.0] @@ -22,7 +40,9 @@ function test_globally_constant( ) itp = NDInterpolation(u, itp_dims; cache) - @test all(x -> isapprox(x, 2.0; atol = 1.0e-10), eval_grid(itp)) + + grid = eval_grid(itp) + @test all(x -> isapprox(x, 2.0; atol = 1.0e-10), grid) if test_derivatives @test all( x -> isapprox(x, 0.0; atol = 1.0e-10), eval_grid(itp, derivative_orders = (1, 0)) @@ -37,7 +57,7 @@ function test_globally_constant( ID(t1, args1...; t_eval = t1[1:(end - 1)] + diff(t1) / 2, kwargs1...), ID(t2, args2...; t_eval = t2[1:(end - 1)] + diff(t2) / 2, kwargs2...), ) - itp = NDInterpolation(u, itp_dims) + itp = NDInterpolation(u, itp_dims; cache) @test all(x -> isapprox(x, 2.0; atol = 1.0e-10), eval_grid(itp)) return if test_derivatives @test all( @@ -49,15 +69,17 @@ function test_globally_constant( end end -function test_analytic(itp::NDInterpolation{N_in}, f) where {N_in} +function test_analytic(itp::NDInterpolation{<:Any, N_in}, f) where {N_in} + used_interp_dims = remove(NoInterpolationDimension, itp.interp_dims) # Evaluation in data points - ts = ntuple(dim_in -> itp.interp_dims[dim_in].t, N_in) + ts = map(d -> d.t, used_interp_dims) for t in Iterators.product(ts...) @test itp(t) ≈ f(t...) end # Evaluation between data points ts_ = ntuple(dim_in -> ts[dim_in][1:(end - 1)] + diff(ts[dim_in]) / 2, N_in) + t = first(Iterators.product(ts_...)) for t in Iterators.product(ts_...) @test itp(t) ≈ f(t...) end @@ -68,7 +90,6 @@ end test_globally_constant(LinearInterpolationDimension) f(t1, t2) = 3.0 + 2.3t1 - 4.7t2 - Random.seed!(1) t1 = cumsum(rand(10)) t2 = cumsum(rand(10)) @@ -80,14 +101,24 @@ end u = f.(t1, t2') itp = NDInterpolation(u, itp_dims) test_analytic(itp, f) + + itp_dims = ( + NoInterpolationDimension(), + LinearInterpolationDimension(t2; t_eval = t2 ./ 2), + ) + itp = NDInterpolation(u, itp_dims) + eval_grid(itp) end @testset "BSpline Interpolation" begin - test_globally_constant( - BSplineInterpolationDimension, args1 = (2,), args2 = (3,), - kwargs1 = (:max_derivative_order_eval => 1,), - kwargs2 = (:max_derivative_order_eval => 1,) - ) + ID = BSplineInterpolationDimension + args1 = (2,) + args2 = (3,) + kwargs1 = (:max_derivative_order_eval => 1,) + kwargs2 = (:max_derivative_order_eval => 1,) + cache = EmptyCache() + test_derivatives = true + test_globally_constant(ID; args1, args2, kwargs1, kwargs2, cache, test_derivatives) f(t1, t2, t3) = t1^2 + t2^2 + t3^2 @@ -105,13 +136,14 @@ end @testset "NURBS Interpolation" begin Random.seed!(10) - test_globally_constant( - BSplineInterpolationDimension; args1 = (3,), args2 = (1,), - kwargs1 = (:max_derivative_order_eval => 1,), - kwargs2 = (:max_derivative_order_eval => 1,), - cache = NURBSWeights(rand(7, 5)), - test_derivatives = false - ) + ID = BSplineInterpolationDimension + args1 = (3,) + args2 = (1,) + kwargs1 = (:max_derivative_order_eval => 1,) + kwargs2 = (:max_derivative_order_eval => 1,) + cache = NURBSWeights(rand(7, 5)) + test_derivatives = false + test_globally_constant(ID; args1, args2, kwargs1, kwargs2, cache, test_derivatives) ## Circle representation # Knots @@ -148,3 +180,17 @@ end @test allunique(points_on_circle[2:end]) @test all(point -> point[1]^2 + point[2]^2 ≈ 1, points_on_circle) end + +@testset "Mixed Interpolation" begin + t2 = cumsum(rand(4)) + t3 = collect(0:(π / 2):(2π)) + u = ones(3, 4, 5) + itp_dims = ( + NoInterpolationDimension(), + LinearInterpolationDimension(t2; t_eval = t2), + BSplineInterpolationDimension(t3, 1; t_eval = t3), + ) + itp = NDInterpolation(u, itp_dims) + @test isapprox(itp(1.0, 4.0), [1.0, 1.0, 1.0]) + eval_grid(itp) +end