From acaca2139120d1756d84fefeb1ea903998f2a60b Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Fri, 18 Jul 2025 14:56:23 -0400 Subject: [PATCH 1/4] 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 2ca55cf6b8..06a226cdbe 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 b9dd04e62dfd224445a3108dfa16baadda72fef8 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Fri, 18 Jul 2025 15:02:53 -0400 Subject: [PATCH 2/4] 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 06a226cdbe..5f120f1c62 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 e115f0daaea9ae13dec9154d57563bd00a3e7b25 Mon Sep 17 00:00:00 2001 From: oscarddssmith Date: Thu, 7 Aug 2025 13:06:48 -0400 Subject: [PATCH 3/4] 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 5f120f1c62..e588912a42 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 afa99d6ca1b4529c6b28618bd18a4e5f5177879a Mon Sep 17 00:00:00 2001 From: oscarddssmith Date: Thu, 7 Aug 2025 22:43:07 -0600 Subject: [PATCH 4/4] 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 e588912a42..941867905d 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 DiffEqBase.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