Merge pull request #1069 from SciML/dual_func

Handle dual detection on SciMLFunctions

Author: ChrisRackauckas

Project.toml

@@ -125,9 +125,10 @@ SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [targets]
-test = ["Distributed", "Measurements", "MonteCarloMeasurements", "Unitful", "LabelledArrays", "ForwardDiff", "SparseArrays", "InteractiveUtils", "Plots", "Pkg", "Random", "StaticArrays", "SafeTestsets", "Statistics", "Test", "Distributions", "Aqua", "BenchmarkTools"]
+test = ["Distributed", "Measurements", "MonteCarloMeasurements", "Unitful", "LabelledArrays", "SymbolicIndexingInterface", "ForwardDiff", "SparseArrays", "InteractiveUtils", "Plots", "Pkg", "Random", "StaticArrays", "SafeTestsets", "Statistics", "Test", "Distributions", "Aqua", "BenchmarkTools"]

src/forwarddiff.jl

@@ -155,6 +155,15 @@ const FORWARDDIFF_AUTODETECTION_FAILURE_MESSAGE = """
                                 end
                                 ```
 
+                                To opt a type out of the dual checking, define an overload
+                                that returns Any. For example:
+
+                                ```julia
+                                function DiffEqBase.anyeltypedual(::YourType, ::Type{Val{counter}}) where {counter}
+                                    Any
+                                end
+                                ```
+
                                 If you have defined this on a common type which should
                                 be more generally supported, please open a pull request
                                 adding this dispatch. If you need help defining this dispatch,
@@ -338,6 +347,8 @@ function anyeltypedual(x::NamedTuple, ::Type{Val{counter}} = Val{0}) where {coun
     anyeltypedual(values(x))
 end
 
+DiffEqBase.anyeltypedual(f::SciMLBase.AbstractSciMLFunction, ::Type{Val{counter}}) where {counter} = Any
+
 @inline promote_u0(::Nothing, p, t0) = nothing
 
 @inline function promote_u0(u0, p, t0)

test/downstream/default_linsolve_structure.jl

@@ -6,23 +6,23 @@ jp_diag = Diagonal(zeros(2))
 fun = ODEFunction(f; jac = jac, jac_prototype = jp_diag)
 prob = ODEProblem(fun, ones(2), (1.0, 10.0))
 sol = solve(prob, Rosenbrock23())
-@test sol[end] ≈ [10.0, 10.0]
+@test sol.u[end] ≈ [10.0, 10.0]
 @test length(sol) < 60
 
 sol = solve(prob, Rosenbrock23(autodiff = false))
-@test sol[end] ≈ [10.0, 10.0]
+@test sol.u[end] ≈ [10.0, 10.0]
 @test length(sol) < 60
 
 jp = Tridiagonal(jp_diag)
 fun = ODEFunction(f; jac = jac, jac_prototype = jp)
 prob = ODEProblem(fun, ones(2), (1.0, 10.0))
 
 sol = solve(prob, Rosenbrock23())
-@test sol[end] ≈ [10.0, 10.0]
+@test sol.u[end] ≈ [10.0, 10.0]
 @test length(sol) < 60
 
 sol = solve(prob, Rosenbrock23(autodiff = false))
-@test sol[end] ≈ [10.0, 10.0]
+@test sol.u[end] ≈ [10.0, 10.0]
 @test length(sol) < 60
 
 #=
@@ -43,7 +43,7 @@ sol = solve(prob,Rosenbrock23())
     fun = ODEFunction(f; jac = jac, jac_prototype = jp)
     prob = ODEProblem(fun, ones(2), (1.0, 10.0))
     sol = solve(prob, Rosenbrock23(autodiff = false))
-    @test sol[end] ≈ [10.0, 10.0]
+    @test sol.u[end] ≈ [10.0, 10.0]
     @test length(sol) < 60
 end
 
@@ -52,6 +52,6 @@ end
     fun = ODEFunction(f; jac = jac, jac_prototype = jp)
     prob = ODEProblem(fun, ones(2), (1.0, 10.0))
     sol = solve(prob, Rosenbrock23(autodiff = false))
-    @test sol[end] ≈ [10.0, 10.0]
+    @test sol.u[end] ≈ [10.0, 10.0]
     @test length(sol) < 60
 end

test/downstream/ensemble_analysis.jl

@@ -25,7 +25,7 @@ vecarr = timeseries_steps_quantile(sim, 0.5)
 m_series, v_series = timeseries_steps_meanvar(sim)
 summ = EnsembleSummary(sim)
 
-m4, v4 = m_series[3], v_series[3]
+m4, v4 = m_series.u[3], v_series.u[3]
 covar_mat = timeseries_steps_meancov(sim)[3, 3]
 @test m ≈ m4
 @test v ≈ v4
@@ -54,7 +54,7 @@ m_series = timeseries_point_median(sim, 0:(1 // 2^(3)):1)
 m_series = timeseries_point_quantile(sim, 0.5, 0:(1 // 2^(3)):1)
 m_series, v_series = timeseries_point_meanvar(sim, 0:(1 // 2^(3)):1)
 summ = EnsembleSummary(sim, 0:(1 // 2^(3)):1)
-m5, v5 = m_series[5], v_series[5]
+m5, v5 = m_series.u[5], v_series.u[5]
 @test m ≈ m5
 @test v ≈ v5
 m6, m7, v6 = timeseries_point_meancov(sim, 0:(1 // 2^(3)):1, 0:(1 // 2^(3)):1)[5, 5]
@@ -83,7 +83,7 @@ vecarr = timeseries_steps_quantile(sim, 0.5)
 m_series, v_series = timeseries_steps_meanvar(sim)
 summ = EnsembleSummary(sim)
 
-m4, v4 = m_series[3], v_series[3]
+m4, v4 = m_series.u[3], v_series.u[3]
 covar_mat = timeseries_steps_meancov(sim)[3, 3]
 @test m ≈ m4
 @test v ≈ v4
@@ -112,7 +112,7 @@ m_series = timeseries_point_median(sim, 0:(1 // 2^(3)):1)
 m_series = timeseries_point_quantile(sim, 0.5, 0:(1 // 2^(3)):1)
 m_series, v_series = timeseries_point_meanvar(sim, 0:(1 // 2^(3)):1)
 summ = EnsembleSummary(sim, 0:(1 // 2^(3)):1)
-m5, v5 = m_series[5], v_series[5]
+m5, v5 = m_series.u[5], v_series.u[5]
 @test m ≈ m5
 @test v ≈ v5
 m6, m7, v6 = timeseries_point_meancov(sim, 0:(1 // 2^(3)):1, 0:(1 // 2^(3)):1)[5, 5]

test/downstream/tables.jl

@@ -1,11 +1,11 @@
-using OrdinaryDiffEq, DataFrames, Test
+using OrdinaryDiffEq, DataFrames, Test, SymbolicIndexingInterface
 f_2dlinear = (du, u, p, t) -> du .= 1.01u;
 prob = ODEProblem(f_2dlinear, rand(2, 2), (0.0, 1.0));
 sol1 = solve(prob, Euler(); dt = 1 // 2^(4));
 df = DataFrame(sol1)
 @test names(df) == ["timestamp", "value1", "value2", "value3", "value4"]
 
-prob = ODEProblem(ODEFunction(f_2dlinear, syms = [:a, :b, :c, :d]), rand(2, 2), (0.0, 1.0));
+prob = ODEProblem(ODEFunction(f_2dlinear, sys = SymbolicIndexingInterface.SymbolCache([:a, :b, :c, :d], [], :t)), rand(2, 2), (0.0, 1.0));
 sol2 = solve(prob, Euler(); dt = 1 // 2^(4));
 df = DataFrame(sol2)
 @test names(df) == ["timestamp", "a", "b", "c", "d"]

test/forwarddiff_dual_detection.jl

@@ -347,3 +347,4 @@ prob = ODEProblem{false}(f, u0, tspan)
 foo = SciMLBase.build_solution(
     prob, DiffEqBase.InternalEuler.FwdEulerAlg(), [u0, u0], [0.0, 1.0])
 DiffEqBase.anyeltypedual((; x = foo))
+DiffEqBase.anyeltypedual((; x = foo, y = prob.f))