Skip to content

Commit de6eb96

Browse files
ErikQQYErikQQY
committed
Put DualAbstractNonlinearProblem solving in subpackages
1 parent ce721be commit de6eb96

File tree

11 files changed

+205
-24
lines changed

11 files changed

+205
-24
lines changed

‎lib/NonlinearSolveBase/ext/NonlinearSolveBaseForwardDiffExt.jl‎

Lines changed: 5 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,12 @@ using SciMLBase: SciMLBase, AbstractNonlinearProblem, IntervalNonlinearProblem,
1212

1313
using NonlinearSolveBase: NonlinearSolveBase, ImmutableNonlinearProblem,
1414
AbstractNonlinearSolveAlgorithm, Utils, InternalAPI,
15-
AbstractNonlinearSolveCache
15+
AbstractNonlinearSolveCache, NonlinearSolvePolyAlgorithm
1616

1717
const DI = DifferentiationInterface
1818

19-
const ALL_SOLVER_TYPES = [
20-
Nothing, AbstractNonlinearSolveAlgorithm
19+
const GENERAL_SOLVER_TYPES = [
20+
Nothing, AbstractNonlinearSolveAlgorithm, NonlinearSolvePolyAlgorithm
2121
]
2222

2323
const DualNonlinearProblem = NonlinearProblem{
@@ -121,7 +121,7 @@ function NonlinearSolveBase.nonlinearsolve_dual_solution(
121121
return map(((uᵢ, pᵢ),) -> Dual{T, V, P}(uᵢ, pᵢ), zip(u, Utils.restructure(u, partials)))
122122
end
123123

124-
for algType in ALL_SOLVER_TYPES
124+
for algType in GENERAL_SOLVER_TYPES
125125
@eval function SciMLBase.__solve(
126126
prob::DualAbstractNonlinearProblem, alg::$(algType), args...; kwargs...
127127
)
@@ -157,7 +157,7 @@ function InternalAPI.reinit!(
157157
return cache
158158
end
159159

160-
for algType in ALL_SOLVER_TYPES
160+
for algType in GENERAL_SOLVER_TYPES
161161
@eval function SciMLBase.__init(
162162
prob::DualAbstractNonlinearProblem, alg::$(algType), args...; kwargs...
163163
)
@@ -200,13 +200,4 @@ nodual_value(x) = x
200200
nodual_value(x::Dual) = ForwardDiff.value(x)
201201
nodual_value(x::AbstractArray{<:Dual}) = map(ForwardDiff.value, x)
202202

203-
"""
204-
pickchunksize(x) = pickchunksize(length(x))
205-
pickchunksize(x::Int)
206-
207-
Determine the chunk size for ForwardDiff and PolyesterForwardDiff based on the input length.
208-
"""
209-
@inline pickchunksize(x) = pickchunksize(length(x))
210-
@inline pickchunksize(x::Int) = ForwardDiff.pickchunksize(x)
211-
212203
end

‎lib/NonlinearSolveBase/src/NonlinearSolveBase.jl‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,4 +84,6 @@ export DescentResult, SteepestDescent, NewtonDescent, DampedNewtonDescent, Dogle
8484

8585
export NonlinearSolvePolyAlgorithm
8686

87+
export pickchunksize
88+
8789
end

‎lib/NonlinearSolveBase/src/common_defaults.jl‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,3 +45,12 @@ function get_tolerance(::Union{StaticArray, Number}, ::Nothing, ::Type{T}) where
4545
# Rational numbers can throw an error if used inside GPU Kernels
4646
return T(real(oneunit(T)) * (eps(real(one(T)))^(real(T)(0.8))))
4747
end
48+
49+
"""
50+
pickchunksize(x) = pickchunksize(length(x))
51+
pickchunksize(x::Int)
52+
53+
Determine the chunk size for ForwardDiff and PolyesterForwardDiff based on the input length.
54+
"""
55+
@inline pickchunksize(x) = pickchunksize(length(x))
56+
@inline pickchunksize(x::Int) = ForwardDiff.pickchunksize(x)

‎lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ using SciMLBase: SciMLBase, AbstractNonlinearProblem, NLStats, ReturnCode,
2929
using SciMLJacobianOperators: VecJacOperator, JacVecOperator, StatefulJacobianOperator
3030

3131
using FiniteDiff: FiniteDiff # Default Finite Difference Method
32-
using ForwardDiff: ForwardDiff # Default Forward Mode AD
32+
using ForwardDiff: ForwardDiff, Dual # Default Forward Mode AD
3333

3434
include("raphson.jl")
3535
include("gauss_newton.jl")
@@ -41,6 +41,8 @@ include("poly_algs.jl")
4141

4242
include("solve.jl")
4343

44+
include("forward_diff.jl")
45+
4446
@setup_workload begin
4547
nonlinear_functions = (
4648
(NonlinearFunction{false, NoSpecialize}((u, p) -> u .* u .- p), 0.1),

‎lib/NonlinearSolveFirstOrder/src/forward_diff.jl‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
const DualNonlinearProblem = NonlinearProblem{
2+
<:Union{Number, <:AbstractArray}, iip,
3+
<:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}}
4+
} where {iip, T, V, P}
5+
const DualNonlinearLeastSquaresProblem = NonlinearLeastSquaresProblem{
6+
<:Union{Number, <:AbstractArray}, iip,
7+
<:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}}
8+
} where {iip, T, V, P}
9+
const DualAbstractNonlinearProblem = Union{
10+
DualNonlinearProblem, DualNonlinearLeastSquaresProblem
11+
}
12+
13+
function SciMLBase.__init(
14+
prob::DualAbstractNonlinearProblem, alg::GeneralizedFirstOrderAlgorithm, args...; kwargs...
15+
)
16+
p = NonlinearSolveBase.nodual_value(prob.p)
17+
newprob = SciMLBase.remake(prob; u0 = NonlinearSolveBase.nodual_value(prob.u0), p)
18+
cache = init(newprob, alg, args...; kwargs...)
19+
return NonlinearSolveForwardDiffCache(
20+
cache, newprob, alg, prob.p, p, ForwardDiff.partials(prob.p)
21+
)
22+
end
23+
24+
function SciMLBase.__solve(
25+
prob::DualAbstractNonlinearProblem, alg::GeneralizedFirstOrderAlgorithm, args...; kwargs...
26+
)
27+
sol, partials = NonlinearSolveBase.nonlinearsolve_forwarddiff_solve(
28+
prob, alg, args...; kwargs...
29+
)
30+
dual_soln = NonlinearSolveBase.nonlinearsolve_dual_solution(sol.u, partials, prob.p)
31+
return SciMLBase.build_solution(
32+
prob, alg, dual_soln, sol.resid; sol.retcode, sol.stats, sol.original
33+
)
34+
end

‎lib/NonlinearSolveQuasiNewton/Project.toml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,12 @@ SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
1818
SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
1919
StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
2020

21+
[weakdeps]
22+
ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
23+
24+
[extensions]
25+
NonlinearSolveQuasiNewtonForwardDiffExt = "ForwardDiff"
26+
2127
[compat]
2228
ADTypes = "1.9.0"
2329
Aqua = "0.8"

‎lib/NonlinearSolveQuasiNewton/ext/NonlinearSolveQuasiNewtonForwardDiffExt.jl‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
module NonlinearSolveQuasiNewtonForwardDiffExt
2+
3+
using CommonSolve: CommonSolve, solve
4+
using ForwardDiff: ForwardDiff, Dual
5+
using SciMLBase: SciMLBase, AbstractNonlinearProblem, IntervalNonlinearProblem,
6+
NonlinearProblem, NonlinearLeastSquaresProblem, remake
7+
8+
using NonlinearSolveBase: NonlinearSolveBase
9+
10+
using NonlinearSolveQuasiNewton: QuasiNewtonAlgorithm
11+
12+
const DualNonlinearProblem = NonlinearProblem{
13+
<:Union{Number, <:AbstractArray}, iip,
14+
<:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}}
15+
} where {iip, T, V, P}
16+
const DualNonlinearLeastSquaresProblem = NonlinearLeastSquaresProblem{
17+
<:Union{Number, <:AbstractArray}, iip,
18+
<:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}}
19+
} where {iip, T, V, P}
20+
const DualAbstractNonlinearProblem = Union{
21+
DualNonlinearProblem, DualNonlinearLeastSquaresProblem
22+
}
23+
24+
function SciMLBase.__solve(
25+
prob::DualAbstractNonlinearProblem, alg::QuasiNewtonAlgorithm, args...; kwargs...
26+
)
27+
sol, partials = NonlinearSolveBase.nonlinearsolve_forwarddiff_solve(
28+
prob, alg, args...; kwargs...
29+
)
30+
dual_soln = NonlinearSolveBase.nonlinearsolve_dual_solution(sol.u, partials, prob.p)
31+
return SciMLBase.build_solution(
32+
prob, alg, dual_soln, sol.resid; sol.retcode, sol.stats, sol.original
33+
)
34+
end
35+
36+
function SciMLBase.__init(
37+
prob::DualAbstractNonlinearProblem, alg::QuasiNewtonAlgorithm, args...; kwargs...
38+
)
39+
p = nodual_value(prob.p)
40+
newprob = SciMLBase.remake(prob; u0 = nodual_value(prob.u0), p)
41+
cache = init(newprob, alg, args...; kwargs...)
42+
return NonlinearSolveForwardDiffCache(
43+
cache, newprob, alg, prob.p, p, ForwardDiff.partials(prob.p)
44+
)
45+
end
46+
47+
end

