diff --git a/Project.toml b/Project.toml index 989a5133..f94950fb 100644 --- a/Project.toml +++ b/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"] diff --git a/ext/DataInterpolationsReactantExt.jl b/ext/DataInterpolationsReactantExt.jl new file mode 100644 index 00000000..5ab3ba73 --- /dev/null +++ b/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 diff --git a/test/reactant_tests.jl b/test/reactant_tests.jl new file mode 100644 index 00000000..d7df8f44 --- /dev/null +++ b/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 diff --git a/test/runtests.jl b/test/runtests.jl index 7a65a8f7..2e4be922 100644 --- a/test/runtests.jl +++ b/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