diff --git a/Project.toml b/Project.toml index ed0b27f95..ac1676c87 100644 --- a/Project.toml +++ b/Project.toml @@ -9,6 +9,7 @@ DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56" FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index cae730bc3..25f4a6e3f 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -15,6 +15,7 @@ import ArrayInterface import LinearSolve using DiffEqBase using SparseDiffTools +using LineSearches @reexport using SciMLBase using SciMLBase: NLStats @@ -65,7 +66,8 @@ PrecompileTools.@compile_workload begin end export RadiusUpdateSchemes - +export LineSearches export NewtonRaphson, TrustRegion, LevenbergMarquardt + end # module diff --git a/src/jacobian.jl b/src/jacobian.jl index 8296069e0..1574e4b13 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -183,3 +183,19 @@ function jacobian_autodiff(f, x::AbstractArray, nonlinfun, alg) jac_prototype = jac_prototype, chunksize = chunk_size), num_of_chunks) end + +function simple_jacobian(cache, x::Number) + @unpack f, p = cache + g = Base.Fix2(f, p) + ForwardDiff.derivative(g, x) +end + +function simple_jacobian(cache, x::AbstractArray{<:Number}) + @unpack f, fu, p, prob = cache + if !SciMLBase.isinplace(prob) + g = Base.Fix2(f, p) + return ForwardDiff.jacobian(g, x) + else + return ForwardDiff.jacobian((fu, x) -> f(fu, x, p), fu, x) + end +end diff --git a/src/raphson.jl b/src/raphson.jl index 24e5799fd..a082e6bb1 100644 --- a/src/raphson.jl +++ b/src/raphson.jl @@ -34,6 +34,7 @@ for large-scale and numerically-difficult nonlinear systems. - `diff_type`: the type of finite differencing used if `autodiff = false`. Defaults to `Val{:forward}` for forward finite differences. For more details on the choices, see the [FiniteDiff.jl](https://github.com/JuliaDiff/FiniteDiff.jl) documentation. + - `linesearch`: the line search algorithm used. Defaults to no line search being used. - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the linear solves within the Newton method. Defaults to `nothing`, which means it uses the LinearSolve.jl default algorithm choice. For more information on available algorithm @@ -48,22 +49,23 @@ for large-scale and numerically-difficult nonlinear systems. Currently, the linear solver and chunk size choice only applies to in-place defined `NonlinearProblem`s. That is expected to change in the future. """ -struct NewtonRaphson{CS, AD, FDT, L, P, ST, CJ} <: +struct NewtonRaphson{CS, AD, FDT, L, P, ST, CJ, LS} <: AbstractNewtonAlgorithm{CS, AD, FDT, ST, CJ} linsolve::L precs::P + linesearch::LS end function NewtonRaphson(; chunk_size = Val{0}(), autodiff = Val{true}(), standardtag = Val{true}(), concrete_jac = nothing, - diff_type = Val{:forward}, linsolve = nothing, precs = DEFAULT_PRECS) + diff_type = Val{:forward}, linesearch = LineSearches.Static(), linsolve = nothing, precs = DEFAULT_PRECS) NewtonRaphson{_unwrap_val(chunk_size), _unwrap_val(autodiff), diff_type, typeof(linsolve), typeof(precs), _unwrap_val(standardtag), - _unwrap_val(concrete_jac)}(linsolve, - precs) + _unwrap_val(concrete_jac), typeof(linesearch)}(linsolve, + precs, linesearch) end -mutable struct NewtonRaphsonCache{iip, fType, algType, uType, duType, resType, pType, +mutable struct NewtonRaphsonCache{iip, fType, algType, uType, duType, resType, pType, foType, goType, INType, tolType, probType, ufType, L, jType, JC} f::fType @@ -71,6 +73,9 @@ mutable struct NewtonRaphsonCache{iip, fType, algType, uType, duType, resType, p u::uType fu::resType p::pType + α::Real + f_o::foType # stores value of objective + g_o::goType # stores grad of objective at x uf::ufType linsolve::L J::jType @@ -85,7 +90,7 @@ mutable struct NewtonRaphsonCache{iip, fType, algType, uType, duType, resType, p stats::NLStats function NewtonRaphsonCache{iip}(f::fType, alg::algType, u::uType, fu::resType, - p::pType, uf::ufType, linsolve::L, J::jType, + p::pType, α::Real, f_o::foType, g_o::goType, uf::ufType, linsolve::L, J::jType, du1::duType, jac_config::JC, force_stop::Bool, maxiters::Int, internalnorm::INType, @@ -93,11 +98,11 @@ mutable struct NewtonRaphsonCache{iip, fType, algType, uType, duType, resType, p prob::probType, stats::NLStats) where { iip, fType, algType, uType, - duType, resType, pType, INType, + duType, resType, pType, foType, goType, INType, tolType, probType, ufType, L, jType, JC} - new{iip, fType, algType, uType, duType, resType, pType, INType, tolType, - probType, ufType, L, jType, JC}(f, alg, u, fu, p, + new{iip, fType, algType, uType, duType, resType, pType, foType, goType, INType, tolType, + probType, ufType, L, jType, JC}(f, alg, u, fu, p, α, f_o, g_o, uf, linsolve, J, du1, jac_config, force_stop, maxiters, internalnorm, retcode, abstol, prob, stats) @@ -105,7 +110,7 @@ mutable struct NewtonRaphsonCache{iip, fType, algType, uType, duType, resType, p end function jacobian_caches(alg::NewtonRaphson, f, u, p, ::Val{false}) - JacobianWrapper(f, p), nothing, nothing, nothing, nothing + JacobianWrapper(f, p), nothing, nothing, zero(u), nothing end function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::NewtonRaphson, @@ -122,6 +127,9 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::NewtonRaphson end f = prob.f p = prob.p + α = 1.0 #line search coefficient + f_o = 0.0 + g_o = zero(u) if iip fu = zero(u) f(fu, u, p) @@ -130,7 +138,7 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::NewtonRaphson end uf, linsolve, J, du1, jac_config = jacobian_caches(alg, f, u, p, Val(iip)) - return NewtonRaphsonCache{iip}(f, alg, u, fu, p, uf, linsolve, J, du1, jac_config, + return NewtonRaphsonCache{iip}(f, alg, u, fu, p, α, f_o, g_o, uf, linsolve, J, du1, jac_config, false, maxiters, internalnorm, ReturnCode.Default, abstol, prob, NLStats(1, 0, 0, 0, 0)) end @@ -139,13 +147,13 @@ function perform_step!(cache::NewtonRaphsonCache{true}) @unpack u, fu, f, p, alg = cache @unpack J, linsolve, du1 = cache jacobian!(J, cache) - # u = u - J \ fu linres = dolinsolve(alg.precs, linsolve, A = J, b = _vec(fu), linu = _vec(du1), p = p, reltol = cache.abstol) cache.linsolve = linres.cache - @. u = u - du1 - f(fu, u, p) + perform_linesearch!(cache) + @. cache.u = cache.u - cache.α * cache.du1 + f(cache.fu, cache.u, p) if cache.internalnorm(cache.fu) < cache.abstol cache.force_stop = true @@ -160,7 +168,9 @@ end function perform_step!(cache::NewtonRaphsonCache{false}) @unpack u, fu, f, p = cache J = jacobian(cache, f) - cache.u = u - J \ fu + cache.du1 = J \ fu + perform_linesearch!(cache) + cache.u = cache.u - cache.α * cache.du1 cache.fu = f(cache.u, p) if iszero(cache.fu) || cache.internalnorm(cache.fu) < cache.abstol cache.force_stop = true @@ -172,6 +182,44 @@ function perform_step!(cache::NewtonRaphsonCache{false}) return nothing end +function objective_linesearch(cache::NewtonRaphsonCache) ## returns the objective functions required in LS + @unpack f, p = cache + + function fo(x) + return dot(value_f(cache, x), value_f(cache, x)) / 2 + end + + function g!(g_o, x) + mul!(g_o, simple_jacobian(cache, x)', value_f(cache, x)) + g_o + end + + function fg!(g_o, x) + g!(g_o, x) + fo(x) + end + return fo, g!, fg! +end + +function perform_linesearch!(cache::NewtonRaphsonCache) + @unpack u, fu, du1, alg, α = cache + fo, g!, fg! = objective_linesearch(cache) + cache.f_o = fo(u) + g!(cache.g_o, u) + @unpack f_o, g_o = cache + ϕ(α) = fo(u .- α .* du1) + + function dϕ(α) + g!(g_o, u .- α .* du1) + return dot(g_o, -du1) + end + + function ϕdϕ(α) + return (fg!(g_o, u .- α .* du1), dot(g_o, -du1)) + end + cache.α, cache.f_o = alg.linesearch(ϕ, dϕ, ϕdϕ, 1.0, f_o, dot(du1, -g_o)) +end + function SciMLBase.solve!(cache::NewtonRaphsonCache) while !cache.force_stop && cache.stats.nsteps < cache.maxiters perform_step!(cache) @@ -207,3 +255,18 @@ function SciMLBase.reinit!(cache::NewtonRaphsonCache{iip}, u0 = cache.u; p = cac cache.retcode = ReturnCode.Default return cache end + +## some helper func ### + +function value_f(cache::NewtonRaphsonCache, x) + iip = SciMLBase.isinplace(cache.prob) + @unpack f, p = cache + if iip + res = zero(x) + f(res, x, p) + else + res = f(x, p) + end + res +end + diff --git a/src/utils.jl b/src/utils.jl index 9d72e230f..eacf06727 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -157,3 +157,6 @@ function rfunc(r::R, c2::R, M::R, γ1::R, γ2::R, β::R) where {R <: Real} # R-f return (1 - γ1 - β) * (exp(r - c2) + β / (1 - γ1 - β)) end end + + +