Skip to content

Commit

Permalink
handle complex numbers correctly by default
Browse files Browse the repository at this point in the history
  • Loading branch information
@avik-pal
avik-pal committed Dec 19, 2023
1 parent aa282b7 commit 133b1d5
Show file tree
Hide file tree
Showing 4 changed files with 145 additions and 59 deletions.
6 changes: 4 additions & 2 deletions Project.toml
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "NonlinearSolve"
uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
authors = ["SciML"]
version = "3.1.1"
version = "3.1.2"

[deps]
ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
Expand Down Expand Up @@ -71,6 +71,7 @@ MaybeInplace = "0.1.1"
NLsolve = "4.5"
NaNMath = "1"
NonlinearProblemLibrary = "0.1.1"
OrdinaryDiffEq = "6"
Pkg = "1"
PrecompileTools = "1.2"
Printf = "<0.0.1, 1"
Expand Down Expand Up @@ -106,6 +107,7 @@ MINPACK = "4854310b-de5a-5eb6-a2a5-c1dee2bd17f9"
NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
NonlinearProblemLibrary = "b7050fa9-e91f-4b37-bcee-a89a063da141"
OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
Expand All @@ -117,4 +119,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

[targets]
test = ["Aqua", "Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools", "NonlinearProblemLibrary", "LeastSquaresOptim", "FastLevenbergMarquardt", "NaNMath", "BandedMatrices", "DiffEqBase", "StableRNGs", "MINPACK", "NLsolve"]
test = ["Aqua", "Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools", "NonlinearProblemLibrary", "LeastSquaresOptim", "FastLevenbergMarquardt", "NaNMath", "BandedMatrices", "DiffEqBase", "StableRNGs", "MINPACK", "NLsolve", "OrdinaryDiffEq"]
166 changes: 110 additions & 56 deletions src/default.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -165,8 +165,8 @@ function SciMLBase.reinit!(cache::NonlinearSolvePolyAlgorithmCache, args...; kwa
end

"""
RobustMultiNewton(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS,
autodiff = nothing)
RobustMultiNewton(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing,
precs = DEFAULT_PRECS, autodiff = nothing)
A polyalgorithm focused on robustness. It uses a mixture of Newton methods with different
globalizing techniques (trust region updates, line searches, etc.) in order to find a
Expand All @@ -176,6 +176,11 @@ Basically, if this algorithm fails, then "most" good ways of solving your proble
you may need to think about reformulating the model (either there is an issue with the model,
or more precision / more stable linear solver choice is required).
### Arguments
- `T`: The eltype of the initial guess. It is only used to check if some of the algorithms
are compatible with the problem type. Defaults to `Float64`.
### Keyword Arguments
- `autodiff`: determines the backend used for the Jacobian. Note that this argument is
Expand All @@ -196,28 +201,38 @@ or more precision / more stable linear solver choice is required).
algorithms, consult the
[LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/).
"""
function RobustMultiNewton(; concrete_jac = nothing, linsolve = nothing,
precs = DEFAULT_PRECS, autodiff = nothing)
algs = (TrustRegion(; concrete_jac, linsolve, precs),
TrustRegion(; concrete_jac, linsolve, precs, autodiff,
radius_update_scheme = RadiusUpdateSchemes.Bastin),
NewtonRaphson(; concrete_jac, linsolve, precs, linesearch = BackTracking(),
autodiff),
TrustRegion(; concrete_jac, linsolve, precs,
radius_update_scheme = RadiusUpdateSchemes.NLsolve, autodiff),
TrustRegion(; concrete_jac, linsolve, precs,
radius_update_scheme = RadiusUpdateSchemes.Fan, autodiff))
function RobustMultiNewton(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing,
precs = DEFAULT_PRECS, autodiff = nothing) where {T}
if __is_complex(T)
# Let's atleast have something here for complex numbers
algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),)
else
algs = (TrustRegion(; concrete_jac, linsolve, precs),
TrustRegion(; concrete_jac, linsolve, precs, autodiff,
radius_update_scheme = RadiusUpdateSchemes.Bastin),
NewtonRaphson(; concrete_jac, linsolve, precs, linesearch = BackTracking(),
autodiff),
TrustRegion(; concrete_jac, linsolve, precs,
radius_update_scheme = RadiusUpdateSchemes.NLsolve, autodiff),
TrustRegion(; concrete_jac, linsolve, precs,
radius_update_scheme = RadiusUpdateSchemes.Fan, autodiff))
end
return NonlinearSolvePolyAlgorithm(algs, Val(:NLS))
end

