Skip to content

Commit 196adbb

Browse files
avik-palavik-pal
committed
refactor: move algorithms to First Order sub-package
1 parent 21ffb22 commit 196adbb

File tree

9 files changed

+380
-290
lines changed

9 files changed

+380
-290
lines changed

‎lib/NonlinearSolveFirstOrder/Project.toml

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,15 +3,53 @@ uuid = "5959db7a-ea39-4486-b5fe-2dd0bf03d60d"
33
authors = ["Avik Pal <[email protected]> and contributors"]
44
version = "1.0.0"
55

6+
[deps]
7+
ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
8+
ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
9+
CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
10+
ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
11+
DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
12+
FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
13+
ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
14+
LineSearch = "87fe0de2-c867-4266-b59a-2f0a94fc965b"
15+
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
16+
LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
17+
MaybeInplace = "bb5d69b7-63fc-4a16-80bd-7e42200c7bdb"
18+
NonlinearSolveBase = "be0214bd-f91f-a760-ac4e-3421ce2b2da0"
19+
PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
20+
Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
21+
SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
22+
SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
23+
Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
24+
StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
25+
626
[compat]
27+
ADTypes = "1.9.0"
728
Aqua = "0.8"
29+
ArrayInterface = "7.16.0"
30+
CommonSolve = "0.2.4"
31+
ConcreteStructs = "0.2.3"
32+
DiffEqBase = "6.155.3"
833
ExplicitImports = "1.5"
34+
FiniteDiff = "2.26.0"
35+
ForwardDiff = "0.10.36"
936
Hwloc = "3"
1037
InteractiveUtils = "<0.0.1, 1"
38+
LineSearch = "0.1.4"
39+
LinearAlgebra = "1.11.0"
40+
LinearSolve = "2.36.1"
41+
MaybeInplace = "0.1.4"
1142
NonlinearProblemLibrary = "0.1.2"
43+
NonlinearSolveBase = "1.1"
1244
Pkg = "1.10"
45+
PrecompileTools = "1.2"
1346
ReTestItems = "1.24"
47+
Reexport = "1"
48+
SciMLBase = "2.54"
49+
SciMLOperators = "0.3.11"
50+
Setfield = "1.1.1"
1451
StableRNGs = "1"
52+
StaticArraysCore = "1.4.3"
1553
Test = "1.10"
1654
julia = "1.10"
1755

‎lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,23 +3,29 @@ module NonlinearSolveFirstOrder
33
using Reexport: @reexport
44
using PrecompileTools: @compile_workload, @setup_workload
55

6+
using ADTypes: ADTypes
67
using ArrayInterface: ArrayInterface
78
using CommonSolve: CommonSolve
89
using ConcreteStructs: @concrete
9-
using DiffEqBase: DiffEqBase # Needed for `init` / `solve` dispatches
10+
using DiffEqBase: DiffEqBase # Needed for `init` / `solve` dispatches
11+
using FiniteDiff: FiniteDiff # Default Finite Difference Method
12+
using ForwardDiff: ForwardDiff # Default Forward Mode AD
1013
using LinearAlgebra: LinearAlgebra, Diagonal, dot, inv, diag
11-
using LinearSolve: LinearSolve # Trigger Linear Solve extension in NonlinearSolveBase
14+
using LinearSolve: LinearSolve # Trigger Linear Solve extension in NonlinearSolveBase
1215
using MaybeInplace: @bb
1316
using NonlinearSolveBase: NonlinearSolveBase, AbstractNonlinearSolveAlgorithm,
1417
AbstractNonlinearSolveCache, AbstractResetCondition,
1518
AbstractResetConditionCache, AbstractApproximateJacobianStructure,
1619
AbstractJacobianCache, AbstractJacobianInitialization,
1720
AbstractApproximateJacobianUpdateRule, AbstractDescentDirection,
1821
AbstractApproximateJacobianUpdateRuleCache,
22+
AbstractDampingFunction, AbstractDampingFunctionCache,
1923
Utils, InternalAPI, get_timer_output, @static_timeit,
20-
update_trace!, L2_NORM, NewtonDescent
24+
update_trace!, L2_NORM,
25+
NewtonDescent, DampedNewtonDescent
2126
using SciMLBase: SciMLBase, AbstractNonlinearProblem, NLStats, ReturnCode
2227
using SciMLOperators: AbstractSciMLOperator
28+
using Setfield: @set!
2329
using StaticArraysCore: StaticArray, Size, MArray
2430

2531
include("raphson.jl")

‎lib/NonlinearSolveFirstOrder/src/gauss_newton.jl

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,22 @@
1+
"""
2+
GaussNewton(;
3+
concrete_jac = nothing, linsolve = nothing, linesearch = missing, precs = nothing,
4+
autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing
5+
)
16
7+
An advanced GaussNewton implementation with support for efficient handling of sparse
8+
matrices via colored automatic differentiation and preconditioned linear solvers. Designed
9+
for large-scale and numerically-difficult nonlinear systems.
10+
"""
11+
function GaussNewton(;
12+
concrete_jac = nothing, linsolve = nothing, linesearch = missing, precs = nothing,
13+
autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing
14+
)
15+
return GeneralizedFirstOrderAlgorithm(;
16+
linesearch,
17+
descent = NewtonDescent(; linsolve, precs),
18+
autodiff, vjp_autodiff, jvp_autodiff,
19+
concrete_jac,
20+
name = :GaussNewton
21+
)
22+
end

