Skip to content
New issue

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

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

Already on GitHub? Sign in to your account

Add a default for NLLS #277

Merged
merged 5 commits into from
Nov 6, 2023
Merged
Show file tree
Hide file tree
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
Expand Up @@ -16,7 +16,9 @@ GeneralKlement
## Polyalgorithms

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

Expand Down
12 changes: 4 additions & 8 deletions docs/src/solvers/NonlinearLeastSquaresSolvers.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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
Expand Up @@ -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,
Expand Down
8 changes: 5 additions & 3 deletions ext/NonlinearSolveLeastSquaresOptimExt.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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))

Expand Down
3 changes: 2 additions & 1 deletion src/NonlinearSolve.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
55 changes: 49 additions & 6 deletions src/default.jl
Original file line number Diff line number Diff line change
Expand Up @@ -250,14 +250,57 @@
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...)

Check warning on line 300 in src/default.jl

View check run for this annotation

Codecov / codecov/patch

src/default.jl#L299-L300

Added lines #L299 - L300 were not covered by tests
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
Expand Up @@ -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]
Expand Down
9 changes: 5 additions & 4 deletions test/nonlinear_least_squares.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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

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

for prob in nlls_problems, solver in solvers
Expand All @@ -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

Expand All @@ -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
Loading