[RFC] Splitting up NonlinearSolve

[avik-pal]

Similar to the treatment that was given to OrdinaryDiffEq, we want to split up this package into subpackages. This will be much easier to do here, since all the solvers here are extremely modular (that was one of the central points of our paper https://arxiv.org/abs/2403.16341). This is my first proposal:

• NonlinearSolveBase.jl
  • currently houses the termination conditions and AD rules for forward mode
  • This will also house most of what is present in timer_ouputs.jl, abstract_types.jl, and internal.jl internal/. The AD wrappers and such also go here via extensions
• LineSearch.jl (empty repo exists as of now): Let's move https://github.com/SciML/NonlinearSolve.jl/blob/master/src/globalization/line_search.jl to that package. We can iterate on a 2.0 design and move to native solvers later (is that okay @Vaibhavdixit02?) DONE WAITING FOR A MERGE
• Next we do the main solver packages:
  • NonlinearSolveApproximateJacobian.jl -- LBroyden, Broyden, Klement, and other solvers (and utilities) that reduce and are used in https://github.com/SciML/NonlinearSolve.jl/blob/master/src/core/approximate_jacobian.jl.
  • NonlinearSolveFirstOrder.jl -- This is the biggest chunk housing TrustRegion, NewtonRaphson, etc.
  • NonlinearSolveSpectralMethods.jl -- DFSane and contents of https://github.com/SciML/NonlinearSolve.jl/blob/master/src/core/spectral_methods.jl.
  • Where do we put the extension solvers? NonlinearSolve.jl or another subpackage NonlinearSolveExtensionAlgs.jl?
• NonlinearSolve.jl -- We define the default algorithms here. We need to pull in all the packages (except Spectral methods) for this part.
• BracketingNonlinearSolve -- Split out of SimpleNonlinearSolve
• SimpleNonlinearSolve -- migrate to using bracketingnonlinearsolve and nonlinearsolvebase

To cut down more aggressively on compile times we need to finish NonlinearSolveBase.jl and use that in-place of DiffEqBase.jl.

Also this unblocks solvers like #404 that need to rely on TaylorDiff and heavier packages.

cc @oscardssmith @ChrisRackauckas

[avik-pal]

Also as a part of this change should we move SimpleNonlinearSolve.jl into lib instead of a separate repo?

[ChrisRackauckas]

Also as a part of this change should we move SimpleNonlinearSolve.jl into lib instead of a separate repo?

I'd entertain that.

[oscardssmith]

IMO this split-up probably won't help much. The Nonlinear solvers are much simpler than the ODE ones and they use more of the same functionality. As such, I doubt this will lead to significant speed benefits.

[Vaibhavdixit02]

LineSearch.jl (empty repo exists as of now): Let's move https://github.com/SciML/NonlinearSolve.jl/blob/master/src/globalization/line_search.jl to that package. We can iterate on a 2.0 design and move to native solvers later

Sounds perfect, probably the best way to get any activity started there 😅

[ChrisRackauckas]

The 4 are rather different. LineSearch should be split anyways because of its use in optimization. We don't need to load spectral methods at all for any case users don't request it since they aren't in the defaults. And by default we don't need NonlinearSolveApproximateJacobian with the ODE solvers, we just want the first order methods. So we would probably cut out about half of what's loaded, and for everything else it would parallelize the precompilation.

Where do we put the extension solvers? NonlinearSolve.jl or another subpackage NonlinearSolveExtensionAlgs.jl?

I think NonlinearSolveExtensionAlgs.jl makes sense for that so it's a low dep.

IMO this split-up probably won't help much. The Nonlinear solvers are much simpler than the ODE ones and they use more of the same functionality. As such, I doubt this will lead to significant speed benefits.

I'm not sure that's really the case. What we're really get is parallelized precompilation. This package takes longer than OrdinaryDiffEq to precompile and is the biggest precompilation step right now. Splitting it to parallelize 4 way would already be a huge improvement if we just decrease it by 50%. I think we'd get even lower.

The other compile time decrease is to make some internal functions non-algorithm specific. That is another avenue to investigate, but I'd try to cut out the 3 portions and only depend on first order downstream first and see what we get. I wouldn't be surprised at something that's a good 3x in installation and precompilation time improvement, which would be a massive improvement.

[avik-pal]

Also we should add a simple ASV benchmark script to all the repos to track load times. For example https://github.com/LuxDL/Lux.jl/blob/main/bench/asv.jl is very quick to run and adds comments like LuxDL/Lux.jl#848 (comment) with the load time

[ParamThakkar123]

Should I take this one ? If nobody has taken it yet ?

[ParamThakkar123]

Should I take this one ? If nobody has taken it yet ?

@ChrisRackauckas @avik-pal ??

[avik-pal]

We received an email from someone who wanted to take this on.

[ParamThakkar123]

I sent an email to sciml today for getting a claim on this project an hour ago

[ParamThakkar123]

Did the email come from [email protected]??

[ParamThakkar123]

@avik-pal ??

[ChrisRackauckas]

Let's work this out here. @Suavesito-Olimpiada put a claim in first so that has priority. However, it's unclear right now if that's for NonlinearSolve.jl or BoundaryValueDiffEq.jl. And it looks like @avik-pal already got started on this one himself #458 independently of that, so let's try to spread things out so we aren't overlapping 😅. I think we should hear from:

1. @Suavesito-Olimpiada which one did you intend to claim?
2. @avik-pal were you planning to do this one yourself? Or were you just starting the PR?

[ParamThakkar123]

Okay. No problem. I will put my step back from here 😅

[avik-pal]

@avik-pal were you planning to do this one yourself? Or were you just starting the PR?

I was just setting up an initial sketch and PoC of the benefits of the split.

[Suavesito-Olimpiada]

I intended to claim both (as per here), but, started with this repo, as per the answer to my questions in email. 😅

I saw the PR in #458. From what I understand, that is independent of the plan exposed here by Avik.

[tansongchen]

@avik-pal How is this splitting going? If I want to revive #404 , should I make a package or an extension?