From 6f0d159da648d9ee4253d90744de00bf8232a746 Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Mon, 17 Jan 2022 23:30:32 -0500 Subject: [PATCH 1/4] specialize Newton on static arrays MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ```julia using StaticArrays using NonlinearSolve function f(x, _) F1 = (x[1] + 3) * (x[2]^3 - 7) + 18 F2 = sin(x[2] * exp(x[1]) - 1) SA[F1,F2] end function f!(F, x) F[1] = (x[1] + 3) * (x[2]^3 - 7) + 18 F[2] = sin(x[2] * exp(x[1]) - 1) nothing end x0 = [0.1; 1.2] x0s = SVector{size(x0)...}(x0) x0m = MVector{size(x0)...}(x0) prob = NonlinearProblem{false}(f, x0s) using NLsolve, BenchmarkTools @btime sol = solve(prob,NewtonRaphson()); # 320.000 ns (2 allocations: 128 bytes) @btime nlsolve(f!, x0m); # 1.460 μs (35 allocations: 1.36 KiB) ``` --- src/scalar.jl | 15 ++++++++++----- src/utils.jl | 3 +++ 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/src/scalar.jl b/src/scalar.jl index 5b3c25966..b5193988a 100644 --- a/src/scalar.jl +++ b/src/scalar.jl @@ -1,12 +1,17 @@ -function SciMLBase.solve(prob::NonlinearProblem{<:Number}, alg::NewtonRaphson, args...; xatol = nothing, xrtol = nothing, maxiters = 1000, kwargs...) +function SciMLBase.solve(prob::NonlinearProblem{<:Union{Number,SVector}}, alg::NewtonRaphson, args...; xatol = nothing, xrtol = nothing, maxiters = 1000, kwargs...) f = Base.Fix2(prob.f, prob.p) x = float(prob.u0) fx = float(prob.u0) T = typeof(x) - atol = xatol !== nothing ? xatol : oneunit(T) * (eps(one(T)))^(4//5) - rtol = xrtol !== nothing ? xrtol : eps(one(T))^(4//5) + atol = xatol !== nothing ? xatol : oneunit(eltype(T)) * (eps(one(eltype(T))))^(4//5) + rtol = xrtol !== nothing ? xrtol : eps(one(eltype(T)))^(4//5) + + if typeof(x) <: Number + xo = oftype(one(eltype(x)), Inf) + else + xo = map(x->oftype(one(eltype(x)), Inf),x) + end - xo = oftype(x, Inf) for i in 1:maxiters if alg_autodiff(alg) fx, dfx = value_derivative(f, x) @@ -52,7 +57,7 @@ end function SciMLBase.solve(prob::NonlinearProblem{<:Number, iip, <:Dual{T,V,P}}, alg::NewtonRaphson, args...; kwargs...) where {iip, T, V, P} sol, partials = scalar_nlsolve_ad(prob, alg, args...; kwargs...) return SciMLBase.build_solution(prob, alg, Dual{T,V,P}(sol.u, partials), sol.resid; retcode=sol.retcode) - + end function SciMLBase.solve(prob::NonlinearProblem{<:Number, iip, <:AbstractArray{<:Dual{T,V,P}}}, alg::NewtonRaphson, args...; kwargs...) where {iip, T, V, P} sol, partials = scalar_nlsolve_ad(prob, alg, args...; kwargs...) diff --git a/src/utils.jl b/src/utils.jl index 4c0ec9d2b..1e3d2e4cc 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -224,6 +224,9 @@ function value_derivative(f::F, x::R) where {F,R} ForwardDiff.value(out), ForwardDiff.extract_derivative(T, out) end +# Todo: improve this dispatch +value_derivative(f::F, x::SVector) where F = f(x),ForwardDiff.jacobian(f, x) + value(x) = x value(x::Dual) = ForwardDiff.value(x) value(x::AbstractArray{<:Dual}) = map(ForwardDiff.value, x) From 711bc5382dadb07d1cae043ccf72ae414d473e4e Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Tue, 18 Jan 2022 05:49:35 -0500 Subject: [PATCH 2/4] Fix forwarddiff overloads --- Project.toml | 3 +- src/scalar.jl | 7 +- test/basictests.jl | 181 ++++++++++++++++++++++++++++++++++++++++++ test/runtests.jl | 192 +++------------------------------------------ 4 files changed, 197 insertions(+), 186 deletions(-) create mode 100644 test/basictests.jl diff --git a/Project.toml b/Project.toml index 336bc6739..08e2eb3bc 100644 --- a/Project.toml +++ b/Project.toml @@ -34,7 +34,8 @@ julia = "1.6" [extras] BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["BenchmarkTools", "Test", "ForwardDiff"] +test = ["BenchmarkTools", "SafeTestsets", "Test", "ForwardDiff"] diff --git a/src/scalar.jl b/src/scalar.jl index b5193988a..f97a9e312 100644 --- a/src/scalar.jl +++ b/src/scalar.jl @@ -15,6 +15,9 @@ function SciMLBase.solve(prob::NonlinearProblem{<:Union{Number,SVector}}, alg::N for i in 1:maxiters if alg_autodiff(alg) fx, dfx = value_derivative(f, x) + elseif x isa AbstractArray + fx = f(x) + dfx = FiniteDiff.finite_difference_jacobian(f, x, alg.diff_type, eltype(x), fx) else fx = f(x) dfx = FiniteDiff.finite_difference_derivative(f, x, alg.diff_type, eltype(x), fx) @@ -54,12 +57,12 @@ function scalar_nlsolve_ad(prob, alg, args...; kwargs...) return sol, partials end -function SciMLBase.solve(prob::NonlinearProblem{<:Number, iip, <:Dual{T,V,P}}, alg::NewtonRaphson, args...; kwargs...) where {iip, T, V, P} +function SciMLBase.solve(prob::NonlinearProblem{<:Union{Number,SVector}, iip, <:Dual{T,V,P}}, alg::NewtonRaphson, args...; kwargs...) where {iip, T, V, P} sol, partials = scalar_nlsolve_ad(prob, alg, args...; kwargs...) return SciMLBase.build_solution(prob, alg, Dual{T,V,P}(sol.u, partials), sol.resid; retcode=sol.retcode) end -function SciMLBase.solve(prob::NonlinearProblem{<:Number, iip, <:AbstractArray{<:Dual{T,V,P}}}, alg::NewtonRaphson, args...; kwargs...) where {iip, T, V, P} +function SciMLBase.solve(prob::NonlinearProblem{<:Union{Number,SVector}, iip, <:AbstractArray{<:Dual{T,V,P}}}, alg::NewtonRaphson, args...; kwargs...) where {iip, T, V, P} sol, partials = scalar_nlsolve_ad(prob, alg, args...; kwargs...) return SciMLBase.build_solution(prob, alg, Dual{T,V,P}(sol.u, partials), sol.resid; retcode=sol.retcode) end diff --git a/test/basictests.jl b/test/basictests.jl new file mode 100644 index 000000000..0e12f32c9 --- /dev/null +++ b/test/basictests.jl @@ -0,0 +1,181 @@ +using NonlinearSolve +using StaticArrays +using BenchmarkTools +using Test + +function benchmark_immutable(f, u0) + probN = NonlinearProblem{false}(f, u0) + solver = init(probN, NewtonRaphson(), tol = 1e-9) + sol = solve!(solver) +end + +function benchmark_mutable(f, u0) + probN = NonlinearProblem{false}(f, u0) + solver = init(probN, NewtonRaphson(), tol = 1e-9) + sol = (reinit!(solver, probN); solve!(solver)) +end + +function benchmark_scalar(f, u0) + probN = NonlinearProblem{false}(f, u0) + sol = (solve(probN, NewtonRaphson())) +end + +function ff(u,p) + u .* u .- 2 +end +const cu0 = @SVector[1.0, 1.0] +function sf(u,p) + u * u - 2 +end +const csu0 = 1.0 + +sol = benchmark_immutable(ff, cu0) +@test sol.retcode === Symbol(NonlinearSolve.DEFAULT) +@test all(sol.u .* sol.u .- 2 .< 1e-9) +sol = benchmark_mutable(ff, cu0) +@test sol.retcode === Symbol(NonlinearSolve.DEFAULT) +@test all(sol.u .* sol.u .- 2 .< 1e-9) +sol = benchmark_scalar(sf, csu0) +@test sol.retcode === Symbol(NonlinearSolve.DEFAULT) +@test sol.u * sol.u - 2 < 1e-9 + +@test (@ballocated benchmark_immutable(ff, cu0)) == 0 +@test (@ballocated benchmark_mutable(ff, cu0)) < 200 +@test (@ballocated benchmark_scalar(sf, csu0)) == 0 + +# AD Tests +using ForwardDiff + +# Immutable +f, u0 = (u, p) -> u .* u .- p, @SVector[1.0, 1.0] + +g = function (p) + probN = NonlinearProblem{false}(f, u0, p) + sol = solve(probN, NewtonRaphson(), tol = 1e-9) + return sol.u[end] +end + +for p in 1.0:0.1:100.0 + @test g(p) ≈ sqrt(p) + @test ForwardDiff.derivative(g, p) ≈ 1/(2*sqrt(p)) +end + +# Scalar +f, u0 = (u, p) -> u * u - p, 1.0 + +# NewtonRaphson +g = function (p) + probN = NonlinearProblem{false}(f, oftype(p, u0), p) + sol = solve(probN, NewtonRaphson()) + return sol.u +end + +@test ForwardDiff.derivative(g, 1.0) ≈ 0.5 + +for p in 1.1:0.1:100.0 + @test g(p) ≈ sqrt(p) + @test ForwardDiff.derivative(g, p) ≈ 1/(2*sqrt(p)) +end + +u0 = (1.0, 20.0) +# Falsi +g = function (p) + probN = NonlinearProblem{false}(f, typeof(p).(u0), p) + sol = solve(probN, Falsi()) + return sol.left +end + +for p in 1.1:0.1:100.0 + @test g(p) ≈ sqrt(p) + @test ForwardDiff.derivative(g, p) ≈ 1/(2*sqrt(p)) +end + +f, u0 = (u, p) -> p[1] * u * u - p[2], (1.0, 100.0) +t = (p) -> [sqrt(p[2] / p[1])] +p = [0.9, 50.0] +for alg in [Bisection(), Falsi()] + global g, p + g = function (p) + probN = NonlinearProblem{false}(f, u0, p) + sol = solve(probN, Bisection()) + return [sol.left] + end + + @test g(p) ≈ [sqrt(p[2] / p[1])] + @test ForwardDiff.jacobian(g, p) ≈ ForwardDiff.jacobian(t, p) +end + +gnewton = function (p) + probN = NonlinearProblem{false}(f, 0.5, p) + sol = solve(probN, NewtonRaphson()) + return [sol.u] +end +@test gnewton(p) ≈ [sqrt(p[2] / p[1])] +@test ForwardDiff.jacobian(gnewton, p) ≈ ForwardDiff.jacobian(t, p) + +# Error Checks + +f, u0 = (u, p) -> u .* u .- 2.0, @SVector[1.0, 1.0] +probN = NonlinearProblem(f, u0) + +@test solve(probN, NewtonRaphson()).u[end] ≈ sqrt(2.0) +@test solve(probN, NewtonRaphson(); immutable = false).u[end] ≈ sqrt(2.0) +@test solve(probN, NewtonRaphson(;autodiff=false)).u[end] ≈ sqrt(2.0) +@test solve(probN, NewtonRaphson(;autodiff=false)).u[end] ≈ sqrt(2.0) + +for u0 in [1.0, [1, 1.0]] + local f, probN, sol + f = (u, p) -> u .* u .- 2.0 + probN = NonlinearProblem(f, u0) + sol = sqrt(2) * u0 + + @test solve(probN, NewtonRaphson()).u ≈ sol + @test solve(probN, NewtonRaphson()).u ≈ sol + @test solve(probN, NewtonRaphson(;autodiff=false)).u ≈ sol +end + +# Bisection Tests +f, u0 = (u, p) -> u .* u .- 2.0, (1.0, 2.0) +probB = NonlinearProblem(f, u0) + +# Falsi +solver = init(probB, Falsi()) +sol = solve!(solver) +@test sol.left ≈ sqrt(2.0) + +# this should call the fast scalar overload +@test solve(probB, Bisection()).left ≈ sqrt(2.0) + +# these should call the iterator version +solver = init(probB, Bisection()) +@test solver isa NonlinearSolve.BracketingImmutableSolver +@test solve!(solver).left ≈ sqrt(2.0) + +# Garuntee Tests for Bisection +f = function (u, p) + if u < 2.0 + return u - 2.0 + elseif u > 3.0 + return u - 3.0 + else + return 0.0 + end +end +probB = NonlinearProblem(f, (0.0, 4.0)) + +solver = init(probB, Bisection(;exact_left = true)) +sol = solve!(solver) +@test f(sol.left, nothing) < 0.0 +@test f(nextfloat(sol.left), nothing) >= 0.0 + +solver = init(probB, Bisection(;exact_right = true)) +sol = solve!(solver) +@test f(sol.right, nothing) > 0.0 +@test f(prevfloat(sol.right), nothing) <= 0.0 + +solver = init(probB, Bisection(;exact_left = true, exact_right = true); immutable = false) +sol = solve!(solver) +@test f(sol.left, nothing) < 0.0 +@test f(nextfloat(sol.left), nothing) >= 0.0 +@test f(sol.right, nothing) > 0.0 +@test f(prevfloat(sol.right), nothing) <= 0.0 diff --git a/test/runtests.jl b/test/runtests.jl index 7b3367904..b578e43d6 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,187 +1,13 @@ -using NonlinearSolve -using StaticArrays -using BenchmarkTools -using Test +using Pkg +using SafeTestsets +const LONGER_TESTS = false -function benchmark_immutable(f, u0) - probN = NonlinearProblem{false}(f, u0) - solver = init(probN, NewtonRaphson(), tol = 1e-9) - sol = solve!(solver) -end - -function benchmark_mutable(f, u0) - probN = NonlinearProblem{false}(f, u0) - solver = init(probN, NewtonRaphson(), immutable = false, tol = 1e-9) - sol = (reinit!(solver, probN); solve!(solver)) -end - -function benchmark_scalar(f, u0) - probN = NonlinearProblem{false}(f, u0) - sol = (solve(probN, NewtonRaphson())) -end - -function ff(u,p) - u .* u .- 2 -end -const cu0 = @SVector[1.0, 1.0] -function sf(u,p) - u * u - 2 -end -const csu0 = 1.0 - -sol = benchmark_immutable(ff, cu0) -@test sol.retcode === Symbol(NonlinearSolve.DEFAULT) -@test all(sol.u .* sol.u .- 2 .< 1e-9) -sol = benchmark_mutable(ff, cu0) -@test sol.retcode === Symbol(NonlinearSolve.DEFAULT) -@test all(sol.u .* sol.u .- 2 .< 1e-9) -sol = benchmark_scalar(sf, csu0) -@test sol.retcode === Symbol(NonlinearSolve.DEFAULT) -@test sol.u * sol.u - 2 < 1e-9 - -@test (@ballocated benchmark_immutable(ff, cu0)) == 0 -@test (@ballocated benchmark_mutable(ff, cu0)) < 200 -@test (@ballocated benchmark_scalar(sf, csu0)) == 0 - -# AD Tests -using ForwardDiff - -# Immutable -f, u0 = (u, p) -> u .* u .- p, @SVector[1.0, 1.0] - -g = function (p) - probN = NonlinearProblem{false}(f, u0, p) - sol = solve(probN, NewtonRaphson(), immutable = true, tol = 1e-9) - return sol.u[end] -end - -for p in 1.0:0.1:100.0 - @test g(p) ≈ sqrt(p) - @test ForwardDiff.derivative(g, p) ≈ 1/(2*sqrt(p)) -end - -# Scalar -f, u0 = (u, p) -> u * u - p, 1.0 - -# NewtonRaphson -g = function (p) - probN = NonlinearProblem{false}(f, oftype(p, u0), p) - sol = solve(probN, NewtonRaphson()) - return sol.u -end - -@test ForwardDiff.derivative(g, 1.0) ≈ 0.5 - -for p in 1.1:0.1:100.0 - @test g(p) ≈ sqrt(p) - @test ForwardDiff.derivative(g, p) ≈ 1/(2*sqrt(p)) -end +const GROUP = get(ENV, "GROUP", "All") +const is_APPVEYOR = Sys.iswindows() && haskey(ENV,"APPVEYOR") -u0 = (1.0, 20.0) -# Falsi -g = function (p) - probN = NonlinearProblem{false}(f, typeof(p).(u0), p) - sol = solve(probN, Falsi()) - return sol.left -end - -for p in 1.1:0.1:100.0 - @test g(p) ≈ sqrt(p) - @test ForwardDiff.derivative(g, p) ≈ 1/(2*sqrt(p)) -end - -f, u0 = (u, p) -> p[1] * u * u - p[2], (1.0, 100.0) -t = (p) -> [sqrt(p[2] / p[1])] -p = [0.9, 50.0] -for alg in [Bisection(), Falsi()] - global g, p - g = function (p) - probN = NonlinearProblem{false}(f, u0, p) - sol = solve(probN, Bisection()) - return [sol.left] - end - - @test g(p) ≈ [sqrt(p[2] / p[1])] - @test ForwardDiff.jacobian(g, p) ≈ ForwardDiff.jacobian(t, p) -end - -gnewton = function (p) - probN = NonlinearProblem{false}(f, 0.5, p) - sol = solve(probN, NewtonRaphson()) - return [sol.u] -end -@test gnewton(p) ≈ [sqrt(p[2] / p[1])] -@test ForwardDiff.jacobian(gnewton, p) ≈ ForwardDiff.jacobian(t, p) - -# Error Checks - -f, u0 = (u, p) -> u .* u .- 2.0, @SVector[1.0, 1.0] -probN = NonlinearProblem(f, u0) - -@test solve(probN, NewtonRaphson()).u[end] ≈ sqrt(2.0) -@test solve(probN, NewtonRaphson(); immutable = false).u[end] ≈ sqrt(2.0) -@test solve(probN, NewtonRaphson(;autodiff=false)).u[end] ≈ sqrt(2.0) -@test solve(probN, NewtonRaphson(;autodiff=false); immutable = false).u[end] ≈ sqrt(2.0) +@time begin -for u0 in [1.0, [1, 1.0]] - local f, probN, sol - f = (u, p) -> u .* u .- 2.0 - probN = NonlinearProblem(f, u0) - sol = sqrt(2) * u0 - - @test solve(probN, NewtonRaphson()).u ≈ sol - @test solve(probN, NewtonRaphson(); immutable = false).u ≈ sol - @test solve(probN, NewtonRaphson(;autodiff=false)).u ≈ sol - @test solve(probN, NewtonRaphson(;autodiff=false); immutable = false).u ≈ sol +if GROUP == "All" || GROUP == "Interface" + #@time @safetestset "Linear Solver Tests" begin include("interface/linear_solver_test.jl") end + @time @safetestset "Basic Tests + Some AD" begin include("basictests.jl") end end - -# Bisection Tests -f, u0 = (u, p) -> u .* u .- 2.0, (1.0, 2.0) -probB = NonlinearProblem(f, u0) - -# Falsi -solver = init(probB, Falsi()) -sol = solve!(solver) -@test sol.left ≈ sqrt(2.0) - -# this should call the fast scalar overload -@test solve(probB, Bisection()).left ≈ sqrt(2.0) - -# these should call the iterator version -solver = init(probB, Bisection()) -@test solver isa NonlinearSolve.BracketingImmutableSolver -# Question: Do we need BracketingImmutableSolver? We have fast scalar overload and -# Bracketing solvers work only for scalars. - -solver = init(probB, Bisection(); immutable = false) -# @test solver isa NonlinearSolve.BracketingSolver -@test solve!(solver).left ≈ sqrt(2.0) - -# Garuntee Tests for Bisection -f = function (u, p) - if u < 2.0 - return u - 2.0 - elseif u > 3.0 - return u - 3.0 - else - return 0.0 - end -end -probB = NonlinearProblem(f, (0.0, 4.0)) - -solver = init(probB, Bisection(;exact_left = true); immutable = false) -sol = solve!(solver) -@test f(sol.left, nothing) < 0.0 -@test f(nextfloat(sol.left), nothing) >= 0.0 - -solver = init(probB, Bisection(;exact_right = true); immutable = false) -sol = solve!(solver) -@test f(sol.right, nothing) > 0.0 -@test f(prevfloat(sol.right), nothing) <= 0.0 - -solver = init(probB, Bisection(;exact_left = true, exact_right = true); immutable = false) -sol = solve!(solver) -@test f(sol.left, nothing) < 0.0 -@test f(nextfloat(sol.left), nothing) >= 0.0 -@test f(sol.right, nothing) > 0.0 -@test f(prevfloat(sol.right), nothing) <= 0.0 From 82bb7da7c8a112dc0b6d36d8dd0be4d3b390ef11 Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Tue, 18 Jan 2022 06:01:27 -0500 Subject: [PATCH 3/4] fix test setup --- Project.toml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 08e2eb3bc..cd2a09d1b 100644 --- a/Project.toml +++ b/Project.toml @@ -34,8 +34,9 @@ julia = "1.6" [extras] BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["BenchmarkTools", "SafeTestsets", "Test", "ForwardDiff"] +test = ["BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff"] From 47bcc55378b8c901a742852111454e019904f94f Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Tue, 18 Jan 2022 06:10:08 -0500 Subject: [PATCH 4/4] fix test setup --- test/basictests.jl | 2 +- test/runtests.jl | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/test/basictests.jl b/test/basictests.jl index 0e12f32c9..bdcb4b457 100644 --- a/test/basictests.jl +++ b/test/basictests.jl @@ -50,7 +50,7 @@ using ForwardDiff f, u0 = (u, p) -> u .* u .- p, @SVector[1.0, 1.0] g = function (p) - probN = NonlinearProblem{false}(f, u0, p) + probN = NonlinearProblem{false}(f, csu0, p) sol = solve(probN, NewtonRaphson(), tol = 1e-9) return sol.u[end] end diff --git a/test/runtests.jl b/test/runtests.jl index b578e43d6..58afd0c94 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -11,3 +11,5 @@ if GROUP == "All" || GROUP == "Interface" #@time @safetestset "Linear Solver Tests" begin include("interface/linear_solver_test.jl") end @time @safetestset "Basic Tests + Some AD" begin include("basictests.jl") end end + +end