DataInterpolations support

Project.toml

@@ -8,22 +8,28 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
-SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 
 [weakdeps]
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0"
 
 [extensions]
 SparseConnectivityTracerNNlibExt = "NNlib"
 SparseConnectivityTracerSpecialFunctionsExt = "SpecialFunctions"
+SparseConnectivityTracerDataInterpolationsExt = "DataInterpolations"
+
+[extras]
+DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 
 [compat]
 ADTypes = "1"
+DataInterpolations = "6"
 DocStringExtensions = "0.9"
 FillArrays = "1"
 LinearAlgebra = "<0.0.1, 1"

docs/src/dev/adding_overloads.md

@@ -20,7 +20,7 @@ to improve the performance of your functions or to work around some of SCT's [li
     This page of the documentation shows you how.
 
 !!! tip "Copy one of our package extensions"
-    The easiest way to add overloads is to copy one of our [package extensions](https://github.com/adrhill/SparseConnectivityTracer.jl/tree/main/ext) and to modify it.
+    The easiest way to add overloads is to copy one of our package extensions, [e.g. our NNlib extension](https://github.com/adrhill/SparseConnectivityTracer.jl/blob/main/ext/SparseConnectivityTracerNNlibExt.jl), and to modify it.
     Please upstream your additions by opening a pull request! We will help you out to get your feature merged.
 
 ## Operator classification

ext/SparseConnectivityTracerDataInterpolationsExt.jl

@@ -0,0 +1,173 @@
+# WARNING: If you are following the "Adding Overloads" guide's advice to copy an existing package extension, 
+# copy another, less complicated one!
+module SparseConnectivityTracerDataInterpolationsExt
+
+if isdefined(Base, :get_extension)
+    using SparseConnectivityTracer: AbstractTracer, Dual, primal, tracer
+    using SparseConnectivityTracer: GradientTracer, gradient_tracer_1_to_1
+    using SparseConnectivityTracer: HessianTracer, hessian_tracer_1_to_1
+    using SparseConnectivityTracer: Fill # from FillArrays.jl
+    import DataInterpolations:
+        AbstractInterpolation,
+        LinearInterpolation,
+        QuadraticInterpolation,
+        LagrangeInterpolation,
+        AkimaInterpolation,
+        ConstantInterpolation,
+        QuadraticSpline,
+        CubicSpline,
+        BSplineInterpolation,
+        BSplineApprox
+    # TODO: support when Julia 1.6 is dropped
+    # CubicHermiteSpline,
+    # PCHIPInterpolation,
+    # QuinticHermiteSpline
+else
+    using ..SparseConnectivityTracer: AbstractTracer, Dual, primal, tracer
+    using ..SparseConnectivityTracer: GradientTracer, gradient_tracer_1_to_1
+    using ..SparseConnectivityTracer: HessianTracer, hessian_tracer_1_to_1
+    using ..SparseConnectivityTracer: Fill # from FillArrays.jl
+    import ..DataInterpolations:
+        AbstractInterpolation,
+        LinearInterpolation,
+        QuadraticInterpolation,
+        LagrangeInterpolation,
+        AkimaInterpolation,
+        ConstantInterpolation,
+        QuadraticSpline,
+        CubicSpline,
+        BSplineInterpolation,
+        BSplineApprox
+    # TODO: support when Julia 1.6 is dropped
+    # CubicHermiteSpline,
+    # PCHIPInterpolation,
+    # QuinticHermiteSpline
+end
+
+#========================#
+# General interpolations #
+#========================#
+
+# We assume that with the exception of ConstantInterpolation and LinearInterpolation,
+# all interpolations have a non-zero second derivative at some point in the input domain.
+
+for I in (
+    :QuadraticInterpolation,
+    :LagrangeInterpolation,
+    :AkimaInterpolation,
+    :QuadraticSpline,
+    :CubicSpline,
+    :BSplineInterpolation,
+    :BSplineApprox,
+    # TODO: support when Julia 1.6 is dropped
+    # :CubicHermiteSpline,
+    # :QuinticHermiteSpline,
+)
+    # 1D Interpolations (uType<:AbstractVector)
+    @eval function (interp::$(I){uType})(t::GradientTracer) where {uType<:AbstractVector}
+        return gradient_tracer_1_to_1(t, false)
+    end
+    @eval function (interp::$(I){uType})(t::HessianTracer) where {uType<:AbstractVector}
+        return hessian_tracer_1_to_1(t, false, false)
+    end
+
+    # ND Interpolations (uType<:AbstractMatrix)
+    @eval function (interp::$(I){uType})(t::GradientTracer) where {uType<:AbstractMatrix}
+        t = gradient_tracer_1_to_1(t, false)
+        nstates = size(interp.u, 1)
+        return Fill(t, nstates)
+    end
+    @eval function (interp::$(I){uType})(t::HessianTracer) where {uType<:AbstractMatrix}
+        t = hessian_tracer_1_to_1(t, false, false)
+        nstates = size(interp.u, 1)
+        return Fill(t, nstates)
+    end
+end
+
+# Some Interpolations require custom overloads on `Dual` due to mutation of caches.
+for I in (
+    :LagrangeInterpolation,
+    :BSplineInterpolation,
+    :BSplineApprox,
+    # TODO: support when Julia 1.6 is dropped
+    # :CubicHermiteSpline,
+    # :QuinticHermiteSpline,
+)
+    @eval function (interp::$(I){uType})(d::Dual) where {uType<:AbstractVector}
+        p = interp(primal(d))
+        t = interp(tracer(d))
+        return Dual(p, t)
+    end
+
+    @eval function (interp::$(I){uType})(d::Dual) where {uType<:AbstractMatrix}
+        p = interp(primal(d))
+        t = interp(tracer(d))
+        return Dual.(p, t)
+    end
+end
+
+#=======================#
+# ConstantInterpolation #
+#=======================#
+
+# 1D Interpolations (uType<:AbstractVector)
+function (interp::ConstantInterpolation{uType})(
+    t::GradientTracer
+) where {uType<:AbstractVector}
+    return gradient_tracer_1_to_1(t, true)
+end
+function (interp::ConstantInterpolation{uType})(
+    t::HessianTracer
+) where {uType<:AbstractVector}
+    return hessian_tracer_1_to_1(t, true, true)
+end
+
+# ND Interpolations (uType<:AbstractMatrix)
+function (interp::ConstantInterpolation{uType})(
+    t::GradientTracer
+) where {uType<:AbstractMatrix}
+    t = gradient_tracer_1_to_1(t, true)
+    nstates = size(interp.u, 1)
+    return Fill(t, nstates)
+end
+function (interp::ConstantInterpolation{uType})(
+    t::HessianTracer
+) where {uType<:AbstractMatrix}
+    t = hessian_tracer_1_to_1(t, true, true)
+    nstates = size(interp.u, 1)
+    return Fill(t, nstates)
+end
+
+#=====================#
+# LinearInterpolation #
+#=====================#
+
+# 1D Interpolations (uType<:AbstractVector)
+function (interp::LinearInterpolation{uType})(
+    t::GradientTracer
+) where {uType<:AbstractVector}
+    return gradient_tracer_1_to_1(t, false)
+end
+function (interp::LinearInterpolation{uType})(
+    t::HessianTracer
+) where {uType<:AbstractVector}
+    return hessian_tracer_1_to_1(t, false, true)
+end
+
+# ND Interpolations (uType<:AbstractMatrix)
+function (interp::LinearInterpolation{uType})(
+    t::GradientTracer
+) where {uType<:AbstractMatrix}
+    t = gradient_tracer_1_to_1(t, false)
+    nstates = size(interp.u, 1)
+    return Fill(t, nstates)
+end
+function (interp::LinearInterpolation{uType})(
+    t::HessianTracer
+) where {uType<:AbstractMatrix}
+    t = hessian_tracer_1_to_1(t, false, true)
+    nstates = size(interp.u, 1)
+    return Fill(t, nstates)
+end
+
+end # module SparseConnectivityTracerDataInterpolationsExt

src/SparseConnectivityTracer.jl

@@ -50,6 +50,9 @@ function __init__()
         @require NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" include(
             "../ext/SparseConnectivityTracerNNlibExt.jl"
         )
+        @require DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0" include(
+            "../ext/SparseConnectivityTracerDataInterpolationsExt.jl"
+        )
     end
 end
 

test/Project.toml

@@ -5,6 +5,7 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
+DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"