Merge pull request #2729 from SciML/nonlinearsolve_inference

Fix inference with NonlinearSolveAlg

Author: ChrisRackauckas

lib/OrdinaryDiffEqBDF/Project.toml

@@ -63,6 +63,7 @@ DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
+NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 ODEProblemLibrary = "fdc4e326-1af4-4b90-96e7-779fcce2daa5"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
@@ -72,4 +73,4 @@ JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 
 [targets]
-test = ["DiffEqDevTools", "ForwardDiff", "Random", "SafeTestsets", "Test", "ODEProblemLibrary", "StaticArrays", "Enzyme", "LinearSolve", "JET", "Aqua"]
+test = ["DiffEqDevTools", "ForwardDiff", "Random", "SafeTestsets", "Test", "ODEProblemLibrary", "StaticArrays", "Enzyme", "LinearSolve", "JET", "Aqua", "NonlinearSolve"]

lib/OrdinaryDiffEqBDF/src/OrdinaryDiffEqBDF.jl

@@ -32,7 +32,8 @@ import OrdinaryDiffEqCore
 using OrdinaryDiffEqDifferentiation: UJacobianWrapper
 using OrdinaryDiffEqNonlinearSolve: NLNewton, du_alias_or_new, build_nlsolver,
                                     nlsolve!, nlsolvefail, isnewton, markfirststage!,
-                                    set_new_W!, DIRK, compute_step!, COEFFICIENT_MULTISTEP
+                                    set_new_W!, DIRK, compute_step!, COEFFICIENT_MULTISTEP,
+                                    NonlinearSolveAlg
 import ADTypes
 import ADTypes: AutoForwardDiff, AutoFiniteDiff, AbstractADType
 

lib/OrdinaryDiffEqBDF/test/inference_tests.jl

@@ -0,0 +1,13 @@
+using OrdinaryDiffEqBDF, NonlinearSolve, Test
+
+prob = ODEProblem((du,u,p,t) -> du .= u, zeros(1), (0.0,1.0))
+nlalg = FBDF(autodiff=false, nlsolve = OrdinaryDiffEqBDF.NonlinearSolveAlg(TrustRegion(autodiff = AutoFiniteDiff())))
+basicalg = FBDF(autodiff=false)
+basicalgad = FBDF()
+
+nlsolver = @inferred OrdinaryDiffEqBDF.build_nlsolver(basicalg, prob.u0, prob.u0, prob.p, 0.0, 0.0, prob.f, prob.u0, Float64,
+        Float64, Float64, 0.0, 0.0, Val(true))
+nlsolver = @inferred OrdinaryDiffEqBDF.build_nlsolver(nlalg, prob.u0, prob.u0, prob.p, 0.0, 0.0, prob.f, prob.u0, Float64,
+        Float64, Float64, 0.0, 0.0, Val(true))
+nlsolver = @test_throws Any @inferred OrdinaryDiffEqBDF.build_nlsolver(basicalgad, prob.u0, prob.u0, prob.p, 0.0, 0.0, prob.f, prob.u0, Float64,
+        Float64, Float64, 0.0, 0.0, Val(true))

lib/OrdinaryDiffEqBDF/test/runtests.jl

@@ -5,6 +5,7 @@ using SafeTestsets
 @time @safetestset "DAE Event Tests" include("dae_event.jl")
 @time @safetestset "DAE Initialization Tests" include("dae_initialization_tests.jl")
 
+@time @safetestset "BDF Inference Tests" include("inference_tests.jl")
 @time @safetestset "BDF Convergence Tests" include("bdf_convergence_tests.jl")
 @time @safetestset "BDF Regression Tests" include("bdf_regression_tests.jl")
 

lib/OrdinaryDiffEqNonlinearSolve/src/utils.jl

@@ -153,6 +153,17 @@ function build_nlsolver(alg, u, uprev, p, t, dt, f::F, rate_prototype,
         uBottomEltypeNoUnits, tTypeNoUnits, γ, c, α, iip)
 end
 
