Add tests

Author: avik-pal

Project.toml

@@ -50,7 +50,7 @@ FiniteDiff = "2"
 ForwardDiff = "0.10.3"
 LeastSquaresOptim = "0.8"
 LineSearches = "7"
-LinearAlgebra = "1.9"
+LinearAlgebra = "<0.0.1, 1"
 LinearSolve = "2.12"
 NonlinearProblemLibrary = "0.1"
 PrecompileTools = "1"
@@ -59,7 +59,7 @@ Reexport = "0.2, 1"
 SciMLBase = "2.8.2"
 SciMLOperators = "0.3"
 SimpleNonlinearSolve = "0.1.23"
-SparseArrays = "1.9"
+SparseArrays = "<0.0.1, 1"
 SparseDiffTools = "2.12"
 StaticArraysCore = "1.4"
 UnPack = "1.0"

src/gaussnewton.jl

@@ -29,23 +29,30 @@ for large-scale and numerically-difficult nonlinear least squares problems.
   - `linesearch`: the line search algorithm to use. Defaults to [`LineSearch()`](@ref),
     which means that no line search is performed. Algorithms from `LineSearches.jl` can be
     used here directly, and they will be converted to the correct `LineSearch`.
+  - `vjp_autodiff`: Automatic Differentiation Backend used for vector-jacobian products.
+    This is applicable if the linear solver doesn't require a concrete jacobian, for eg.,
+    Krylov Methods. Defaults to `nothing`, which means if the problem is out of place and
+    `Zygote` is loaded then, we use `AutoZygote`. In all other, cases `FiniteDiff` is used.
 """
 @concrete struct GaussNewton{CJ, AD} <: AbstractNewtonAlgorithm{CJ, AD}
     ad::AD
     linsolve
     precs
     linesearch
+    vjp_autodiff
 end
 
 function set_ad(alg::GaussNewton{CJ}, ad) where {CJ}
-    return GaussNewton{CJ}(ad, alg.linsolve, alg.precs, alg.linesearch)
+    return GaussNewton{CJ}(ad, alg.linsolve, alg.precs, alg.linesearch, alg.vjp_autodiff)
 end
 
 function GaussNewton(; concrete_jac = nothing, linsolve = nothing,
-        linesearch = LineSearch(), precs = DEFAULT_PRECS, adkwargs...)
+        linesearch = LineSearch(), precs = DEFAULT_PRECS, vjp_autodiff = nothing,
+        adkwargs...)
     ad = default_adargs_to_adtype(; adkwargs...)
     linesearch = linesearch isa LineSearch ? linesearch : LineSearch(; method = linesearch)
-    return GaussNewton{_unwrap_val(concrete_jac)}(ad, linsolve, precs, linesearch)
+    return GaussNewton{_unwrap_val(concrete_jac)}(ad, linsolve, precs, linesearch,
+        vjp_autodiff)
 end
 
 @concrete mutable struct GaussNewtonCache{iip} <: AbstractNonlinearSolveCache{iip}

src/jacobian.jl

@@ -2,6 +2,9 @@
     JᵀJ
     Jᵀ
 end
+
+SciMLBase.isinplace(JᵀJ::KrylovJᵀJ) = isinplace(JᵀJ.Jᵀ)
+
 sparsity_detection_alg(_, _) = NoSparsityDetection()
 function sparsity_detection_alg(f, ad::AbstractSparseADType)
     if f.sparsity === nothing
@@ -67,12 +70,10 @@ function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f::F, u, p, ::Val
         jac_cache = nothing
     end
 
-    # FIXME: To properly support needsJᵀJ without Jacobian, we need to implement
-    #        a reverse diff operation with the seed being `Jx`, this is not yet implemented
-    J = if !(linsolve_needs_jac || alg_wants_jac || needsJᵀJ)
+    J = if !(linsolve_needs_jac || alg_wants_jac)
         if f.jvp === nothing
             # We don't need to construct the Jacobian
-            JacVec(uf, u; autodiff = __get_nonsparse_ad(alg.ad))
+            JacVec(uf, u; fu, autodiff = __get_nonsparse_ad(alg.ad))
         else
             if iip
                 jvp = (_, u, v) -> (du = similar(fu); f.jvp(du, v, u, p); du)
@@ -96,9 +97,9 @@ function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f::F, u, p, ::Val
     du = _mutable_zero(u)
 
     if needsJᵀJ
-        # TODO: Pass in `jac_transpose_autodiff`
-        JᵀJ, Jᵀfu = __init_JᵀJ(J, _vec(fu), uf, u;
-            jac_autodiff = __get_nonsparse_ad(alg.ad))
+        JᵀJ, Jᵀfu = __init_JᵀJ(J, _vec(fu), uf, u; f,
+            vjp_autodiff = __get_nonsparse_ad(_getproperty(alg, Val(:vjp_autodiff))),
+            jvp_autodiff = __get_nonsparse_ad(alg.ad))
     end
 
     if linsolve_init
@@ -141,26 +142,29 @@ function __init_JᵀJ(J::StaticArray, fu, args...; kwargs...)
     JᵀJ = MArray{Tuple{size(J, 2), size(J, 2)}, eltype(J)}(undef)
     return JᵀJ, J' * fu
 end
-function __init_JᵀJ(J::FunctionOperator, fu, uf, u, args...;
-        jac_transpose_autodiff = nothing, jac_autodiff = nothing, kwargs...)
-    autodiff = __concrete_jac_transpose_autodiff(jac_transpose_autodiff, jac_autodiff, uf)
-    Jᵀ = VecJac(uf, u; autodiff)
+function __init_JᵀJ(J::FunctionOperator, fu, uf, u, args...; f = nothing,
+        vjp_autodiff = nothing, jvp_autodiff = nothing, kwargs...)
+    # FIXME: Proper fix to this requires the FunctionOperator patch
+    if f !== nothing && f.vjp !== nothing
+        @warn "Currently we don't make use of user provided `jvp`. This is planned to be \
+               fixed in the near future."
+    end
+    autodiff = __concrete_vjp_autodiff(vjp_autodiff, jvp_autodiff, uf)
+    Jᵀ = VecJac(uf, u; fu, autodiff)
     JᵀJ_op = SciMLOperators.cache_operator(Jᵀ * J, u)
     JᵀJ = KrylovJᵀJ(JᵀJ_op, Jᵀ)
     Jᵀfu = Jᵀ * fu
     return JᵀJ, Jᵀfu
 end
 
-SciMLBase.isinplace(JᵀJ::KrylovJᵀJ) = isinplace(JᵀJ.Jᵀ)
-
-function __concrete_jac_transpose_autodiff(jac_transpose_autodiff, jac_autodiff, uf)
-    if jac_transpose_autodiff === nothing
+function __concrete_vjp_autodiff(vjp_autodiff, jvp_autodiff, uf)
+    if vjp_autodiff === nothing
         if isinplace(uf)
             # VecJac can be only FiniteDiff
             return AutoFiniteDiff()
         else
             # Short circuit if we see that FiniteDiff was used for J computation
-            jac_autodiff isa AutoFiniteDiff && return jac_autodiff
+            jvp_autodiff isa AutoFiniteDiff && return jvp_autodiff
             # Check if Zygote is loaded then use Zygote else use FiniteDiff
             if haskey(Base.loaded_modules,
                 Base.PkgId(Base.UUID("e88e6eb3-aa80-5325-afca-941959d7151f"), "Zygote"))
@@ -170,7 +174,13 @@ function __concrete_jac_transpose_autodiff(jac_transpose_autodiff, jac_autodiff,
             end
         end
     else
-        return __get_nonsparse_ad(jac_transpose_autodiff)
+        ad = __get_nonsparse_ad(vjp_autodiff)
+        if isinplace(uf) && ad isa AutoZygote
+            @warn "Attempting to use Zygote.jl for linesearch on an in-place problem. \
+                Falling back to finite differencing."
+            return AutoFiniteDiff()
+        end
+        return ad
     end
 end
 

src/linesearch.jl

@@ -131,6 +131,10 @@ function LineSearchesJLCache(ls::LineSearch, f::F, u, p, fu1, IIP::Val{iip}) whe
     end
 
     function g!(u, fu)
+        if f.jvp !== nothing
+            @warn "Currently we don't make use of user provided `jvp` in linesearch. This \
+                   is planned to be fixed in the near future." maxlog=1
+        end
         op = VecJac(SciMLBase.JacobianWrapper(f, p), u; fu = fu1, autodiff)
         if iip
             mul!(g₀, op, fu)

src/trustRegion.jl

@@ -141,9 +141,12 @@ for large-scale and numerically-difficult nonlinear systems.
     `expand_threshold < r` (with `r` defined in `shrink_threshold`). Defaults to `2.0`.
   - `max_shrink_times`: the maximum number of times to shrink the trust region radius in a
     row, `max_shrink_times` is exceeded, the algorithm returns. Defaults to `32`.
+  - `vjp_autodiff`: Automatic Differentiation Backend used for vector-jacobian products.
+    This is applicable if the linear solver doesn't require a concrete jacobian, for eg.,
+    Krylov Methods. Defaults to `nothing`, which means if the problem is out of place and
+    `Zygote` is loaded then, we use `AutoZygote`. In all other, cases `FiniteDiff` is used.
 """
