diff --git a/ext/NonlinearSolveEnlsipExt.jl b/ext/NonlinearSolveEnlsipExt.jl new file mode 100644 index 000000000..0e9123808 --- /dev/null +++ b/ext/NonlinearSolveEnlsipExt.jl @@ -0,0 +1,67 @@ +module NonlinearSolveEnlsipExt + +using FastClosures: @closure +using NonlinearSolve: NonlinearSolve, EnlsipJL +using SciMLBase: SciMLBase, NonlinearLeastSquaresProblem, ReturnCode +using Enlsip: Enlsip + +function SciMLBase.__solve(prob::NonlinearLeastSquaresProblem, alg::EnlsipJL, args...; + abstol = nothing, reltol = nothing, maxiters = 1000, + alias_u0::Bool = false, maxtime = nothing, show_trace::Val{ST} = Val(false), + termination_condition = nothing, kwargs...) where {ST} + NonlinearSolve.__test_termination_condition(termination_condition, :EnlsipJL) + + f, u0, resid = NonlinearSolve.__construct_extension_f( + prob; alias_u0, can_handle_oop = Val(true), force_oop = Val(true)) + + f_aug = @closure u -> begin + u_ = view(u, 1:(length(u) - 1)) + r = f(u_) + return vcat(r, u[length(u)]) + end + + eq_cons = u -> [u[end]] + + abstol = NonlinearSolve.DEFAULT_TOLERANCE(abstol, eltype(u0)) + reltol = NonlinearSolve.DEFAULT_TOLERANCE(reltol, eltype(u0)) + + maxtime = maxtime === nothing ? typemax(eltype(u0)) : maxtime + + jac_fn = NonlinearSolve.__construct_extension_jac( + prob, alg, u0, resid; alg.autodiff, can_handle_oop = Val(true)) + + n = length(u0) + 1 + m = length(resid) + 1 + + jac_fn_aug = @closure u -> begin + u_ = view(u, 1:(length(u) - 1)) + J = jac_fn(u_) + J_full = similar(u, (m, n)) + J_full[1:(m - 1), 1:(n - 1)] .= J + fill!(J_full[1:(m - 1), n], false) + fill!(J_full[m, 1:(n - 1)], false) + return J_full + end + + u0 = vcat(u0, 0.0) + u_low = [eltype(u0)(ifelse(i == 1, -Inf, 0)) for i in 1:length(u0)] + u_up = [eltype(u0)(ifelse(i == 1, Inf, 0)) for i in 1:length(u0)] + + model = Enlsip.CnlsModel( + f_aug, n, m; starting_point = u0, jacobian_residuals = jac_fn_aug, + x_low = u_low, x_upp = u_up, nb_eqcons = 1, eq_constraints = eq_cons) + Enlsip.solve!(model; max_iter = maxiters, time_limit = maxtime, silent = !ST, + abs_tol = abstol, rel_tol = reltol, x_tol = reltol) + + sol_u = Enlsip.solution(model) + resid = Enlsip.sum_sq_residuals(model) + + status = Enlsip.status(model) + retcode = status === :found_first_order_stationary_point ? ReturnCode.Success : + status === :maximum_iterations_exceeded ? ReturnCode.MaxIters : + status === :time_limit_exceeded ? ReturnCode.MaxTime : ReturnCode.Failure + + return SciMLBase.build_solution(prob, alg, sol_u, resid; retcode, original = model) +end + +end diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index b92aac597..8db1ea0d2 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -162,8 +162,8 @@ export NonlinearSolvePolyAlgorithm, RobustMultiNewton, FastShortcutNonlinearPoly FastShortcutNLLSPolyalg # Extension Algorithms -export LeastSquaresOptimJL, FastLevenbergMarquardtJL, CMINPACK, NLsolveJL, NLSolversJL, - FixedPointAccelerationJL, SpeedMappingJL, SIAMFANLEquationsJL +export LeastSquaresOptimJL, FastLevenbergMarquardtJL, CMINPACK, EnlsipJL, NLsolveJL, + NLSolversJL, FixedPointAccelerationJL, SpeedMappingJL, SIAMFANLEquationsJL # Advanced Algorithms -- Without Bells and Whistles export GeneralizedFirstOrderAlgorithm, ApproximateJacobianSolveAlgorithm, GeneralizedDFSane diff --git a/src/algorithms/extension_algs.jl b/src/algorithms/extension_algs.jl index ea9ab79ab..795e3797f 100644 --- a/src/algorithms/extension_algs.jl +++ b/src/algorithms/extension_algs.jl @@ -484,3 +484,36 @@ function SIAMFANLEquationsJL(; method = :newton, delta = 1e-3, linsolve = nothin end return SIAMFANLEquationsJL(method, delta, linsolve, m, beta, autodiff) end + +""" + EnlsipJL(; autodiff = nothing) + +Wrapper over [Enlsip.jl](https://github.com/UncertainLab/Enlsip.jl) for solving Nonlinear +Least Squares Problems. + +### 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 + `nothing` which means that a default is selected according to the problem specification! + +!!! note + + This algorithm is only available if `Enlsip.jl` is installed. + +!!! warning + + Enlsip is designed for constrained NLLS problems. However, since we don't support + constraints in NonlinearSolve.jl currently, we add a dummy constraint to the problem + before calling Enlsip. +""" +@concrete struct EnlsipJL <: AbstractNonlinearSolveExtensionAlgorithm + autodiff +end + +function EnlsipJL(; autodiff = nothing) + if Base.get_extension(@__MODULE__, :NonlinearSolveEnlsipExt) === nothing + error("EnlsipJL requires Enlsip.jl to be loaded") + end + return EnlsipJL(autodiff) +end diff --git a/test/misc/qa_tests.jl b/test/misc/qa_tests.jl index a3b370b5b..cab3f2ee0 100644 --- a/test/misc/qa_tests.jl +++ b/test/misc/qa_tests.jl @@ -16,7 +16,7 @@ end @testitem "Explicit Imports" tags=[:misc] begin using NonlinearSolve, ADTypes, SimpleNonlinearSolve, SciMLBase - import BandedMatrices, FastLevenbergMarquardt, FixedPointAcceleration, + import BandedMatrices, Enlsip, FastLevenbergMarquardt, FixedPointAcceleration, LeastSquaresOptim, MINPACK, NLsolve, NLSolvers, SIAMFANLEquations, SpeedMapping, Symbolics, Zygote diff --git a/test/wrappers/nlls_tests.jl b/test/wrappers/nlls_tests.jl index 53cea758d..1d004df81 100644 --- a/test/wrappers/nlls_tests.jl +++ b/test/wrappers/nlls_tests.jl @@ -1,7 +1,7 @@ @testsetup module WrapperNLLSSetup using Reexport @reexport using LinearAlgebra, StableRNGs, StaticArrays, Random, ForwardDiff, Zygote -import FastLevenbergMarquardt, LeastSquaresOptim, MINPACK +import FastLevenbergMarquardt, LeastSquaresOptim, MINPACK, Enlsip true_function(x, θ) = @. θ[1] * exp(θ[2] * x) * cos(θ[3] * x + θ[4]) true_function(y, x, θ) = (@. y = θ[1] * exp(θ[2] * x) * cos(θ[3] * x + θ[4])) @@ -46,7 +46,7 @@ end end end -@testitem "FastLevenbergMarquardt.jl + CMINPACK: Jacobian Provided" setup=[WrapperNLLSSetup] tags=[:wrappers] begin +@testitem "FastLevenbergMarquardt.jl + CMINPACK + Einsip: Jacobian Provided" setup=[WrapperNLLSSetup] tags=[:wrappers] begin function jac!(J, θ, p) resid = zeros(length(p)) ForwardDiff.jacobian!(J, (resid, θ) -> loss_function(resid, θ, p), resid, θ) @@ -71,13 +71,14 @@ end solvers = Any[FastLevenbergMarquardtJL(linsolve) for linsolve in (:cholesky, :qr)] Sys.isapple() || push!(solvers, CMINPACK()) + push!(solvers, EnlsipJL()) for solver in solvers, prob in probs sol = solve(prob, solver; maxiters = 10000, abstol = 1e-8) @test maximum(abs, sol.resid) < 1e-6 end end -@testitem "FastLevenbergMarquardt.jl + CMINPACK: Jacobian Not Provided" setup=[WrapperNLLSSetup] tags=[:wrappers] begin +@testitem "FastLevenbergMarquardt.jl + CMINPACK + Einsip: Jacobian Not Provided" setup=[WrapperNLLSSetup] tags=[:wrappers] begin probs = [ NonlinearLeastSquaresProblem( NonlinearFunction{true}(loss_function; resid_prototype = zero(y_target)), @@ -92,6 +93,9 @@ end autodiff in (nothing, AutoForwardDiff(), AutoFiniteDiff())]) Sys.isapple() || append!(solvers, [CMINPACK(; method) for method in (:auto, :lm, :lmdif)]) + append!(solvers, + [EnlsipJL(; autodiff) + for autodiff in (nothing, AutoForwardDiff(), AutoFiniteDiff())]) for solver in solvers, prob in probs sol = solve(prob, solver; maxiters = 10000, abstol = 1e-8)