Skip to content

Commit

Permalink
test: comprehensive testing of root finding
Browse files Browse the repository at this point in the history
  • Loading branch information
avik-pal committed Oct 7, 2024
1 parent 61ea504 commit 7f15b30
Show file tree
Hide file tree
Showing 9 changed files with 195 additions and 17 deletions.
22 changes: 20 additions & 2 deletions docs/src/release_notes.md
Original file line number Diff line number Diff line change
@@ -1,14 +1,32 @@
# 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.
- Compat for `SimpleNonlinearSolve` has been bumped to `v1`.
- 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`
Expand Down
4 changes: 3 additions & 1 deletion lib/NonlinearSolveBase/src/termination_conditions.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down
10 changes: 5 additions & 5 deletions lib/SimpleNonlinearSolve/src/SimpleNonlinearSolve.jl
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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
Expand Down
3 changes: 2 additions & 1 deletion lib/SimpleNonlinearSolve/src/halley.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
2 changes: 1 addition & 1 deletion lib/SimpleNonlinearSolve/src/lbroyden.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
11 changes: 7 additions & 4 deletions lib/SimpleNonlinearSolve/test/core/23_test_problems_tests.jl
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
@testsnippet RobustnessTestSnippet begin
using NonlinearProblemLibrary, NonlinearSolveBase, LinearAlgebra
using NonlinearProblemLibrary, NonlinearSolveBase, LinearAlgebra, ADTypes

problems = NonlinearProblemLibrary.problems
dicts = NonlinearProblemLibrary.dicts
Expand Down Expand Up @@ -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]] = []
Expand All @@ -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()
Expand All @@ -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]
Expand Down
2 changes: 1 addition & 1 deletion lib/SimpleNonlinearSolve/test/core/exotic_type_tests.jl
Original file line number Diff line number Diff line change
@@ -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)
Expand Down
154 changes: 154 additions & 0 deletions lib/SimpleNonlinearSolve/test/core/rootfind_tests.jl
Original file line number Diff line number Diff line change
@@ -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
4 changes: 2 additions & 2 deletions lib/SimpleNonlinearSolve/test/gpu/cuda_tests.jl
Original file line number Diff line number Diff line change
@@ -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)
Expand Down Expand Up @@ -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()
Expand Down

0 comments on commit 7f15b30

Please sign in to comment.