-@concrete struct TrustRegion{CJ, AD, MTR} <:
-                 AbstractNewtonAlgorithm{CJ, AD}
+@concrete struct TrustRegion{CJ, AD, MTR} <: AbstractNewtonAlgorithm{CJ, AD}
     ad::AD
     linsolve
     precs
@@ -156,25 +159,27 @@ for large-scale and numerically-difficult nonlinear systems.
     shrink_factor::MTR
     expand_factor::MTR
     max_shrink_times::Int
+    vjp_autodiff
 end
 
 function set_ad(alg::TrustRegion{CJ}, ad) where {CJ}
     return TrustRegion{CJ}(ad, alg.linsolve, alg.precs, alg.radius_update_scheme,
         alg.max_trust_radius, alg.initial_trust_radius, alg.step_threshold,
         alg.shrink_threshold, alg.expand_threshold, alg.shrink_factor, alg.expand_factor,
-        alg.max_shrink_times)
+        alg.max_shrink_times, alg.vjp_autodiff)
 end
 
 function TrustRegion(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS,
         radius_update_scheme::RadiusUpdateSchemes.T = RadiusUpdateSchemes.Simple, #defaults to conventional radius update
         max_trust_radius::Real = 0 // 1, initial_trust_radius::Real = 0 // 1,
         step_threshold::Real = 1 // 10000, shrink_threshold::Real = 1 // 4,
         expand_threshold::Real = 3 // 4, shrink_factor::Real = 1 // 4,
-        expand_factor::Real = 2 // 1, max_shrink_times::Int = 32, adkwargs...)
+        expand_factor::Real = 2 // 1, max_shrink_times::Int = 32, vjp_autodiff = nothing,
+        adkwargs...)
     ad = default_adargs_to_adtype(; adkwargs...)
     return TrustRegion{_unwrap_val(concrete_jac)}(ad, linsolve, precs, radius_update_scheme,
         max_trust_radius, initial_trust_radius, step_threshold, shrink_threshold,
-        expand_threshold, shrink_factor, expand_factor, max_shrink_times)
+        expand_threshold, shrink_factor, expand_factor, max_shrink_times, vjp_autodiff)
 end
 
 @concrete mutable struct TrustRegionCache{iip, trustType, floatType} <:
