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

feat: NonlinearSolveBase + BracketingNonlinearSolve + SimpleNonlinearSolve #458

Merged
merged 92 commits into from
Oct 23, 2024

Conversation

avik-pal
Copy link
Member

@avik-pal avik-pal commented Aug 23, 2024

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 -- Remove LibGit2 dep JuliaDocs/DocStringExtensions.jl#167
    2. DifferentiationInterface -- SparseArrays Reduce load time? JuliaDiff/DifferentiationInterface.jl#416
    3. FiniteDiff -- again SparseArrays
    4. MaybeInplace -- again SparseArrays (fixed by refactor: move SparseArrays into an extension MaybeInplace.jl#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

  • NonlinearSolveBase
    • Base Package
    • Testing
    • AD selection preferences
    • ForwardAD Primitives
      • Rootfinding
      • NLLS
  • BracketingNonlinearSolve
    • All core algorithms
    • Add tests
    • ForwardAD support
  • SimpleNonlinearSolve
  • CI scripts
    • scripts for the new subpackages
      • NonlinearSolveBase
      • SimpleNonlinearSolve
      • BracketingNonlinearSolve
    • reenable NonlinearSolve testing
  • NonlinearSolve -- update to use these subpackages
    • AD support
    • update to use __init and __solve
    • NonlinearSolveBase
      • Termination Condtiions
      • AD selection
  • Add QA testing for all the repos
    • SimpleNonlinearSolve
    • BracketingNonlinearSolve
    • NonlinearSolveBase

@avik-pal avik-pal marked this pull request as draft August 23, 2024 00:35
@avik-pal
Copy link
Member Author

avik-pal commented Aug 23, 2024

@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)

@avik-pal avik-pal changed the title feat: setup core packages NonlinearSolveBase feat: NonlinearSolveBase + BracketingNonlinearSolve + SimpleNonlinearSolve Aug 23, 2024
@ChrisRackauckas
Copy link
Member

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
Copy link
Member

You also want the diffeqbase.jl functionwrappers though.

@avik-pal
Copy link
Member Author

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
Copy link
Member

@avik-pal avik-pal force-pushed the ap/split_start branch 2 times, most recently from 2768ae8 to 734b2f8 Compare September 17, 2024 17:18
@avik-pal
Copy link
Member Author

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
Copy link
Member Author

avik-pal commented Sep 17, 2024

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 😅

@avik-pal avik-pal force-pushed the ap/split_start branch 2 times, most recently from 07cc36c to 1d75d7f Compare September 25, 2024 02:43
@avik-pal avik-pal changed the base branch from master to ap/linesearch September 25, 2024 02:43
@avik-pal avik-pal force-pushed the ap/split_start branch 2 times, most recently from 5ebe24c to 30487a1 Compare September 26, 2024 09:02
@avik-pal avik-pal force-pushed the ap/linesearch branch 2 times, most recently from 40c4df9 to 173dd01 Compare September 26, 2024 09:15
@avik-pal avik-pal force-pushed the ap/linesearch branch 4 times, most recently from 947fb85 to b2d061a Compare October 3, 2024 23:55
Base automatically changed from ap/linesearch to master October 4, 2024 16:11
@avik-pal avik-pal force-pushed the ap/split_start branch 2 times, most recently from dcf2c7f to d6167c9 Compare October 4, 2024 17:40
@avik-pal avik-pal force-pushed the ap/split_start branch 2 times, most recently from f3a7ccc to f36989a Compare October 22, 2024 22:54
@avik-pal
Copy link
Member Author

This PR is now in a mergeable state finally! Was able to keep breaking changes to a minimal

  - Use of termination conditions from `DiffEqBase` has been removed. Use the termination
    conditions from `NonlinearSolveBase` instead.
  - If no autodiff is provided, we now choose from a list of autodiffs based on the packages
    loaded. For example, if `Enzyme` is loaded, we will default to that (for reverse mode).
    In general, we don't guarantee the exact autodiff selected if `autodiff` is not provided
    (i.e. `nothing`).

@avik-pal avik-pal marked this pull request as ready for review October 22, 2024 23:05
@avik-pal avik-pal force-pushed the ap/split_start branch 2 times, most recently from da057e7 to 7f57584 Compare October 22, 2024 23:55
@avik-pal
Copy link
Member Author

I will merge and get registrations going. We are not releasing the new NonlinearSolve right now anyways till #481 is finished.

@avik-pal avik-pal merged commit 21b02bd into master Oct 23, 2024
52 of 75 checks passed
@avik-pal avik-pal deleted the ap/split_start branch October 23, 2024 12:22
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.

3 participants