SciML · avik-pal · May 24, 2024 · May 24, 2024
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "NonlinearSolve"
 uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 authors = ["SciML"]
-version = "3.12.0"
+version = "3.12.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -32,6 +32,7 @@ TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 
 [weakdeps]
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
+Enlsip = "d5306a6b-d590-428d-a53a-eb3bb2d36f2d"
 FastLevenbergMarquardt = "7a0df574-e128-4d35-8cbd-3d84502bf7ce"
 FixedPointAcceleration = "817d07cb-a79a-5c30-9a31-890123675176"
 LeastSquaresOptim = "0fc2ff8b-aaa3-5acd-a817-1944a5e08891"
@@ -45,6 +46,7 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
 NonlinearSolveBandedMatricesExt = "BandedMatrices"
+NonlinearSolveEnlsipExt = "Enlsip"
 NonlinearSolveFastLevenbergMarquardtExt = "FastLevenbergMarquardt"
 NonlinearSolveFixedPointAccelerationExt = "FixedPointAcceleration"
 NonlinearSolveLeastSquaresOptimExt = "LeastSquaresOptim"
@@ -65,7 +67,9 @@ BenchmarkTools = "1.4"
 CUDA = "5.2"
 ConcreteStructs = "0.2.3"
 DiffEqBase = "6.149.0"
+Enlsip = "0.9"
 Enzyme = "0.12"
+ExplicitImports = "1.4.4"
 FastBroadcast = "0.2.8"
 FastClosures = "0.3.2"
 FastLevenbergMarquardt = "0.1"
@@ -115,7 +119,9 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+Enlsip = "d5306a6b-d590-428d-a53a-eb3bb2d36f2d"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
 FastLevenbergMarquardt = "7a0df574-e128-4d35-8cbd-3d84502bf7ce"
 FixedPointAcceleration = "817d07cb-a79a-5c30-9a31-890123675176"
 LeastSquaresOptim = "0fc2ff8b-aaa3-5acd-a817-1944a5e08891"
@@ -139,4 +145,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "BandedMatrices", "BenchmarkTools", "CUDA", "Enzyme", "FastLevenbergMarquardt", "FixedPointAcceleration", "LeastSquaresOptim", "MINPACK", "ModelingToolkit", "NLSolvers", "NLsolve", "NaNMath", "NonlinearProblemLibrary", "OrdinaryDiffEq", "Pkg", "Random", "ReTestItems", "SIAMFANLEquations", "SpeedMapping", "StableRNGs", "StaticArrays", "Sundials", "Symbolics", "Test", "Zygote"]
+test = ["Aqua", "BandedMatrices", "BenchmarkTools", "CUDA", "Enlsip", "Enzyme", "ExplicitImports", "FastLevenbergMarquardt", "FixedPointAcceleration", "LeastSquaresOptim", "MINPACK", "ModelingToolkit", "NLSolvers", "NLsolve", "NaNMath", "NonlinearProblemLibrary", "OrdinaryDiffEq", "Pkg", "Random", "ReTestItems", "SIAMFANLEquations", "SpeedMapping", "StableRNGs", "StaticArrays", "Sundials", "Symbolics", "Test", "Zygote"]
diff --git a/ext/NonlinearSolveBandedMatricesExt.jl b/ext/NonlinearSolveBandedMatricesExt.jl
@@ -1,6 +1,9 @@
 module NonlinearSolveBandedMatricesExt
 
-using BandedMatrices, LinearAlgebra, NonlinearSolve, SparseArrays
+using BandedMatrices: BandedMatrix
+using LinearAlgebra: Diagonal
+using NonlinearSolve: NonlinearSolve
+using SparseArrays: sparse
 
 # This is used if we vcat a Banded Jacobian with a Diagonal Matrix in Levenberg
 @inline NonlinearSolve._vcat(B::BandedMatrix, D::Diagonal) = vcat(sparse(B), D)

diff --git a/ext/NonlinearSolveEnlsipExt.jl 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/ext/NonlinearSolveFastLevenbergMarquardtExt.jl b/ext/NonlinearSolveFastLevenbergMarquardtExt.jl
@@ -1,10 +1,13 @@
 module NonlinearSolveFastLevenbergMarquardtExt
 
