Add Reactant extension for branchless interpolation

Project.toml

@@ -17,6 +17,7 @@ FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 RegularizationTools = "29dad682-9a27-4bc3-9c72-016788665182"
 SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
@@ -25,6 +26,7 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 [extensions]
 DataInterpolationsChainRulesCoreExt = "ChainRulesCore"
 DataInterpolationsOptimExt = "Optim"
+DataInterpolationsReactantExt = "Reactant"
 DataInterpolationsRegularizationToolsExt = "RegularizationTools"
 DataInterpolationsSparseConnectivityTracerExt = ["SparseConnectivityTracer", "FillArrays"]
 DataInterpolationsSymbolicsExt = "Symbolics"
@@ -47,6 +49,7 @@ Mooncake = "0.4"
 Optim = "1.6, 2"
 PrettyTables = "2.4, 3"
 QuadGK = "2.9.1"
+Reactant = "0.2"
 RecipesBase = "1.3"
 Reexport = "1"
 RegularizationTools = "0.6"
@@ -71,6 +74,7 @@ Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 RegularizationTools = "29dad682-9a27-4bc3-9c72-016788665182"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
@@ -82,4 +86,4 @@ Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "AllocCheck", "BenchmarkTools", "SafeTestsets", "ChainRulesCore", "Mooncake", "Optim", "RegularizationTools", "Test", "StableRNGs", "FiniteDifferences", "QuadGK", "ForwardDiff", "StaticArrays", "Symbolics", "Unitful", "Zygote", "SparseConnectivityTracer"]
+test = ["Aqua", "AllocCheck", "BenchmarkTools", "SafeTestsets", "ChainRulesCore", "Mooncake", "Optim", "Reactant", "RegularizationTools", "Test", "StableRNGs", "FiniteDifferences", "QuadGK", "ForwardDiff", "StaticArrays", "Symbolics", "Unitful", "Zygote", "SparseConnectivityTracer"]

ext/DataInterpolationsReactantExt.jl

