diff --git a/lib/OrdinaryDiffEqNonlinearSolve/src/OrdinaryDiffEqNonlinearSolve.jl b/lib/OrdinaryDiffEqNonlinearSolve/src/OrdinaryDiffEqNonlinearSolve.jl index 3a345be247..fc7b23e67c 100644 --- a/lib/OrdinaryDiffEqNonlinearSolve/src/OrdinaryDiffEqNonlinearSolve.jl +++ b/lib/OrdinaryDiffEqNonlinearSolve/src/OrdinaryDiffEqNonlinearSolve.jl @@ -14,7 +14,8 @@ import DiffEqBase: OrdinaryDiffEqTag, calculate_residuals, calculate_residuals!, BrownFullBasicInit, ShampineCollocationInit import ConstructionBase import PreallocationTools: DiffCache, get_tmp -using SimpleNonlinearSolve: SimpleTrustRegion, SimpleGaussNewton +using SimpleNonlinearSolve: SimpleTrustRegion, SimpleGaussNewton, + AbstractSimpleNonlinearSolveAlgorithm using NonlinearSolve: FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg, NewtonRaphson, step! using MuladdMacro: @muladd @@ -56,7 +57,7 @@ import OrdinaryDiffEqCore: _initialize_dae!, import OrdinaryDiffEqDifferentiation: update_W!, is_always_new, build_uf, build_J_W, WOperator, StaticWOperator, wrapprecs, build_jac_config, dolinsolve, - resize_jac_config!, jacobian2W!, jacobian! + resize_jac_config!, jacobian2W!, jacobian!, jacobian import StaticArraysCore: StaticArray diff --git a/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl b/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl index 149049e6ec..12c10218b3 100644 --- a/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl +++ b/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl @@ -41,11 +41,29 @@ function initialize!( cache.tstep = integrator.t + nlsolver.c * dt (; ustep, tstep, k, invγdt) = cache - if SciMLBase.has_stats(integrator) + if SciMLBase.has_stats(integrator) && hasfield(typeof(cache.cache), :stats) integrator.stats.nf += cache.cache.stats.nf integrator.stats.njacs += cache.cache.stats.njacs integrator.stats.nsolve += cache.cache.stats.nsolve end + + if cache.W !== nothing + dtgamma = method === DIRK ? γ * dt : γ * dt / α + W_γdt = cache.W_γdt + first_call = iszero(W_γdt) + should_update = first_call || alg.always_new || + nlsolver.status === Divergence || + abs(inv(dtgamma) / inv(W_γdt) - 1) > alg.new_W_dt_cutoff + if should_update + _update_nlsolvealg_W_oop!(cache, integrator, dtgamma, tstep) + cache.W_γdt = dtgamma + integrator.stats.nw += 1 + cache.new_W = true + else + cache.new_W = false + end + end + if f isa DAEFunction nlp_params = (tmp, α, tstep, invγdt, p, dt, uprev, f) else @@ -153,6 +171,22 @@ function _update_nlsolvealg_W!(nlcache, integrator, dtgamma, tstep, new_jac = tr return nothing end +function _update_nlsolvealg_W_oop!(nlcache, integrator, dtgamma, tstep) + (; f, p, uprev, alg) = integrator + mass_matrix = f.mass_matrix + uf = nlcache.uf + J_new = if SciMLBase.has_jac(f) + f.jac(uprev, p, tstep) + else + uf.f = nlsolve_f(f, alg) + uf.p = p + uf.t = tstep + jacobian(uf, uprev, integrator) + end + nlcache.W[] = StaticWOperator(J_new - mass_matrix * inv(dtgamma)) + return nothing +end + ## compute_step! @muladd function compute_step!(nlsolver::NLSolver{<:NonlinearSolveAlg, false}, integrator) @@ -701,7 +735,8 @@ function Base.resize!(nlcache::NonlinearSolveCache, ::AbstractNLSolver, integrat nlcache.du1 === nothing || resize!(nlcache.du1, i) nlcache.weight === nothing || resize!(nlcache.weight, i) nlcache.jac_config === nothing || resize_jac_config!(nlcache, integrator) - nlcache.W === nothing || resize_J_W!(nlcache, integrator, i) + # Ref{StaticWOperator} = OOP path, fixed-size, not resizable. + nlcache.W === nothing || nlcache.W isa Ref || resize_J_W!(nlcache, integrator, i) nlcache.W_γdt = zero(nlcache.W_γdt) nlcache.new_W = true return nothing diff --git a/lib/OrdinaryDiffEqNonlinearSolve/src/utils.jl b/lib/OrdinaryDiffEqNonlinearSolve/src/utils.jl index 01ed522bca..08f0b45ac7 100644 --- a/lib/OrdinaryDiffEqNonlinearSolve/src/utils.jl +++ b/lib/OrdinaryDiffEqNonlinearSolve/src/utils.jl @@ -499,21 +499,38 @@ function build_nlsolver( γ = tTypeNoUnits(γ) α = tTypeNoUnits(α) dt = tTypeNoUnits(dt) + use_w_reuse = !isdae && f.nlstep_data === nothing && + W isa StaticWOperator && + nlalg.alg isa AbstractSimpleNonlinearSolveAlgorithm nlf = isdae ? oopdaenlf : oopodenlf nlp_params = if isdae (tmp, α, tstep, invγdt, p, dt, uprev, f) else (tmp, γ, α, tstep, invγdt, DIRK, p, dt, f) end - prob = NonlinearProblem( - NonlinearFunction{false, SciMLBase.FullSpecialize}(nlf), - copy(ztmp), nlp_params - ) + W_ref = use_w_reuse ? Ref{typeof(W)}(W) : nothing + prob = if use_w_reuse + nlf_jac = let W_ref = W_ref + (z, p) -> W_ref[].W + end + NonlinearProblem( + NonlinearFunction{false, SciMLBase.FullSpecialize}( + nlf; jac = nlf_jac, jac_prototype = W.W + ), + copy(ztmp), nlp_params + ) + else + NonlinearProblem( + NonlinearFunction{false, SciMLBase.FullSpecialize}(nlf), + copy(ztmp), nlp_params + ) + end inner_alg = _nlalg_with_linsolve(nlalg.alg, alg.linsolve) cache = init(prob, inner_alg, verbose = verbose.nonlinear_verbosity) nlcache = NonlinearSolveCache( nothing, tstep, nothing, nothing, invγdt, prob, cache, - nothing, nothing, nothing, nothing, nothing, nothing, + nothing, W_ref, use_w_reuse ? uf : nothing, + nothing, nothing, nothing, zero(tstep), true ) else diff --git a/lib/OrdinaryDiffEqNonlinearSolve/test/linear_nonlinear_tests.jl b/lib/OrdinaryDiffEqNonlinearSolve/test/linear_nonlinear_tests.jl index 049ee7b034..f16aae0e42 100644 --- a/lib/OrdinaryDiffEqNonlinearSolve/test/linear_nonlinear_tests.jl +++ b/lib/OrdinaryDiffEqNonlinearSolve/test/linear_nonlinear_tests.jl @@ -1,7 +1,8 @@ using OrdinaryDiffEqSDIRK, OrdinaryDiffEqBDF, OrdinaryDiffEqRosenbrock, Test, Random, - LinearAlgebra, LinearSolve, ADTypes + LinearAlgebra, LinearSolve, ADTypes, SciMLBase using OrdinaryDiffEqNonlinearSolve: NonlinearSolveAlg using NonlinearSolve: NewtonRaphson +using SimpleNonlinearSolve: SimpleNewtonRaphson Random.seed!(123) A = 0.01 * rand(3, 3) @@ -208,3 +209,20 @@ let integ = init( step!(integ) @test !iszero(integ.cache.nlsolver.cache.weight) end + +using StaticArrays +let + vdp_static(u, p, t) = SVector(u[2], p[1] * ((1 - u[1]^2) * u[2] - u[1])) + prob = ODEProblem( + vdp_static, SVector(2.0, 0.0), (0.0, 1.0), SVector(1.0e3) + ) + sol = solve( + prob, + TRBDF2( + nlsolve = NonlinearSolveAlg(SimpleNewtonRaphson(; autodiff = AutoForwardDiff())), + concrete_jac = true + ); reltol = 1.0e-8, abstol = 1.0e-10 + ) + @test SciMLBase.successful_retcode(sol.retcode) + @test sol.stats.njacs > 0 +end