From 7f15b30cb55acd4f69171c4c21867dc37e83ccb6 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 7 Oct 2024 17:13:19 -0400 Subject: [PATCH] test: comprehensive testing of root finding --- docs/src/release_notes.md | 22 ++- .../src/termination_conditions.jl | 4 +- .../src/SimpleNonlinearSolve.jl | 10 +- lib/SimpleNonlinearSolve/src/halley.jl | 3 +- lib/SimpleNonlinearSolve/src/lbroyden.jl | 2 +- .../test/core/23_test_problems_tests.jl | 11 +- .../test/core/exotic_type_tests.jl | 2 +- .../test/core/rootfind_tests.jl | 154 ++++++++++++++++++ .../test/gpu/cuda_tests.jl | 4 +- 9 files changed, 195 insertions(+), 17 deletions(-) diff --git a/docs/src/release_notes.md b/docs/src/release_notes.md index 9d517267e..0208b1e66 100644 --- a/docs/src/release_notes.md +++ b/docs/src/release_notes.md @@ -1,6 +1,24 @@ # Release Notes -## Breaking Changes in `NonlinearSolve.jl` v3 +## Oct '24 + +### Breaking Changes in `NonlinearSolve.jl` v4 + +### Breaking Changes in `SimpleNonlinearSolve.jl` v2 + + - `Auto*` structs are no longer exported. Load `ADTypes` to access them. + - Use of termination conditions from `DiffEqBase` has been removed. Use the termination + conditions from `NonlinearSolveBase` instead. + - We no longer export the entire `SciMLBase`. Instead selected functionality relevant to + `SimpleNonlinearSolve` has been exported. + - 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. In general, we + don't guarantee the exact autodiff selected if `autodiff` is not provided (i.e. + `nothing`). + +## Dec '23 + +### Breaking Changes in `NonlinearSolve.jl` v3 - `GeneralBroyden` and `GeneralKlement` have been renamed to `Broyden` and `Klement` respectively. @@ -8,7 +26,7 @@ - The old style of specifying autodiff with `chunksize`, `standardtag`, etc. has been deprecated in favor of directly specifying the autodiff type, like `AutoForwardDiff`. -## Breaking Changes in `SimpleNonlinearSolve.jl` v1 +### Breaking Changes in `SimpleNonlinearSolve.jl` v1 - Batched solvers have been removed in favor of `BatchedArrays.jl`. Stay tuned for detailed tutorials on how to use `BatchedArrays.jl` with `NonlinearSolve` & `SimpleNonlinearSolve` diff --git a/lib/NonlinearSolveBase/src/termination_conditions.jl b/lib/NonlinearSolveBase/src/termination_conditions.jl index 7978c19b8..3e957f139 100644 --- a/lib/NonlinearSolveBase/src/termination_conditions.jl +++ b/lib/NonlinearSolveBase/src/termination_conditions.jl @@ -26,7 +26,7 @@ function update_u!!(cache::NonlinearTerminationModeCache, u) if cache.u isa AbstractArray && ArrayInterface.can_setindex(cache.u) copyto!(cache.u, u) else - cache.u .= u + cache.u = u end end @@ -60,6 +60,8 @@ function SciMLBase.init( else u_diff_cache = u_unaliased end + best_value = initial_objective + max_stalled_steps = mode.max_stalled_steps else initial_objective = nothing objectives_trace = nothing diff --git a/lib/SimpleNonlinearSolve/src/SimpleNonlinearSolve.jl b/lib/SimpleNonlinearSolve/src/SimpleNonlinearSolve.jl index 4ed53295e..081201a24 100644 --- a/lib/SimpleNonlinearSolve/src/SimpleNonlinearSolve.jl +++ b/lib/SimpleNonlinearSolve/src/SimpleNonlinearSolve.jl @@ -1,19 +1,19 @@ module SimpleNonlinearSolve using Accessors: @reset -using CommonSolve: CommonSolve, solve +using CommonSolve: CommonSolve, solve, init, solve! using ConcreteStructs: @concrete using FastClosures: @closure using LineSearch: LiFukushimaLineSearch using LinearAlgebra: LinearAlgebra, dot using MaybeInplace: @bb, setindex_trait, CannotSetindex, CanSetindex using PrecompileTools: @compile_workload, @setup_workload -using SciMLBase: AbstractNonlinearAlgorithm, NonlinearProblem, NonlinearLeastSquaresProblem, - IntervalNonlinearProblem, ReturnCode +using SciMLBase: SciMLBase, AbstractNonlinearAlgorithm, NonlinearFunction, NonlinearProblem, + NonlinearLeastSquaresProblem, IntervalNonlinearProblem, ReturnCode, remake using StaticArraysCore: StaticArray, SArray, SVector, MArray # AD Dependencies -using ADTypes: ADTypes +using ADTypes: ADTypes, AutoForwardDiff using DifferentiationInterface: DifferentiationInterface using FiniteDiff: FiniteDiff using ForwardDiff: ForwardDiff @@ -121,7 +121,7 @@ export IntervalNonlinearProblem export Alefeld, Bisection, Brent, Falsi, ITP, Ridder -export NonlinearProblem, NonlinearLeastSquaresProblem +export NonlinearFunction, NonlinearProblem, NonlinearLeastSquaresProblem export SimpleBroyden, SimpleKlement, SimpleLimitedMemoryBroyden export SimpleDFSane diff --git a/lib/SimpleNonlinearSolve/src/halley.jl b/lib/SimpleNonlinearSolve/src/halley.jl index 6b8948248..30eb1a821 100644 --- a/lib/SimpleNonlinearSolve/src/halley.jl +++ b/lib/SimpleNonlinearSolve/src/halley.jl @@ -35,7 +35,8 @@ function SciMLBase.__solve( abstol, reltol, tc_cache = NonlinearSolveBase.init_termination_cache( prob, abstol, reltol, fx, x, termination_condition, Val(:simple)) - autodiff = NonlinearSolveBase.select_jacobian_autodiff(prob, alg.autodiff) + # The way we write the 2nd order derivatives, we know Enzyme won't work there + autodiff = alg.autodiff === nothing ? AutoForwardDiff() : alg.autodiff @bb xo = copy(x) diff --git a/lib/SimpleNonlinearSolve/src/lbroyden.jl b/lib/SimpleNonlinearSolve/src/lbroyden.jl index ce39ca10d..3a17e0936 100644 --- a/lib/SimpleNonlinearSolve/src/lbroyden.jl +++ b/lib/SimpleNonlinearSolve/src/lbroyden.jl @@ -38,7 +38,7 @@ function SciMLBase.__solve( args...; termination_condition = nothing, kwargs...) if prob.u0 isa SArray if termination_condition === nothing || - termination_condition isa AbsNormTerminationMode + termination_condition isa NonlinearSolveBase.AbsNormTerminationMode return internal_static_solve( prob, alg, args...; termination_condition, kwargs...) end diff --git a/lib/SimpleNonlinearSolve/test/core/23_test_problems_tests.jl b/lib/SimpleNonlinearSolve/test/core/23_test_problems_tests.jl index 7bb28c649..6dd85b94b 100644 --- a/lib/SimpleNonlinearSolve/test/core/23_test_problems_tests.jl +++ b/lib/SimpleNonlinearSolve/test/core/23_test_problems_tests.jl @@ -1,5 +1,5 @@ @testsnippet RobustnessTestSnippet begin - using NonlinearProblemLibrary, NonlinearSolveBase, LinearAlgebra + using NonlinearProblemLibrary, NonlinearSolveBase, LinearAlgebra, ADTypes problems = NonlinearProblemLibrary.problems dicts = NonlinearProblemLibrary.dicts @@ -40,7 +40,7 @@ end @testitem "23 Test Problems: SimpleNewtonRaphson" setup=[RobustnessTestSnippet] tags=[:core] begin - alg_ops = (SimpleNewtonRaphson(),) + alg_ops = (SimpleNewtonRaphson(; autodiff = AutoForwardDiff()),) broken_tests = Dict(alg => Int[] for alg in alg_ops) broken_tests[alg_ops[1]] = [] @@ -49,7 +49,7 @@ end end @testitem "23 Test Problems: SimpleHalley" setup=[RobustnessTestSnippet] tags=[:core] begin - alg_ops = (SimpleHalley(),) + alg_ops = (SimpleHalley(; autodiff = AutoForwardDiff()),) broken_tests = Dict(alg => Int[] for alg in alg_ops) if Sys.isapple() @@ -62,7 +62,10 @@ end end @testitem "23 Test Problems: SimpleTrustRegion" setup=[RobustnessTestSnippet] tags=[:core] begin - alg_ops = (SimpleTrustRegion(), SimpleTrustRegion(; nlsolve_update_rule = Val(true))) + alg_ops = ( + SimpleTrustRegion(; autodiff = AutoForwardDiff()), + SimpleTrustRegion(; nlsolve_update_rule = Val(true), autodiff = AutoForwardDiff()) + ) broken_tests = Dict(alg => Int[] for alg in alg_ops) broken_tests[alg_ops[1]] = [3, 15, 16, 21] diff --git a/lib/SimpleNonlinearSolve/test/core/exotic_type_tests.jl b/lib/SimpleNonlinearSolve/test/core/exotic_type_tests.jl index 7e7223ea0..6227348a2 100644 --- a/lib/SimpleNonlinearSolve/test/core/exotic_type_tests.jl +++ b/lib/SimpleNonlinearSolve/test/core/exotic_type_tests.jl @@ -1,5 +1,5 @@ @testitem "BigFloat Support" tags=[:core] begin - using SimpleNonlinearSolve, LinearAlgebra + using SimpleNonlinearSolve, LinearAlgebra, ADTypes, SciMLBase fn_iip = NonlinearFunction{true}((du, u, p) -> du .= u .* u .- p) fn_oop = NonlinearFunction{false}((u, p) -> u .* u .- p) diff --git a/lib/SimpleNonlinearSolve/test/core/rootfind_tests.jl b/lib/SimpleNonlinearSolve/test/core/rootfind_tests.jl index 8b1378917..c99b670cb 100644 --- a/lib/SimpleNonlinearSolve/test/core/rootfind_tests.jl +++ b/lib/SimpleNonlinearSolve/test/core/rootfind_tests.jl @@ -1 +1,155 @@ +@testsnippet RootfindTestSnippet begin + using StaticArrays, Random, LinearAlgebra, ForwardDiff, NonlinearSolveBase, SciMLBase + using ADTypes, PolyesterForwardDiff, Enzyme, ReverseDiff + quadratic_f(u, p) = u .* u .- p + quadratic_f!(du, u, p) = (du .= u .* u .- p) + + function newton_fails(u, p) + return 0.010000000000000002 .+ + 10.000000000000002 ./ (1 .+ + (0.21640425613334457 .+ + 216.40425613334457 ./ (1 .+ + (0.21640425613334457 .+ + 216.40425613334457 ./ (1 .+ 0.0006250000000000001(u .^ 2.0))) .^ 2.0)) .^ + 2.0) .- 0.0011552453009332421u .- p + end + + const TERMINATION_CONDITIONS = [ + NormTerminationMode(Base.Fix1(maximum, abs)), + RelTerminationMode(), + RelNormTerminationMode(Base.Fix1(maximum, abs)), + RelNormSafeTerminationMode(Base.Fix1(maximum, abs)), + RelNormSafeBestTerminationMode(Base.Fix1(maximum, abs)), + AbsTerminationMode(), + AbsNormTerminationMode(Base.Fix1(maximum, abs)), + AbsNormSafeTerminationMode(Base.Fix1(maximum, abs)), + AbsNormSafeBestTerminationMode(Base.Fix1(maximum, abs)) + ] + + function run_nlsolve_oop(f::F, u0, p = 2.0; solver) where {F} + return solve(NonlinearProblem{false}(f, u0, p), solver; abstol = 1e-9) + end + function run_nlsolve_iip(f!::F, u0, p = 2.0; solver) where {F} + return solve(NonlinearProblem{true}(f!, u0, p), solver; abstol = 1e-9) + end +end + +@testitem "First Order Methods" setup=[RootfindTestSnippet] tags=[:core] begin + @testset for alg in ( + SimpleNewtonRaphson, + SimpleTrustRegion, + (; kwargs...) -> SimpleTrustRegion(; kwargs..., nlsolve_update_rule = Val(true)) + ) + @testset for autodiff in ( + AutoForwardDiff(), + AutoFiniteDiff(), + AutoReverseDiff(), + AutoEnzyme(), + nothing + ) + @testset "[OOP] u0: $(typeof(u0))" for u0 in ( + [1.0, 1.0], @SVector[1.0, 1.0], 1.0) + sol = run_nlsolve_oop(quadratic_f, u0; solver = alg(; autodiff)) + @test SciMLBase.successful_retcode(sol) + @test maximum(abs, quadratic_f(sol.u, 2.0)) < 1e-9 + end + + @testset "[IIP] u0: $(typeof(u0))" for u0 in ([1.0, 1.0],) + sol = run_nlsolve_iip(quadratic_f!, u0; solver = alg(; autodiff)) + @test SciMLBase.successful_retcode(sol) + @test maximum(abs, quadratic_f(sol.u, 2.0)) < 1e-9 + end + + @testset "Termination Condition: $(nameof(typeof(termination_condition))) u0: $(nameof(typeof(u0)))" for termination_condition in TERMINATION_CONDITIONS, + u0 in (1.0, [1.0, 1.0], @SVector[1.0, 1.0]) + + probN = NonlinearProblem(quadratic_f, u0, 2.0) + @test all(solve( + probN, alg(; autodiff = AutoForwardDiff()); termination_condition).u .≈ + sqrt(2.0)) + end + end + end +end + +@testitem "Second Order Methods" setup=[RootfindTestSnippet] tags=[:core] begin + @testset for alg in ( + SimpleHalley, + ) + @testset for autodiff in ( + AutoForwardDiff(), + AutoFiniteDiff(), + AutoReverseDiff(), + nothing + ) + @testset "[OOP] u0: $(typeof(u0))" for u0 in ( + [1.0, 1.0], @SVector[1.0, 1.0], 1.0) + sol = run_nlsolve_oop(quadratic_f, u0; solver = alg(; autodiff)) + @test SciMLBase.successful_retcode(sol) + @test maximum(abs, quadratic_f(sol.u, 2.0)) < 1e-9 + end + end + + @testset "Termination Condition: $(nameof(typeof(termination_condition))) u0: $(nameof(typeof(u0)))" for termination_condition in TERMINATION_CONDITIONS, + u0 in (1.0, [1.0, 1.0], @SVector[1.0, 1.0]) + + probN = NonlinearProblem(quadratic_f, u0, 2.0) + @test all(solve( + probN, alg(; autodiff = AutoForwardDiff()); termination_condition).u .≈ + sqrt(2.0)) + end + end +end + +@testitem "Derivative Free Methods" setup=[RootfindTestSnippet] tags=[:core] begin + @testset "$(nameof(typeof(alg)))" for alg in ( + SimpleBroyden(), + SimpleKlement(), + SimpleDFSane(), + SimpleLimitedMemoryBroyden(), + SimpleBroyden(; linesearch = Val(true)), + SimpleLimitedMemoryBroyden(; linesearch = Val(true)) + ) + @testset "[OOP] u0: $(typeof(u0))" for u0 in ([1.0, 1.0], @SVector[1.0, 1.0], 1.0) + sol = run_nlsolve_oop(quadratic_f, u0; solver = alg) + @test SciMLBase.successful_retcode(sol) + @test maximum(abs, quadratic_f(sol.u, 2.0)) < 1e-9 + end + + @testset "[IIP] u0: $(typeof(u0))" for u0 in ([1.0, 1.0],) + sol = run_nlsolve_iip(quadratic_f!, u0; solver = alg) + @test SciMLBase.successful_retcode(sol) + @test maximum(abs, quadratic_f(sol.u, 2.0)) < 1e-9 + end + + @testset "Termination Condition: $(nameof(typeof(termination_condition))) u0: $(nameof(typeof(u0)))" for termination_condition in TERMINATION_CONDITIONS, + u0 in (1.0, [1.0, 1.0], @SVector[1.0, 1.0]) + + probN = NonlinearProblem(quadratic_f, u0, 2.0) + @test all(solve(probN, alg; termination_condition).u .≈ sqrt(2.0)) + end + end +end + +@testitem "Newton Fails" setup=[RootfindTestSnippet] tags=[:core] begin + u0 = [-10.0, -1.0, 1.0, 2.0, 3.0, 4.0, 10.0] + p = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + + @testset "$(nameof(typeof(alg)))" for alg in ( + SimpleDFSane(), + SimpleTrustRegion(), + SimpleHalley(), + SimpleTrustRegion(; nlsolve_update_rule = Val(true)) + ) + sol = run_nlsolve_oop(newton_fails, u0, p; solver = alg) + @test SciMLBase.successful_retcode(sol) + @test maximum(abs, newton_fails(sol.u, p)) < 1e-9 + end +end + +@testitem "Kwargs Propagation" setup=[RootfindTestSnippet] tags=[:core] begin + prob = NonlinearProblem(quadratic_f, ones(4), 2.0; maxiters = 2) + sol = solve(prob, SimpleNewtonRaphson()) + @test sol.retcode === ReturnCode.MaxIters +end diff --git a/lib/SimpleNonlinearSolve/test/gpu/cuda_tests.jl b/lib/SimpleNonlinearSolve/test/gpu/cuda_tests.jl index d51ce353f..09fa52304 100644 --- a/lib/SimpleNonlinearSolve/test/gpu/cuda_tests.jl +++ b/lib/SimpleNonlinearSolve/test/gpu/cuda_tests.jl @@ -1,5 +1,5 @@ @testitem "Solving on CUDA" tags=[:cuda] begin - using StaticArrays, CUDA, SimpleNonlinearSolve + using StaticArrays, CUDA, SimpleNonlinearSolve, ADTypes if CUDA.functional() CUDA.allowscalar(false) @@ -47,7 +47,7 @@ end @testitem "CUDA Kernel Launch Test" tags=[:cuda] begin - using StaticArrays, CUDA, SimpleNonlinearSolve + using StaticArrays, CUDA, SimpleNonlinearSolve, ADTypes using NonlinearSolveBase: ImmutableNonlinearProblem if CUDA.functional()