"""
FastShortcutNonlinearPolyalg(; concrete_jac = nothing, linsolve = nothing,
precs = DEFAULT_PRECS, must_use_jacobian::Val = Val(false),
prefer_simplenonlinearsolve::Val{SA} = Val(false), autodiff = nothing)
FastShortcutNonlinearPolyalg(::Type{T} = Float64; concrete_jac = nothing,
linsolve = nothing, precs = DEFAULT_PRECS, must_use_jacobian::Val = Val(false),
prefer_simplenonlinearsolve::Val{SA} = Val(false), autodiff = nothing) where {T}
A polyalgorithm focused on balancing speed and robustness. It first tries less robust methods
for more performance and then tries more robust techniques if the faster ones fail.
### Arguments
- `T`: The eltype of the initial guess. It is only used to check if some of the algorithms
are compatible with the problem type. Defaults to `Float64`.
### Keyword Arguments
- `autodiff`: determines the backend used for the Jacobian. Note that this argument is
Expand All @@ -238,53 +253,76 @@ for more performance and then tries more robust techniques if the faster ones fa
algorithms, consult the
[LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/).
"""
function FastShortcutNonlinearPolyalg(; concrete_jac = nothing, linsolve = nothing,
precs = DEFAULT_PRECS, must_use_jacobian::Val{JAC} = Val(false),
function FastShortcutNonlinearPolyalg(::Type{T} = Float64; concrete_jac = nothing,
linsolve = nothing, precs = DEFAULT_PRECS, must_use_jacobian::Val{JAC} = Val(false),
prefer_simplenonlinearsolve::Val{SA} = Val(false),
autodiff = nothing) where {JAC, SA}
autodiff = nothing) where {T, JAC, SA}
if JAC
algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),
NewtonRaphson(; concrete_jac, linsolve, precs, linesearch = BackTracking(),
autodiff),
TrustRegion(; concrete_jac, linsolve, precs, autodiff),
TrustRegion(; concrete_jac, linsolve, precs,
radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff))
else
# SimpleNewtonRaphson and SimpleTrustRegion are not robust to singular Jacobians
# and thus are not included in the polyalgorithm
if SA
algs = (SimpleBroyden(),
Broyden(; init_jacobian = Val(:true_jacobian)),
SimpleKlement(),
NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),
NewtonRaphson(; concrete_jac, linsolve, precs, linesearch = BackTracking(),
autodiff),
NewtonRaphson(; concrete_jac, linsolve, precs, linesearch = BackTracking(),
autodiff),
TrustRegion(; concrete_jac, linsolve, precs,
radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff))
if __is_complex(T)
algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),)

Check warning on line 262 in src/default.jl

View check run for this annotation

Codecov / codecov/patch

src/default.jl#L261-L262 

Added lines #L261 - L262 were not covered by tests
else
algs = (Broyden(),
Broyden(; init_jacobian = Val(:true_jacobian)),
Klement(; linsolve, precs),
NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),
algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),

Check warning on line 264 in src/default.jl

View check run for this annotation

Codecov / codecov/patch

src/default.jl#L264 

