SciML / NonlinearSolve.jl

High-performance and differentiation-enabled nonlinear solvers (Newton methods), bracketed rootfinding (bisection, Falsi), with sparsity and Newton-Krylov support.
https://docs.sciml.ai/NonlinearSolve/stable/
MIT License
232 stars 40 forks source link

feat: `NonlinearSolveBase` + `BracketingNonlinearSolve` + `SimpleNonlinearSolve` #458

Open avik-pal opened 1 month ago

avik-pal commented 1 month ago

First step in cutting down on load times. This is mostly independent of #456

  1. Move away from DiffEqBase to NonlinearSolveBase in the Simple packages.
  2. Split SimpleNonlinearSolve into (SimpleNonlinearSolve has longer load times than Lux which sounds extremely wrong)
    1. SimpleNonlinearSolve -- retains its current functionality but drops DiffEqBase and loads NonlinearSolveBase and BracketingNonlinearSolve
    2. BracketingNonlinearSolve
  3. What does NonlinearSolveBase have?
    1. Termination Conditions that were kind of incorrectly put (by me) into DiffEqBase
    2. Forward AD overloads via extensions
    3. Helper functions for norm that is used across NonlinearSolve
    4. Moved the definition of ImmutableNonlinearProblem from SimpleNonlinearSolve
  4. What are the other sources of bad load time?
    1. DocStringExtensions -- https://github.com/JuliaDocs/DocStringExtensions.jl/pull/167
    2. DifferentiationInterface -- SparseArrays https://github.com/gdalle/DifferentiationInterface.jl/issues/416
    3. FiniteDiff -- again SparseArrays
    4. MaybeInplace -- again SparseArrays (fixed by https://github.com/SciML/MaybeInplace.jl/pull/11)
    5. MLStyle via SciMLBase -- I have been told this is being looked into.
  5. Pending downstream changes:
    1. DiffEqBase loads NonlinearSolveBase and adds get_concrete_problem

There's a bunch of stuff in NonlinearSolve that can eventually be moved into base, but that is something to do later on.

Tasklist

avik-pal commented 1 month ago
@time begin
    using SimpleNonlinearSolve
    prob_brack = IntervalNonlinearProblem{false}(
        (u, p) -> u^2 - p, (0.0, 25.0), 5)
    solve(prob_brack, ITP())
end
# 1.574422 seconds (1.17 M allocations: 81.038 MiB, 2.57% gc time, 6.58% compilation time)

@time begin
    using BracketingNonlinearSolve
    prob_brack = IntervalNonlinearProblem{false}(
        (u, p) -> u^2 - p, (0.0, 25.0), 5)
    solve(prob_brack, ITP())
end
# 0.504557 seconds (429.92 k allocations: 30.137 MiB, 6.16% gc time, 25.34% compilation time)
ChrisRackauckas commented 1 month ago

Move away from DiffEqBase to NonlinearSolveBase in the Simple packages.

Should we move the DiffEqBase/solve.jl to SciMLBase? We probably should and make that a bit more generic, and then reuse it in more contexts anyways.

ChrisRackauckas commented 1 month ago

You also want the diffeqbase.jl functionwrappers though.

avik-pal commented 1 month ago

Should we move the DiffEqBase/solve.jl to SciMLBase? We probably should and make that a bit more generic, and then reuse it in more contexts anyways.

I think so. They don't seem to be diffeq specific.

You also want the diffeqbase.jl functionwrappers though.

Not sure what you are referring to

ChrisRackauckas commented 1 month ago

This part: https://github.com/SciML/DiffEqBase.jl/blob/master/src/norecompile.jl

avik-pal commented 2 weeks ago

Ok bracketing solvers are now done. We are still not as fast as loading Roots.jl (which is like 0.009567s) but this is as fast as we can get here. All the remaining timings are coming from upstream packages:

julia> @time_imports using BracketingNonlinearSolve
      0.2 ms  ConcreteStructs
      0.2 ms  CommonSolve
      0.8 ms  ArrayInterface
      0.2 ms  FastClosures
      0.1 ms  Compat → CompatLinearAlgebraExt
               ┌ 0.0 ms DocStringExtensions.__init__() 
     92.6 ms  DocStringExtensions 99.23% compilation time
     10.7 ms  RecipesBase
      0.6 ms  StaticArraysCore
      0.3 ms  ArrayInterfaceStaticArraysCoreExt
      0.5 ms  RuntimeGeneratedFunctions
               ┌ 0.0 ms InverseFunctions.__init__() 
      1.0 ms  InverseFunctions
      0.2 ms  ConstructionBase
      0.1 ms  CompositionsBase
      0.1 ms  CompositionsBaseInverseFunctionsExt
      0.2 ms  ConstructionBaseLinearAlgebraExt
               ┌ 0.0 ms Accessors.__init__() 
     11.7 ms  Accessors
      2.5 ms  SymbolicIndexingInterface
      0.1 ms  InverseFunctionsDatesExt
      0.2 ms  AccessorsDatesExt
      0.2 ms  Adapt
      0.3 ms  GPUArraysCore
      0.1 ms  ArrayInterfaceGPUArraysCoreExt
     18.2 ms  RecursiveArrayTools
      0.3 ms  SciMLStructures
      5.7 ms  FunctionWrappers
      0.2 ms  FunctionWrappersWrappers
      0.2 ms  EnumX
      2.2 ms  ADTypes
     81.4 ms  MLStyle
      5.7 ms  Expronicon
      0.2 ms  Reexport
      8.2 ms  SciMLOperators
      0.2 ms  SciMLOperatorsStaticArraysCoreExt
               ┌ 0.0 ms SciMLBase.__init__() 
     69.2 ms  SciMLBase
      1.5 ms  NonlinearSolveBase
      6.3 ms  BracketingNonlinearSolve
avik-pal commented 2 weeks ago

Also I am not sure how to do the SimpleNonlinearSolve move without breaking older versions. If we change the path to SimpleNonlinearSolve to this repo in the registry, then older versions don't exist in the git history right? IIUC it would still exist in the pkg server, but do we want to do that :sweat_smile: