From 24f83d8be056a367ce942cfd5b733b8c9213febb Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Sat, 25 May 2024 16:29:04 -0700 Subject: [PATCH 1/4] Use DifferentiationInterface --- Project.toml | 6 +- ...leNonlinearSolvePolyesterForwardDiffExt.jl | 20 -- src/SimpleNonlinearSolve.jl | 6 +- src/nlsolve/halley.jl | 14 +- src/nlsolve/raphson.jl | 11 +- src/nlsolve/trustRegion.jl | 13 +- src/utils.jl | 238 ++++-------------- 7 files changed, 78 insertions(+), 230 deletions(-) delete mode 100644 ext/SimpleNonlinearSolvePolyesterForwardDiffExt.jl diff --git a/Project.toml b/Project.toml index 7023d27..5284b3e 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "SimpleNonlinearSolve" uuid = "727e6d20-b764-4bd8-a329-72de5adea6c7" authors = ["SciML"] -version = "1.8.1" +version = "1.9.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" @@ -9,6 +9,7 @@ ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471" DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5" +DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a" FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" @@ -21,7 +22,6 @@ StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" [weakdeps] ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" -PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" @@ -29,7 +29,6 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [extensions] SimpleNonlinearSolveChainRulesCoreExt = "ChainRulesCore" -SimpleNonlinearSolvePolyesterForwardDiffExt = "PolyesterForwardDiff" SimpleNonlinearSolveReverseDiffExt = "ReverseDiff" SimpleNonlinearSolveStaticArraysExt = "StaticArrays" SimpleNonlinearSolveTrackerExt = "Tracker" @@ -45,6 +44,7 @@ ChainRulesCore = "1.22" ConcreteStructs = "0.2.3" DiffEqBase = "6.149" DiffResults = "1.1" +DifferentiationInterface = "0.4" ExplicitImports = "1.5.0" FastClosures = "0.3.2" FiniteDiff = "2.22" diff --git a/ext/SimpleNonlinearSolvePolyesterForwardDiffExt.jl b/ext/SimpleNonlinearSolvePolyesterForwardDiffExt.jl deleted file mode 100644 index ac898ac..0000000 --- a/ext/SimpleNonlinearSolvePolyesterForwardDiffExt.jl +++ /dev/null @@ -1,20 +0,0 @@ -module SimpleNonlinearSolvePolyesterForwardDiffExt - -using PolyesterForwardDiff: PolyesterForwardDiff -using SimpleNonlinearSolve: SimpleNonlinearSolve - -@inline SimpleNonlinearSolve.__is_extension_loaded(::Val{:PolyesterForwardDiff}) = true - -@inline function SimpleNonlinearSolve.__polyester_forwarddiff_jacobian!( - f!::F, y, J, x, chunksize) where {F} - PolyesterForwardDiff.threaded_jacobian!(f!, y, J, x, chunksize) - return J -end - -@inline function SimpleNonlinearSolve.__polyester_forwarddiff_jacobian!( - f::F, J, x, chunksize) where {F} - PolyesterForwardDiff.threaded_jacobian!(f, J, x, chunksize) - return J -end - -end diff --git a/src/SimpleNonlinearSolve.jl b/src/SimpleNonlinearSolve.jl index 1ff66b9..a024f06 100644 --- a/src/SimpleNonlinearSolve.jl +++ b/src/SimpleNonlinearSolve.jl @@ -3,13 +3,15 @@ module SimpleNonlinearSolve using PrecompileTools: @compile_workload, @setup_workload, @recompile_invalidations @recompile_invalidations begin - using ADTypes: ADTypes, AutoFiniteDiff, AutoForwardDiff, AutoPolyesterForwardDiff + using ADTypes: ADTypes, AbstractADType, AutoFiniteDiff, AutoForwardDiff, + AutoPolyesterForwardDiff using ArrayInterface: ArrayInterface using ConcreteStructs: @concrete using DiffEqBase: DiffEqBase, AbstractNonlinearTerminationMode, AbstractSafeNonlinearTerminationMode, AbstractSafeBestNonlinearTerminationMode, AbsNormTerminationMode, NONLINEARSOLVE_DEFAULT_NORM + using DifferentiationInterface: DifferentiationInterface using DiffResults: DiffResults using FastClosures: @closure using FiniteDiff: FiniteDiff @@ -25,6 +27,8 @@ using PrecompileTools: @compile_workload, @setup_workload, @recompile_invalidati using StaticArraysCore: StaticArray, SVector, SMatrix, SArray, MArray, Size end +const DI = DifferentiationInterface + @reexport using SciMLBase abstract type AbstractSimpleNonlinearSolveAlgorithm <: AbstractNonlinearAlgorithm end diff --git a/src/nlsolve/halley.jl b/src/nlsolve/halley.jl index e550258..5ff5eea 100644 --- a/src/nlsolve/halley.jl +++ b/src/nlsolve/halley.jl @@ -12,8 +12,9 @@ A low-overhead implementation of Halley's Method. ### Keyword Arguments - - `autodiff`: determines the backend used for the Hessian. Defaults to `nothing`. Valid - choices are `AutoForwardDiff()` or `AutoFiniteDiff()`. + - `autodiff`: determines the backend used for the Hessian. Defaults to `nothing` (i.e. + automatic backend selection). Valid choices include backends from + `DifferentiationInterface.jl`. !!! warning @@ -26,13 +27,11 @@ end function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleHalley, args...; abstol = nothing, reltol = nothing, maxiters = 1000, alias_u0 = false, termination_condition = nothing, kwargs...) - isinplace(prob) && - error("SimpleHalley currently only supports out-of-place nonlinear problems") - x = __maybe_unaliased(prob.u0, alias_u0) fx = _get_fx(prob, x) T = eltype(x) + f = __fixed_parameter_function(prob) autodiff = __get_concrete_autodiff(prob, alg.autodiff) abstol, reltol, tc_cache = init_termination_cache( prob, abstol, reltol, fx, x, termination_condition) @@ -51,7 +50,7 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleHalley, args...; for i in 1:maxiters # Hessian Computation is unfortunately type unstable - fx, dfx, d2fx = compute_jacobian_and_hessian(autodiff, prob, fx, x) + fx, dfx, d2fx = compute_jacobian_and_hessian(autodiff, prob, f, fx, x) setindex_trait(x) === CannotSetindex() && (A = dfx) # Factorize Once and Reuse @@ -78,9 +77,8 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleHalley, args...; cᵢ = _restructure(cᵢ, cᵢ_) if i == 1 - if iszero(fx) + iszero(fx) && return build_solution(prob, alg, x, fx; retcode = ReturnCode.Success) - end else # Termination Checks tc_sol = check_termination(tc_cache, fx, x, xo, prob, alg) diff --git a/src/nlsolve/raphson.jl b/src/nlsolve/raphson.jl index 7b419ce..ca6864b 100644 --- a/src/nlsolve/raphson.jl +++ b/src/nlsolve/raphson.jl @@ -13,9 +13,9 @@ and static array problems. ### Keyword Arguments - - `autodiff`: determines the backend used for the Jacobian. Defaults to - `nothing`. Valid choices are `AutoPolyesterForwardDiff()`, `AutoForwardDiff()` or - `AutoFiniteDiff()`. + - `autodiff`: determines the backend used for the Jacobian. Defaults to `nothing` (i.e. + automatic backend selection). Valid choices include jacobian backends from + `DifferentiationInterface.jl`. """ @kwdef @concrete struct SimpleNewtonRaphson <: AbstractNewtonAlgorithm autodiff = nothing @@ -30,13 +30,14 @@ function SciMLBase.__solve(prob::Union{NonlinearProblem, NonlinearLeastSquaresPr fx = _get_fx(prob, x) autodiff = __get_concrete_autodiff(prob, alg.autodiff) @bb xo = copy(x) - J, jac_cache = jacobian_cache(autodiff, prob.f, fx, x, prob.p) + f = __fixed_parameter_function(prob) + J, jac_cache = jacobian_cache(autodiff, prob, f, fx, x) abstol, reltol, tc_cache = init_termination_cache( prob, abstol, reltol, fx, x, termination_condition) for i in 1:maxiters - fx, dfx = value_and_jacobian(autodiff, prob.f, fx, x, prob.p, jac_cache; J) + fx, dfx = value_and_jacobian(autodiff, prob, f, fx, x, jac_cache; J) if i == 1 iszero(fx) && build_solution(prob, alg, x, fx; retcode = ReturnCode.Success) diff --git a/src/nlsolve/trustRegion.jl b/src/nlsolve/trustRegion.jl index a19cf2c..6ff263e 100644 --- a/src/nlsolve/trustRegion.jl +++ b/src/nlsolve/trustRegion.jl @@ -10,9 +10,9 @@ scalar and static array problems. ### Keyword Arguments - - `autodiff`: determines the backend used for the Jacobian. Defaults to - `nothing`. Valid choices are `AutoPolyesterForwardDiff()`, `AutoForwardDiff()` or - `AutoFiniteDiff()`. + - `autodiff`: determines the backend used for the Jacobian. Defaults to `nothing` (i.e. + automatic backend selection). Valid choices include jacobian backends from + `DifferentiationInterface.jl`. - `max_trust_radius`: the maximum radius of the trust region. Defaults to `max(norm(f(u0)), maximum(u0) - minimum(u0))`. - `initial_trust_radius`: the initial trust region radius. Defaults to @@ -85,8 +85,9 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleTrustRegion, args. fx = _get_fx(prob, x) norm_fx = norm(fx) @bb xo = copy(x) - J, jac_cache = jacobian_cache(autodiff, prob.f, fx, x, prob.p) - fx, ∇f = value_and_jacobian(autodiff, prob.f, fx, x, prob.p, jac_cache; J) + f = __fixed_parameter_function(prob) + J, jac_cache = jacobian_cache(autodiff, prob, f, fx, x) + fx, ∇f = value_and_jacobian(autodiff, prob, f, fx, x, jac_cache; J) abstol, reltol, tc_cache = init_termination_cache( prob, abstol, reltol, fx, x, termination_condition) @@ -144,7 +145,7 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleTrustRegion, args. # Take the step. @bb @. xo = x - fx, ∇f = value_and_jacobian(autodiff, prob.f, fx, x, prob.p, jac_cache; J) + fx, ∇f = value_and_jacobian(autodiff, prob, f, fx, x, jac_cache; J) # Update the trust region radius. if !_unwrap_val(alg.nlsolve_update_rule) && r > η₃ diff --git a/src/utils.jl b/src/utils.jl index 333e54d..ae380d7 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -1,9 +1,4 @@ -struct SimpleNonlinearSolveTag end - -function ForwardDiff.checktag(::Type{<:ForwardDiff.Tag{<:SimpleNonlinearSolveTag, <:T}}, - f::F, x::AbstractArray{T}) where {T, F} - return true -end +struct HasAnalyticJacobian end """ __prevfloat_tdir(x, x0, x1) @@ -26,203 +21,63 @@ Return the maximum of `a` and `b` if `x1 > x0`, otherwise return the minimum. """ __max_tdir(a, b, x0, x1) = ifelse(x1 > x0, max(a, b), min(a, b)) -__standard_tag(::Nothing, x) = ForwardDiff.Tag(SimpleNonlinearSolveTag(), eltype(x)) -__standard_tag(tag::ForwardDiff.Tag, _) = tag -__standard_tag(tag, x) = ForwardDiff.Tag(tag, eltype(x)) - -__pick_forwarddiff_chunk(x) = ForwardDiff.Chunk(length(x)) -function __pick_forwarddiff_chunk(x::StaticArray) - L = prod(Size(x)) - if L ≤ ForwardDiff.DEFAULT_CHUNK_THRESHOLD - return ForwardDiff.Chunk{L}() - else - return ForwardDiff.Chunk{ForwardDiff.DEFAULT_CHUNK_THRESHOLD}() - end -end - -function __get_jacobian_config(ad::AutoForwardDiff{CS}, f::F, x) where {F, CS} - ck = (CS === nothing || CS ≤ 0) ? __pick_forwarddiff_chunk(x) : ForwardDiff.Chunk{CS}() - tag = __standard_tag(ad.tag, x) - return __forwarddiff_jacobian_config(f, x, ck, tag) -end -function __get_jacobian_config(ad::AutoForwardDiff{CS}, f!::F, y, x) where {F, CS} - ck = (CS === nothing || CS ≤ 0) ? __pick_forwarddiff_chunk(x) : ForwardDiff.Chunk{CS}() - tag = __standard_tag(ad.tag, x) - return ForwardDiff.JacobianConfig(f!, y, x, ck, tag) -end - -function __forwarddiff_jacobian_config(f::F, x, ck::ForwardDiff.Chunk, tag) where {F} - return ForwardDiff.JacobianConfig(f, x, ck, tag) -end -function __forwarddiff_jacobian_config( - f::F, x::SArray, ck::ForwardDiff.Chunk{N}, tag) where {F, N} - seeds = ForwardDiff.construct_seeds(ForwardDiff.Partials{N, eltype(x)}) - duals = ForwardDiff.Dual{typeof(tag), eltype(x), N}.(x) - return ForwardDiff.JacobianConfig{typeof(tag), eltype(x), N, typeof(duals)}( - seeds, duals) +function __fixed_parameter_function(prob::NonlinearProblem) + isinplace(prob) && return @closure (du, u) -> prob.f(du, u, prob.p) + return Base.Fix2(prob.f, prob.p) end -function __get_jacobian_config(ad::AutoPolyesterForwardDiff{CS}, args...) where {CS} - x = last(args) - return (CS === nothing || CS ≤ 0) ? __pick_forwarddiff_chunk(x) : - ForwardDiff.Chunk{CS}() -end +function value_and_jacobian( + ad, prob::NonlinearProblem, f::F, y, x, cache; J = nothing) where {F} + x isa Number && return DI.value_and_derivative(f, ad, x, cache) -""" - value_and_jacobian(ad, f, y, x, p, cache; J = nothing) - -Compute `f(x), d/dx f(x)` in the most efficient way based on `ad`. None of the arguments -except `cache` (& `J` if not nothing) are mutated. -""" -function value_and_jacobian(ad, f::F, y, x::X, p, cache; J = nothing) where {F, X} - if isinplace(f) - _f = (du, u) -> f(du, u, p) - if SciMLBase.has_jac(f) - f.jac(J, x, p) - _f(y, x) - return y, J - elseif ad isa AutoForwardDiff - res = DiffResults.DiffResult(y, J) - ForwardDiff.jacobian!(res, _f, y, x, cache) - return DiffResults.value(res), DiffResults.jacobian(res) - elseif ad isa AutoFiniteDiff - FiniteDiff.finite_difference_jacobian!(J, _f, x, cache) - _f(y, x) - return y, J - elseif ad isa AutoPolyesterForwardDiff - __polyester_forwarddiff_jacobian!(_f, y, J, x, cache) - return y, J + if isinplace(prob) + if cache isa HasAnalyticJacobian + prob.f.jac(J, x, p) + f(y, x) else - throw(ArgumentError("Unsupported AD method: $(ad)")) + DI.jacobian!(f, y, J, ad, x, cache) end + return y, J else - _f = Base.Fix2(f, p) - if SciMLBase.has_jac(f) - return _f(x), f.jac(x, p) - elseif ad isa AutoForwardDiff - if ArrayInterface.can_setindex(x) - res = DiffResults.DiffResult(y, J) - ForwardDiff.jacobian!(res, _f, x, cache) - return DiffResults.value(res), DiffResults.jacobian(res) - else - J_fd = ForwardDiff.jacobian(_f, x, cache) - return _f(x), J_fd - end - elseif ad isa AutoFiniteDiff - J_fd = FiniteDiff.finite_difference_jacobian(_f, x, cache) - return _f(x), J_fd - elseif ad isa AutoPolyesterForwardDiff - __polyester_forwarddiff_jacobian!(_f, J, x, cache) - return _f(x), J - else - throw(ArgumentError("Unsupported AD method: $(ad)")) - end + cache isa HasAnalyticJacobian && return f(x), prob.f.jac(x, prob.p) + J === nothing && return DI.value_and_jacobian(f, ad, x, cache) + y, _ = DI.value_and_jacobian!(f, J, ad, x, cache) + return y, J end end -# Declare functions -function __polyester_forwarddiff_jacobian! end - -function value_and_jacobian(ad, f::F, y, x::Number, p, cache; J = nothing) where {F} - if SciMLBase.has_jac(f) - return f(x, p), f.jac(x, p) - elseif ad isa AutoForwardDiff - T = typeof(__standard_tag(ad.tag, x)) - out = f(ForwardDiff.Dual{T}(x, one(x)), p) - return ForwardDiff.value(out), ForwardDiff.extract_derivative(T, out) - elseif ad isa AutoPolyesterForwardDiff - # Just use ForwardDiff - T = typeof(__standard_tag(nothing, x)) - out = f(ForwardDiff.Dual{T}(x, one(x)), p) - return ForwardDiff.value(out), ForwardDiff.extract_derivative(T, out) - elseif ad isa AutoFiniteDiff - _f = Base.Fix2(f, p) - return _f(x), FiniteDiff.finite_difference_derivative(_f, x, ad.fdtype) - else - throw(ArgumentError("Unsupported AD method: $(ad)")) - end -end +function jacobian_cache(ad, prob::NonlinearProblem, f::F, y, x) where {F} + x isa Number && return (nothing, DI.prepare_derivative(f, ad, x)) -""" - jacobian_cache(ad, f, y, x, p) --> J, cache - -Returns a Jacobian Matrix and a cache for the Jacobian computation. -""" -function jacobian_cache(ad, f::F, y, x::X, p) where {F, X <: AbstractArray} - if isinplace(f) - _f = (du, u) -> f(du, u, p) + if isinplace(prob) J = similar(y, length(y), length(x)) - if SciMLBase.has_jac(f) - return J, nothing - elseif ad isa AutoForwardDiff || ad isa AutoPolyesterForwardDiff - return J, __get_jacobian_config(ad, _f, y, x) - elseif ad isa AutoFiniteDiff - return J, FiniteDiff.JacobianCache(copy(x), copy(y), copy(y), ad.fdtype) - else - throw(ArgumentError("Unsupported AD method: $(ad)")) - end + SciMLBase.has_jac(prob.f) && return J, HasAnalyticJacobian() + return J, DI.prepare_jacobian(f, y, ad, x) else - _f = Base.Fix2(f, p) - if SciMLBase.has_jac(f) - return nothing, nothing - elseif ad isa AutoForwardDiff - J = ArrayInterface.can_setindex(x) ? similar(y, length(y), length(x)) : nothing - return J, __get_jacobian_config(ad, _f, x) - elseif ad isa AutoPolyesterForwardDiff - @assert ArrayInterface.can_setindex(x) "PolyesterForwardDiff requires mutable inputs. Use AutoForwardDiff instead." - J = similar(y, length(y), length(x)) - return J, __get_jacobian_config(ad, _f, x) - elseif ad isa AutoFiniteDiff - return nothing, FiniteDiff.JacobianCache(copy(x), copy(y), copy(y), ad.fdtype) - else - throw(ArgumentError("Unsupported AD method: $(ad)")) - end + SciMLBase.has_jac(prob.f) && return nothing, HasAnalyticJacobian() + J = ArrayInterface.can_setindex(x) ? similar(y, length(y), length(x)) : nothing + return J, DI.prepare_jacobian(f, ad, x) end end -jacobian_cache(ad, f::F, y, x::Number, p) where {F} = nothing, nothing - -function compute_jacobian_and_hessian(ad::AutoForwardDiff, prob, _, x::Number) - fx = prob.f(x, prob.p) - J_fn = Base.Fix1(ForwardDiff.derivative, Base.Fix2(prob.f, prob.p)) - dfx = J_fn(x) - d2fx = ForwardDiff.derivative(J_fn, x) - return fx, dfx, d2fx -end - -function compute_jacobian_and_hessian(ad::AutoForwardDiff, prob, fx, x) - if isinplace(prob) - error("Inplace version for Nested ForwardDiff Not Implemented Yet!") - else - f = Base.Fix2(prob.f, prob.p) - fx = f(x) - J_fn = Base.Fix1(ForwardDiff.jacobian, f) - dfx = J_fn(x) - d2fx = ForwardDiff.jacobian(J_fn, x) - return fx, dfx, d2fx +function compute_jacobian_and_hessian(ad, prob::NonlinearProblem, f::F, y, x) where {F} + if x isa Number + df = @closure x -> DI.derivative(f, ad, x) + return f(x), df(x), DI.derivative(df, ad, x) end -end - -function compute_jacobian_and_hessian(ad::AutoFiniteDiff, prob, _, x::Number) - fx = prob.f(x, prob.p) - J_fn = x -> FiniteDiff.finite_difference_derivative( - Base.Fix2(prob.f, prob.p), x, ad.fdtype) - dfx = J_fn(x) - d2fx = FiniteDiff.finite_difference_derivative(J_fn, x, ad.fdtype) - return fx, dfx, d2fx -end -function compute_jacobian_and_hessian(ad::AutoFiniteDiff, prob, fx, x) if isinplace(prob) - error("Inplace version for Nested FiniteDiff Not Implemented Yet!") - else - f = Base.Fix2(prob.f, prob.p) - fx = f(x) - J_fn = x -> FiniteDiff.finite_difference_jacobian(f, x, ad.fdtype) - dfx = J_fn(x) - d2fx = FiniteDiff.finite_difference_jacobian(J_fn, x, ad.fdtype) - return fx, dfx, d2fx + df = @closure x -> begin + res = similar(y, promote_type(eltype(y), eltype(x))) + return DI.jacobian(f, res, ad, x) + end + J, H = DI.value_and_jacobian(df, ad, x) + f(y, x) + return y, J, H end + + df = @closure x -> DI.jacobian(f, ad, x) + return f(x), df(x), DI.jacobian(df, ad, x) end __init_identity_jacobian(u::Number, fu, α = true) = oftype(u, α) @@ -360,10 +215,19 @@ end end # Decide which AD backend to use -@inline __get_concrete_autodiff(prob, ad::ADTypes.AbstractADType; kwargs...) = ad +@inline __get_concrete_autodiff(prob, ad::AbstractADType; kwargs...) = ad +@inline function __get_concrete_autodiff(prob, ad::AutoForwardDiff{nothing}; kwargs...) + return AutoForwardDiff(; chunksize = ForwardDiff.pickchunksize(length(prob.u0)), ad.tag) +end +@inline function __get_concrete_autodiff( + prob, ad::AutoPolyesterForwardDiff{nothing}; kwargs...) + return AutoPolyesterForwardDiff(; + chunksize = ForwardDiff.pickchunksize(length(prob.u0)), ad.tag) +end @inline function __get_concrete_autodiff(prob, ::Nothing; kwargs...) - return ifelse( - ForwardDiff.can_dual(eltype(prob.u0)), AutoForwardDiff(), AutoFiniteDiff()) + return ifelse(ForwardDiff.can_dual(eltype(prob.u0)), + AutoForwardDiff(; chunksize = ForwardDiff.pickchunksize(length(prob.u0))), + AutoFiniteDiff()) end @inline __reshape(x::Number, args...) = x From ae0bf1017dcce1612186886662890738d0fe1605 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Sat, 25 May 2024 17:31:06 -0700 Subject: [PATCH 2/4] Clean up AD --- Project.toml | 2 +- src/SimpleNonlinearSolve.jl | 8 ++-- src/ad.jl | 92 ++++++++++++++----------------------- src/nlsolve/halley.jl | 6 +-- src/nlsolve/lbroyden.jl | 6 +-- src/utils.jl | 32 +++++++++---- 6 files changed, 68 insertions(+), 78 deletions(-) diff --git a/Project.toml b/Project.toml index 5284b3e..a7aa31f 100644 --- a/Project.toml +++ b/Project.toml @@ -47,7 +47,7 @@ DiffResults = "1.1" DifferentiationInterface = "0.4" ExplicitImports = "1.5.0" FastClosures = "0.3.2" -FiniteDiff = "2.22" +FiniteDiff = "2.23.1" ForwardDiff = "0.10.36" LinearAlgebra = "1.10" LinearSolve = "2.30" diff --git a/src/SimpleNonlinearSolve.jl b/src/SimpleNonlinearSolve.jl index a024f06..04a32c9 100644 --- a/src/SimpleNonlinearSolve.jl +++ b/src/SimpleNonlinearSolve.jl @@ -20,10 +20,10 @@ using PrecompileTools: @compile_workload, @setup_workload, @recompile_invalidati mul!, norm, transpose using MaybeInplace: @bb, setindex_trait, CanSetindex, CannotSetindex using Reexport: @reexport - using SciMLBase: SciMLBase, IntervalNonlinearProblem, NonlinearFunction, - NonlinearLeastSquaresProblem, NonlinearProblem, ReturnCode, init, - remake, solve, AbstractNonlinearAlgorithm, build_solution, isinplace, - _unwrap_val + using SciMLBase: SciMLBase, AbstractNonlinearProblem, IntervalNonlinearProblem, + NonlinearFunction, NonlinearLeastSquaresProblem, NonlinearProblem, + ReturnCode, init, remake, solve, AbstractNonlinearAlgorithm, + build_solution, isinplace, _unwrap_val using StaticArraysCore: StaticArray, SVector, SMatrix, SArray, MArray, Size end diff --git a/src/ad.jl b/src/ad.jl index f42651b..ffae7cc 100644 --- a/src/ad.jl +++ b/src/ad.jl @@ -1,25 +1,15 @@ -function SciMLBase.solve( - prob::NonlinearProblem{<:Union{Number, <:AbstractArray}, iip, - <:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}}}, - alg::AbstractSimpleNonlinearSolveAlgorithm, - args...; - kwargs...) where {T, V, P, iip} - sol, partials = __nlsolve_ad(prob, alg, args...; kwargs...) - dual_soln = __nlsolve_dual_soln(sol.u, partials, prob.p) - return SciMLBase.build_solution( - prob, alg, dual_soln, sol.resid; sol.retcode, sol.stats, sol.original) -end - -function SciMLBase.solve( - prob::NonlinearLeastSquaresProblem{ - <:AbstractArray, iip, <:Union{<:AbstractArray{<:Dual{T, V, P}}}}, - alg::AbstractSimpleNonlinearSolveAlgorithm, - args...; - kwargs...) where {T, V, P, iip} - sol, partials = __nlsolve_ad(prob, alg, args...; kwargs...) - dual_soln = __nlsolve_dual_soln(sol.u, partials, prob.p) - return SciMLBase.build_solution( - prob, alg, dual_soln, sol.resid; sol.retcode, sol.stats, sol.original) +for pType in (NonlinearProblem, NonlinearLeastSquaresProblem) + @eval function SciMLBase.solve( + prob::$(pType){<:Union{Number, <:AbstractArray}, iip, + <:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}}}, + alg::AbstractSimpleNonlinearSolveAlgorithm, + args...; + kwargs...) where {T, V, P, iip} + sol, partials = __nlsolve_ad(prob, alg, args...; kwargs...) + dual_soln = __nlsolve_dual_soln(sol.u, partials, prob.p) + return SciMLBase.build_solution( + prob, alg, dual_soln, sol.resid; sol.retcode, sol.stats, sol.original) + end end for algType in (Bisection, Brent, Alefeld, Falsi, ITP, Ridder) @@ -47,8 +37,7 @@ function __nlsolve_ad( tspan = value.(prob.tspan) newprob = IntervalNonlinearProblem(prob.f, tspan, p; prob.kwargs...) else - u0 = value(prob.u0) - newprob = NonlinearProblem(prob.f, u0, p; prob.kwargs...) + newprob = remake(prob; p, u0 = value(prob.u0)) end sol = solve(newprob, alg, args...; kwargs...) @@ -73,12 +62,8 @@ function __nlsolve_ad( end function __nlsolve_ad(prob::NonlinearLeastSquaresProblem, alg, args...; kwargs...) - p = value(prob.p) - u0 = value(prob.u0) - newprob = NonlinearLeastSquaresProblem(prob.f, u0, p; prob.kwargs...) - + newprob = remake(prob; p = value(prob.p), u0 = value(prob.u0)) sol = solve(newprob, alg, args...; kwargs...) - uu = sol.u # First check for custom `vjp` then custom `Jacobian` and if nothing is provided use @@ -86,7 +71,7 @@ function __nlsolve_ad(prob::NonlinearLeastSquaresProblem, alg, args...; kwargs.. if SciMLBase.has_vjp(prob.f) if isinplace(prob) _F = @closure (du, u, p) -> begin - resid = similar(du, length(sol.resid)) + resid = __similar(du, length(sol.resid)) prob.f(resid, u, p) prob.f.vjp(du, resid, u, p) du .*= 2 @@ -101,9 +86,9 @@ function __nlsolve_ad(prob::NonlinearLeastSquaresProblem, alg, args...; kwargs.. elseif SciMLBase.has_jac(prob.f) if isinplace(prob) _F = @closure (du, u, p) -> begin - J = similar(du, length(sol.resid), length(u)) + J = __similar(du, length(sol.resid), length(u)) prob.f.jac(J, u, p) - resid = similar(du, length(sol.resid)) + resid = __similar(du, length(sol.resid)) prob.f(resid, u, p) mul!(reshape(du, 1, :), vec(resid)', J, 2, false) return nothing @@ -116,35 +101,30 @@ function __nlsolve_ad(prob::NonlinearLeastSquaresProblem, alg, args...; kwargs.. else if isinplace(prob) _F = @closure (du, u, p) -> begin - resid = similar(du, length(sol.resid)) - res = DiffResults.DiffResult( - resid, similar(du, length(sol.resid), length(u))) _f = @closure (du, u) -> prob.f(du, u, p) - ForwardDiff.jacobian!(res, _f, resid, u) - mul!(reshape(du, 1, :), vec(DiffResults.value(res))', - DiffResults.jacobian(res), 2, false) + resid = __similar(du, length(sol.resid)) + v, J = DI.value_and_jacobian(_f, resid, AutoForwardDiff(), u) + mul!(reshape(du, 1, :), vec(v)', J, 2, false) return nothing end else # For small problems, nesting ForwardDiff is actually quite fast + _f = Base.Fix2(prob.f, newprob.p) if __is_extension_loaded(Val(:Zygote)) && (length(uu) + length(sol.resid) ≥ 50) - _F = @closure (u, p) -> __zygote_compute_nlls_vjp(prob.f, u, p) + # TODO: Remove once DI has the value_and_pullback_split defined + _F = @closure (u, p) -> __zygote_compute_nlls_vjp(_f, u, p) else _F = @closure (u, p) -> begin - T = promote_type(eltype(u), eltype(p)) - res = DiffResults.DiffResult(similar(u, T, size(sol.resid)), - similar(u, T, length(sol.resid), length(u))) - ForwardDiff.jacobian!(res, Base.Fix2(prob.f, p), u) - return reshape( - 2 .* vec(DiffResults.value(res))' * DiffResults.jacobian(res), - size(u)) + _f = Base.Fix2(prob.f, p) + v, J = DI.value_and_jacobian(_f, AutoForwardDiff(), u) + return reshape(2 .* vec(v)' * J, size(u)) end end end end - f_p = __nlsolve_∂f_∂p(prob, _F, uu, p) - f_x = __nlsolve_∂f_∂u(prob, _F, uu, p) + f_p = __nlsolve_∂f_∂p(prob, _F, uu, newprob.p) + f_x = __nlsolve_∂f_∂u(prob, _F, uu, newprob.p) z_arr = -f_x \ f_p @@ -152,7 +132,7 @@ function __nlsolve_ad(prob::NonlinearLeastSquaresProblem, alg, args...; kwargs.. sumfun = ((z, p),) -> map(zᵢ -> zᵢ * ForwardDiff.partials(p), z) if uu isa Number partials = sum(sumfun, zip(z_arr, pp)) - elseif p isa Number + elseif pp isa Number partials = sumfun((z_arr, pp)) else partials = sum(sumfun, zip(eachcol(z_arr), pp)) @@ -164,7 +144,7 @@ end @inline function __nlsolve_∂f_∂p(prob, f::F, u, p) where {F} if isinplace(prob) __f = p -> begin - du = similar(u, promote_type(eltype(u), eltype(p))) + du = __similar(u, promote_type(eltype(u), eltype(p))) f(du, u, p) return du end @@ -182,16 +162,12 @@ end @inline function __nlsolve_∂f_∂u(prob, f::F, u, p) where {F} if isinplace(prob) - du = similar(u) - __f = (du, u) -> f(du, u, p) - ForwardDiff.jacobian(__f, du, u) + __f = @closure (du, u) -> f(du, u, p) + return ForwardDiff.jacobian(__f, __similar(u), u) else __f = Base.Fix2(f, p) - if u isa Number - return ForwardDiff.derivative(__f, u) - else - return ForwardDiff.jacobian(__f, u) - end + u isa Number && return ForwardDiff.derivative(__f, u) + return ForwardDiff.jacobian(__f, u) end end diff --git a/src/nlsolve/halley.jl b/src/nlsolve/halley.jl index 5ff5eea..d3fe1cd 100644 --- a/src/nlsolve/halley.jl +++ b/src/nlsolve/halley.jl @@ -39,9 +39,9 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleHalley, args...; @bb xo = copy(x) if setindex_trait(x) === CanSetindex() - A = similar(x, length(x), length(x)) - Aaᵢ = similar(x, length(x)) - cᵢ = similar(x) + A = __similar(x, length(x), length(x)) + Aaᵢ = __similar(x, length(x)) + cᵢ = __similar(x) else A = x Aaᵢ = x diff --git a/src/nlsolve/lbroyden.jl b/src/nlsolve/lbroyden.jl index b34d4cd..1eeda4f 100644 --- a/src/nlsolve/lbroyden.jl +++ b/src/nlsolve/lbroyden.jl @@ -272,7 +272,7 @@ end return :(return SVector{$N, $T}(($(getcalls...)))) end -__lbroyden_threshold_cache(x, ::Val{threshold}) where {threshold} = similar(x, threshold) +__lbroyden_threshold_cache(x, ::Val{threshold}) where {threshold} = __similar(x, threshold) function __lbroyden_threshold_cache(x::StaticArray, ::Val{threshold}) where {threshold} return zeros(MArray{Tuple{threshold}, eltype(x)}) end @@ -298,7 +298,7 @@ end end end function __init_low_rank_jacobian(u, fu, ::Val{threshold}) where {threshold} - Vᵀ = similar(u, threshold, length(u)) - U = similar(u, length(fu), threshold) + Vᵀ = __similar(u, threshold, length(u)) + U = __similar(u, length(fu), threshold) return U, Vᵀ end diff --git a/src/utils.jl b/src/utils.jl index ae380d7..0bf027f 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -21,13 +21,13 @@ Return the maximum of `a` and `b` if `x1 > x0`, otherwise return the minimum. """ __max_tdir(a, b, x0, x1) = ifelse(x1 > x0, max(a, b), min(a, b)) -function __fixed_parameter_function(prob::NonlinearProblem) +function __fixed_parameter_function(prob::AbstractNonlinearProblem) isinplace(prob) && return @closure (du, u) -> prob.f(du, u, prob.p) return Base.Fix2(prob.f, prob.p) end function value_and_jacobian( - ad, prob::NonlinearProblem, f::F, y, x, cache; J = nothing) where {F} + ad, prob::AbstractNonlinearProblem, f::F, y, x, cache; J = nothing) where {F} x isa Number && return DI.value_and_derivative(f, ad, x, cache) if isinplace(prob) @@ -46,21 +46,22 @@ function value_and_jacobian( end end -function jacobian_cache(ad, prob::NonlinearProblem, f::F, y, x) where {F} +function jacobian_cache(ad, prob::AbstractNonlinearProblem, f::F, y, x) where {F} x isa Number && return (nothing, DI.prepare_derivative(f, ad, x)) if isinplace(prob) - J = similar(y, length(y), length(x)) + J = __similar(y, length(y), length(x)) SciMLBase.has_jac(prob.f) && return J, HasAnalyticJacobian() return J, DI.prepare_jacobian(f, y, ad, x) else SciMLBase.has_jac(prob.f) && return nothing, HasAnalyticJacobian() - J = ArrayInterface.can_setindex(x) ? similar(y, length(y), length(x)) : nothing + J = ArrayInterface.can_setindex(x) ? __similar(y, length(y), length(x)) : nothing return J, DI.prepare_jacobian(f, ad, x) end end -function compute_jacobian_and_hessian(ad, prob::NonlinearProblem, f::F, y, x) where {F} +function compute_jacobian_and_hessian( + ad, prob::AbstractNonlinearProblem, f::F, y, x) where {F} if x isa Number df = @closure x -> DI.derivative(f, ad, x) return f(x), df(x), DI.derivative(df, ad, x) @@ -68,7 +69,7 @@ function compute_jacobian_and_hessian(ad, prob::NonlinearProblem, f::F, y, x) wh if isinplace(prob) df = @closure x -> begin - res = similar(y, promote_type(eltype(y), eltype(x))) + res = __similar(y, promote_type(eltype(y), eltype(x))) return DI.jacobian(f, res, ad, x) end J, H = DI.value_and_jacobian(df, ad, x) @@ -83,7 +84,7 @@ end __init_identity_jacobian(u::Number, fu, α = true) = oftype(u, α) __init_identity_jacobian!!(J::Number) = one(J) function __init_identity_jacobian(u, fu, α = true) - J = similar(u, promote_type(eltype(u), eltype(fu)), length(fu), length(u)) + J = __similar(u, promote_type(eltype(u), eltype(fu)), length(fu), length(u)) fill!(J, zero(eltype(J))) J[diagind(J)] .= eltype(J)(α) return J @@ -129,7 +130,7 @@ end T = eltype(x) return T.(f.resid_prototype) else - fx = similar(x) + fx = __similar(x) f(fx, x, p) return fx end @@ -242,3 +243,16 @@ end # Extension function __zygote_compute_nlls_vjp end + +function __similar(x, args...; kwargs...) + y = similar(x, args...; kwargs...) + return __init_bigfloat_array!!(y) +end + +function __init_bigfloat_array!!(x) + if ArrayInterface.can_setindex(x) + eltype(x) <: BigFloat && fill!(x, BigFloat(0)) + return x + end + return x +end From fb2f20dabe4bbe60686e0de4909ae5470254c251 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Sat, 25 May 2024 18:29:32 -0700 Subject: [PATCH 3/4] Remove the static arrays special casing --- Project.toml | 10 +++++----- ext/SimpleNonlinearSolveStaticArraysExt.jl | 7 ------- src/SimpleNonlinearSolve.jl | 1 + src/ad.jl | 6 ++++-- src/nlsolve/dfsane.jl | 15 +++++++++------ src/utils.jl | 14 +++++++------- test/core/rootfind_tests.jl | 9 ++++----- 7 files changed, 30 insertions(+), 32 deletions(-) delete mode 100644 ext/SimpleNonlinearSolveStaticArraysExt.jl diff --git a/Project.toml b/Project.toml index a7aa31f..1d9cc13 100644 --- a/Project.toml +++ b/Project.toml @@ -18,19 +18,18 @@ MaybeInplace = "bb5d69b7-63fc-4a16-80bd-7e42200c7bdb" PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" +Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46" StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" [weakdeps] ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" -StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [extensions] SimpleNonlinearSolveChainRulesCoreExt = "ChainRulesCore" SimpleNonlinearSolveReverseDiffExt = "ReverseDiff" -SimpleNonlinearSolveStaticArraysExt = "StaticArrays" SimpleNonlinearSolveTrackerExt = "Tracker" SimpleNonlinearSolveZygoteExt = "Zygote" @@ -40,7 +39,7 @@ AllocCheck = "0.1.1" Aqua = "0.8" ArrayInterface = "7.9" CUDA = "5.2" -ChainRulesCore = "1.22" +ChainRulesCore = "1.23" ConcreteStructs = "0.2.3" DiffEqBase = "6.149" DiffResults = "1.1" @@ -59,13 +58,14 @@ PrecompileTools = "1.2" Random = "1.10" ReTestItems = "1.23" Reexport = "1.2" -ReverseDiff = "1.15" +ReverseDiff = "1.15.3" SciMLBase = "2.37.0" SciMLSensitivity = "7.58" +Setfield = "1.1.1" StaticArrays = "1.9" StaticArraysCore = "1.4.2" Test = "1.10" -Tracker = "0.2.32" +Tracker = "0.2.33" Zygote = "0.6.69" julia = "1.10" diff --git a/ext/SimpleNonlinearSolveStaticArraysExt.jl b/ext/SimpleNonlinearSolveStaticArraysExt.jl deleted file mode 100644 index c865084..0000000 --- a/ext/SimpleNonlinearSolveStaticArraysExt.jl +++ /dev/null @@ -1,7 +0,0 @@ -module SimpleNonlinearSolveStaticArraysExt - -using SimpleNonlinearSolve: SimpleNonlinearSolve - -@inline SimpleNonlinearSolve.__is_extension_loaded(::Val{:StaticArrays}) = true - -end diff --git a/src/SimpleNonlinearSolve.jl b/src/SimpleNonlinearSolve.jl index 04a32c9..eba6d99 100644 --- a/src/SimpleNonlinearSolve.jl +++ b/src/SimpleNonlinearSolve.jl @@ -24,6 +24,7 @@ using PrecompileTools: @compile_workload, @setup_workload, @recompile_invalidati NonlinearFunction, NonlinearLeastSquaresProblem, NonlinearProblem, ReturnCode, init, remake, solve, AbstractNonlinearAlgorithm, build_solution, isinplace, _unwrap_val + using Setfield: @set! using StaticArraysCore: StaticArray, SVector, SMatrix, SArray, MArray, Size end diff --git a/src/ad.jl b/src/ad.jl index ffae7cc..bb5afea 100644 --- a/src/ad.jl +++ b/src/ad.jl @@ -109,10 +109,12 @@ function __nlsolve_ad(prob::NonlinearLeastSquaresProblem, alg, args...; kwargs.. end else # For small problems, nesting ForwardDiff is actually quite fast - _f = Base.Fix2(prob.f, newprob.p) if __is_extension_loaded(Val(:Zygote)) && (length(uu) + length(sol.resid) ≥ 50) # TODO: Remove once DI has the value_and_pullback_split defined - _F = @closure (u, p) -> __zygote_compute_nlls_vjp(_f, u, p) + _F = @closure (u, p) -> begin + _f = Base.Fix2(prob.f, p) + return __zygote_compute_nlls_vjp(_f, u, p) + end else _F = @closure (u, p) -> begin _f = Base.Fix2(prob.f, p) diff --git a/src/nlsolve/dfsane.jl b/src/nlsolve/dfsane.jl index 7dd1522..835ee4b 100644 --- a/src/nlsolve/dfsane.jl +++ b/src/nlsolve/dfsane.jl @@ -77,12 +77,8 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleDFSane{M}, args... α_1 = one(T) f_1 = fx_norm - history_f_k = if x isa SArray || - (x isa Number && __is_extension_loaded(Val(:StaticArrays))) - ones(SVector{M, T}) * fx_norm - else - fill(fx_norm, M) - end + history_f_k = x isa SArray ? ones(SVector{M, T}) * fx_norm : + __history_vec(fx_norm, Val(M)) # Generate the cache @bb x_cache = similar(x) @@ -150,6 +146,8 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleDFSane{M}, args... # Store function value if history_f_k isa SVector history_f_k = Base.setindex(history_f_k, fx_norm_new, mod1(k, M)) + elseif history_f_k isa NTuple + @set! history_f_k[mod1(k, M)] = fx_norm_new else history_f_k[mod1(k, M)] = fx_norm_new end @@ -158,3 +156,8 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleDFSane{M}, args... return build_solution(prob, alg, x, fx; retcode = ReturnCode.MaxIters) end + +@inline @generated function __history_vec(fx_norm, ::Val{M}) where {M} + M ≥ 11 && return :(fill(fx_norm, M)) # Julia can't specialize here + return :(ntuple(Returns(fx_norm), $(M))) +end diff --git a/src/utils.jl b/src/utils.jl index 0bf027f..2e76a4b 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -32,16 +32,15 @@ function value_and_jacobian( if isinplace(prob) if cache isa HasAnalyticJacobian - prob.f.jac(J, x, p) + prob.f.jac(J, x, prob.p) f(y, x) - else - DI.jacobian!(f, y, J, ad, x, cache) + return y, J end - return y, J + return DI.value_and_jacobian!(f, y, J, ad, x, cache) else cache isa HasAnalyticJacobian && return f(x), prob.f.jac(x, prob.p) J === nothing && return DI.value_and_jacobian(f, ad, x, cache) - y, _ = DI.value_and_jacobian!(f, J, ad, x, cache) + y, J = DI.value_and_jacobian!(f, J, ad, x, cache) return y, J end end @@ -63,8 +62,9 @@ end function compute_jacobian_and_hessian( ad, prob::AbstractNonlinearProblem, f::F, y, x) where {F} if x isa Number - df = @closure x -> DI.derivative(f, ad, x) - return f(x), df(x), DI.derivative(df, ad, x) + H = DI.second_derivative(f, ad, x) + v, J = DI.value_and_derivative(f, ad, x) + return v, J, H end if isinplace(prob) diff --git a/test/core/rootfind_tests.jl b/test/core/rootfind_tests.jl index 1ef0757..6165314 100644 --- a/test/core/rootfind_tests.jl +++ b/test/core/rootfind_tests.jl @@ -34,7 +34,6 @@ end export quadratic_f, quadratic_f!, quadratic_f2, newton_fails, TERMINATION_CONDITIONS, benchmark_nlsolve_oop, benchmark_nlsolve_iip - end @testitem "First Order Methods" setup=[RootfindingTesting] tags=[:core] begin @@ -42,7 +41,7 @@ end SimpleTrustRegion, (args...; kwargs...) -> SimpleTrustRegion( args...; nlsolve_update_rule = Val(true), kwargs...)) - @testset "AutoDiff: $(nameof(typeof(autodiff))))" for autodiff in ( + @testset "AutoDiff: $(nameof(typeof(autodiff)))" for autodiff in ( AutoFiniteDiff(), AutoForwardDiff(), AutoPolyesterForwardDiff()) @testset "[OOP] u0: $(typeof(u0))" for u0 in ( [1.0, 1.0], @SVector[1.0, 1.0], 1.0) @@ -59,7 +58,7 @@ end end end - @testset "Termination condition: $(termination_condition) u0: $(nameof(typeof(u0)))" for termination_condition in TERMINATION_CONDITIONS, + @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) @@ -79,7 +78,7 @@ end end end - @testset "Termination condition: $(termination_condition) u0: $(nameof(typeof(u0)))" for termination_condition in TERMINATION_CONDITIONS, + @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) @@ -104,7 +103,7 @@ end @test all(abs.(sol.u .* sol.u .- 2) .< 1e-9) end - @testset "Termination condition: $(termination_condition) u0: $(nameof(typeof(u0)))" for termination_condition in TERMINATION_CONDITIONS, + @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) From 29ae93938cf908767ff658d246d6ec2ef6d3dcd4 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Sat, 25 May 2024 20:07:01 -0700 Subject: [PATCH 4/4] Add Halley to 23 test problems --- test/core/23_test_problems_tests.jl | 23 ++++++++++++++++++----- test/core/rootfind_tests.jl | 6 ++++++ 2 files changed, 24 insertions(+), 5 deletions(-) diff --git a/test/core/23_test_problems_tests.jl b/test/core/23_test_problems_tests.jl index 9625c68..e455609 100644 --- a/test/core/23_test_problems_tests.jl +++ b/test/core/23_test_problems_tests.jl @@ -41,7 +41,7 @@ end export problems, dicts, test_on_library end -@testitem "SimpleNewtonRaphson" setup=[RobustnessTesting] tags=[:core] begin +@testitem "23 Test Problems: SimpleNewtonRaphson" setup=[RobustnessTesting] tags=[:core] begin alg_ops = (SimpleNewtonRaphson(),) broken_tests = Dict(alg => Int[] for alg in alg_ops) @@ -50,7 +50,20 @@ end test_on_library(problems, dicts, alg_ops, broken_tests) end -@testitem "SimpleTrustRegion" setup=[RobustnessTesting] tags=[:core] begin +@testitem "23 Test Problems: SimpleHalley" setup=[RobustnessTesting] tags=[:core] begin + alg_ops = (SimpleHalley(),) + + broken_tests = Dict(alg => Int[] for alg in alg_ops) + if Sys.isapple() + broken_tests[alg_ops[1]] = [1, 5, 11, 15, 16, 18] + else + broken_tests[alg_ops[1]] = [1, 5, 15, 16, 18] + end + + test_on_library(problems, dicts, alg_ops, broken_tests) +end + +@testitem "23 Test Problems: SimpleTrustRegion" setup=[RobustnessTesting] tags=[:core] begin alg_ops = (SimpleTrustRegion(), SimpleTrustRegion(; nlsolve_update_rule = Val(true))) broken_tests = Dict(alg => Int[] for alg in alg_ops) @@ -60,7 +73,7 @@ end test_on_library(problems, dicts, alg_ops, broken_tests) end -@testitem "SimpleDFSane" setup=[RobustnessTesting] tags=[:core] begin +@testitem "23 Test Problems: SimpleDFSane" setup=[RobustnessTesting] tags=[:core] begin alg_ops = (SimpleDFSane(),) broken_tests = Dict(alg => Int[] for alg in alg_ops) @@ -73,7 +86,7 @@ end test_on_library(problems, dicts, alg_ops, broken_tests) end -@testitem "SimpleBroyden" retries=5 setup=[RobustnessTesting] tags=[:core] begin +@testitem "23 Test Problems: SimpleBroyden" retries=5 setup=[RobustnessTesting] tags=[:core] begin alg_ops = (SimpleBroyden(),) broken_tests = Dict(alg => Int[] for alg in alg_ops) @@ -82,7 +95,7 @@ end test_on_library(problems, dicts, alg_ops, broken_tests) end -@testitem "SimpleKlement" setup=[RobustnessTesting] tags=[:core] begin +@testitem "23 Test Problems: SimpleKlement" setup=[RobustnessTesting] tags=[:core] begin alg_ops = (SimpleKlement(),) broken_tests = Dict(alg => Int[] for alg in alg_ops) diff --git a/test/core/rootfind_tests.jl b/test/core/rootfind_tests.jl index 6165314..eecdeca 100644 --- a/test/core/rootfind_tests.jl +++ b/test/core/rootfind_tests.jl @@ -76,6 +76,12 @@ end @test SciMLBase.successful_retcode(sol) @test all(abs.(sol.u .* sol.u .- 2) .< 1e-9) end + + @testset "[IIP] u0: $(nameof(typeof(u0)))" for u0 in ([1.0, 1.0],) + sol = benchmark_nlsolve_iip(quadratic_f!, u0; solver = SimpleHalley(; autodiff)) + @test SciMLBase.successful_retcode(sol) + @test all(abs.(sol.u .* sol.u .- 2) .< 1e-9) + end end @testset "Termination condition: $(nameof(typeof(termination_condition))) u0: $(nameof(typeof(u0)))" for termination_condition in TERMINATION_CONDITIONS,