@@ -422,7 +427,7 @@ function retrospective_step!(cache::TrustRegionCache)
     @unpack H, g, du = cache
 
     return -(get_loss(fu_prev) - get_loss(fu)) /
-           (dot(du, g) + dot(du, H, du) / 2)
+           (dot(_vec(du), _vec(g)) + __lr_mul(Val(isinplace(cache)), H, _vec(du)) / 2)
 end
 
 function trust_region_step!(cache::TrustRegionCache)

src/utils.jl

@@ -305,6 +305,12 @@ end
 __issingular(x::AbstractMatrix{T}) where {T} = cond(x) > inv(sqrt(eps(real(T))))
 __issingular(x) = false ## If SciMLOperator and such
 
+# Safe getproperty
+@generated function _getproperty(s::S, ::Val{X}) where {S, X}
+    hasfield(S, X) && return :(s.$X)
+    return :(nothing)
+end
+
 # If factorization is LU then perform that and update the linsolve cache
 # else check if the matrix is singular
 function _try_factorize_and_check_singular!(linsolve, X)

test/basictests.jl

@@ -142,41 +142,52 @@ end
 # --- TrustRegion tests ---
 
 @testset "TrustRegion" begin
-    function benchmark_nlsolve_oop(f, u0, p = 2.0; radius_update_scheme, kwargs...)
+    function benchmark_nlsolve_oop(f, u0, p = 2.0; radius_update_scheme, linsolve = nothing,
+            vjp_autodiff = nothing, kwargs...)
         prob = NonlinearProblem{false}(f, u0, p)