+function daenlf(ztmp, z, p)
+    tmp, ustep, γ, α, tstep, k, invγdt, _p, dt, f = p
+    _compute_rhs!(tmp, ztmp, ustep, γ, α, tstep, k, invγdt, _p, dt, f, z)[1]
+end
+
+function odenlf(ztmp, z, p)
+    tmp, ustep, γ, α, tstep, k, invγdt, method, _p, dt, f = p
+    _compute_rhs!(
+    tmp, ztmp, ustep, γ, α, tstep, k, invγdt, method, _p, dt, f, z)[1]
+end
+
 function build_nlsolver(
         alg, nlalg::Union{NLFunctional, NLAnderson, NLNewton, NonlinearSolveAlg},
         u, uprev, p, t, dt,
@@ -215,19 +226,11 @@ function build_nlsolver(
         if nlalg isa NonlinearSolveAlg
             α = tTypeNoUnits(α)
             dt = tTypeNoUnits(dt)
-            if isdae
-                nlf = (ztmp, z, p) -> begin
-                    tmp, ustep, γ, α, tstep, k, invγdt, _p, dt, f = p
-                    _compute_rhs!(tmp, ztmp, ustep, γ, α, tstep, k, invγdt, _p, dt, f, z)[1]
-                end
-                nlp_params = (tmp, ustep, γ, α, tstep, k, invγdt, p, dt, f)
+            nlf = isdae ? daenlf : odenlf
+            nlp_params = if isdae
+                (tmp, ustep, γ, α, tstep, k, invγdt, p, dt, f)
             else
-                nlf = (ztmp, z, p) -> begin
-                    tmp, ustep, γ, α, tstep, k, invγdt, method, _p, dt, f = p
-                    _compute_rhs!(
-                        tmp, ztmp, ustep, γ, α, tstep, k, invγdt, method, _p, dt, f, z)[1]
-                end
-                nlp_params = (tmp, ustep, γ, α, tstep, k, invγdt, DIRK, p, dt, f)
+                (tmp, ustep, γ, α, tstep, k, invγdt, DIRK, p, dt, f)
             end
             prob = NonlinearProblem(NonlinearFunction{true}(nlf), ztmp, nlp_params)
             cache = init(prob, nlalg.alg)
@@ -262,6 +265,16 @@ function build_nlsolver(
         Divergence, nlcache)
 end
 
+function oopdaenlf(z, p)
+    tmp, α, tstep, invγdt, _p, dt, uprev, f = p
+    _compute_rhs(tmp, α, tstep, invγdt, p, dt, uprev, f, z)[1]
+end
+
+function oopodenlf(z, p)
+    tmp, γ, α, tstep, invγdt, method, _p, dt, f = p
+    _compute_rhs(tmp, γ, α, tstep, invγdt, method, _p, dt, f, z)[1]
+end
+
 function build_nlsolver(
         alg, nlalg::Union{NLFunctional, NLAnderson, NLNewton, NonlinearSolveAlg},
         u, uprev, p,
@@ -300,18 +313,11 @@ function build_nlsolver(
         if nlalg isa NonlinearSolveAlg
             α = tTypeNoUnits(α)
             dt = tTypeNoUnits(dt)
-            if isdae
-                nlf = (z, p) -> begin
-                    tmp, α, tstep, invγdt, _p, dt, uprev, f = p
-                    _compute_rhs(tmp, α, tstep, invγdt, p, dt, uprev, f, z)[1]
-                end
-                nlp_params = (tmp, α, tstep, invγdt, p, dt, uprev, f)
+            nlf = isdae ? oopdaenlf : oopodenlf
+            nlp_params = if isdae
+                (tmp, α, tstep, invγdt, p, dt, uprev, f)
             else
-                nlf = (z, p) -> begin
-                    tmp, γ, α, tstep, invγdt, method, _p, dt, f = p
-                    _compute_rhs(tmp, γ, α, tstep, invγdt, method, _p, dt, f, z)[1]
-                end
-                nlp_params = (tmp, γ, α, tstep, invγdt, DIRK, p, dt, f)
+                (tmp, γ, α, tstep, invγdt, DIRK, p, dt, f)
             end
             prob = NonlinearProblem(NonlinearFunction{false}(nlf), copy(ztmp), nlp_params)
             cache = init(prob, nlalg.alg)