Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Switch to DifferentiationInterface in Implicit Solvers #2545

Draft
wants to merge 116 commits into
base: master
Choose a base branch
from

Conversation

jClugstor
Copy link
Contributor

@jClugstor jClugstor commented Nov 27, 2024

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

This is at a point where we can do stuff like this:

using OrdinaryDiffEqCore
using OrdinaryDiffEqSDIRK
using ADTypes
using EnzymeCore

function lorenz!(du, u, p, t)
    du[1] = 10.0 * (u[2] - u[1])
    du[2] = u[1] * (28.0 - u[3]) - u[2]
    du[3] = u[1] * u[2] - (8 / 3) * u[3]
end

u0 = [1.0; 0.0; 0.0]
tspan = (0.0, 100.0)
prob = ODEProblem{true, SciMLBase.NoSpecialize}(lorenz!, u0, tspan)
sol = solve(prob, ImplicitEuler(autodiff=AutoEnzyme(function_annotation = EnzymeCore.Const)))

and it actually uses sparsity detection and greedy jacobian coloring plus Enzyme to compute the Jacobians.

Some things I'm unsure about:

  1. The current behavior is to use Jacobian coloring and SparseDiffTools by default. In order to keep that up, we have to wrap any ADType given in an AutoSparse unless it's already an AutoSparse. This does change the ADType that the user entered to be wrapped in an AutoSparse, which feels weird to me. Maybe there should be an option to just directly use the ADType entered, but by default we wrap it into an AutoSparse? I'm not sure.
  2. The biggest issue is that the way the sparsity detectors work with DI is by using operator overloading (both TracerSparsityDetector and SymbolicsSparsityDetector do), but that's an issue when using AutoSpecialilzation, because of the FunctionWrappers. The solution I found was to just unwrap the function in the preparation process. I'm not sure what performance implication this will have, but I don't think it should do much, since the preparation should be run just once.
  3. There's still pieces in here that use raw SparseDiffTools, (build_J_W) that I haven't looked in to how to convert to DI yet.
  4. I may need to fix some of the versions.
  5. There are some places that are getting sparse things where it's not expected.

@jClugstor jClugstor changed the title Switch to DifferentiationInterface in Implicit Solver Switch to DifferentiationInterface in Implicit Solvers Nov 27, 2024
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.

2 participants