-        return solve(prob, TrustRegion(; radius_update_scheme); abstol = 1e-9, kwargs...)
+        return solve(prob, TrustRegion(; radius_update_scheme, linsolve, vjp_autodiff);
+            abstol = 1e-9, kwargs...)
     end
 
-    function benchmark_nlsolve_iip(f, u0, p = 2.0; radius_update_scheme, kwargs...)
+    function benchmark_nlsolve_iip(f, u0, p = 2.0; radius_update_scheme, linsolve = nothing,
+            vjp_autodiff = nothing, kwargs...)
         prob = NonlinearProblem{true}(f, u0, p)
-        return solve(prob, TrustRegion(; radius_update_scheme); abstol = 1e-9, kwargs...)
+        return solve(prob, TrustRegion(; radius_update_scheme, linsolve, vjp_autodiff);
+            abstol = 1e-9, kwargs...)
     end
 
     radius_update_schemes = [RadiusUpdateSchemes.Simple, RadiusUpdateSchemes.NocedalWright,
         RadiusUpdateSchemes.NLsolve, RadiusUpdateSchemes.Hei, RadiusUpdateSchemes.Yuan,
         RadiusUpdateSchemes.Fan, RadiusUpdateSchemes.Bastin]
     u0s = ([1.0, 1.0], @SVector[1.0, 1.0], 1.0)
+    linear_solvers = [nothing, LUFactorization(), KrylovJL_GMRES()]
 
-    @testset "[OOP] u0: $(typeof(u0)) radius_update_scheme: $(radius_update_scheme)" for u0 in u0s,
-        radius_update_scheme in radius_update_schemes
+    @testset "[OOP] u0: $(typeof(u0)) radius_update_scheme: $(radius_update_scheme) linear_solver: $(linsolve)" for u0 in u0s,
+        radius_update_scheme in radius_update_schemes, linsolve in linear_solvers
+
+        !(u0 isa Array) && linsolve !== nothing && continue
+
+        abstol = ifelse(linsolve isa KrylovJL, 1e-6, 1e-9)
 
-        sol = benchmark_nlsolve_oop(quadratic_f, u0; radius_update_scheme)
+        sol = benchmark_nlsolve_oop(quadratic_f, u0; radius_update_scheme, linsolve, abstol)
         @test SciMLBase.successful_retcode(sol)
-        @test all(abs.(sol.u .* sol.u .- 2) .< 1e-9)
+        @test all(abs.(sol.u .* sol.u .- 2) .< abstol)
 
         cache = init(NonlinearProblem{false}(quadratic_f, u0, 2.0),
-            TrustRegion(; radius_update_scheme); abstol = 1e-9)
+            TrustRegion(; radius_update_scheme, linsolve); abstol)
         @test (@ballocated solve!($cache)) < 200
     end
 
-    @testset "[IIP] u0: $(typeof(u0)) radius_update_scheme: $(radius_update_scheme)" for u0 in ([
-            1.0, 1.0],), radius_update_scheme in radius_update_schemes
-        sol = benchmark_nlsolve_iip(quadratic_f!, u0; radius_update_scheme)
+    @testset "[IIP] u0: $(typeof(u0)) radius_update_scheme: $(radius_update_scheme) linear_solver: $(linsolve)" for u0 in ([
+            1.0, 1.0],), radius_update_scheme in radius_update_schemes, linsolve in linear_solvers
+        abstol = ifelse(linsolve isa KrylovJL, 1e-6, 1e-9)
+        sol = benchmark_nlsolve_iip(quadratic_f!, u0; radius_update_scheme, linsolve,
+            abstol)
         @test SciMLBase.successful_retcode(sol)
