Add Enlsip Extension

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

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

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

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
 

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)