‎lib/NonlinearSolveSpectralMethods/Project.toml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,12 @@ PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
1414
Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
1515
SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
1616

17+
[weakdeps]
18+
ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
19+
20+
[extensions]
21+
NonlinearSolveSpectralMethodsForwardDiffExt = "ForwardDiff"
22+
1723
[compat]
1824
Aqua = "0.8"
1925
BenchmarkTools = "1.5.0"

‎lib/NonlinearSolveSpectralMethods/ext/NonlinearSolveSpectralMethodsForwardDiffExt.jl‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
module NonlinearSolveSpectralMethodsForwardDiffExt
2+
3+
using CommonSolve: CommonSolve, solve
4+
using ForwardDiff: ForwardDiff, Dual
5+
using SciMLBase: SciMLBase, AbstractNonlinearProblem, IntervalNonlinearProblem,
6+
NonlinearProblem, NonlinearLeastSquaresProblem, remake
7+
8+
using NonlinearSolveBase: NonlinearSolveBase
9+
10+
using NonlinearSolveSpectralMethods: GeneralizedDFSane
11+
12+
const DualNonlinearProblem = NonlinearProblem{
13+
<:Union{Number, <:AbstractArray}, iip,
14+
<:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}}
15+
} where {iip, T, V, P}
16+
const DualNonlinearLeastSquaresProblem = NonlinearLeastSquaresProblem{
17+
<:Union{Number, <:AbstractArray}, iip,
18+
<:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}}
19+
} where {iip, T, V, P}
20+
const DualAbstractNonlinearProblem = Union{
21+
DualNonlinearProblem, DualNonlinearLeastSquaresProblem
22+
}
23+
24+
function SciMLBase.__solve(
25+
prob::DualAbstractNonlinearProblem, alg::GeneralizedDFSane, args...; kwargs...
26+
)
27+
sol, partials = NonlinearSolveBase.nonlinearsolve_forwarddiff_solve(
28+
prob, alg, args...; kwargs...
29+
)
30+
dual_soln = NonlinearSolveBase.nonlinearsolve_dual_solution(sol.u, partials, prob.p)
31+
return SciMLBase.build_solution(
32+
prob, alg, dual_soln, sol.resid; sol.retcode, sol.stats, sol.original
33+
)
34+
end
35+
36+
function SciMLBase.__init(
37+
prob::DualAbstractNonlinearProblem, alg::GeneralizedDFSane, args...; kwargs...
38+
)
39+
p = nodual_value(prob.p)
40+
newprob = SciMLBase.remake(prob; u0 = nodual_value(prob.u0), p)
41+
cache = init(newprob, alg, args...; kwargs...)
42+
return NonlinearSolveForwardDiffCache(
43+
cache, newprob, alg, prob.p, p, ForwardDiff.partials(prob.p)
44+
)
45+
end
46+
47+
end

