diff --git a/lib/ModelingToolkitBase/src/problems/odeproblem.jl b/lib/ModelingToolkitBase/src/problems/odeproblem.jl index 5e16204dbc..9d3138e0fc 100644 --- a/lib/ModelingToolkitBase/src/problems/odeproblem.jl +++ b/lib/ModelingToolkitBase/src/problems/odeproblem.jl @@ -1,11 +1,15 @@ """ - generate_ODENLStepData(sys, u0, p, mm, nlstep_compile, nlstep_scc) + generate_ODENLStepData(sys, u0, p, mm, nlstep_compile, nlstep_scc; jac = false) Generate the NLStep data for implicit ODE solvers. This is a stub that throws an error if called without ModelingToolkit loaded. The actual implementation is provided by ModelingToolkit when it is loaded. + +When `jac = true`, the analytic Jacobian of the teared inner nonlinear system is +generated symbolically and attached to `nlprob.f.jac`, so NonlinearSolve does not +have to recompute it via AD/FD on every Newton iteration. """ -function generate_ODENLStepData(sys, u0, p, mm, nlstep_compile, nlstep_scc) +function generate_ODENLStepData(sys, u0, p, mm, nlstep_compile, nlstep_scc; jac = false) error( """ `nlstep=true` requires ModelingToolkit.jl to be loaded. @@ -65,7 +69,7 @@ Base.@nospecializeinfer @fallback_iip_specialize function SciMLBase.ODEFunction{ _M = concrete_massmatrix(M; sparse, u0) if nlstep - ode_nlstep = generate_ODENLStepData(sys, u0, p, M, nlstep_compile, nlstep_scc) + ode_nlstep = generate_ODENLStepData(sys, u0, p, M, nlstep_compile, nlstep_scc; jac) else ode_nlstep = nothing end diff --git a/src/systems/solver_nlprob.jl b/src/systems/solver_nlprob.jl index 9bd0d107db..f96ad0c58a 100644 --- a/src/systems/solver_nlprob.jl +++ b/src/systems/solver_nlprob.jl @@ -1,6 +1,7 @@ function MTKBase.generate_ODENLStepData( sys::System, u0, p, mm = calculate_massmatrix(sys), - nlstep_compile::Bool = true, nlstep_scc::Bool = false + nlstep_compile::Bool = true, nlstep_scc::Bool = false; + jac::Bool = false ) nlsys, outer_tmp, inner_tmp = inner_nlsystem(sys, mm, nlstep_compile) state = ProblemState(; u = u0, p) @@ -15,10 +16,13 @@ function MTKBase.generate_ODENLStepData( haskey(op, v) && continue op[v] = getsym(sys, v)(state) end + # Forward the outer ODEFunction's `jac` flag so the inner teared nonlinear + # problem also carries an analytic Jacobian. This lets NonlinearSolve.jl + # skip the per-iteration AD/FD Jacobian computation on the inner Newton. nlprob = if nlstep_scc - SCCNonlinearProblem(nlsys, op; build_initializeprob = false) + SCCNonlinearProblem(nlsys, op; build_initializeprob = false, jac) else - NonlinearProblem(nlsys, op; build_initializeprob = false) + NonlinearProblem(nlsys, op; build_initializeprob = false, jac) end subsetidxs = [findfirst(isequal(y), unknowns(sys)) for y in unknowns(nlsys)]