From 5c87095e9f9da20ff40c1c56360a41db61f4892e Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Fri, 18 Jul 2025 14:56:23 -0400 Subject: [PATCH 1/5] Reuse integrators in nlstep and NonlinearSolveAlg integrations This is a pretty major performance boost --- lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl | 5 ++--- test/modelingtoolkit/nlstep_tests.jl | 8 ++++---- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl b/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl index 3e593b7766..981f687da4 100644 --- a/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl +++ b/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl @@ -76,15 +76,14 @@ function initialize!(nlsolver::NLSolver{<:NonlinearSolveAlg, true}, nlstep_data.set_outer_tmp(nlstep_data.nlprob, zero(z)) 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 diff --git a/test/modelingtoolkit/nlstep_tests.jl b/test/modelingtoolkit/nlstep_tests.jl index 132accd5f8..6d3d96a88b 100644 --- a/test/modelingtoolkit/nlstep_tests.jl +++ b/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,7 +60,7 @@ 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) @@ -68,7 +68,7 @@ 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)); From a5d70bad7f8d4d6704b9fab56a798e702fa0eccc Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Fri, 18 Jul 2025 15:02:53 -0400 Subject: [PATCH 2/5] Reduce loop allocations --- lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl b/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl index 981f687da4..996e190621 100644 --- a/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl +++ b/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl @@ -66,14 +66,15 @@ function initialize!(nlsolver::NLSolver{<:NonlinearSolveAlg, true}, nlstep_data = f.nlstep_data if nlstep_data !== nothing + ztmp .= 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, ztmp) 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, ztmp) end nlstep_data.nlprob.u0 .= @view z[nlstep_data.u0perm] SciMLBase.reinit!(cache.cache, nlstep_data.nlprob.u0, p=nlstep_data.nlprob.p) From 1a837c3f92be936eff367c4a8f38eb528a3e5909 Mon Sep 17 00:00:00 2001 From: oscarddssmith Date: Thu, 7 Aug 2025 13:06:48 -0400 Subject: [PATCH 3/5] switch tmp and inplace probmap --- lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl b/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl index 996e190621..a36a05043e 100644 --- a/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl +++ b/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl @@ -66,21 +66,21 @@ function initialize!(nlsolver::NLSolver{<:NonlinearSolveAlg, true}, nlstep_data = f.nlstep_data if nlstep_data !== nothing - ztmp .= 0 + 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, ztmp) + 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, ztmp) + nlstep_data.set_outer_tmp(nlstep_data.nlprob, atmp) end nlstep_data.nlprob.u0 .= @view z[nlstep_data.u0perm] 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) + nlp_params = (tmp, atmp, ustep, γ, α, tstep, k, invγdt, p, dt, f) else nlp_params = (tmp, ustep, γ, α, tstep, k, invγdt, method, p, dt, f) 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]), From 3d1a6bad13ad0ba19435de8572dd81045f736c40 Mon Sep 17 00:00:00 2001 From: oscarddssmith Date: Thu, 7 Aug 2025 22:43:07 -0600 Subject: [PATCH 4/5] it works --- lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl b/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl index a36a05043e..d12eae3f20 100644 --- a/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl +++ b/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 @@ -80,7 +80,7 @@ function initialize!(nlsolver::NLSolver{<:NonlinearSolveAlg, true}, SciMLBase.reinit!(cache.cache, nlstep_data.nlprob.u0, p=nlstep_data.nlprob.p) else if f isa DAEFunction - nlp_params = (tmp, atmp, ustep, γ, α, tstep, k, invγdt, p, dt, f) + 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 From 9ae0860dfe99f7000e5f1ad681477d779320fea1 Mon Sep 17 00:00:00 2001 From: Oscar Smith Date: Thu, 2 Oct 2025 12:32:33 -0400 Subject: [PATCH 5/5] bump compat --- lib/OrdinaryDiffEqNonlinearSolve/Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/OrdinaryDiffEqNonlinearSolve/Project.toml b/lib/OrdinaryDiffEqNonlinearSolve/Project.toml index 1e621a32c7..d12ca68235 100644 --- a/lib/OrdinaryDiffEqNonlinearSolve/Project.toml +++ b/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"