@@ -0,0 +1,191 @@
+module DataInterpolationsReactantExt
+
+using DataInterpolations: DataInterpolations,
+    AbstractInterpolation,
+    LinearInterpolation,
+    QuadraticInterpolation,
+    ConstantInterpolation,
+    ExtrapolationType,
+    get_parameters,
+    linear_interpolation_parameters
+import DataInterpolations: _interpolate
+using Reactant: Reactant, TracedRNumber
+
+# =============================================================================
+# Branchless interpolation for Reactant traced numbers.
+#
+# When Reactant traces through an ODE function, the time `t` becomes a
+# TracedRNumber. Standard DataInterpolations methods use `if t < ...` which
+# fails because TracedRNumber{Bool} can't be used in boolean context.
+#
+# Strategy: evaluate the interpolation formula for all segments and use
+# `ifelse` (which Reactant lowers to stablehlo.select) to pick the right one.
+# This is O(n) in knot count but fully traceable.
+# =============================================================================
+
+# ===========================================================================
+# Helper: get slope for a given segment index (concrete integer)
+# ===========================================================================
+
+function _get_slope(A::LinearInterpolation, idx::Int)
+    if A.cache_parameters
+        return A.p.slope[idx]
+    else
+        return linear_interpolation_parameters(A.u, A.t, idx)
+    end
+end
+
+# ===========================================================================
+# LinearInterpolation — branchless
+# ===========================================================================
+
+function _interpolate(A::LinearInterpolation{<:AbstractVector}, t::TracedRNumber)
+    n = length(A.t)
+
+    # Evaluate first segment
+    u1 = oftype(t, A.u[1])
+    slope1 = oftype(t, _get_slope(A, 1))
+    t1 = oftype(t, A.t[1])
+    result = u1 + slope1 * (t - t1)
+
+    # Cascade through remaining segments: if t >= A.t[i], use segment i
+    for i in 2:(n - 1)
+        ui = oftype(t, A.u[i])
+        slope_i = oftype(t, _get_slope(A, i))
+        ti = oftype(t, A.t[i])
+        seg_val = ui + slope_i * (t - ti)
+        result = ifelse(t >= ti, seg_val, result)
+    end
+
+    # Handle extrapolation
+    left_val = _traced_left_extrap_val(A, t)
+    right_val = _traced_right_extrap_val(A, t)
+    tmin = oftype(t, first(A.t))
+    tmax = oftype(t, last(A.t))
+
+    result = ifelse(t < tmin, left_val, result)
+    result = ifelse(t > tmax, right_val, result)
+
+    return result
+end
+
+function _traced_left_extrap_val(A::LinearInterpolation, t)
+    ext = A.extrapolation_left
+    if ext == ExtrapolationType.Constant
+        return oftype(t, first(A.u))
+    else
+        # Linear, Extension, or fallback: extend with first segment's slope
+        slope = oftype(t, _get_slope(A, 1))
+        return oftype(t, first(A.u)) + slope * (t - oftype(t, first(A.t)))
+    end
+end
+
+function _traced_right_extrap_val(A::LinearInterpolation, t)
+    ext = A.extrapolation_right
+    n = length(A.t)
+    if ext == ExtrapolationType.Constant
+        return oftype(t, last(A.u))
+    else
+        slope = oftype(t, _get_slope(A, n - 1))
+        return oftype(t, last(A.u)) + slope * (t - oftype(t, last(A.t)))
+    end
+end
+
+# ===========================================================================
+# ConstantInterpolation — branchless
+# ===========================================================================
+
+function _interpolate(A::ConstantInterpolation{<:AbstractVector}, t::TracedRNumber)
+    n = length(A.t)
+
+    if A.dir === :left
+        # :left — use value at the largest t[i] <= t
+        result = oftype(t, A.u[1])
+        for i in 2:n
+            ti = oftype(t, A.t[i])
+            result = ifelse(t >= ti, oftype(t, A.u[i]), result)
+        end
+    else
+        # :right — use value at the smallest t[i] >= t
+        result = oftype(t, A.u[n])
+        for i in (n - 1):-1:1
+            ti = oftype(t, A.t[i])
+            result = ifelse(t <= ti, oftype(t, A.u[i]), result)
+        end
+    end
+
+    # Handle extrapolation
+    tmin = oftype(t, first(A.t))
+    tmax = oftype(t, last(A.t))
+
+    ext_left = A.extrapolation_left
+    left_val = if ext_left == ExtrapolationType.Constant || ext_left == ExtrapolationType.Extension
+        oftype(t, first(A.u))
+    else
+        oftype(t, first(A.u))
+    end
+
+    ext_right = A.extrapolation_right
+    right_val = if ext_right == ExtrapolationType.Constant || ext_right == ExtrapolationType.Extension
+        oftype(t, last(A.u))
+    else
+        oftype(t, last(A.u))
+    end
+
+    result = ifelse(t < tmin, left_val, result)
+    result = ifelse(t > tmax, right_val, result)
+
+    return result
+end
+
+# ===========================================================================
+# QuadraticInterpolation — branchless
+# ===========================================================================
+
+function _interpolate(A::QuadraticInterpolation{<:AbstractVector}, t::TracedRNumber)
+    n = length(A.t)
+
+    # First segment
+    α1, β1 = get_parameters(A, 1)
+    Δt = t - oftype(t, A.t[1])
+    result = oftype(t, A.u[1]) + Δt * (oftype(t, α1) * Δt + oftype(t, β1))
+
+    # Remaining segments
+    for i in 2:(n - 2)
+        αi, βi = get_parameters(A, i)
+        Δti = t - oftype(t, A.t[i])
+        seg_val = oftype(t, A.u[i]) + Δti * (oftype(t, αi) * Δti + oftype(t, βi))
+        ti = oftype(t, A.t[i])
+        result = ifelse(t >= ti, seg_val, result)
+    end
+
+    # Handle extrapolation
+    tmin = oftype(t, first(A.t))
+    tmax = oftype(t, last(A.t))
+
+    ext_left = A.extrapolation_left
+    left_val = if ext_left == ExtrapolationType.Constant
+        oftype(t, first(A.u))
+    else
+        α, β = get_parameters(A, 1)
+        Δt_left = t - oftype(t, A.t[1])
+        oftype(t, A.u[1]) + Δt_left * (oftype(t, α) * Δt_left + oftype(t, β))
+    end
+
+    ext_right = A.extrapolation_right
+    last_idx = max(n - 2, 1)
+    right_val = if ext_right == ExtrapolationType.Constant
+        oftype(t, last(A.u))
+    else
+        α, β = get_parameters(A, last_idx)
+        Δt_right = t - oftype(t, A.t[last_idx])
+        oftype(t, A.u[last_idx]) + Δt_right * (oftype(t, α) * Δt_right + oftype(t, β))
+    end
+
+    result = ifelse(t < tmin, left_val, result)
+    result = ifelse(t > tmax, right_val, result)
+
+    return result
+end
+
+end # module