-using ArrayInterface, NonlinearSolve, SciMLBase
-import ConcreteStructs: @concrete
-import FastClosures: @closure
-import FastLevenbergMarquardt as FastLM
-import StaticArraysCore: SArray
+using ArrayInterface: ArrayInterface
+using FastClosures: @closure
+using FastLevenbergMarquardt: FastLevenbergMarquardt
+using NonlinearSolve: NonlinearSolve, FastLevenbergMarquardtJL
+using SciMLBase: SciMLBase, NonlinearLeastSquaresProblem, NonlinearProblem, ReturnCode
+using StaticArraysCore: SArray
+
+const FastLM = FastLevenbergMarquardt
 
 @inline function _fast_lm_solver(::FastLevenbergMarquardtJL{linsolve}, x) where {linsolve}
     if linsolve === :cholesky

diff --git a/ext/NonlinearSolveFixedPointAccelerationExt.jl b/ext/NonlinearSolveFixedPointAccelerationExt.jl
@@ -1,6 +1,8 @@
 module NonlinearSolveFixedPointAccelerationExt
 
-using NonlinearSolve, FixedPointAcceleration, SciMLBase
+using NonlinearSolve: NonlinearSolve, FixedPointAccelerationJL
+using SciMLBase: SciMLBase, NonlinearProblem, ReturnCode
+using FixedPointAcceleration: FixedPointAcceleration, fixed_point
 
 function SciMLBase.__solve(prob::NonlinearProblem, alg::FixedPointAccelerationJL, args...;
         abstol = nothing, maxiters = 1000, alias_u0::Bool = false,

diff --git a/ext/NonlinearSolveLeastSquaresOptimExt.jl b/ext/NonlinearSolveLeastSquaresOptimExt.jl
@@ -1,8 +1,11 @@
 module NonlinearSolveLeastSquaresOptimExt
 
-using NonlinearSolve, SciMLBase
-import ConcreteStructs: @concrete
-import LeastSquaresOptim as LSO
+using ConcreteStructs: @concrete
+using LeastSquaresOptim: LeastSquaresOptim
+using NonlinearSolve: NonlinearSolve, LeastSquaresOptimJL, TraceMinimal
+using SciMLBase: SciMLBase, NonlinearLeastSquaresProblem, NonlinearProblem, ReturnCode
+
+const LSO = LeastSquaresOptim
 
 @inline function _lso_solver(::LeastSquaresOptimJL{alg, ls}) where {alg, ls}
     linsolve = ls === :qr ? LSO.QR() :

diff --git a/ext/NonlinearSolveMINPACKExt.jl b/ext/NonlinearSolveMINPACKExt.jl
@@ -1,7 +1,9 @@
 module NonlinearSolveMINPACKExt
 
-using MINPACK, NonlinearSolve, SciMLBase
-import FastClosures: @closure
+using MINPACK: MINPACK
+using NonlinearSolve: NonlinearSolve, CMINPACK
+using SciMLBase: SciMLBase, NonlinearLeastSquaresProblem, NonlinearProblem, ReturnCode
+using FastClosures: @closure
 
 function SciMLBase.__solve(
         prob::Union{NonlinearLeastSquaresProblem, NonlinearProblem}, alg::CMINPACK,

diff --git a/ext/NonlinearSolveNLSolversExt.jl b/ext/NonlinearSolveNLSolversExt.jl
@@ -1,7 +1,13 @@
 module NonlinearSolveNLSolversExt
 
-using ADTypes, FastClosures, NonlinearSolve, NLSolvers, SciMLBase, LinearAlgebra
-using FiniteDiff, ForwardDiff
+using ADTypes: ADTypes, AutoFiniteDiff, AutoForwardDiff, AutoPolyesterForwardDiff
+using FastClosures: @closure
+using FiniteDiff: FiniteDiff
+using ForwardDiff: ForwardDiff
+using LinearAlgebra: norm
+using NLSolvers: NLSolvers, NEqOptions, NEqProblem
+using NonlinearSolve: NonlinearSolve, NLSolversJL
+using SciMLBase: SciMLBase, NonlinearProblem, ReturnCode
 
 function SciMLBase.__solve(prob::NonlinearProblem, alg::NLSolversJL, args...;
         abstol = nothing, reltol = nothing, maxiters = 1000,

diff --git a/ext/NonlinearSolveNLsolveExt.jl b/ext/NonlinearSolveNLsolveExt.jl
@@ -1,6 +1,8 @@
 module NonlinearSolveNLsolveExt
 
-using NonlinearSolve, NLsolve, SciMLBase
+using NonlinearSolve: NonlinearSolve, NLsolveJL, TraceMinimal
+using NLsolve: NLsolve, OnceDifferentiable, nlsolve
+using SciMLBase: SciMLBase, NonlinearProblem, ReturnCode
 
 function SciMLBase.__solve(
         prob::NonlinearProblem, alg::NLsolveJL, args...; abstol = nothing,

diff --git a/ext/NonlinearSolveSIAMFANLEquationsExt.jl b/ext/NonlinearSolveSIAMFANLEquationsExt.jl
@@ -1,7 +1,10 @@
 module NonlinearSolveSIAMFANLEquationsExt
 
-using NonlinearSolve, SIAMFANLEquations, SciMLBase
-import FastClosures: @closure
+using FastClosures: @closure
+using NonlinearSolve: NonlinearSolve, SIAMFANLEquationsJL
+using SciMLBase: SciMLBase, NonlinearProblem, ReturnCode
+using SIAMFANLEquations: SIAMFANLEquations, aasol, nsol, nsoli, nsolsc, ptcsol, ptcsoli,
+                         ptcsolsc, secant
 
 @inline function __siam_fanl_equations_retcode_mapping(sol)
     if sol.errcode == 0

diff --git a/ext/NonlinearSolveSpeedMappingExt.jl b/ext/NonlinearSolveSpeedMappingExt.jl
@@ -1,6 +1,8 @@
 module NonlinearSolveSpeedMappingExt
 
-using NonlinearSolve, SciMLBase, SpeedMapping
+using NonlinearSolve: NonlinearSolve, SpeedMappingJL
+using SciMLBase: SciMLBase, NonlinearProblem, ReturnCode
+using SpeedMapping: speedmapping
 
 function SciMLBase.__solve(prob::NonlinearProblem, alg::SpeedMappingJL, args...;
         abstol = nothing, maxiters = 1000, alias_u0::Bool = false,

diff --git a/ext/NonlinearSolveSymbolicsExt.jl b/ext/NonlinearSolveSymbolicsExt.jl
@@ -1,6 +1,6 @@
 module NonlinearSolveSymbolicsExt
 
-import NonlinearSolve, Symbolics
+using NonlinearSolve: NonlinearSolve
 
 NonlinearSolve.is_extension_loaded(::Val{:Symbolics}) = true
 

diff --git a/ext/NonlinearSolveZygoteExt.jl b/ext/NonlinearSolveZygoteExt.jl
@@ -1,6 +1,6 @@
 module NonlinearSolveZygoteExt
 
-import NonlinearSolve, Zygote
+using NonlinearSolve: NonlinearSolve
 
 NonlinearSolve.is_extension_loaded(::Val{:Zygote}) = true
 

diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl
@@ -4,37 +4,58 @@ if isdefined(Base, :Experimental) && isdefined(Base.Experimental, Symbol("@max_m
     @eval Base.Experimental.@max_methods 1
 end
 
-import Reexport: @reexport
-import PrecompileTools: @recompile_invalidations, @compile_workload, @setup_workload
+using Reexport: @reexport
+using PrecompileTools: @recompile_invalidations, @compile_workload, @setup_workload
 
 @recompile_invalidations begin
-    using ADTypes, ConcreteStructs, DiffEqBase, FastBroadcast, FastClosures, LazyArrays,
-          LinearAlgebra, LinearSolve, MaybeInplace, Preferences, Printf, SciMLBase,
-          SimpleNonlinearSolve, SparseArrays, SparseDiffTools
-
-    import ArrayInterface: ArrayInterface, undefmatrix, can_setindex, restructure,
-                           fast_scalar_indexing, ismutable
-    import DiffEqBase: AbstractNonlinearTerminationMode,
-                       AbstractSafeNonlinearTerminationMode,
-                       AbstractSafeBestNonlinearTerminationMode,
-                       NonlinearSafeTerminationReturnCode, get_termination_mode
-    import FiniteDiff
-    import ForwardDiff
-    import ForwardDiff: Dual
-    import LineSearches
-    import LinearSolve: ComposePreconditioner, InvPreconditioner, needs_concrete_A
-    import RecursiveArrayTools: recursivecopy!, recursivefill!
-
-    import SciMLBase: AbstractNonlinearAlgorithm, JacobianWrapper, AbstractNonlinearProblem,
-                      AbstractSciMLOperator, NLStats, _unwrap_val, has_jac, isinplace
-    import SparseDiffTools: AbstractSparsityDetection
-    import StaticArraysCore: StaticArray, SVector, SArray, MArray, Size, SMatrix, MMatrix
-    import SymbolicIndexingInterface: SymbolicIndexingInterface, ParameterIndexingProxy,
-                                      symbolic_container, parameter_values, state_values,
-                                      getu
+    using ADTypes: AutoFiniteDiff, AutoForwardDiff, AutoPolyesterForwardDiff, AutoZygote,
+                   AutoEnzyme, AutoSparse
+    # FIXME: deprecated, remove in future
+    using ADTypes: AutoSparseFiniteDiff, AutoSparseForwardDiff,
+                   AutoSparsePolyesterForwardDiff, AutoSparseZygote
+
+    using ArrayInterface: ArrayInterface, can_setindex, restructure, fast_scalar_indexing,
+                          ismutable
+    using ConcreteStructs: @concrete
+    using DiffEqBase: DiffEqBase, AbstractNonlinearTerminationMode,
+                      AbstractSafeBestNonlinearTerminationMode, AbsNormTerminationMode,
+                      AbsSafeBestTerminationMode, AbsSafeTerminationMode,
+                      AbsTerminationMode, NormTerminationMode, RelNormTerminationMode,
+                      RelSafeBestTerminationMode, RelSafeTerminationMode,
+                      RelTerminationMode, SimpleNonlinearSolveTerminationMode,
+                      SteadyStateDiffEqTerminationMode
+    using FastBroadcast: @..
+    using FastClosures: @closure
+    using FiniteDiff: FiniteDiff
+    using ForwardDiff: ForwardDiff, Dual
+    using LazyArrays: LazyArrays, ApplyArray, cache
+    using LinearAlgebra: LinearAlgebra, ColumnNorm, Diagonal, I, LowerTriangular, Symmetric,
+                         UpperTriangular, axpy!, cond, diag, diagind, dot, issuccess,
+                         istril, istriu, lu, mul!, norm, pinv, tril!, triu!
+    using LineSearches: LineSearches
+    using LinearSolve: LinearSolve, LUFactorization, QRFactorization, ComposePreconditioner,
+                       InvPreconditioner, needs_concrete_A
+    using MaybeInplace: @bb
+    using Printf: @printf
+    using Preferences: Preferences, @load_preference, @set_preferences!
+    using RecursiveArrayTools: recursivecopy!, recursivefill!
+    using SciMLBase: AbstractNonlinearAlgorithm, JacobianWrapper, AbstractNonlinearProblem,
+                     AbstractSciMLOperator, _unwrap_val, has_jac, isinplace
+    using SparseArrays: AbstractSparseMatrix, SparseMatrixCSC
+    using SparseDiffTools: SparseDiffTools, AbstractSparsityDetection,
+                           ApproximateJacobianSparsity, JacPrototypeSparsityDetection,
+                           NoSparsityDetection, PrecomputedJacobianColorvec,
+                           SymbolicsSparsityDetection, auto_jacvec, auto_jacvec!,
+                           auto_vecjac, init_jacobian, num_jacvec, num_jacvec!, num_vecjac,
+                           num_vecjac!, sparse_jacobian, sparse_jacobian!,
+                           sparse_jacobian_cache
+    using StaticArraysCore: StaticArray, SVector, SArray, MArray, Size, SMatrix
+    using SymbolicIndexingInterface: SymbolicIndexingInterface, ParameterIndexingProxy,
+                                     symbolic_container, parameter_values, state_values,
+                                     getu
 end
 
-@reexport using ADTypes, SciMLBase, SimpleNonlinearSolve
+@reexport using SciMLBase, SimpleNonlinearSolve
 
 # Type-Inference Friendly Check for Extension Loading
 is_extension_loaded(::Val) = false
@@ -141,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
@@ -166,4 +187,11 @@ export SteadyStateDiffEqTerminationMode, SimpleNonlinearSolveTerminationMode,
 # Tracing Functionality
 export TraceAll, TraceMinimal, TraceWithJacobianConditionNumber
 
+# Reexport ADTypes
+export AutoFiniteDiff, AutoForwardDiff, AutoPolyesterForwardDiff, AutoZygote, AutoEnzyme,
+       AutoSparse
+# FIXME: deprecated, remove in future
+export AutoSparseFiniteDiff, AutoSparseForwardDiff, AutoSparsePolyesterForwardDiff,
+       AutoSparseZygote
+
 end # module
diff --git a/src/abstract_types.jl b/src/abstract_types.jl
@@ -311,12 +311,12 @@ function returns_norm_form_damping(f::F) where {F}
 end
 
 """
-    AbstractNonlinearSolveOperator <: SciMLBase.AbstractSciMLOperator
+    AbstractNonlinearSolveOperator <: AbstractSciMLOperator
 
 NonlinearSolve.jl houses a few custom operators. These will eventually be moved out but till
 then this serves as the abstract type for them.
 """
-abstract type AbstractNonlinearSolveOperator{T} <: SciMLBase.AbstractSciMLOperator{T} end
+abstract type AbstractNonlinearSolveOperator{T} <: AbstractSciMLOperator{T} end
 
 # Approximate Jacobian Algorithms
 """