‎lib/NonlinearSolveFirstOrder/src/pseudo_transient.jl

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,90 @@
1+
"""
2+
PseudoTransient(;
3+
concrete_jac = nothing, linesearch = missing, alpha_initial = 1e-3,
4+
linsolve = nothing, precs = nothing,
5+
autodiff = nothing, jvp_autodiff = nothing, vjp_autodiff = nothing
6+
)
17
8+
An implementation of PseudoTransient Method [coffey2003pseudotransient](@cite) that is used
9+
to solve steady state problems in an accelerated manner. It uses an adaptive time-stepping
10+
to integrate an initial value of nonlinear problem until sufficient accuracy in the desired
11+
steady-state is achieved to switch over to Newton's method and gain a rapid convergence.
12+
This implementation specifically uses "switched evolution relaxation"
13+
[kelley1998convergence](@cite) SER method.
14+
15+
### Keyword Arguments
16+
17+
- `alpha_initial` : the initial pseudo time step. It defaults to `1e-3`. If it is small,
18+
you are going to need more iterations to converge but it can be more stable.
19+
"""
20+
function PseudoTransient(;
21+
concrete_jac = nothing, linesearch = missing, alpha_initial = 1e-3,
22+
linsolve = nothing, precs = nothing,
23+
autodiff = nothing, jvp_autodiff = nothing, vjp_autodiff = nothing
24+
)
25+
return GeneralizedFirstOrderAlgorithm(;
26+
linesearch,
27+
descent = DampedNewtonDescent(;
28+
linsolve, precs, initial_damping = alpha_initial,
29+
damping_fn = SwitchedEvolutionRelaxation()
30+
),
31+
autodiff,
32+
jvp_autodiff,
33+
vjp_autodiff,
34+
concrete_jac,
35+
name = :PseudoTransient
36+
)
37+
end
38+
39+
"""
40+
SwitchedEvolutionRelaxation()
41+
42+
Method for updating the damping parameter in the [`PseudoTransient`](@ref) method based on
43+
"switched evolution relaxation" [kelley1998convergence](@cite) SER method.
44+
"""
45+
struct SwitchedEvolutionRelaxation <: AbstractDampingFunction end
46+
47+
"""
48+
SwitchedEvolutionRelaxationCache <: AbstractDampingFunctionCache
49+
50+
Cache for the [`SwitchedEvolutionRelaxation`](@ref) method.
51+
"""
52+
@concrete mutable struct SwitchedEvolutionRelaxationCache <: AbstractDampingFunctionCache
53+
res_norm
54+
α⁻¹
55+
internalnorm
56+
end
57+
58+
function NonlinearSolveBase.requires_normal_form_jacobian(::Union{
59+
SwitchedEvolutionRelaxation, SwitchedEvolutionRelaxationCache})
60+
return false
61+
end
62+
63+
function NonlinearSolveBase.requires_normal_form_rhs(::Union{
64+
SwitchedEvolutionRelaxation, SwitchedEvolutionRelaxationCache})
65+
return false
66+
end
67+
68+
function InternalAPI.init(
69+
prob::AbstractNonlinearProblem, f::SwitchedEvolutionRelaxation,
70+
initial_damping, J, fu, u, args...;
71+
internalnorm::F = L2_NORM, kwargs...
72+
) where {F}
73+
T = promote_type(eltype(u), eltype(fu))
74+
return SwitchedEvolutionRelaxationCache(
75+
internalnorm(fu),
76+
T(inv(initial_damping)),
77+
internalnorm
78+
)
79+
end
80+
81+
(damping::SwitchedEvolutionRelaxationCache)(::Nothing) = damping.α⁻¹
82+
83+
function InternalAPI.solve!(
84+
damping::SwitchedEvolutionRelaxationCache, J, fu, args...; kwargs...
85+
)
86+
res_norm = damping.internalnorm(fu)
87+
damping.α⁻¹ *= res_norm / damping.res_norm
88+
damping.res_norm = res_norm
89+
return damping.α⁻¹
90+
end

‎lib/NonlinearSolveFirstOrder/src/raphson.jl

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,22 @@
1+
"""
2+
NewtonRaphson(;
3+
concrete_jac = nothing, linsolve = nothing, linesearch = missing, precs = nothing,
4+
autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing
5+
)
16
7+
An advanced NewtonRaphson implementation with support for efficient handling of sparse
8+
matrices via colored automatic differentiation and preconditioned linear solvers. Designed
9+
for large-scale and numerically-difficult nonlinear systems.
10+
"""
11+
function NewtonRaphson(;
12+
concrete_jac = nothing, linsolve = nothing, linesearch = missing, precs = nothing,
13+
autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing
14+
)
15+
return GeneralizedFirstOrderAlgorithm(;
16+
linesearch,
17+
descent = NewtonDescent(; linsolve, precs),
18+
autodiff, vjp_autodiff, jvp_autodiff,
19+
concrete_jac,
20+
name = :NewtonRaphson
21+
)
22+
end

0 commit comments

Comments
 (0)