test/reactant_tests.jl

@@ -0,0 +1,65 @@
+using DataInterpolations
+using Reactant
+using Test
+
+function test_reactant_interpolation(interp, t_vals; atol = 1.0e-12)
+    for t_val in t_vals
+        expected = interp(t_val)
+        f = Reactant.compile(
+            (t) -> interp(t),
+            (Reactant.to_rarray(t_val; track_numbers = true),)
+        )
+        result = Float64(f(Reactant.to_rarray(t_val; track_numbers = true)))
+        @test isapprox(result, expected; atol = atol)
+    end
+    return
+end
+
+@testset "LinearInterpolation" begin
+    @testset "basic" begin
+        interp = LinearInterpolation([0.0, 1.0, 4.0], [0.0, 0.5, 1.0])
+        test_reactant_interpolation(interp, [0.0, 0.25, 0.5, 0.75, 1.0])
+    end
+
+    @testset "Extension extrapolation" begin
+        interp = LinearInterpolation(
+            [0.0, 2.0], [0.0, 1.0]; extrapolation = ExtrapolationType.Extension
+        )
+        test_reactant_interpolation(interp, [-0.5, 0.0, 0.5, 1.0, 1.5])
+    end
+
+    @testset "Constant extrapolation" begin
+        interp = LinearInterpolation(
+            [0.0, 2.0], [0.0, 1.0];
+            extrapolation_left = ExtrapolationType.Constant,
+            extrapolation_right = ExtrapolationType.Constant
+        )
+        test_reactant_interpolation(interp, [-0.5, 0.0, 0.5, 1.0, 1.5])
+    end
+
+    @testset "many knots" begin
+        u = sin.(range(0, 2π; length = 20))
+        t = collect(range(0.0, 1.0; length = 20))
+        interp = LinearInterpolation(u, t)
+        test_reactant_interpolation(interp, range(0.0, 1.0; length = 50))
+    end
+end
+
+@testset "ConstantInterpolation" begin
+    @testset "dir=:left" begin
+        interp = ConstantInterpolation([1.0, 3.0, 5.0], [0.0, 0.5, 1.0])
+        test_reactant_interpolation(interp, [0.0, 0.25, 0.5, 0.75, 1.0])
+    end
+
+    @testset "dir=:right" begin
+        interp = ConstantInterpolation(
+            [1.0, 3.0, 5.0], [0.0, 0.5, 1.0]; dir = :right
+        )
+        test_reactant_interpolation(interp, [0.0, 0.25, 0.5, 0.75, 1.0])
+    end
+end
+
+@testset "QuadraticInterpolation" begin
+    interp = QuadraticInterpolation([0.0, 1.0, 0.0], [0.0, 0.5, 1.0])
+    test_reactant_interpolation(interp, [0.0, 0.25, 0.5, 0.75, 1.0]; atol = 1.0e-10)
+end

test/runtests.jl

@@ -27,6 +27,7 @@ const GROUP = get(ENV, "GROUP", "All")
             @safetestset "SparseConnectivityTracer Tests" include("sparseconnectivitytracer_tests.jl")
             @safetestset "Zygote support Tests" include("zygote_tests.jl")
             @safetestset "Mooncake support Tests" include("mooncake_tests.jl")
+            @safetestset "Reactant support Tests" include("reactant_tests.jl")
         end
     end