Added line #L264 was not covered by tests
NewtonRaphson(; concrete_jac, linsolve, precs, linesearch = BackTracking(),
autodiff),
TrustRegion(; concrete_jac, linsolve, precs, autodiff),
TrustRegion(; concrete_jac, linsolve, precs,
radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff))
end
else
# SimpleNewtonRaphson and SimpleTrustRegion are not robust to singular Jacobians
# and thus are not included in the polyalgorithm
if SA
if __is_complex(T)
algs = (SimpleBroyden(),

Check warning on line 276 in src/default.jl

View check run for this annotation

Codecov / codecov/patch

src/default.jl#L276 

Added line #L276 was not covered by tests
Broyden(; init_jacobian = Val(:true_jacobian)),
SimpleKlement(),
NewtonRaphson(; concrete_jac, linsolve, precs, autodiff))
else
algs = (SimpleBroyden(),
Broyden(; init_jacobian = Val(:true_jacobian)),
SimpleKlement(),
NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),
NewtonRaphson(; concrete_jac, linsolve, precs,
linesearch = BackTracking(), autodiff),
NewtonRaphson(; concrete_jac, linsolve, precs,
linesearch = BackTracking(), autodiff),
TrustRegion(; concrete_jac, linsolve, precs,
radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff))
end
else
if __is_complex(T)
algs = (Broyden(),
Broyden(; init_jacobian = Val(:true_jacobian)),
Klement(; linsolve, precs),
NewtonRaphson(; concrete_jac, linsolve, precs, autodiff))
else
algs = (Broyden(),
Broyden(; init_jacobian = Val(:true_jacobian)),
Klement(; linsolve, precs),
NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),
NewtonRaphson(; concrete_jac, linsolve, precs,
linesearch = BackTracking(), autodiff),
TrustRegion(; concrete_jac, linsolve, precs, autodiff),
TrustRegion(; concrete_jac, linsolve, precs,
radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff))
end
end
end
return NonlinearSolvePolyAlgorithm(algs, Val(:NLS))
end

"""
FastShortcutNLLSPolyalg(; concrete_jac = nothing, linsolve = nothing,
precs = DEFAULT_PRECS, kwargs...)
FastShortcutNLLSPolyalg(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing,
precs = DEFAULT_PRECS, kwargs...)
A polyalgorithm focused on balancing speed and robustness. It first tries less robust methods
for more performance and then tries more robust techniques if the faster ones fail.
### Arguments
- `T`: The eltype of the initial guess. It is only used to check if some of the algorithms
are compatible with the problem type. Defaults to `Float64`.
### Keyword Arguments
- `autodiff`: determines the backend used for the Jacobian. Note that this argument is
Expand All @@ -305,42 +343,58 @@ for more performance and then tries more robust techniques if the faster ones fa
algorithms, consult the
[LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/).
"""
function FastShortcutNLLSPolyalg(; concrete_jac = nothing, linsolve = nothing,
precs = DEFAULT_PRECS, kwargs...)
algs = (GaussNewton(; concrete_jac, linsolve, precs, kwargs...),
GaussNewton(; concrete_jac, linsolve, precs, linesearch = BackTracking(),
kwargs...),
LevenbergMarquardt(; concrete_jac, linsolve, precs, kwargs...))
function FastShortcutNLLSPolyalg(::Type{T} = Float64; concrete_jac = nothing,
linsolve = nothing, precs = DEFAULT_PRECS, kwargs...) where {T}
if __is_complex(T)
algs = (GaussNewton(; concrete_jac, linsolve, precs, kwargs...),

Check warning on line 349 in src/default.jl

View check run for this annotation

Codecov / codecov/patch

src/default.jl#L349 

Added line #L349 was not covered by tests
LevenbergMarquardt(; concrete_jac, linsolve, precs, kwargs...))
else
algs = (GaussNewton(; concrete_jac, linsolve, precs, kwargs...),
TrustRegion(; concrete_jac, linsolve, precs, kwargs...),
GaussNewton(; concrete_jac, linsolve, precs, linesearch = BackTracking(),
kwargs...),
TrustRegion(; concrete_jac, linsolve, precs,
radius_update_scheme = RadiusUpdateSchemes.Bastin, kwargs...),
LevenbergMarquardt(; concrete_jac, linsolve, precs, kwargs...))
end
return NonlinearSolvePolyAlgorithm(algs, Val(:NLLS))
end

## Defaults

## TODO: In the long run we want to use an `Assumptions` API like LinearSolve to specify
## the conditioning of the Jacobian and such

## TODO: Currently some of the algorithms like LineSearches / TrustRegion don't support
## complex numbers. We should use the `DiffEqBase` trait for this once all of the
## NonlinearSolve algorithms support it. For now we just do a check and remove the
## unsupported ones from default

## Defaults to a fast and robust poly algorithm in most cases. If the user went through
## the trouble of specifying a custom jacobian function, we should use algorithms that
## can use that!

