Merge pull request #2730 from oscardssmith/os/fix-FBDF-ad

fix autodiff for FBDF

Author: oscardssmith

lib/OrdinaryDiffEqBDF/src/bdf_caches.jl

@@ -360,7 +360,7 @@ function alg_cache(alg::QNDF{MO}, u, rate_prototype, ::Type{uEltypeNoUnits},
         ::Val{false}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits
 } where {MO}
     max_order = MO
-    γ, c = one(uEltypeNoUnits), 1
+    γ, c = 1//1, 1
     nlsolver = build_nlsolver(alg, u, uprev, p, t, dt, f, rate_prototype, uEltypeNoUnits,
         uBottomEltypeNoUnits, tTypeNoUnits, γ, c, Val(false))
     dtprev = one(dt)
@@ -426,7 +426,7 @@ function alg_cache(alg::QNDF{MO}, u, rate_prototype, ::Type{uEltypeNoUnits},
         ::Val{true}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits
 } where {MO}
     max_order = MO
-    γ, c = one(eltype(alg.kappa)), 1
+    γ, c = 1//1, 1
     nlsolver = build_nlsolver(alg, u, uprev, p, t, dt, f, rate_prototype, uEltypeNoUnits,
         uBottomEltypeNoUnits, tTypeNoUnits, γ, c, Val(true))
     fsalfirst = zero(rate_prototype)
@@ -541,7 +541,7 @@ function alg_cache(alg::FBDF{MO}, u, rate_prototype, ::Type{uEltypeNoUnits},
         dt, reltol, p, calck,
         ::Val{false}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits
 } where {MO}
-    γ, c = one(uEltypeNoUnits), 1
+    γ, c = 1//1, 1
     max_order = MO
     nlsolver = build_nlsolver(alg, u, uprev, p, t, dt, f, rate_prototype, uEltypeNoUnits,
         uBottomEltypeNoUnits, tTypeNoUnits, γ, c, Val(false))
@@ -616,7 +616,7 @@ function alg_cache(alg::FBDF{MO}, u, rate_prototype, ::Type{uEltypeNoUnits},
         dt, reltol, p, calck,
         ::Val{true}) where {MO, uEltypeNoUnits, uBottomEltypeNoUnits,
         tTypeNoUnits}
-    γ, c = one(uEltypeNoUnits), 1
+    γ, c = 1//1, 1
     fsalfirst = zero(rate_prototype)
     max_order = MO
     nlsolver = build_nlsolver(alg, u, uprev, p, t, dt, f, rate_prototype, uEltypeNoUnits,

lib/OrdinaryDiffEqBDF/test/bdf_regression_tests.jl

@@ -1,10 +1,10 @@
-using OrdinaryDiffEqBDF, Test
+using OrdinaryDiffEqBDF, ForwardDiff, Test
 
-foop = (u, p, t) -> u
-proboop = ODEProblem(foop, ones(2), (0.0, 1000.0))
+foop = (u, p, t) -> u * p
+proboop = ODEProblem(foop, ones(2), (0.0, 1000.0), 1.0)
 
-fiip = (du, u, p, t) -> du .= u
-probiip = ODEProblem(fiip, ones(2), (0.0, 1000.0))
+fiip = (du, u, p, t) -> du .= u .* p
+probiip = ODEProblem(fiip, ones(2), (0.0, 1000.0), 1.0)
 
 @testset "FBDF reinit" begin
     for prob in [proboop, probiip]
@@ -18,3 +18,19 @@ probiip = ODEProblem(fiip, ones(2), (0.0, 1000.0))
         @test integ.sol.t[end] >= 700
     end
 end
+
+function ad_helper(alg, prob)
+    function costoop(p)
+        _oprob = remake(prob; p)
+        sol = solve(_oprob, alg, saveat=1:10)
+        return sum(sol)
+    end
+end
+
+@testset "parameter autodiff" begin
+    for prob in [proboop, probiip]
+        for alg in [FBDF(), QNDF()]
+            ForwardDiff.derivative(ad_helper(alg, prob), 1.0)
+        end
+    end
+end

lib/OrdinaryDiffEqBDF/test/dae_ad_tests.jl

@@ -14,6 +14,16 @@ function f(du, u, p, t)
         +p[1] * u[1] - p[2] * u[2]^2 - p[3] * u[2] * u[3] - du[2],
         u[1] + u[2] + u[3] - 1.0]
 end
+function f_ode(du, u, p, t)
+    du .= [-p[1] * u[1] + p[3] * u[2] * u[3],
+        +p[1] * u[1] - p[2] * u[2]^2 - p[3] * u[2] * u[3],
+        u[1] + u[2] + u[3] - 1.0]
+end
+function f_ode(u, p, t)
+    [-p[1] * u[1] + p[3] * u[2] * u[3],
+        +p[1] * u[1] - p[2] * u[2]^2 - p[3] * u[2] * u[3],
+        u[1] + u[2] + u[3] - 1.0]
+end
 p = [0.04, 3e7, 1e4]
 u₀ = [1.0, 0, 0]
 du₀ = [-0.04, 0.04, 0.0]
@@ -22,16 +32,18 @@ differential_vars = [true, true, false]
 M = Diagonal([1.0, 1.0, 0.0])
 prob = DAEProblem(f, du₀, u₀, tspan, p, differential_vars = differential_vars)
 prob_oop = DAEProblem{false}(f, du₀, u₀, tspan, p, differential_vars = differential_vars)
-f_mm = ODEFunction{true, SciMLBase.AutoSpecialize}(f, mass_matrix = M)
+f_mm = ODEFunction{true}(f_ode, mass_matrix = M)
 prob_mm = ODEProblem(f_mm, u₀, tspan, p)
+f_mm_oop = ODEFunction{false}(f_ode, mass_matrix = M)
+prob_mm_oop = ODEProblem(f_mm_oop, u₀, tspan, p)
 @test_broken sol1 = @inferred solve(prob, DFBDF(autodiff=afd_cs3), dt = 1e-5, abstol = 1e-8, reltol = 1e-8)
 @test_broken sol2 = @inferred solve(prob_oop, DFBDF(autodiff=afd_cs3), dt = 1e-5, abstol = 1e-8, reltol = 1e-8)
 @test_broken sol3 = @inferred solve(prob_mm, FBDF(autodiff=afd_cs3), 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
 @testset "Inplace: $(isinplace(_prob)), DAEProblem: $(_prob isa DAEProblem), BrownBasic: $(initalg isa BrownFullBasicInit), Autodiff: $autodiff" for _prob in [
-        prob, prob_oop, prob_mm],
+        prob, prob_oop, prob_mm, prob_mm_oop],
     initalg in [BrownFullBasicInit(), ShampineCollocationInit()], autodiff in [afd_cs3, AutoFiniteDiff()]
 
     alg = (_prob isa DAEProblem) ? DFBDF(; autodiff) : FBDF(; autodiff)
@@ -40,9 +52,5 @@ prob_mm = ODEProblem(f_mm, u₀, tspan, p)
             reltol = 1e-14, initializealg = initalg)
         sum(sol)
     end
-    if _prob isa DAEProblem
-        @test ForwardDiff.gradient(f, [0.04, 3e7, 1e4])≈[0, 0, 0] atol=1e-8
-    else
-        @test_broken ForwardDiff.gradient(f, [0.04, 3e7, 1e4])≈[0, 0, 0] atol=1e-8
-    end
+    @test ForwardDiff.gradient(f, [0.04, 3e7, 1e4])≈[0, 0, 0] atol=1e-8
 end