Merge pull request #2760 from SciML/reinit

Reuse integrators in nlstep and NonlinearSolveAlg integrations

Author: ChrisRackauckas

lib/OrdinaryDiffEqNonlinearSolve/Project.toml

@@ -37,7 +37,7 @@ AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 
 [compat]
-NonlinearSolve = "4.10"
+NonlinearSolve = "4.12"
 ForwardDiff = "0.10.38, 1"
 Test = "<0.0.1, 1"
 FastBroadcast = "0.3"

lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl

@@ -56,7 +56,7 @@ function initialize!(nlsolver::NLSolver{<:NonlinearSolveAlg, true},
     cache.invγdt = inv(dt * nlsolver.γ)
     cache.tstep = integrator.t + nlsolver.c * dt
 
-    @unpack ustep, tstep, k, invγdt = cache
+    @unpack ustep, atmp, tstep, k, invγdt = cache
 
     if SciMLBase.has_stats(integrator)
         integrator.stats.nf += cache.cache.stats.nf
@@ -66,25 +66,25 @@ function initialize!(nlsolver::NLSolver{<:NonlinearSolveAlg, true},
 
     nlstep_data = f.nlstep_data
     if nlstep_data !== nothing
+        atmp .= 0
         if method === COEFFICIENT_MULTISTEP
             nlstep_data.set_γ_c(nlstep_data.nlprob, (one(t), one(t), α * invγdt, tstep))
-            nlstep_data.set_inner_tmp(nlstep_data.nlprob, zero(z))
+            nlstep_data.set_inner_tmp(nlstep_data.nlprob, atmp)
             nlstep_data.set_outer_tmp(nlstep_data.nlprob, tmp)
         else
             nlstep_data.set_γ_c(nlstep_data.nlprob, (dt, γ, one(t), tstep))
             nlstep_data.set_inner_tmp(nlstep_data.nlprob, tmp)
-            nlstep_data.set_outer_tmp(nlstep_data.nlprob, zero(z))
+            nlstep_data.set_outer_tmp(nlstep_data.nlprob, atmp)
         end
         nlstep_data.nlprob.u0 .= @view z[nlstep_data.u0perm]
-        cache.cache = init(nlstep_data.nlprob, alg.alg)
+        SciMLBase.reinit!(cache.cache, nlstep_data.nlprob.u0, p=nlstep_data.nlprob.p)
     else
         if f isa DAEFunction
             nlp_params = (tmp, ztmp, ustep, γ, α, tstep, k, invγdt, p, dt, f)
         else
             nlp_params = (tmp, ustep, γ, α, tstep, k, invγdt, method, p, dt, f)
         end
-        new_prob = remake(cache.prob, p = nlp_params, u0 = z)
-        cache.cache = init(new_prob, alg.alg)
+        SciMLBase.reinit!(cache.cache, z, p=nlp_params)
     end
     nothing
 end
@@ -127,7 +127,7 @@ end
             nlcache.prob, nlcache.alg, nlcache.u, nlcache.fu;
             nlcache.retcode, nlcache.stats, nlcache.trace
         )
-        ztmp .= nlstep_data.nlprobmap(nlstepsol)
+        nlstep_data.nlprobmap(ztmp, nlstepsol)
         ustep = compute_ustep!(ustep, tmp, γ, z, method)
         calculate_residuals!(@view(atmp[nlstep_data.u0perm]), nlcache.fu, 
                              @view(uprev[nlstep_data.u0perm]), 

test/modelingtoolkit/nlstep_tests.jl

@@ -10,7 +10,7 @@ using Test
 eqs = [D(y₁) ~ -k₁ * y₁ + k₃ * y₂ * y₃,
     D(y₂) ~ k₁ * y₁ - k₂ * y₂^2 - k₃ * y₂ * y₃,
     D(y₃) ~ k₂ * y₂^2]
-@mtkbuild rober = ODESystem(eqs, t)
+@mtkcompile rober = ODESystem(eqs, t)
 prob = ODEProblem(rober, [[y₁, y₂, y₃] .=> [1.0; 0.0; 0.0]; [k₁, k₂, k₃] .=> (0.04, 3e7, 1e4)], (0.0, 1e5), jac = true)
 prob2 = ODEProblem(rober, [[y₁, y₂, y₃] .=> [1.0; 0.0; 0.0]; [k₁, k₂, k₃] .=> (0.04, 3e7, 1e4)], (0.0, 1e5), jac = true, nlstep = true)
 
@@ -29,7 +29,7 @@ sol1 = solve(prob, TRBDF2(autodiff=AutoFiniteDiff(), nlsolve = nlalg));
 sol2 = solve(prob2, TRBDF2(autodiff=AutoFiniteDiff(), nlsolve = nlalg));
 
 @test sol1.t != sol2.t
-@test sol1 != sol2
+@test sol1.u != sol2.u
 @test sol1(sol1.t) ≈ sol2(sol1.t) atol=1e-3
 
 testprob = ODEProblem(rober, [[y₁, y₂, y₃] .=> [1.0; 0.0; 0.0]; [k₁, k₂, k₃] .=> (0.04, 3e7, 1e4)], (0.0, 1.0), nlstep = true)
@@ -60,15 +60,15 @@ sim = analyticless_test_convergence(dts, testprob, FBDF(autodiff=AutoFiniteDiff(
 eqs_nonaut = [D(y₁) ~ -k₁ * y₁ + (t+1) * k₃ * y₂ * y₃,
     D(y₂) ~ k₁ * y₁ - (t+1) * k₂ * y₂^2 - (t+1) * k₃ * y₂ * y₃,
     D(y₃) ~ (t+1) * k₂ * y₂^2]
-@mtkbuild rober_nonaut = ODESystem(eqs_nonaut, t)
+@mtkcompile rober_nonaut = ODESystem(eqs_nonaut, t)
 prob = ODEProblem(rober_nonaut, [[y₁, y₂, y₃] .=> [1.0; 0.0; 0.0]; [k₁, k₂, k₃] .=> (0.04, 3e7, 1e4)], (0.0, 1e5), jac = true)
 prob2 = ODEProblem(rober_nonaut, [[y₁, y₂, y₃] .=> [1.0; 0.0; 0.0]; [k₁, k₂, k₃] .=> (0.04, 3e7, 1e4)], (0.0, 1e5), jac = true, nlstep = true)
 
 sol1 = solve(prob, FBDF(autodiff=AutoFiniteDiff(), nlsolve = nlalg));
 sol2 = solve(prob2, FBDF(autodiff=AutoFiniteDiff(), nlsolve = nlalg));
 
 @test sol1.t != sol2.t
-@test sol1 != sol2
+@test sol1.u != sol2.u
 @test sol1(sol1.t) ≈ sol2(sol1.t) atol=1e-3
 
 sol1 = solve(prob, TRBDF2(autodiff=AutoFiniteDiff(), nlsolve = nlalg));