function SciMLBase.__init(prob::NonlinearProblem, ::Nothing, args...; kwargs...)
must_use_jacobian = Val(prob.f.jac !== nothing)
return SciMLBase.__init(prob, FastShortcutNonlinearPolyalg(; must_use_jacobian),
return SciMLBase.__init(prob,
FastShortcutNonlinearPolyalg(eltype(prob.u0); must_use_jacobian),
args...; kwargs...)
end

function SciMLBase.__solve(prob::NonlinearProblem, ::Nothing, args...; kwargs...)
must_use_jacobian = Val(prob.f.jac !== nothing)
prefer_simplenonlinearsolve = Val(prob.u0 isa SArray)
return SciMLBase.__solve(prob,
FastShortcutNonlinearPolyalg(; must_use_jacobian,
FastShortcutNonlinearPolyalg(eltype(prob.u0); must_use_jacobian,
prefer_simplenonlinearsolve), args...; kwargs...)
end

function SciMLBase.__init(prob::NonlinearLeastSquaresProblem, ::Nothing, args...; kwargs...)
return SciMLBase.__init(prob, FastShortcutNLLSPolyalg(), args...; kwargs...)
return SciMLBase.__init(prob, FastShortcutNLLSPolyalg(eltype(prob.u0)), args...;

Check warning on line 392 in src/default.jl

View check run for this annotation

Codecov / codecov/patch

src/default.jl#L392 

Added line #L392 was not covered by tests
kwargs...)
end

function SciMLBase.__solve(prob::NonlinearLeastSquaresProblem, ::Nothing, args...;
kwargs...)
return SciMLBase.__solve(prob, FastShortcutNLLSPolyalg(), args...; kwargs...)
return SciMLBase.__solve(prob, FastShortcutNLLSPolyalg(eltype(prob.u0)), args...;
kwargs...)
end
5 changes: 5 additions & 0 deletions src/utils.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -494,3 +494,8 @@ end
@inline __diag(x::AbstractMatrix) = diag(x)
@inline __diag(x::AbstractVector) = x
@inline __diag(x::Number) = x

@inline __is_complex(::Type{ComplexF64}) = true
@inline __is_complex(::Type{ComplexF32}) = true
@inline __is_complex(::Type{Complex}) = true

Check warning on line 500 in src/utils.jl

View check run for this annotation

Codecov / codecov/patch

src/utils.jl#L499-L500 

Added lines #L499 - L500 were not covered by tests
@inline __is_complex(::Type{T}) where {T} = false
27 changes: 26 additions & 1 deletion test/polyalgs.jl
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
using NonlinearSolve, Test, NaNMath
using NonlinearSolve, Test, NaNMath, OrdinaryDiffEq

f(u, p) = u .* u .- 2
u0 = [1.0, 1.0]
Expand Down Expand Up @@ -58,3 +58,28 @@ p = 2.0
prob = NonlinearProblem(ff_interval, u0, p)
sol = solve(prob; abstol = 1e-9)
@test SciMLBase.successful_retcode(sol)

# Shooting Problem: Taken from BoundaryValueDiffEq.jl
# Testing for Complex Valued Root Finding. For Complex valued inputs we drop some of the
# algorithms which dont support those.
function ode_func!(du, u, p, t)
du[1] = u[2]
du[2] = -u[1]
return nothing
end

function objective_function!(resid, u0, p)
odeprob = ODEProblem{true}(ode_func!, u0, (0.0, 100.0), p)
sol = solve(odeprob, Tsit5(), abstol = 1e-9, reltol = 1e-9, verbose = false)
resid[1] = sol(0.0)[1]
resid[2] = sol(100.0)[1] - 1.0
return nothing
end

prob = NonlinearProblem{true}(objective_function!, [0.0, 1.0] .+ 1im)
sol = solve(prob; abstol = 1e-10)
@test SciMLBase.successful_retcode(sol)
# This test is not meant to return success but test that all the default solvers can handle
# complex valued problems
@test_nowarn solve(prob; abstol = 1e-19, maxiters = 10)
@test_nowarn solve(prob, RobustMultiNewton(eltype(prob.u0)); abstol = 1e-19, maxiters = 10)

0 comments on commit 133b1d5

Please sign in to comment.