refactor: add output dimensionality as a parameter in AbstractInterpo…

…lations
SciML · Oct 9, 2024 · 1df1996 · 1df1996
1 parent e73a150
commit 1df1996
Show file tree

Hide file tree

Showing 6 changed files with 108 additions and 73 deletions.
diff --git a/src/DataInterpolations.jl b/src/DataInterpolations.jl
@@ -2,7 +2,7 @@ module DataInterpolations
 
 ### Interface Functionality
 
-abstract type AbstractInterpolation{T} end
+abstract type AbstractInterpolation{T, N} end
 
 using LinearAlgebra, RecipesBase
 using PrettyTables
@@ -92,8 +92,8 @@ export LinearInterpolation, QuadraticInterpolation, LagrangeInterpolation,
 
 # added for RegularizationSmooth, JJS 11/27/21
 ### Regularization data smoothing and interpolation
-struct RegularizationSmooth{uType, tType, T, T2, ITP <: AbstractInterpolation{T}} <:
-       AbstractInterpolation{T}
+struct RegularizationSmooth{uType, tType, T, T2, N, ITP <: AbstractInterpolation{T, N}} <:
+       AbstractInterpolation{T, N}
     u::uType
     û::uType
     t::tType
@@ -116,7 +116,8 @@ struct RegularizationSmooth{uType, tType, T, T2, ITP <: AbstractInterpolation{T}
             alg,
             Aitp,
             extrapolate)
-        new{typeof(u), typeof(t), eltype(u), typeof(λ), typeof(Aitp)}(
+        N = get_output_dim(u)
+        new{typeof(u), typeof(t), eltype(u), typeof(λ), N, typeof(Aitp)}(
             u,
             û,
             t,
@@ -143,8 +144,9 @@ struct CurvefitCache{
     lbType,
     algType,
     pminType,
-    T
-} <: AbstractInterpolation{T}
+    T,
+    N
+} <: AbstractInterpolation{T, N}
     u::uType
     t::tType
     m::mType        # model type
@@ -155,9 +157,10 @@ struct CurvefitCache{
     pmin::pminType  # optimized params
     extrapolate::Bool
     function CurvefitCache(u, t, m, p0, ub, lb, alg, pmin, extrapolate)
+        N = get_output_dim(u)
         new{typeof(u), typeof(t), typeof(m),
             typeof(p0), typeof(ub), typeof(lb),
-            typeof(alg), typeof(pmin), eltype(u)}(u,
+            typeof(alg), typeof(pmin), eltype(u), N}(u,
             t,
             m,
             p0,

diff --git a/src/derivatives.jl b/src/derivatives.jl
@@ -151,14 +151,14 @@ function _derivative(A::BSplineInterpolation{<:AbstractVector{<:Number}}, t::Num
     n = length(A.t)
     scale = (A.p[idx + 1] - A.p[idx]) / (A.t[idx + 1] - A.t[idx])
     t_ = A.p[idx] + (t - A.t[idx]) * scale
-    N = t isa ForwardDiff.Dual ? zeros(eltype(t), n) : A.N
-    spline_coefficients!(N, A.d - 1, A.k, t_)
+    sc = t isa ForwardDiff.Dual ? zeros(eltype(t), n) : A.sc
+    spline_coefficients!(sc, A.d - 1, A.k, t_)
     ducum = zero(eltype(A.u))
     if t == A.t[1]
         ducum = (A.c[2] - A.c[1]) / (A.k[A.d + 2])
     else
         for i in 1:(n - 1)
-            ducum += N[i + 1] * (A.c[i + 1] - A.c[i]) / (A.k[i + A.d + 1] - A.k[i + 1])
+            ducum += sc[i + 1] * (A.c[i + 1] - A.c[i]) / (A.k[i + A.d + 1] - A.k[i + 1])
         end
     end
     ducum * A.d * scale
@@ -172,14 +172,14 @@ function _derivative(A::BSplineApprox{<:AbstractVector{<:Number}}, t::Number, ig
     idx = get_idx(A, t, iguess)
     scale = (A.p[idx + 1] - A.p[idx]) / (A.t[idx + 1] - A.t[idx])
     t_ = A.p[idx] + (t - A.t[idx]) * scale
-    N = t isa ForwardDiff.Dual ? zeros(eltype(t), A.h) : A.N
-    spline_coefficients!(N, A.d - 1, A.k, t_)
+    sc = t isa ForwardDiff.Dual ? zeros(eltype(t), A.h) : A.sc
+    spline_coefficients!(sc, A.d - 1, A.k, t_)
     ducum = zero(eltype(A.u))
     if t == A.t[1]
         ducum = (A.c[2] - A.c[1]) / (A.k[A.d + 2])
     else
         for i in 1:(A.h - 1)
-            ducum += N[i + 1] * (A.c[i + 1] - A.c[i]) / (A.k[i + A.d + 1] - A.k[i + 1])
+            ducum += sc[i + 1] * (A.c[i + 1] - A.c[i]) / (A.k[i + A.d + 1] - A.k[i + 1])
         end
     end
     ducum * A.d * scale

diff --git a/src/integral_inverses.jl b/src/integral_inverses.jl
@@ -1,4 +1,4 @@
-abstract type AbstractIntegralInverseInterpolation{T} <: AbstractInterpolation{T} end
+abstract type AbstractIntegralInverseInterpolation{T, N} <: AbstractInterpolation{T, N} end
 
 """
     invert_integral(A::AbstractInterpolation)::AbstractIntegralInverseInterpolation
@@ -33,15 +33,16 @@ Can be easily constructed with `invert_integral(A::LinearInterpolation{<:Abstrac
   - `t` : Given by `A.I` (the cumulative integral of `A`)
   - `A` : The `LinearInterpolation` object
 """
-struct LinearInterpolationIntInv{uType, tType, itpType, T} <:
-       AbstractIntegralInverseInterpolation{T}
+struct LinearInterpolationIntInv{uType, tType, itpType, T, N} <:
+       AbstractIntegralInverseInterpolation{T, N}
     u::uType
     t::tType
     extrapolate::Bool
     iguesser::Guesser{tType}
     itp::itpType
     function LinearInterpolationIntInv(u, t, A)
-        new{typeof(u), typeof(t), typeof(A), eltype(u)}(
+        N = get_output_dim(u)
+        new{typeof(u), typeof(t), typeof(A), eltype(u), N}(
             u, t, A.extrapolate, Guesser(t), A)
     end
 end
@@ -79,15 +80,16 @@ Can be easily constructed with `invert_integral(A::ConstantInterpolation{<:Abstr
   - `t` : Given by `A.I` (the cumulative integral of `A`)
   - `A` : The `ConstantInterpolation` object
 """
-struct ConstantInterpolationIntInv{uType, tType, itpType, T} <:
-       AbstractIntegralInverseInterpolation{T}
+struct ConstantInterpolationIntInv{uType, tType, itpType, T, N} <:
+       AbstractIntegralInverseInterpolation{T, N}
     u::uType
     t::tType
     extrapolate::Bool
     iguesser::Guesser{tType}
     itp::itpType
     function ConstantInterpolationIntInv(u, t, A)
-        new{typeof(u), typeof(t), typeof(A), eltype(u)}(
+        N = get_output_dim(u)
+        new{typeof(u), typeof(t), typeof(A), eltype(u), N}(
             u, t, A.extrapolate, Guesser(t), A
         )
     end