From 57e4eca42e0c0d609dd42fce13d221c7e310ee40 Mon Sep 17 00:00:00 2001 From: yonatanwesen Date: Thu, 21 Sep 2023 23:58:34 -0400 Subject: [PATCH 1/3] reworking pt using avik's cleanup --- src/NonlinearSolve.jl | 5 +- src/jacobian.jl | 8 ++- src/pseudotransient.jl | 150 +++++++++++++++++++++++++++++++++++++++++ test/basictests.jl | 57 ++++++++++++++++ 4 files changed, 217 insertions(+), 3 deletions(-) create mode 100644 src/pseudotransient.jl diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 615f96c03..962551893 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -39,6 +39,7 @@ include("linesearch.jl") include("raphson.jl") include("trustRegion.jl") include("levenberg.jl") +include("pseudotransient.jl") include("jacobian.jl") include("ad.jl") @@ -49,7 +50,7 @@ PrecompileTools.@compile_workload begin prob = NonlinearProblem{false}((u, p) -> u .* u .- p, T(0.1), T(2)) precompile_algs = if VERSION ≥ v"1.7" - (NewtonRaphson(), TrustRegion(), LevenbergMarquardt()) + (NewtonRaphson(), TrustRegion(), LevenbergMarquardt(), PseudoTransient()) else (NewtonRaphson(),) end @@ -68,7 +69,7 @@ end export RadiusUpdateSchemes -export NewtonRaphson, TrustRegion, LevenbergMarquardt +export NewtonRaphson, TrustRegion, LevenbergMarquardt, PseudoTransient export LineSearch diff --git a/src/jacobian.jl b/src/jacobian.jl index 01c030462..1ff29d803 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -85,7 +85,13 @@ function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f, u, p, ::Val{ii end du = _mutable_zero(u) - linprob = LinearProblem(J, _vec(fu); u0 = _vec(du)) + if alg isa PseudoTransient + alpha = convert(eltype(u), alg.alpha_initial) + J_new = J - (1 / alpha) * I + linprob = LinearProblem(J_new, _vec(fu); u0 = _vec(du)) + else + linprob = LinearProblem(J, _vec(fu); u0 = _vec(du)) + end weight = similar(u) recursivefill!(weight, true) diff --git a/src/pseudotransient.jl b/src/pseudotransient.jl new file mode 100644 index 000000000..8213850d3 --- /dev/null +++ b/src/pseudotransient.jl @@ -0,0 +1,150 @@ +@concrete struct PseudoTransient{CJ, AD} <: AbstractNewtonAlgorithm{CJ, AD} + ad::AD + linsolve + precs + alpha_initial +end + +concrete_jac(::PseudoTransient{CJ}) where {CJ} = CJ + +function PseudoTransient(; concrete_jac = nothing, linsolve = nothing, + precs = DEFAULT_PRECS, alpha_initial = 1e-3, adkwargs...) + ad = default_adargs_to_adtype(; adkwargs...) + return PseudoTransient{_unwrap_val(concrete_jac)}(ad, linsolve, precs, alpha_initial) +end + +@concrete mutable struct PseudoTransientCache{iip} + f + alg + u + fu1 + fu2 + du + p + alpha + res_norm + uf + linsolve + J + jac_cache + force_stop + maxiters::Int + internalnorm + retcode::ReturnCode.T + abstol + prob + stats::NLStats +end + +isinplace(::PseudoTransientCache{iip}) where {iip} = iip + +function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::PseudoTransient, args...; + alias_u0 = false, maxiters = 1000, abstol = 1e-6, internalnorm = DEFAULT_NORM, + kwargs...) where {uType, iip} + @unpack f, u0, p = prob + u = alias_u0 ? u0 : deepcopy(u0) + if iip + fu1 = f.resid_prototype === nothing ? zero(u) : f.resid_prototype + f(fu1, u, p) + else + fu1 = _mutable(f(u, p)) + end + uf, linsolve, J, fu2, jac_cache, du = jacobian_caches(alg, f, u, p, Val(iip)) + alpha = convert(eltype(u), alg.alpha_initial) + res_norm = internalnorm(fu1) + + return PseudoTransientCache{iip}(f, alg, u, fu1, fu2, du, p, alpha, res_norm, uf, + linsolve, J, + jac_cache, false, maxiters, internalnorm, ReturnCode.Default, abstol, prob, + NLStats(1, 0, 0, 0, 0)) +end + +function perform_step!(cache::PseudoTransientCache{true}) + @unpack u, fu1, f, p, alg, J, linsolve, du, alpha = cache + jacobian!!(J, cache) + J_new = J - (1 / alpha) * I + + # u = u - J \ fu + linres = dolinsolve(alg.precs, linsolve; A = J_new, b = _vec(fu1), linu = _vec(du), + p, reltol = cache.abstol) + cache.linsolve = linres.cache + @. u = u - du + f(fu1, u, p) + + new_norm = cache.internalnorm(fu1) + cache.alpha *= cache.res_norm / new_norm + cache.res_norm = new_norm + + new_norm < cache.abstol && (cache.force_stop = true) + cache.stats.nf += 1 + cache.stats.njacs += 1 + cache.stats.nsolve += 1 + cache.stats.nfactors += 1 + return nothing +end + +function perform_step!(cache::PseudoTransientCache{false}) + @unpack u, fu1, f, p, alg, linsolve, alpha = cache + + cache.J = jacobian!!(cache.J, cache) + # u = u - J \ fu + if linsolve === nothing + cache.du = fu1 / cache.J + else + linres = dolinsolve(alg.precs, linsolve; A = cache.J - (1 / alpha) * I, + b = _vec(fu1), + linu = _vec(cache.du), p, reltol = cache.abstol) + cache.linsolve = linres.cache + end + cache.u = @. u - cache.du # `u` might not support mutation + cache.fu1 = f(cache.u, p) + + new_norm = cache.internalnorm(fu1) + cache.alpha *= cache.res_norm / new_norm + cache.res_norm = new_norm + + new_norm < cache.abstol && (cache.force_stop = true) + cache.stats.nf += 1 + cache.stats.njacs += 1 + cache.stats.nsolve += 1 + cache.stats.nfactors += 1 + return nothing +end + +function SciMLBase.solve!(cache::PseudoTransientCache) + while !cache.force_stop && cache.stats.nsteps < cache.maxiters + perform_step!(cache) + cache.stats.nsteps += 1 + end + + if cache.stats.nsteps == cache.maxiters + cache.retcode = ReturnCode.MaxIters + else + cache.retcode = ReturnCode.Success + end + + return SciMLBase.build_solution(cache.prob, cache.alg, cache.u, cache.fu1; + cache.retcode, cache.stats) +end + +function SciMLBase.reinit!(cache::PseudoTransientCache{iip}, u0 = cache.u; p = cache.p, + abstol = cache.abstol, maxiters = cache.maxiters) where {iip} + cache.p = p + if iip + recursivecopy!(cache.u, u0) + cache.f(cache.fu1, cache.u, p) + else + # don't have alias_u0 but cache.u is never mutated for OOP problems so it doesn't matter + cache.u = u0 + cache.fu1 = cache.f(cache.u, p) + end + cache.alpha = convert(eltype(u), 1e-3) + cache.res_norm = internalnorm(cache.fu1) + cache.abstol = abstol + cache.maxiters = maxiters + cache.stats.nf = 1 + cache.stats.nsteps = 1 + cache.force_stop = false + cache.retcode = ReturnCode.Default + return cache +end diff --git a/test/basictests.jl b/test/basictests.jl index 54e63e93d..396d504fc 100644 --- a/test/basictests.jl +++ b/test/basictests.jl @@ -402,3 +402,60 @@ end end end end + +# --- PseudoTransient tests --- + +@testset "PseudoTransient" begin + + #iip + function test_f!(du, u, p) + du[1] = 2 - 2u[1] + du[2] = u[1] - 4u[2] + return du + end + + #oop + simple_test(u, p) = 0.9f0 .* u .- u + + #test jacobian free PseudoTransient method + function f!(res, u, p) + @. res = u * u - p + end + + @testset "[IIP] u0: $(typeof(u0))" for u0 in (zeros(2),) + probN = NonlinearProblem{true}(test_f!, u0) + sol = solve(probN, PseudoTransient()) + @test sol.retcode == ReturnCode.Success + du = zeros(2) + test_f!(du, sol.u, 4.0) + @test du≈[0, 0] atol=1e-7 + end + + @testset "[OOP] u0: $(typeof(u0))" for u0 in ([1.0f0],) + probN = NonlinearProblem{false}(simple_test, u0) + sol = solve(probN, PseudoTransient(alpha_initial = 1.0)) + @test sol.retcode == ReturnCode.Success + @test abs(sol.u[1]) <= 1.0f-4 + end + + @testset "Jacobian Free PT" begin + u0 = [1.0, 1.0] + p = 2.0 + probN = NonlinearProblem(f!, u0, p) + linsolve = LinearSolve.KrylovJL_GMRES() + sol = solve(probN, + PseudoTransient(alpha_initial = 10.0, linsolve = linsolve), + reltol = 1e-9) + @test sol.retcode == ReturnCode.Success + end + + @testset "NewtonRaphson Fails" begin + p = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + u0 = [-10.0, -1.0, 1.0, 2.0, 3.0, 4.0, 10.0] + probN = NonlinearProblem{false}(newton_fails, u0, p) + sol = solve(probN, PseudoTransient(alpha_initial = 1.0), abstol = 1e-10) + @test all(abs.(newton_fails(sol.u, p)) .< 1e-10) + end +end + +# Test that `PseudoTransient` passes a test that `NewtonRaphson` fails on. From 6b2a1ec8c081d0f0af35092b6d8351715c9ef014 Mon Sep 17 00:00:00 2001 From: yonatanwesen Date: Sun, 24 Sep 2023 21:02:45 -0400 Subject: [PATCH 2/3] forget the 1/alpha term when linsolve is nothing --- src/pseudotransient.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pseudotransient.jl b/src/pseudotransient.jl index 8213850d3..19fe94062 100644 --- a/src/pseudotransient.jl +++ b/src/pseudotransient.jl @@ -89,7 +89,7 @@ function perform_step!(cache::PseudoTransientCache{false}) cache.J = jacobian!!(cache.J, cache) # u = u - J \ fu if linsolve === nothing - cache.du = fu1 / cache.J + cache.du = fu1 / (cache.J - (1 / alpha) * I) else linres = dolinsolve(alg.precs, linsolve; A = cache.J - (1 / alpha) * I, b = _vec(fu1), From 9b6ddf1c31e9b4a6230af766314f38c59267efbd Mon Sep 17 00:00:00 2001 From: yonatanwesen Date: Wed, 4 Oct 2023 13:42:05 -0400 Subject: [PATCH 3/3] added linsolve kwargs so that simple gmres stops --- src/pseudotransient.jl | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/pseudotransient.jl b/src/pseudotransient.jl index 19fe94062..c3308dfb6 100644 --- a/src/pseudotransient.jl +++ b/src/pseudotransient.jl @@ -40,6 +40,7 @@ isinplace(::PseudoTransientCache{iip}) where {iip} = iip function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::PseudoTransient, args...; alias_u0 = false, maxiters = 1000, abstol = 1e-6, internalnorm = DEFAULT_NORM, + linsolve_kwargs = (;), kwargs...) where {uType, iip} @unpack f, u0, p = prob u = alias_u0 ? u0 : deepcopy(u0) @@ -49,7 +50,12 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::PseudoTransie else fu1 = _mutable(f(u, p)) end - uf, linsolve, J, fu2, jac_cache, du = jacobian_caches(alg, f, u, p, Val(iip)) + uf, linsolve, J, fu2, jac_cache, du = jacobian_caches(alg, + f, + u, + p, + Val(iip); + linsolve_kwargs) alpha = convert(eltype(u), alg.alpha_initial) res_norm = internalnorm(fu1) @@ -102,7 +108,6 @@ function perform_step!(cache::PseudoTransientCache{false}) new_norm = cache.internalnorm(fu1) cache.alpha *= cache.res_norm / new_norm cache.res_norm = new_norm - new_norm < cache.abstol && (cache.force_stop = true) cache.stats.nf += 1 cache.stats.njacs += 1