Skip to content

ODENLStepData: forward outer ODEFunction's flag to inner nlprob#4609

Open
singhharsh1708 wants to merge 1 commit into
SciML:masterfrom
singhharsh1708:feat/nlstep-analytic-jac
Open

ODENLStepData: forward outer ODEFunction's flag to inner nlprob#4609
singhharsh1708 wants to merge 1 commit into
SciML:masterfrom
singhharsh1708:feat/nlstep-analytic-jac

Conversation

@singhharsh1708

@singhharsh1708 singhharsh1708 commented Jun 9, 2026

Copy link
Copy Markdown

Closes #4570.

When a user writes ODEProblem(sys, ...; jac = true, nlstep = true), the outer ODEFunction gets an analytic Jacobian but the inner teared NonlinearProblem built by generate_ODENLStepData did not — its f.jac was nothing.

This threads the jac kwarg through generate_ODENLStepData and into the inner NonlinearProblem / SCCNonlinearProblem constructor. With jac = true, MTK now symbolically differentiates the teared inner system once at build time and attaches the result to nlprob.f.jac.

Repro

using ModelingToolkit, SciMLBase
using ModelingToolkit: t_nounits as t, D_nounits as D

@parameters k₁ k₂ k₃
@variables y₁(t) y₂(t) y₃(t)
@mtkcompile rober = System([
D(y₁) ~ -k₁y₁ + k₃y₂y₃,
D(y₂) ~ k₁
y₁ - k₂y₂^2 - k₃y₂y₃,
D(y₃) ~ k₂
y₂^2,
], 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, nlstep = true)

SciMLBase.has_jac(prob.f.nlstep_data.nlprob.f) # false on master, true with this PR
Verified locally
nlprob.f.jac populated when jac = true, untouched (nothing) when jac = false
The attached function is a real MTK GeneratedFunctionWrapper
Solutions still converge with the same retcode and within tolerance of the master result
Motivation
This is a prerequisite for OrdinaryDiffEq.jl-side W-matrix reuse on the nlstep path (SciML/OrdinaryDiffEq.jl#3696). Standalone performance impact is small/mixed (the inner teared problem for benchmark-scale Robertson is 1D, so FD jac is already nearly free); the real wins are downstream when paired with #3696 and on larger MTK systems where the teared inner Jacobian is non-trivial.

Checklist

  • Appropriate tests were added
  • Any code changes were done in a way that does not break public API
  • All documentation related to code changes were updated
  • The new code follows the
    contributor guidelines, in particular the SciML Style Guide and
    COLPRAC.
  • Any new documentation only uses public API

Additional context

Add any other context about the problem here.

@singhharsh1708

Copy link
Copy Markdown
Author

@ChrisRackauckas — small MTK side of the W-reuse story. Threads the existing jac kwarg through generate_ODENLStepData so the inner teared NonlinearProblem gets MTK's symbolic analytic Jacobian when the user writes ODEProblem(...; jac = true, nlstep = true). Prereq for SciML/OrdinaryDiffEq.jl#3696 closure.
Verified locally: has_jac(prob.f.nlstep_data.nlprob.f) flips false → true; same retcode on solve.

@singhharsh1708

Copy link
Copy Markdown
Author

@ChrisRackauckas — friendly bump on this one. Re-verified: CI failures match master's pre-existing red state (same Runic failures on the same files — MTKFMIExt.jl, MTKOrdinaryDiffEqBDFExt.jl etc. — none of which are touched by this PR's 14-line jac kwarg threading). Anything else needed before merge?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

ODENLStepData: set analytic jac on inner nlprob when ODEProblem(jac=true, nlstep=true)

2 participants