‎src/NonlinearSolve.jl‎

Lines changed: 2 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ using LineSearch: BackTracking
1414
using NonlinearSolveBase: NonlinearSolveBase, InternalAPI, AbstractNonlinearSolveAlgorithm,
1515
AbstractNonlinearSolveCache, Utils, L2_NORM,
1616
enable_timer_outputs, disable_timer_outputs,
17-
NonlinearSolvePolyAlgorithm
17+
NonlinearSolvePolyAlgorithm, pickchunksize
1818

1919
using Preferences: set_preferences!
2020
using SciMLBase: SciMLBase, NLStats, ReturnCode, AbstractNonlinearProblem,
@@ -53,14 +53,7 @@ include("extension_algs.jl")
5353

5454
include("default.jl")
5555

56-
const ALL_SOLVER_TYPES = [
57-
Nothing, AbstractNonlinearSolveAlgorithm,
58-
GeneralizedDFSane, GeneralizedFirstOrderAlgorithm, QuasiNewtonAlgorithm,
59-
LeastSquaresOptimJL, FastLevenbergMarquardtJL, NLsolveJL, NLSolversJL,
60-
SpeedMappingJL, FixedPointAccelerationJL, SIAMFANLEquationsJL,
61-
CMINPACK, PETScSNES,
62-
NonlinearSolvePolyAlgorithm
63-
]
56+
include("forward_diff.jl")
6457

6558
@setup_workload begin
6659
nonlinear_functions = (

0 commit comments

Comments
 (0)