-        @test all(abs.(sol.u .* sol.u .- 2) .< 1e-9)
+        @test all(abs.(sol.u .* sol.u .- 2) .< abstol)
 
         cache = init(NonlinearProblem{true}(quadratic_f!, u0, 2.0),
-            TrustRegion(; radius_update_scheme); abstol = 1e-9)
+            TrustRegion(; radius_update_scheme); abstol)
         @test (@ballocated solve!($cache)) ≤ 64
     end
 

test/nonlinear_least_squares.jl

@@ -1,4 +1,4 @@
-using NonlinearSolve, LinearSolve, LinearAlgebra, Test, Random, ForwardDiff
+using NonlinearSolve, LinearSolve, LinearAlgebra, Test, Random, ForwardDiff, Zygote
 import FastLevenbergMarquardt, LeastSquaresOptim
 
 true_function(x, θ) = @. θ[1] * exp(θ[2] * x) * cos(θ[3] * x + θ[4])
@@ -27,9 +27,16 @@ prob_iip = NonlinearLeastSquaresProblem(NonlinearFunction(loss_function;
         resid_prototype = zero(y_target)), θ_init, x)
 
 nlls_problems = [prob_oop, prob_iip]
-solvers = vec(Any[GaussNewton(; linsolve, linesearch)
-                  for linsolve in [nothing, LUFactorization()],
-linesearch in [Static(), BackTracking(), HagerZhang(), StrongWolfe(), MoreThuente()]])
+solvers = []
+for linsolve in [nothing, LUFactorization(), KrylovJL_GMRES()]
+    vjp_autodiffs = linsolve isa KrylovJL ? [nothing, AutoZygote(), AutoFiniteDiff()] :
+                    [nothing]
+    for linesearch in [Static(), BackTracking(), HagerZhang(), StrongWolfe(), MoreThuente()],
+        vjp_autodiff in vjp_autodiffs
+
+        push!(solvers, GaussNewton(; linsolve, linesearch, vjp_autodiff))
+    end
+end
 append!(solvers,
     [
         LevenbergMarquardt(),
@@ -45,6 +52,36 @@ for prob in nlls_problems, solver in solvers
     @test norm(sol.resid) < 1e-6
 end
 
+# This is just for testing that we can use vjp provided by the user
+function vjp(v, θ, p)
+    resid = zeros(length(p))
+    J = ForwardDiff.jacobian((resid, θ) -> loss_function(resid, θ, p), resid, θ)
+    return vec(v' * J)
+end
+
+function vjp!(Jv, v, θ, p)
+    resid = zeros(length(p))
+    J = ForwardDiff.jacobian((resid, θ) -> loss_function(resid, θ, p), resid, θ)
+    mul!(vec(Jv), v', J)
+    return nothing
+end
+
+probs = [
+    NonlinearLeastSquaresProblem(NonlinearFunction{true}(loss_function;
+            resid_prototype = zero(y_target), vjp = vjp!), θ_init, x),
+    NonlinearLeastSquaresProblem(NonlinearFunction{false}(loss_function;
+            resid_prototype = zero(y_target), vjp = vjp), θ_init, x),
+]
+
+for prob in probs, solver in solvers
+    !(solver isa GaussNewton) && continue
+    !(solver.linsolve isa KrylovJL) && continue
+    @test_warn "Currently we don't make use of user provided `jvp`. This is planned to be \
+    fixed in the near future." sol=solve(prob, solver; maxiters = 10000, abstol = 1e-8)
+    sol = solve(prob, solver; maxiters = 10000, abstol = 1e-8)
+    @test norm(sol.resid) < 1e-6
+end
+
 function jac!(J, θ, p)
     resid = zeros(length(p))
     ForwardDiff.jacobian!(J, (resid, θ) -> loss_function(resid, θ, p), resid, θ)