Feat: Handle Adjoints through Initialization

[DhairyaLGandhi]

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

MTK and SciML construct an initialization problem before starting the time stepping to ensure the starting values of the unknowns and parameters adhere to any constraints needed for the system. This PR adds handling for adjoint sensitivities of the NonlinearProblem, NonlinearSquaresProblem, SCCNonlinearProblem etc.

I am opening this to get some feedback regarding how we can accumulate gradients correctly. I have also included a test case for a DAE which I will update to use the values out of SciMLSensitivity.

Add any other context about the problem here.

Currently the gradients get calculated but don't get accumulated, we need to be able to update the gradients for the parameters. Since this is a manual dispatch, the usual graph building in AD is bypassed, and we need to handle this manually. Ideally, we should make it so the cfg itself includes the initialization so we would not have gotten incorrect gradients in the first place 😅 We are also forced to use a LinearProblem instead of \ because it cannot handle singular jacobians.

cc @ChrisRackauckas

[DhairyaLGandhi]

I wanted to ask whether it is preferred to retain \ for the smaller problems, which would be slower with LinearSolve, and if so, we will need to deal with LAPACK errors.

[ChrisRackauckas]

I wanted to ask whether it is preferred to retain \ for the smaller problems, which would be slower with LinearSolve, and if so, we will need to deal with LAPACK errors.

LinearSolve.jl should be faster across the board? It depends a bit on the CPU architecture since it depends on whether it guesses the right LU correctly,

[ChrisRackauckas]

Note that with the latest MTK update, there is now an Initials section where the initial u0 live. That should fix a few things.

[ChrisRackauckas]

this should be before the solve, since you can use the initialization solution from here in the remakes of 397-405 in order to set new u0 and p and thus skip running the initialization a second time.

[DhairyaLGandhi]

How can I indicate to solve to avoid running initialization?

[ChrisRackauckas]

initializealg = NoInit(). Should probably just do CheckInit() for safety but either is fine.

[ChrisRackauckas]

Yeah I don't know about that comment. I think it's just old. (a) \ always allocates because it uses lu instead of lu!, so it's re-allocating the while matrix which is larger than any LinearSolve allocation, and (b) we have since 2023 setup tests on StaticArrays, so the immutable path is non-allocating. I don't think (b) was true when this was written.

[DhairyaLGandhi]

So glad we can remove this branch altogether.

[ChrisRackauckas]

This gradient isn't used? I think this would go into the backpass and if I'm thinking clearly, the resulting return is dp .* igs?

[DhairyaLGandhi]

Not yet. These gradients are currently against the parameters of the initialization problem, not the system exactly. And the mapping between the two is ill defined, so we cannot simply accum

I spoke with @AayushSabharwal about a way to map, it seems initialization_data.intializeprobmap might have some support to return the correctly shaped vector, but there are cases where we cannot know the ordering of dp either.

[DhairyaLGandhi]

There's another subtlety. I am not sure we haven't missed some part of the cfg by manually handling accumulation of gradients. Or any transforms we might need to calculate gradients for. The regular AD graph building typically took care of these details for us, but in this case we would need to worry about incorrect gradients manually

[ChrisRackauckas]

Oh yes, you need to use the initializeprobmap https://github.com/SciML/SciMLBase.jl/blob/master/src/initialization.jl#L268 to map it back to the shape of the initial parameters.

but there are cases where we cannot know the ordering of dp either.

p and dp just need the same ordering, so initializeprobmap should do the trick.

There's another subtlety. I am not sure we haven't missed some part of the cfg by manually handling accumulation of gradients. Or any transforms we might need to calculate gradients for. The regular AD graph building typically took care of these details for us, but in this case we would need to worry about incorrect gradients manually

This is the only change to (u0,p) before solving, so this would account for it, given initializeprobmap is just an index map so an identity function.

[DhairyaLGandhi]

Addressed this occurance in 95ebbf3 to check if this is correct. Will need to work around the global call

[DhairyaLGandhi]

Trying to use the initialization end to end caused gradients against parameters to get dropped. https://github.com/DhairyaLGandhi/SciMLBase.jl/tree/dg/nonlinear is a WIP branch which adds adjoints to the getindex calls, which does capture the expected gradients, but we still end up dropping gradients somewhere in the chain. I am looking into whether we are doing so in the custom adjoints, because that was an issue I had identified for the ODE case.

[ChrisRackauckas]

shouldn't default to ZygoteVJP. Should use the autojacvec of the ODE

[DhairyaLGandhi]

Addressed in 9a8a845

[DhairyaLGandhi]

I could use some understanding of how to handle initialization when MTK analytically solves the problem, and removes all the unknowns. In that case u0, u etc are empty, and we attempt to calculate the jacobian for it in dgdgu_val which seems weird, there are no unknowns to calculate the jacobian for. It would be equally incorrect to use the unknowns from the system itself. The best way to handle this might be to simply add a check for isempty and return. Is that reasonable?

[DhairyaLGandhi]

There seems to be an InitialFailure with DefaultInit

[DhairyaLGandhi]

julia> prob, ps, init = setups[3];

julia> init
OrdinaryDiffEqCore.DefaultInit()

julia> new_sol = solve(prob, Rodas5P(); p = p, sensealg, initializealg = init, abstol = 1e-6, reltol = 1e-3)
┌ Warning: Initialization system is overdetermined. 1 equations for 0 unknowns. Initialization will default to using least squares. `SCCNonlinearProblem` can only be used for initialization of fully determined systems and hence will not be used here. To suppress this warning pass warn_initialize_determined = false. To make this warning into an error, pass fully_determined = true
└ @ ModelingToolkit ~/.julia/packages/ModelingToolkit/aau6A/src/systems/diffeqs/abstractodesystem.jl:1520
retcode: InitialFailure
Interpolation: specialized 4rd order "free" stiffness-aware interpolation
t: 1-element Vector{Float64
}:
 0.0
u: 1-element Vector{Vector{Float64}}:
 [2.0, 0.0, 0.0, 1.0, 0.0]

[ChrisRackauckas]

@AayushSabharwal note that this is the issue of not passing over the initial values and guess values to u0. Without this CheckInit will fail. It's a minor issue but it is something we should handle over time.

[AayushSabharwal]

Yeah I'm working on using the hook in DiffEqBase. It's easy enough to implement, but somehow the returned problem's u0 isn't being respected

[ChrisRackauckas]

dead code?

[DhairyaLGandhi]

I made a mistake pushing cdaa2c7 so have reverted it

[DhairyaLGandhi]

Apologies

[ChrisRackauckas]

?

[DhairyaLGandhi]

Ah this was picked from the history, updated. Thanks