Skip to content

Add a default for NLLS #277

New issue

Have a question about this project? # for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “#”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? # to your account

Merged
merged 5 commits into from
Nov 6, 2023
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions docs/src/api/nonlinearsolve.md
Original file line number Diff line number Diff line change
@@ -16,7 +16,9 @@ GeneralKlement
## Polyalgorithms

```@docs
NonlinearSolvePolyAlgorithm
FastShortcutNonlinearPolyalg
FastShortcutNLLSPolyalg
RobustMultiNewton
```

12 changes: 4 additions & 8 deletions docs/src/solvers/NonlinearLeastSquaresSolvers.md
Original file line number Diff line number Diff line change
@@ -7,7 +7,10 @@ Solves the nonlinear least squares problem defined by `prob` using the algorithm

## Recommended Methods

`LevenbergMarquardt` is a good choice for most problems.
The default method `FastShortcutNLLSPolyalg` is a good choice for most
problems. It is a polyalgorithm that attempts to use a fast algorithm
(`GaussNewton`) and if that fails it falls back to a more robust
algorithm (`LevenbergMarquardt`).

## Full List of Methods

@@ -21,10 +24,3 @@ Solves the nonlinear least squares problem defined by `prob` using the algorithm
problems.
- `SimpleNewtonRaphson()`: Simple Gauss Newton Implementation with `QRFactorization` to
solve a linear least squares problem at each step!

## Example usage

```julia
using NonlinearSolve
sol = solve(prob, LevenbergMarquardt())
```
13 changes: 7 additions & 6 deletions ext/NonlinearSolveFastLevenbergMarquardtExt.jl
Original file line number Diff line number Diff line change
@@ -32,23 +32,24 @@ end
(f::InplaceFunction{false})(fx, x, p) = (fx .= f.f(x, p))

function SciMLBase.__init(prob::NonlinearLeastSquaresProblem,
alg::FastLevenbergMarquardtJL, args...; abstol = 1e-8, reltol = 1e-8,
verbose = false, maxiters = 1000, kwargs...)
alg::FastLevenbergMarquardtJL, args...; alias_u0 = false, abstol = 1e-8,
reltol = 1e-8, verbose = false, maxiters = 1000, kwargs...)
iip = SciMLBase.isinplace(prob)
u0 = alias_u0 ? prob.u0 : deepcopy(prob.u0)

@assert prob.f.jac!==nothing "FastLevenbergMarquardt requires a Jacobian!"

f! = InplaceFunction{iip}(prob.f)
J! = InplaceFunction{iip}(prob.f.jac)

resid_prototype = prob.f.resid_prototype === nothing ?
(!iip ? prob.f(prob.u0, prob.p) : zeros(prob.u0)) :
(!iip ? prob.f(u0, prob.p) : zeros(u0)) :
prob.f.resid_prototype

J = similar(prob.u0, length(resid_prototype), length(prob.u0))
J = similar(u0, length(resid_prototype), length(u0))

solver = _fast_lm_solver(alg, prob.u0)
LM = FastLM.LMWorkspace(prob.u0, resid_prototype, J)
solver = _fast_lm_solver(alg, u0)
LM = FastLM.LMWorkspace(u0, resid_prototype, J)

