Add some inference checks to DAE tests, make type stable

Author: Ickaser

lib/OrdinaryDiffEqBDF/test/dae_ad_tests.jl

@@ -18,8 +18,8 @@ tspan = (0.0, 100000.0)
 differential_vars = [true, true, false]
 prob = DAEProblem(f, du₀, u₀, tspan, p, differential_vars = differential_vars)
 prob_oop = DAEProblem{false}(f, du₀, u₀, tspan, p, differential_vars = differential_vars)
-sol1 = solve(prob, DFBDF(), dt = 1e-5, abstol = 1e-8, reltol = 1e-8)
-sol2 = solve(prob_oop, DFBDF(), dt = 1e-5, abstol = 1e-8, reltol = 1e-8)
+sol1 = @inferred solve(prob, DFBDF(), dt = 1e-5, abstol = 1e-8, reltol = 1e-8)
+sol2 = @inferred solve(prob_oop, DFBDF(), dt = 1e-5, abstol = 1e-8, reltol = 1e-8)
 
 # These tests flex differentiation of the solver and through the initialization
 # To only test the solver part and isolate potential issues, set the initialization to consistent
@@ -29,7 +29,7 @@ sol2 = solve(prob_oop, DFBDF(), dt = 1e-5, abstol = 1e-8, reltol = 1e-8)
 
     alg = DFBDF(; autodiff)
     function f(p)
-        sol = solve(remake(_prob, p = p), alg, abstol = 1e-14,
+        @inferred sol = solve(remake(_prob, p = p), alg, abstol = 1e-14,
             reltol = 1e-14, initializealg = initalg)
         sum(sol)
     end

lib/OrdinaryDiffEqRosenbrock/test/dae_rosenbrock_ad_tests.jl

@@ -1,5 +1,6 @@
 using OrdinaryDiffEqRosenbrock, LinearAlgebra, ForwardDiff, Test
 using OrdinaryDiffEqNonlinearSolve: BrownFullBasicInit, ShampineCollocationInit
+using ADTypes: AutoForwardDiff, AutoFiniteDiff
 
 function rober(du, u, p, t)
     y₁, y₂, y₃ = u
@@ -19,22 +20,23 @@ end
 M = [1.0 0 0
      0 1.0 0
      0 0 0]
-roberf = ODEFunction(rober, mass_matrix = M)
-roberf_oop = ODEFunction{false}(rober, mass_matrix = M)
+# M = Diagonal([1.0, 1.0, 0.0])
+roberf = ODEFunction{true, SciMLBase.AutoSpecialize}(rober, mass_matrix = M)
+roberf_oop = ODEFunction{false, SciMLBase.AutoSpecialize}(rober, mass_matrix = M)
 prob_mm = ODEProblem(roberf, [1.0, 0.0, 0.2], (0.0, 1e5), (0.04, 3e7, 1e4))
 prob_mm_oop = ODEProblem(roberf_oop, [1.0, 0.0, 0.2], (0.0, 1e5), (0.04, 3e7, 1e4))
-sol = solve(prob_mm, Rodas5P(), reltol = 1e-8, abstol = 1e-8)
-sol = solve(prob_mm_oop, Rodas5P(), reltol = 1e-8, abstol = 1e-8)
+sol = @inferred solve(prob_mm, Rodas5P(), reltol = 1e-8, abstol = 1e-8)
+sol = @inferred solve(prob_mm_oop, Rodas5P(), reltol = 1e-8, abstol = 1e-8)
 
 # These tests flex differentiation of the solver and through the initialization
 # To only test the solver part and isolate potential issues, set the initialization to consistent
 @testset "Inplace: $(isinplace(_prob)), DAEProblem: $(_prob isa DAEProblem), BrownBasic: $(initalg isa BrownFullBasicInit), Autodiff: $autodiff" for _prob in [
         prob_mm, prob_mm_oop],
-    initalg in [BrownFullBasicInit(), ShampineCollocationInit()], autodiff in [true, false]
+    initalg in [BrownFullBasicInit(), ShampineCollocationInit()], autodiff in [AutoForwardDiff(chunksize=3), AutoFiniteDiff()]
 
     alg = Rodas5P(; autodiff)
     function f(p)
-        sol = solve(remake(_prob, p = p), alg, abstol = 1e-14,
+        sol = @inferred solve(remake(_prob, p = p), alg, abstol = 1e-14,
             reltol = 1e-14, initializealg = initalg)
         sum(sol)
     end

test/interface/mass_matrix_tests.jl

@@ -1,6 +1,7 @@
 using OrdinaryDiffEq, Test, LinearAlgebra, Statistics
 using OrdinaryDiffEqCore
 using OrdinaryDiffEqNonlinearSolve: NLFunctional, NLAnderson, NLNewton
+using LinearAlgebra: Diagonal
 
 # create mass matrix problems
 function make_mm_probs(mm_A, ::Val{iip}) where {iip}
@@ -194,11 +195,10 @@ end
     u0 = [0.0, 1.0]
     tspan = (0.0, 1.0)
 
-    M = fill(0.0, 2, 2)
-    M[1, 1] = 1.0
+    M = Diagonal([1.0, 0.0])
 
     m_ode_prob = ODEProblem(ODEFunction(f!; mass_matrix = M), u0, tspan)
-    @test_nowarn sol = solve(m_ode_prob, Rosenbrock23())
+    @test_nowarn sol = @inferred solve(m_ode_prob, Rosenbrock23())
 
     M = [0.637947 0.637947
          0.637947 0.637947]
@@ -323,14 +323,14 @@ function dynamics(u, p, t)
 end
 
 x0 = zeros(n, n)
-M = zeros(n * n) |> Diagonal |> Matrix
+M = zeros(n * n) |> Diagonal
 M[1, 1] = true # zero mass matrix breaks rosenbrock
-f = ODEFunction(dynamics!, mass_matrix = M)
+f = ODEFunction{true, SciMLBase.AutoSpecialize}(dynamics!, mass_matrix = M)
 tspan = (0, 10.0)
 prob = ODEProblem(f, x0, tspan)
-foop = ODEFunction(dynamics, mass_matrix = M)
+foop = ODEFunction{false, SciMLBase.AutoSpecialize}(dynamics, mass_matrix = M)
 proboop = ODEProblem(f, x0, tspan)
-sol = solve(prob, Rosenbrock23())
-sol = solve(prob, Rodas4(), initializealg = ShampineCollocationInit())
-sol = solve(proboop, Rodas5())
-sol = solve(proboop, Rodas4(), initializealg = ShampineCollocationInit())
+sol = @inferred solve(prob, Rosenbrock23())
+sol = @inferred solve(prob, Rodas4(), initializealg = ShampineCollocationInit())
+sol = @inferred solve(proboop, Rodas5())
+sol = @inferred solve(proboop, Rodas4(), initializealg = ShampineCollocationInit())