return FastLevenbergMarquardtJLCache(f!, J!, prob, alg, LM, solver,
(; xtol = abstol, ftol = reltol, maxit = maxiters, alg.factor, alg.factoraccept,
8 changes: 5 additions & 3 deletions ext/NonlinearSolveLeastSquaresOptimExt.jl
Original file line number Diff line number Diff line change
@@ -33,17 +33,19 @@ end
(f::FunctionWrapper{false})(du, u) = (du .= f.f(u, f.p))

function SciMLBase.__init(prob::NonlinearLeastSquaresProblem, alg::LeastSquaresOptimJL,
args...; abstol = 1e-8, reltol = 1e-8, verbose = false, maxiters = 1000, kwargs...)
args...; alias_u0 = false, abstol = 1e-8, reltol = 1e-8, verbose = false,
maxiters = 1000, kwargs...)
iip = SciMLBase.isinplace(prob)
u = alias_u0 ? prob.u0 : deepcopy(prob.u0)

f! = FunctionWrapper{iip}(prob.f, prob.p)
g! = prob.f.jac === nothing ? nothing : FunctionWrapper{iip}(prob.f.jac, prob.p)

resid_prototype = prob.f.resid_prototype === nothing ?
(!iip ? prob.f(prob.u0, prob.p) : zeros(prob.u0)) :
(!iip ? prob.f(u, prob.p) : zeros(u)) :
prob.f.resid_prototype

lsoprob = LSO.LeastSquaresProblem(; x = prob.u0, f!, y = resid_prototype, g!,
lsoprob = LSO.LeastSquaresProblem(; x = u, f!, y = resid_prototype, g!,
J = prob.f.jac_prototype, alg.autodiff, output_length = length(resid_prototype))
allocated_prob = LSO.LeastSquaresProblemAllocated(lsoprob, _lso_solver(alg))

3 changes: 2 additions & 1 deletion src/NonlinearSolve.jl
Original file line number Diff line number Diff line change
@@ -147,7 +147,8 @@ export RadiusUpdateSchemes
export NewtonRaphson, TrustRegion, LevenbergMarquardt, DFSane, GaussNewton, PseudoTransient,
GeneralBroyden, GeneralKlement, LimitedMemoryBroyden
export LeastSquaresOptimJL, FastLevenbergMarquardtJL
export NonlinearSolvePolyAlgorithm, RobustMultiNewton, FastShortcutNonlinearPolyalg
export NonlinearSolvePolyAlgorithm,
RobustMultiNewton, FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg

export LineSearch, LiFukushimaLineSearch

55 changes: 49 additions & 6 deletions src/default.jl
Original file line number Diff line number Diff line change
@@ -250,14 +250,57 @@ function FastShortcutNonlinearPolyalg(; concrete_jac = nothing, linsolve = nothi
return NonlinearSolvePolyAlgorithm(algs, Val(:NLS))
end

"""
FastShortcutNLLSPolyalg(; concrete_jac = nothing, linsolve = nothing,
precs = DEFAULT_PRECS, adkwargs...)

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.

### Keyword Arguments

- `autodiff`: determines the backend used for the Jacobian. Note that this argument is
ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to
`AutoForwardDiff()`. Valid choices are types from ADTypes.jl.
- `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used,
then the Jacobian will not be constructed and instead direct Jacobian-vector products
`J*v` are computed using forward-mode automatic differentiation or finite differencing
tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed,
for example for a preconditioner, `concrete_jac = true` can be passed in order to force
the construction of the Jacobian.
- `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
choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/).
- `precs`: the choice of preconditioners for the linear solver. Defaults to using no
preconditioners. For more information on specifying preconditioners for LinearSolve
algorithms, consult the
[LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/).
"""
function FastShortcutNLLSPolyalg(; concrete_jac = nothing, linsolve = nothing,
precs = DEFAULT_PRECS, adkwargs...)
algs = (GaussNewton(; concrete_jac, linsolve, precs, adkwargs...),
GaussNewton(; concrete_jac, linsolve, precs, linesearch = BackTracking(),
adkwargs...),
LevenbergMarquardt(; concrete_jac, linsolve, precs, adkwargs...))
return NonlinearSolvePolyAlgorithm(algs, Val(:NLLS))
end

## Defaults

function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::Nothing, args...;
kwargs...) where {uType, iip}
SciMLBase.__init(prob, FastShortcutNonlinearPolyalg(), args...; kwargs...)
function SciMLBase.__init(prob::NonlinearProblem, ::Nothing, args...; kwargs...)
return SciMLBase.__init(prob, FastShortcutNonlinearPolyalg(), args...; kwargs...)
end

function SciMLBase.__solve(prob::NonlinearProblem, ::Nothing, args...; kwargs...)
return SciMLBase.__solve(prob, FastShortcutNonlinearPolyalg(), args...; kwargs...)
end

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

function SciMLBase.__solve(prob::NonlinearProblem{uType, iip}, alg::Nothing, args...;
kwargs...) where {uType, iip}
SciMLBase.__solve(prob, FastShortcutNonlinearPolyalg(), args...; kwargs...)
function SciMLBase.__solve(prob::NonlinearLeastSquaresProblem, ::Nothing, args...;
kwargs...)
return SciMLBase.__solve(prob, FastShortcutNLLSPolyalg(), args...; kwargs...)
end
2 changes: 1 addition & 1 deletion test/23_test_problems.jl
Original file line number Diff line number Diff line change
@@ -85,7 +85,7 @@ end
# Broyden and Klement Tests are quite flaky and failure seems to be platform dependent
# needs additional investigation before we can enable them
@testset "GeneralBroyden 23 Test Problems" begin
alg_ops = (GeneralBroyden(),)
alg_ops = (GeneralBroyden(; max_resets = 10),)

broken_tests = Dict(alg => Int[] for alg in alg_ops)
broken_tests[alg_ops[1]] = [1, 2, 4, 5, 6, 11, 12, 13, 14]
9 changes: 5 additions & 4 deletions test/nonlinear_least_squares.jl
Original file line number Diff line number Diff line change
@@ -12,12 +12,12 @@ y_target = true_function(x, θ_true)

function loss_function(θ, p)
ŷ = true_function(p, θ)
return abs2.(ŷ .- y_target)
return ŷ .- y_target
end

function loss_function(resid, θ, p)
true_function(resid, p, θ)
resid .= abs2.(resid .- y_target)
resid .= resid .- y_target
return resid
end

@@ -36,6 +36,7 @@ append!(solvers,
LevenbergMarquardt(; linsolve = LUFactorization()),
LeastSquaresOptimJL(:lm),
LeastSquaresOptimJL(:dogleg),
nothing,
])

for prob in nlls_problems, solver in solvers
@@ -45,7 +46,8 @@ for prob in nlls_problems, solver in solvers
end

function jac!(J, θ, p)
ForwardDiff.jacobian!(J, resid -> loss_function(resid, θ, p), θ)
resid = zeros(length(p))
ForwardDiff.jacobian!(J, (resid, θ) -> loss_function(resid, θ, p), resid, θ)
return J
end

@@ -56,6 +58,5 @@ solvers = [FastLevenbergMarquardtJL(:cholesky), FastLevenbergMarquardtJL(:qr)]

for solver in solvers
@time sol = solve(prob, solver; maxiters = 10000, abstol = 1e-8)
@test SciMLBase.successful_retcode(sol)
@test norm(sol.resid) < 1e-6
end